Improving memory utilization of Z2+MoE

[siddharth9820]

Summary

While running experiments, I found that an inordinate amount of memory is used in the gradient upscaling step for high expert to GPU ratios (like 1 or 1/2).

This PR does two things:

• It creates an upper limit on the expert parameter group size (see deepspeed/moe/utils.py).
• Instead of first upscaling all the gradients and running the optimizer once, this PR makes it such that a combined upscaling+optimizer step is executed on each parameter group one by one. Thus at a given time only fp32 gradients of 1 parameter group can exist in GPU memory. (see deepspeed/runtime/zero/stage_1_and_2.py)

Highlight Result

Prior to this PR a 6.7B base model with 16 experts ran OOM on 32 A100 GPUs (40GB).
With the changes, I am able to run the same model with a peak memory utilization of 31.3 GB. Thus at the bare minimum we are saving 21.75% memory for this model.

Train loss curve for reference

Sanity Checks

Train loss curves before and after the changes match

Batch Times

To create the most pathological case, I set global batch size to 8. Yet, there is no penalty in batch times.

Memory Consumption

The memory saved is expected to increase with increasing model sizes.

[siddharth9820]

@tjruwase, please let me know how we would want to offer the user a way to set the max_group_size

[tjruwase]

As discussed, let's add an moe section in ds_config. Perhaps, @awan-10 could help with the design.

[siddharth9820]

@awan-10 Perhaps the moe section can also contain a flag to toggle expert slicing too?