About deepspeed ZeRO-3 training

[Fpw19980724]

Do you have any scripts for deepspeed ZeRO-3 training？

[Alichua]

Hi Tom Holland,

We currently provide DeepSpeed training scripts configured with ZeRO-2 (ds_config_zero2.json).

For our current use case (small-scale models and datasets), ZeRO-2 offers a better balance of speed and stability. Since the model parameters fit comfortably in GPU memory, ZeRO-2 avoids the heavy communication overhead associated with ZeRO-3's parameter sharding, resulting in faster training throughput.

Though, our trainer code is fully compatible with ZeRO-3. If you need to train larger models that exceed single-GPU memory, you can simply switch to a ZeRO-3 config file (e.g., ds_config_zero3.json) without changing the training script. We may include an official ZeRO-3 example in the future for large-scale scenarios.

Best Regards.

[Fpw19980724]

Glad to receive your reply so quickly. It seems there's a problem with my ZeRO-3 configuration that's preventing normal training. By the way, have you tried using other frameworks that can train larger models with longer sequences and more parameters? For example, Megatron etc.
Thanks a lot!

[Alichua]

Thanks for asking,

To answer shortly: Since we focus entirely on small-scale setups, we haven't touched Megatron or run ZeRO-3 in production, so we haven't encountered (or fixed) the specific edge cases that come up with those.

Though TinyEngram's code is theoretically compatible, without real-world testing at that scale, we can't reliably debug those errors for you.

Hence, for ZeRO-3 or Megatron issues, we assume the DeepSpeed community and docs would be much more reliable than our suggestions.

Best,
Chua