PaperlayerView exact revision on arXiv

Source-linked AI summary

Pretraining A Large Language Model using Distributed GPUs: A Memory-Efficient Decentralized Paradigm

Jinrui Zhang, Chaodong Xiao, Aoqi Wu, Xindong Zhang, Lei Zhang

arXiv:2602.11543v3cs.CL

TL;DR

Centralized LLM pretraining requires substantial GPU memory and high-bandwidth communication, while existing decentralized methods still train full models on each node. SPES addresses this gap by training and synchronizing only subsets of MoE experts, with expert-merging warm-up to improve early sparse training. It achieves competitive or matching centralized performance across 2B, 7B, and 9B-scale models while reducing communication and memory demands.

  • Problem

    Existing decentralized LLM training reduces communication requirements but still updates the full model on each node, leaving GPU memory constraints unresolved.

  • Method

    SPES assigns each node a distinct subset of MoE experts, synchronizes trained experts periodically, and merges similar experts during early training.

  • Results

    Up to 33.3% communication cost reduction was reported, while 7B and 9B models matched centralized counterparts under similar data and compute resources.

  • Takeaways & Limitations

    SPES enables MoE LLM pretraining across weakly connected, geographically distributed GPU clusters with lower infrastructure demands and performance on par with comparable centralized baselines.

  • Takeaways & Limitations

    Evaluation was limited to a 9B-parameter model trained on fewer than 500B tokens, so scalability to larger models and longer training remains unvalidated.

Abstract

from arXiv · show

Pretraining large language models (LLMs) typically requires centralized clusters with thousands of high-memory GPUs (e.g., H100/A100). Recent decentralized training methods reduce communication overhead by employing federated optimization; however, they still need to train the entire model on each node, remaining constrained by GPU memory limitations. In this work, we propose SParse Expert Synchronization (SPES), a memory-efficient decentralized framework for pretraining mixture-of-experts (MoE) LLMs. SPES trains only a subset of experts per node, substantially lowering the memory footprint. Each node updates its local experts and periodically synchronizes with other nodes, eliminating full-parameter transmission while ensuring efficient knowledge sharing. To mitigate limited per-expert data utilization under sparse expert updates, we introduce an expert-merging warm-up strategy, where experts exchange knowledge early in training, to rapidly establish foundational capabilities. With SPES, we train a 2B-parameter MoE LLM using 16 standalone 48GB GPUs over internet connections, which achieves competitive performance with centrally trained LLMs under similar computational budgets. We further demonstrate scalability by training a 7B model from scratch and a 9B model upcycled from a dense checkpoint, both of which match prior centralized baselines. Our code is available at https://github.com/zjr2000/SPES.

1. Introduction

Centralized LLM pretraining demands high-memory GPUs and fast interconnects, while existing decentralized approaches reduce bandwidth but still train full models per node. SPES assigns subsets of MoE experts to nodes, adds expert-merging warm-up, and demonstrates competitive multi-scale training with lower resource costs.

  • Motivation: Centralized LLM pretraining is resource-intensive because it requires substantial GPU memory and low-latency, high-bandwidth inter-device communication.These requirements support storing parameters, activations, optimizer states, and gradients and enable model or data parallelism.
  • Motivation: Existing decentralized methods reduce bandwidth through intermittent federated synchronization but still require every node to update the full model.This preserves a substantial per-node memory footprint despite relaxed communication constraints.
  • SPES: SPES trains distinct subsets of MoE experts on different nodes while freezing unassigned experts, reducing per-device memory and communication overhead.Nodes periodically synchronize trained experts rather than transmitting entire model weights.
  • Empirical validation: A 2B-parameter MoE LLM was trained on 16 standalone 48GB NVIDIA L40S GPUs over the internet with performance comparable to centrally trained models under comparable computational budgets.The evaluation also covers 7B and 9B parameter scales, including training from scratch and continual pretraining regimes.
  • Empirical validation: Up to 33.3% communication-cost reduction was reported, while 7B and 9B models matched centralized counterparts trained with similar data and compute resources.The 9B model was upcycled from a strong dense initialization, and the framework used a custom gRPC-based cross-node synchronization protocol.
  • SPES: SPES adds an expert-merging warm-up strategy that periodically aggregates similar experts during early training to strengthen representations under sparse updates.The strategy addresses limited token utilization by allowing experts to exchange knowledge early in training.

