Cannot see effectiveness of Zero 3. Any help is much appreciated.

[hpourmodheji]

Hello Community, DeepSpeed Team and @tjruwase,

I tried to train large bing BERT (w/ ~300M parameters) on 8 GPUs (12GB GeForce GTX 1080 Ti) to see memory reduction with ZeRO 3. It shows linear memory reduction (8x) on the model state (from ~4400MB to ~550MB); however, the max of the memory consumption does not have the same decrease during the training. So what am I maybe missing here that I do not see the effectiveness of ZeRO 3? Any help from you is very much appreciated, as understanding this issue is a road blocker for my research.

Here are the images:

Baseline - GPU Memory Consumption during Training with single GPU:

ZeRO 3 - GPU Memory Consumption during Training on 8 GPUs with ZeRO 3 enabled:

Best,

Hossein

[tjruwase]

@hpourmodheji, what tool are using to generate this memory usage?

Can you please profile memory usage in your code using see_memory_usage(), similar to here? Can you share the logs from that?

[hpourmodheji]

Thanks for your comment, @tjruwase.

I was using tensorboard to generate the memory profiling and visualize it.

I repeated the same experiment using "see_memory_usage()" as you requested. Here are the logs of the investigation:

see_memory_usage_1GPU_No_ZeRO.txt
see_memory_usage_8GPUs_ZeRO3.txt

[hpourmodheji]

Here are some samples of the "see_memory_usage()" results:

Baseline - GPU Memory Consumption during Training with single GPU:

[2022-01-06 16:03:30,144] [INFO] [utils.py:806:see_memory_usage] before forward 5
[2022-01-06 16:03:30,146] [INFO] [utils.py:811:see_memory_usage] MA 4.38 GB Max_MA 7.5 GB CA 9.38 GB Max_CA 9 GB
[2022-01-06 16:03:30,147] [INFO] [utils.py:816:see_memory_usage] CPU Virtual Memory: used = 12.58 GB, percent = 5.0%
[2022-01-06 16:03:30,300] [INFO] [utils.py:806:see_memory_usage] before backward 5
[2022-01-06 16:03:30,301] [INFO] [utils.py:811:see_memory_usage] MA 5.8 GB Max_MA 5.8 GB CA 9.38 GB Max_CA 9 GB
[2022-01-06 16:03:30,302] [INFO] [utils.py:816:see_memory_usage] CPU Virtual Memory: used = 12.6 GB, percent = 5.0%
[2022-01-06 16:03:30,384] [INFO] [utils.py:806:see_memory_usage] before optimizer 5
[2022-01-06 16:03:30,385] [INFO] [utils.py:811:see_memory_usage] MA 5.01 GB Max_MA 5.86 GB CA 9.38 GB Max_CA 9 GB
[2022-01-06 16:03:30,385] [INFO] [utils.py:816:see_memory_usage] CPU Virtual Memory: used = 12.64 GB, percent = 5.0%

ZeRO 3 - GPU Memory Consumption during Training on 8 GPUs with ZeRO 3 enabled:

[2022-01-06 16:09:57,287] [INFO] [utils.py:806:see_memory_usage] before forward 5
[2022-01-06 16:09:57,289] [INFO] [utils.py:811:see_memory_usage] MA 0.55 GB Max_MA 0.64 GB CA 4.57 GB Max_CA 5 GB
[2022-01-06 16:09:57,290] [INFO] [utils.py:816:see_memory_usage] CPU Virtual Memory: used = 44.86 GB, percent = 17.8%
[2022-01-06 16:09:58,444] [INFO] [utils.py:806:see_memory_usage] before backward 5
[2022-01-06 16:09:58,445] [INFO] [utils.py:811:see_memory_usage] MA 1.97 GB Max_MA 1.97 GB CA 4.57 GB Max_CA 5 GB
[2022-01-06 16:09:58,445] [INFO] [utils.py:816:see_memory_usage] CPU Virtual Memory: used = 44.86 GB, percent = 17.8%
[2022-01-06 16:10:00,608] [INFO] [utils.py:806:see_memory_usage] before optimizer 5
[2022-01-06 16:10:00,609] [INFO] [utils.py:811:see_memory_usage] MA 0.63 GB Max_MA 4.04 GB CA 4.57 GB Max_CA 5 GB
[2022-01-06 16:10:00,609] [INFO] [utils.py:816:see_memory_usage] CPU Virtual Memory: used = 44.86 GB, percent = 17.8%

---

Training Stages	Baseline - GPU Memory Consumption during Training with single GPU	ZeRO 3 - GPU Memory Consumption during Training on 8 GPUs with ZeRO 3 enabled	Memory consumption reduction
before forward	MA 4.38 GB	MA 0.55 GB	~8x
before backward	MA 5.8 GB	MA 1.97 GB	~3x
before optimizer	MA 5.01 GB	MA 0.63 G	~8x
before forward	Max_MA 7.5 GB	Max_MA 0.64 GB	~12x
before backward	Max_MA 5.8 GB	Max_MA 1.97 GB	~3x
before optimizer	Max_MA 5.86 GB	Max_MA 4.04 GB	~1.5x

As summarized in the above table:

• by considering memory allocated numbers, we have a 3x memory reduction during training instead of having an 8x reduction.
• by considering maximum memory allocated numbers, we have a 1.5x memory reduction during training instead of having an 8x reduction.

Is this a correct understanding? Using 8 GPUs and ZeRO3, can we achieve an 8x memory consumption reduction in practice?

[tjruwase]

@hpourmodheji, thanks for sharing this table. There are two things to consider.

1. ZeRO-3 provides memory savings for the model and optimizer state, which is what you see as the "before forward" usage. This is because ZeRO-3 partitions those states across the available GPUs. The usage at other points represents consumption by input data, activations, gradients, intermediate buffers, and fragmentation resulting from PyTorch computation. ZeRO3 does not affect the memory consumed by these buffers.

2. PyTorch memory manager maintains a reserve of memory that is not always used for the computation but shows up the memory usage statistics as described here.

So, it is more precise to say that ZeRO3 linearly reduces memory consumption of model and optimizer.

[hpourmodheji]

@tjruwase, thanks for the clarification. Do you have any formulation for memory reduction during training using a ZeRO3 optimizer? Is a 1.5x reduction during training a known/reasonable number?