MAZ3: Memory-Assisted ZeRO-3 for Efficient Collective Communication

Large Language Models (LLMs) have advanced rapidly, but their growing parameter scales and memory demands pose critical challenges for distributed training. Although GPU memory capacity improves steadily, model sizes expand much faster, causing frequent Out-of-Memory (OOM) errors and rising training costs. Existing memory optimization approaches, such as ZeRO-3 and offloading, alleviate per-GPU memory pressure but introduce excessive collective communication, limited computation–communication overlap, and degraded scalability. We present MAZ3, a distributed training framework that mitigates these limitations through three key techniques: (1) Collaborative CPU–GPU memory management, storing full parameters in CPU memory and broadcasting them within nodes to reduce global synchronization; (2) Fine-grained communication–computation overlap, aligning collective operations with model computation to hide latency; (3) Hierarchical aggregation operators, leveraging intra-node NVLink and inter-node NIC channels concurrently to minimize communication overhead. We implement MAZ3 on a multi-GPU cluster and evaluate it with large-scale models. Results show that MAZ3 reduces inter-node model communication (gradients and parameters) by 33%, improves training efficiency by 40.3%, and increases throughput by 67.9% compared with ZeRO-3. Moreover, MAZ3 retains the memory efficiency of ZeRO-3 while approaching the training efficiency and throughput of ZeRO-2 Offload, achieving a balanced trade-off between memory optimization and performance.