DeToNATION: Decoupled Torch Network-Aware Training on Interlinked Online Nodes
This work addresses the computational bottleneck in distributed training for large-scale machine learning models, offering a more efficient method for researchers and practitioners.
The paper tackles the problem of training large neural networks across multiple nodes by introducing FlexDeMo, a hybrid sharded data parallel strategy that reduces inter-node communication by synchronizing only fast-moving gradient components, achieving similar validation loss as full gradient synchronization methods while being substantially faster.
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, DeMo assumes that models fit on a single accelerator. We relax this assumption and introduce FlexDeMo, whereby nodes fully shard model parameters locally between different accelerators, while inter-node communication is reduced by synchronizing only fast-moving components instead of the full gradients -- resulting in a hybrid sharded data parallel training strategy. We further introduce a framework, denoted as DeToNATION, that generalizes DeMo, FlexDeMo, and other popular distributed training schemes such as DiLoCo -- introducing new variations of replication schemes and challenging choices made in DeMo. Our results across language and vision domains show that FlexDeMo attains similar validation loss as hybrid sharded data parallel training employing AdamW and full gradient synchronization, while being substantially faster. FlexDeMo is thus a promising distributed training scheme for the largest machine learning models.