[BUG] "model parallel size" is always 1 with zero3 enabled

[i4never]

Describe the bug
I'm using zero stage3 with optimizer & parameter offloading. The memory used by each gpu should decrease if more gpu is used. (which is not happening). After adding flops_profiler to ds_config, model parallel size remains 1 and params per gpu is not dropping.

• 1 node 1 gpu

world size:                                                   1       
data parallel size:                                           1       
model parallel size:                                          1       
batch size per GPU:                                           2       
params per gpu:                                               83.81 M 
params of model = params per GPU * mp_size:                   83.81 M

• 1 node 4 gpu

world size:                                                   4       
data parallel size:                                           4       
model parallel size:                                          1       
batch size per GPU:                                           2       
params per gpu:                                               83.81 M 
params of model = params per GPU * mp_size:                   83.81 M 

To Reproduce
NA

Expected behavior
Training with bigger batch_size will not cause cuda OOM.
model parallel size printed by flops_profiler should equal to world size.

ds_report output

--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-devel package with yum
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_adam ............... [YES] ...... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.0
 [WARNING]  using untested triton version (2.0.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/data/apps/anaconda3/envs/py38_test1/lib/python3.8/site-packages/torch']
torch version .................... 2.0.1+cu117
deepspeed install path ........... ['/data/apps/anaconda3/envs/py38_test1/lib/python3.8/site-packages/deepspeed']
deepspeed info ................... 0.9.5, unknown, unknown
torch cuda version ............... 11.7
torch hip version ................ None
nvcc version ..................... 11.6
deepspeed wheel compiled w. ...... torch 2.0, cuda 11.7

System info (please complete the following information):

• Ubuntu 18.04
• one machines with x8 A100s
• Python 3.8

Launcher context

python -u -m deepspeed.launcher.launch

Docker context
NO

Additional context
NA

[tjruwase]

@i4never, zero stage 3 is a form of data parallelism, and not model parallelism. DeepSpeed does not implement model parallelism but is compatible with existing forms like tensor slicing and pipeline parallelism.

However, zero stage 3 should reduce per-gpu memory consumption of model parameters and optimizer state. Offloading should also reduce gpu memory consumption by moving parameters and optimizer to CPU memory. Can you share more details of your experience such as logs and repro scripts?

[joshuafc]

@i4never, zero stage 3 is a form of data parallelism, and not model parallelism. DeepSpeed does not implement model parallelism but is compatible with existing forms like tensor slicing and pipeline parallelism.

However, zero stage 3 should reduce per-gpu memory consumption of model parameters and optimizer state. Offloading should also reduce gpu memory consumption by moving parameters and optimizer to CPU memory. Can you share more details of your experience such as logs and repro scripts?

https://www.deepspeed.ai/tutorials/zero/ says "Stage 3: The 16-bit model parameters are partitioned across the processes.".

So the partition of model parameters is not implemented by using Tensor Parallelism or Pipeline Parallelism?

I have been troubled by this doubt for a long time. There is a significant difference in inter-GPU bandwidth between nodes and within a node.

If I have a 4 node cluster. I except the Zero-3 stage will split the model in to 4 block for pipeline parallel, and for each block, the corresponding node will use Tensor Parallel to partition the block parameters. But I can't find a easy sample for how to config zero-3 to specify how to partitioned model parameters.

[tjruwase]

Like I said earlier zero stage 3 is data parallelism, and not tensor or pipeline parallelism. It is a different kind of partitioning from those model parallelism approaches. You might find the following descriptions useful:

1. Comparison of zero data parallelism, tensor parallelism, and pipeline parallelism: https://huggingface.co/blog/bloom-megatron-deepspeed#megatron-deepspeed
2. Animation of zero-infinity, which is an extension of zero stage 3 to use CPU/NVMe memory: https://www.microsoft.com/en-us/research/blog/zero-infinity-and-deepspeed-unlocking-unprecedented-model-scale-for-deep-learning-training/

[noobmaster29]

So using zero stage 3 would not allow large models to be spread across more GPUS? The way it is worded in some articles and tutorials makes it sound like adding more gpus would allow training of overall larger models. Are there any resources you would recommend for model parallelism?

[tjruwase]

@noobmaster29, zero stage 3 does partition large models across available GPUs to enable model size increase. The model partitioning approach is just different from the model parallelism approaches of tensor slicing and pipeline parallelism. Have you taken a look at the above link that compares zero stage 3 with model parallelism? The differences are a bit subtle, so I suggest you play with scaling model sizes with zero stage 3 on a number of GPUs for hands-on experience.

In terms of model parallelism links, here are some great implementations of tensor slicing and pipeline parallelism:

1. Megatron-LM
2. Megratron-DeepSpeed
3. BigScience