DeToNATION: Decoupled Torch Network-Aware Training on Interlinked Online Nodes

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: From, Mogens Henrik, Nielsen, Jacob, Poech, Lukas Galke, Schneider-Kamp, Peter
Format: Preprint
Veröffentlicht: 2025
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866918204621717504
author From, Mogens Henrik
Nielsen, Jacob
Poech, Lukas Galke
Schneider-Kamp, Peter
author_facet From, Mogens Henrik
Nielsen, Jacob
Poech, Lukas Galke
Schneider-Kamp, Peter
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2502_06728
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle DeToNATION: Decoupled Torch Network-Aware Training on Interlinked Online Nodes
From, Mogens Henrik
Nielsen, Jacob
Poech, Lukas Galke
Schneider-Kamp, Peter
Machine Learning
Artificial Intelligence
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.
title DeToNATION: Decoupled Torch Network-Aware Training on Interlinked Online Nodes
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2502.06728