2. Related Work

Prior decentralized LLM training reduces communication through federated optimization, while memory-efficient pretraining mainly relies on sharding and parallelism over tightly coupled accelerators. SPES extends these ideas with cross-node expert sharding and sparse synchronization for heterogeneous, low-bandwidth settings.

  • Decentralized Training: Decentralized LLM pretraining methods such as DiLoCo and Photon use FedAvg to achieve comparable perplexities with substantially lower communication cost than centralized training.Later work improves communication efficiency through new optimizers and architectures tailored to decentralized settings.
  • Memory-Efficient Pretraining: Memory-efficient pretraining methods commonly partition optimizer states, gradients, parameters, or computation using data, pipeline, tensor, and expert parallelism.These approaches primarily target tightly coupled accelerator environments.
  • SPES: SPES distributes MoE experts, gradients, and optimizer states across geographically heterogeneous nodes that communicate only necessary updates.It targets single-GPU nodes and low-bandwidth interconnects where intra-node sharding is infeasible, while complementing existing parallelism paradigms.

3. Memory-Efficient Decentralized Pretraining

SPES partitions MoE expert training across decentralized nodes, reducing per-node memory and communication while periodically sharing expert and shared-module updates. An expert-merging warm-up improves token utilization during sparse training by sharing knowledge among similar experts.

  • Framework overview: SPES partitions expert training across weakly connected nodes and intermittently synchronizes weights, reducing memory usage and communication overhead.Each node trains only its assigned experts while participating in decentralized synchronization.
  • Expert assignment and local training: Each node optimizes its assigned experts and shared parameters while keeping unassigned experts fixed during local updates.The expert parameters are partitioned into disjoint subsets across nodes.
  • Sparse synchronization: During synchronization, shared parameters are aggregated and only updated assigned experts and shared parameters are transmitted.This avoids transmitting the entire model during each synchronization round.
  • Expert-merging warm-up: Sparse expert updates leave tokens routed to frozen experts without gradient contributions, lowering token utilization relative to equivalent centralized training.This motivates the expert-merging warm-up strategy.
  • Expert-merging warm-up: During the initial Tw steps, similar experts are periodically merged with decaying strength α, allowing each expert to benefit from gradients from multiple nodes.Similarity is computed from input projection layers, and merging is restricted to early training to preserve later specialization.
  • Efficiency analysis: With AdamW, SPES reduces per-node memory from 4 × (|ψ| + |Φ|) to 4 × |ψ| + |Φ| + 3 × |Φi|.The reduction follows from storing optimizer states and gradients only for shared parameters and assigned experts.

4. Experiments

