← All projects

Distributed Training with ZeRO — 2D-Parallel Communication from Scratch

Implemented the communication layer of a 2D-parallel training pipeline combining data parallelism and tensor model parallelism (naive and Megatron-style), extended with ZeRO-style optimizer-state sharding, built from scratch on raw MPI collectives.

MPI PyTorch Distributed Training ZeRO
ML SYSTEMS / COMMUNICATION04 / SYSTEM STUDY
Shard the model. Share the work.Two data-parallel groups each contain two tensor-parallel ranks with sharded weights. Horizontal communication combines tensor outputs. Vertical communication all-gathers parameters and reduce-scatters gradients; each rank holds a local optimizer-state shard.DATA GROUP 0DATA GROUP 1R0R1R2R3tensor-parallel collectivestensor-parallel collectivesW / ∇ / AdamW / ∇ / AdamW / ∇ / AdamW / ∇ / Adamall-gather parameters · reduce-scatter gradients
Shard the model. Share the work.Tensor parallel × data parallel

Overview

The communication layer of a 2D-parallel training pipeline, built from scratch directly on raw MPI collectives rather than a higher-level distributed training library. It combines data parallelism with tensor model parallelism — both a naive version and a Megatron-LM-style version — and extends the setup with ZeRO-style optimizer-state sharding to cut memory usage.

Continue to update… soon.