FlexDeMo: Decoupled Momentum Optimization for Hybrid Sharded Data Parallel Training
Training large neural network models requires extensive computational resources, often distributed across several nodes and accelerators. Recent findings suggest that it may be sufficient to only exchange the fast moving components of the gradients, while accumulating momentum locally (Decoupled Momentum, or DeMo). However, when considering larger models that do not fit on a single accelerator, the exchange of gradient information and the integration of DeMo needs to be reconsidered. Here, we propose employing a hybrid sharded data parallel training strategy, FlexDeMo, whereby nodes fully shard model parameters locally between different accelerators, while inter-node communication bandwidth requirements are reduced by synchronizing only fast-moving components instead of the full gradients. This effectively combines previous hybrid sharded strategies with the advantages of decoupled momentum. Our experimental results show that FlexDeMo is on par with hybrid sharded data parallel training employing AdamW and full gradient synchronization in terms of validation loss, demonstrating its viability. Furthermore, FlexDeMo achieves improved training speed compared to full gradient synchronization across nodes. In a bandwidth-constrained 2-node setup, FlexDeMo allows reaching desired levels of validation loss faster than hybrid sharded data parallel training with full gradient synchronization.
Code (0)
등록된 구현이 없습니다.
Methods 이 논문이 사용한 방법론
Similar Papers 제목 키워드 기반
Dion: Distributed Orthonormalized Updates
Recent work has shown that orthonormal matrix updates speed up neural network optimization, improve training stability, and offer better hyperparameter transfer across model sizes. Applying these updates efficiently when…
NOVAK: Unified adaptive optimizer for deep neural networks
This work introduces NOVAK, a modular gradient-based optimization algorithm that integrates adaptive moment estimation, rectified learning-rate scheduling, decoupled weight regularization, multiple variants of Nesterov m…
Scaling Muon for Diffusion Transformers
The matrix-aware optimizer Muon improves large model training by balancing updates across singular directions, yet its scaling behavior and end-to-end efficiency on large Diffusion Transformers (DiTs) remain unclear. We …
SDP4Bit: Toward 4-bit Communication Quantization in Sharded Data Parallelism for LLM Training
Recent years have witnessed a clear trend towards language models with an ever-increasing number of parameters, as well as the growing training overhead and memory usage. Distributed training, particularly through Sharde…
QuantizationLion Secretly Solves Constrained Optimization: As Lyapunov Predicts
Lion (Evolved Sign Momentum), a new optimizer discovered through program search, has shown promising results in training large AI models. It performs comparably or favorably to AdamW but with greater memory efficiency. A…