Experiments evaluate SPES across model scales, training regimes, memory and communication costs, throughput, benchmark performance, and expert-merging ablations. Results show competitive quality with substantially lower resource requirements and scalable performance.

  • Experimental Setup: 2B and 7B models are trained from scratch, while the 9B model is upcycled from a dense initialization.The 7B setup uses four compute nodes with eight NVIDIA A800 GPUs each; the 9B model is initialized from Qwen3-1.7B-Base.
  • Memory Cost Comparison: SPES keeps 2B-model per-GPU memory under 40GB on 16 nodes without sharding, whereas centralized training and DiLoCo require more than 50GB per GPU.SPES achieves this through sparse training, with each node updating only a subset of parameters.
  • Communication Cost Comparison: 65% lower uplink communication results when training a 7B model on four nodes: 9.8GB per SPES node per round versus 28.6GB for DiLoCo and centralized training.SPES synchronizes only shared modules and assigned experts rather than the full parameter set.
  • Training Speed Comparison: 3.67k tokens/s per GPU with SPES at H = 50 compares with 3.79k tokens/s for centralized training on stronger RDMA-connected hardware.Reducing synchronization frequency can further improve SPES throughput.
  • Comparison with Previous Training Paradigms: Across multiple benchmarks, SPES rapidly converges after a slightly slower initial learning curve and ultimately matches or outperforms centralized training and DiLoCo.The controlled comparison uses 1B-parameter MoE models trained on 50B tokens with SlimPajama; Figure 4 evaluates training performance using the OLMo evaluation suite.
  • Performance Comparison with Existing LLMs: SPES-2B and SPES-7B achieve competitive commonsense-reasoning results among models with comparable activated parameter scales, while SPES-9B is competitive with comparable state-of-the-art models using fewer than 500B tokens.The 9B model was stopped early because of resource constraints, while metrics were still improving.

5. Conclusion

SPES provides a decentralized, memory-efficient MoE pretraining paradigm that reduces device and communication demands while matching comparable centralized performance across several scales. Its evaluation remains limited in model size, token count, and task scope.

  • SPES assigns distinct expert subsets to nodes and synchronizes them, reducing per-device memory usage and communication overhead.An expert-merging warm-up strategy improves token utilization per expert and accelerates early convergence.
  • 2B- and 7B-parameter MoE models achieve performance on par with comparable centralized baselines in geographically distributed GPU clusters.The framework also scales to upcycling a 9B model from a dense initialization.
  • Evaluation is limited to a 9B-parameter model trained on fewer than 500B tokens, leaving larger models and longer training durations for future validation.The study also focuses on language understanding rather than multimodal reasoning and generative tasks.

Impact Statement

The statement identifies no specific negative societal consequences requiring discussion and provides a list of supplementary materials. The appendix covers theory, implementation, data, evaluation, additional results, and LLM assistance.

  • The work identifies no specific negative consequences requiring discussion despite potential societal implications.
  • The appendix includes convergence analysis, training hyperparameters, dataset descriptions, sampling ratios, evaluation datasets, and metrics.
  • Additional appendix materials provide extra benchmark results, hyperparameter ablations, and a description of LLM usage in manuscript preparation.

A. Theoretical Analysis of SPES

The theoretical analysis models SPES as block-sparse local optimization with owner-specific expert updates, periodic synchronization, and early expert-merging perturbations. Its assumptions formalize smoothness, stochastic gradients, expert heterogeneity, and bounded merging displacement.

  • SPES updates shared parameters on every node but updates each expert block only on its owner node.The owner partition determines which expert blocks each node may modify.
  • Each node initializes from the global model and performs H local stochastic gradient steps with step size η.A block mask restricts updates to shared parameters and the node’s assigned experts.
  • After local updates, sparse synchronization averages shared parameters and assigns each expert from its owner.
  • During warm-up, expert merging is applied after synchronization while shared parameters remain unchanged.The mixing coefficient α_t is zero after the warm-up period.
  • The analysis assumes smooth local objectives, bounded stochastic gradients, expert-gradient heterogeneity, and bounded merge displacement.Under IID data, the expert-gradient heterogeneity parameter ζ_Φ equals zero.

A.4. Main Convergence Result

