arXiv:2502.06728cs.LGcs.AI2025-02AAAI

提出新型分布式训练框架,显著降低通信开销且保持模型精度。

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

  • 将模型参数分片到多加速器,仅同步梯度中快速变化部分。
  • 在语言与视觉任务上达到与全梯度同步相当的验证损失,速度更快。
  • 适合超大规模模型训练,尤其适用于通信瓶颈严重的场景。

训练大型神经网络需要大量计算资源,通常分布在多个节点和加速器上。最新研究表明,仅交换梯度中的快速变化部分并本地累积动量(解耦动量,DeMo)即可,但该方法假设模型可放入单个加速器。本文放宽此假设,提出 FlexDeMo:各节点在本地将模型参数完全分片于不同加速器之间,同时通过仅同步快速变化组件而非完整梯度来减少节点间通信,形成一种混合分片数据并行训练策略。进一步提出 DeToNATION 框架,统一 DeMo、FlexDeMo 及 DiLoCo 等主流分布式训练方法,引入新的复制策略变体,挑战 DeMo 中的既有选择。跨语言与视觉领域的实验表明,FlexDeMo 在验证损失上与采用 AdamW 和全梯度同步的混合分片数据并行相当,但速度显著提升,是超大规模模型训练的有前景方案。

原文摘要 · Abstract (English)

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.

分布式训练模型并行通信优化大模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。