Research map / Architectures and efficiency / RNNs and normalization
EDiT: A Local-SGD-Based Efficient Distributed Training Method for Large Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes EDiT, a sharded Local SGD method for LLM pretraining with layer-wise forward-pass synchronization, a pseudo-gradient penalty for stability, and an asynchronous variant A-EDiT for heterogeneous clusters.
Problem
Local SGD methods do not handle model sharding for LLMs, incur non-overlapped communication and memory overhead, are unstable on noisy large-scale corpora, and struggle with consistent stragglers and elastic training.
Approach
Combine Local SGD with ZeRO-3/FSDP-style model sharding on a 2D device mesh, synchronize sharded parameters layer-by-layer during the forward pass with prefetching, and filter/weight/clip pseudo-gradients (EMA z-test anomaly elimination, norm-based softmax weighting, clipping); A-EDiT synchronizes on a wall-clock interval so fast workers do more steps.
Key results
- On FineWeb-Edu with Llama 1B, EDiT achieves the best average benchmark accuracy (44.82 vs 44.60 for the mini-batch baseline and 43.96 for DiLoCo).
- EDiT's throughput on Llama 7B is 5.42e4 tokens/s vs 5.14e4 for baseline (about 5% faster), and it is the only Local SGD method besides baseline that runs at 3B and 7B scale without OOM.
- In elastic training simulations on Llama 1B, EDiT yields 4.5% and 2.6% PPL improvements over the baseline in the upscaling and downscaling scenarios.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use EDiT instead of the strongest reported baseline | Llama 1B trained on FineWeb-Edu · FineWeb-Edu public benchmarks (average) | average accuracy: 43.96 -> 44.82 | Helped |
| Use EDiT instead of the strongest reported baseline | Llama 1B trained on the in-house dataset · In-house public benchmarks (average) | average accuracy: 34.62 -> 34.96 | Helped |
| Use A-EDiT instead of EDiT | Llama 1B trained on FineWeb-Edu · FineWeb-Edu public benchmarks (average) | average accuracy: 44.82 -> 43.90 | Hurt |
| Use A-EDiT instead of EDiT | Llama 1B trained on the in-house dataset · In-house public benchmarks (average) | average accuracy: 34.96 -> 34.49 | Hurt |
| Use EDiT instead of Baseline | Llama 1B training · Llama 1B training | throughput (tokens/sec): 2.08e5 -> 2.25e5 | Helped |
| Use A-EDiT instead of EDiT under consistent stragglers | Llama 7B training · Consistent Straggler, Lag 4.5s | TFLOPS: 118.47 -> 223.07 | Helped |
| Use EDiT instead of Baseline under limited bandwidth | Llama 7B training · Limited Bandwidth, Repeat 40 | TFLOPS: 85.18 -> 236.39 | Helped |
| Add anomaly elimination to the pseudo gradient penalty strategy (ablation) | EDiT · In-house dataset | validation PPL | Helped |
| Add weighted averaging to the pseudo gradient penalty strategy (ablation) | EDiT · In-house dataset | validation PPL | Helped |
| Add gradient clipping to the pseudo gradient penalty strategy (ablation) | EDiT · In-house dataset | validation PPL | Helped |
| Add the full pseudo gradient penalty strategy (ablation) | EDiT · In-house dataset | validation PPL | Helped |
Surprising findings
- The optimal inner learning rate for EDiT stays constant (1.5e-4) as worker count grows, whereas the baseline's optimal learning rate increases with workers.
- On the FineWeb-Edu dataset, EDiT's validation perplexity surpasses the mini-batch baseline at 1B scale.
Limitations
- A-EDiT is still negatively affected by stragglers in its performance.
- Elastic training is simulated by halting and restarting training rather than adjusting resources live.
- Large portion of experiments rely on a private in-house dataset that is not released.
Experimental setup
Datasets: FineWeb-Edu in-house pretraining dataset
Benchmarks: MMLU ARC-E ARC-C HellaSwag PIQA CommonSense-QA OpenBookQA WinoGrande
Metrics: validation perplexity training loss accuracy average accuracy throughput (tokens/sec) TFLOPS
Baselines: Baseline mini-batch SGD Post Local SGD DiLoCo CO2 CO2*
Builds on: DiLoCo CO2 Post Local SGD ZeRO FSDP
Closest papers by meaning
- HALoS: Hierarchical Asynchronous Local SGD over Slow Networks for Geo-Distributed Large Language Model Training (ICML 2025 · similarity 0.92)
- Rethinking Memory and Communication Costs for Efficient Data Parallel Training of Large Language Models (NeurIPS 2024 · similarity 0.88)
- Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning (ICLR 2024 · similarity 0.87)
- Enhancing Large Language Model Performance with Gradient-Based Parameter Selection (AAAI 2025 · similarity 0.87)
- Pruning Large Language Models with Semi-Structural Adaptive Sparse Training (AAAI 2025 · similarity 0.87)
- La RoSA: Enhancing LLM Efficiency via Layerwise Rotated Sparse Activation (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.