The convergence analysis bounds SPES under smoothness, stochastic-gradient, heterogeneity, and merging assumptions by separating pre-merge descent from the warm-up perturbation. Shared updates benefit from averaging, whereas owner-only expert updates retain heterogeneity effects.

  • Theorem 1 establishes convergence for SPES under Assumptions 1–4 and a step-size condition, both without and with expert warm-up merging.The result applies for any T ≥ 1 and defines F_inf as the infimum of the global objective.
  • Discussion: The shared block receives 1/N variance reduction from averaging, while owner-only expert updates retain an expert heterogeneity term.
  • Proof strategy: The proof first bounds descent for the pre-merge iterate and then treats merging as a smooth perturbation.
  • Proof strategy: The gradient-error analysis decomposes error into stochastic variance and bias from local drift and expert heterogeneity.Local drift contributes a term of order L^2η^2H^2G^2, while the heterogeneity assumption contributes ζ_Φ^2.
  • Proof strategy: Summing the descent bound across rounds and using F(θ(T,pre)) ≥ F_inf yields the stated convergence inequality.

B. Implementation Details

The experiments use fixed loss coefficients across model scales, with training schedules adjusted by model size and ablation focus.

  • From-scratch 2B and 7B experiments use the listed configurations for the first 70% of training tokens, then halve per-node batch size and set H = 50.
  • The 1B model uses a 50B-token budget, except H and N ablations, which use 100B tokens.
  • Loss coefficients remain fixed across models: cross-entropy 1, load-balancing 0.01, MoE z-loss 0.001, and standard z-loss 1 × 10−

C. Details of Datasets and Sampling Ratio

The study samples training data from multiple open-source corpora, including broad web, scientific, mathematical, code, and multilingual resources, with ratios specified in appendix tables.

  • Training data are sampled from several open-source corpora, with sampling ratios provided in Tables A2 and A3.
  • Ultra-FineWeb combines FineWeb and Chinese FineWeb through verification-based filtering, yielding roughly 1 trillion English and 120 billion Chinese tokens.
  • OLMo-Mix-1124 contains 3.9 trillion tokens, and this work extracts arXiv, OpenWebMath, Algebraic Stack, peS2o, and StarCoder subsets.
  • Nemotron Pretraining Dataset emphasizes math, code, and multilingual question-answering through four specialized components.
  • SlimPajama applies quality filtering and MinHashLSH deduplication, reducing RedPajama from 1.21T to 627B tokens while retaining domain coverage.

D. Evaluation Details

Models are evaluated with standardized lm-evaluation-harness benchmarks covering commonsense reasoning, science, logical reasoning, reading comprehension, and general knowledge.

  • Evaluation uses lm-evaluation-harness version 0.4.7 to facilitate standardized comparisons with prior work.
  • Commonsense and reasoning evaluation includes SIQA, ARC-E, ARC-C, SciQ, PIQA, OpenBookQA, WinoGrande, LogiQA, and BoolQ.
  • Reported metrics include 0-shot accuracy or normalized accuracy, with ARC-C additionally evaluated using 25-shot normalized accuracy.
  • Appendix ablations report averages over eight benchmarks for expert-merging hyperparameters and synchronization-step settings.
  • General knowledge is assessed with MMLU and C-Eval, using 5-shot accuracy for MMLU and 0-shot accuracy for C-Eval.

E. Additional Results

Additional results show competitive performance across benchmarks, stable scaling across node counts, and sensitivity to expert-merging and synchronization schedules.

  • Results on Additional Benchmarks: 26.2 vs. 23.6 on C-Eval and 24.9 vs. 24.6 on MMLU: SPES-7B surpasses MoE++ on C-Eval and remains competitive on MMLU.
  • Results on Additional Benchmarks: SPES-2B performs on par with similarly scaled models using 16 weakly connected nodes.
  • Ablation on Number of Nodes: 50.6 to 49.5: average performance decreases slightly when scaling from 2 to 8 nodes, while remaining competitive across benchmarks.
  • Ablation on Hyperparameters in Expert Merging: 12.5k warmup steps and an interval of Ti = 300 produce the best expert-merging results, whereas shorter, longer, or differently timed merging performs worse.
  • Ablation on Synchronization Steps: H = 50 provides the best balance because larger local-update intervals reduce communication but amplify cross-node model divergence.
Loading 2602.11543v3…