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
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.