mlsanityData to 14 Oct 2026

Mist: Efficient Distributed Training of Large Language Models via Memory-Parallelism Co-Optimization

arXiv 2503.19050 ↗First posted 24 Mar 2025cs.DC cs.AIcited by 7

Not on the map. Why?

Abstract (arXiv)

Various parallelism, such as data, tensor, and pipeline parallelism, along with memory optimizations like activation checkpointing, redundancy elimination, and offloading, have been proposed to accelerate distributed training for Large Language Models. To find the best combination of these techniques, automatic distributed training systems are proposed. However, existing systems only tune a subset of optimizations, due to the lack of overlap awareness, inability to navigate the vast search space, and ignoring the inter-microbatch imbalance, leading to sub-optimal performance. To address these shortcomings, we propose Mist, a memory, overlap, and imbalance-aware automatic distributed training system that comprehensively co-optimizes all memory footprint reduction techniques alongside parallelism. Mist is based on three key ideas: (1) fine-grained overlap-centric scheduling, orchestrating optimizations in an overlapped manner, (2) symbolic-based performance analysis that predicts runtime and memory usage using symbolic expressions for fast tuning, and (3) imbalance-aware hierarchical tuning, decoupling the process into an inter-stage imbalance and overlap aware Mixed Integer Linear Programming problem and an intra-stage Dual-Objective Constrained Optimization problem, and connecting them through Pareto frontier sampling. Our evaluation results show that Mist achieves an average of 1.28× (up to 1.73×) and 1.27× (up to 2.04×) speedup compared to state-of-the-art manual system Megatron-LM and state-of-the-art automatic system Aceso, respectively.

Key facts. Cited by 7 papers in the map (data through 14 Oct 2026). Builds on 47 of its 82 references that are in the map. First posted on arXiv 24 Mar 2025.

Builds on

47 of 82 references.

References · 47

  1. 2020Language Models are Few-Shot Learnerscited by 32,419
  2. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
  3. 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 13,028
  4. 2021Evaluating Large Language Models Trained on Codecited by 8,320
  5. 2016XGBoost: A Scalable Tree Boosting Systemcited by 5,166
  6. 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,703
  7. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
  8. 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,598
  9. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,868
  10. 2020GLU Variants Improve Transformercited by 1,804
  11. 2019Root Mean Square Layer Normalizationcited by 1,480
  12. 2019ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelscited by 1,301
35 more
  1. 2019Energy and Policy Considerations for Deep Learning in NLPcited by 1,254
  2. 2017Mixed Precision Trainingcited by 1,156
  3. 2016Training Deep Nets with Sublinear Memory Costcited by 848
  4. 2022GPT-NeoX-20B: An Open-Source Autoregressive Language Modelcited by 713
  5. 2023PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallelcited by 660
  6. 2021Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LMcited by 632
  7. 2023The Falcon Series of Open Language Modelscited by 535
  8. 2022DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scalecited by 404
  9. 2021ZeRO-Offload: Democratizing Billion-Scale Model Trainingcited by 341
  10. 2022Reducing Activation Recomputation in Large Transformer Modelscited by 302
  11. 2021ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learningcited by 256
  12. 2020Memory-Efficient Pipeline-Parallel DNN Trainingcited by 134
  13. 2020DAPPLE: A Pipelined Data Parallel Approach for Training Large Modelscited by 101
  14. 2019Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerializationcited by 91
  15. 2021TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Modelscited by 75
  16. 2021ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Trainingcited by 61
  17. 2018SuperNeurons: Dynamic GPU Memory Management for Training Deep Neural Networkscited by 55
  18. 2022Galvatron: Efficient Transformer Training over Multiple GPUs Using Automatic Parallelismcited by 53
  19. 2020Dynamic Tensor Rematerializationcited by 46
  20. 2018Supporting Very Large Models using Automatic Dataflow Graph Partitioningcited by 43
  21. 2021Varuna: Scalable, Low-cost Training of Massive Deep Learning Modelscited by 40
  22. 2022GACT: Activation Compressed Training for Generic Network Architecturescited by 31
  23. 2021A Runtime-Based Computational Performance Predictor for Deep Neural Network Trainingcited by 29
  24. 2022Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programscited by 28
  25. 2020torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Modelscited by 26
  26. 2021PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformerscited by 23
  27. 2024GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training with Virtual Memory Stitchingcited by 15
  28. 2022AMP: Automatically Finding Model Parallel Strategies with Heterogeneity Awarenesscited by 13
  29. 2022dPRO: A Generic Profiling and Optimization System for Expediting Distributed DNN Trainingcited by 12
  30. 2022Merak: An Efficient Distributed DNN Training Framework with Automated 3D Parallelism for Giant Foundation Modelscited by 11
  31. 2023Slapo: A Schedule Language for Progressive Optimization of Large Deep Learning Model Trainingcited by 11
  32. 2018Echo: Compiler-based GPU Memory Footprint Reduction for LSTM RNN Trainingcited by 10
  33. 2023Proteus: Simulating the Performance of Distributed DNN Trainingcited by 10
  34. 2023DistSim: A performance model of large-scale hybrid distributed DNN trainingcited by 7
  35. 2022Tempo: Accelerating Transformer-Based Model Training through Memory Footprint Reductioncited by 4

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 3; 2026: 4.

7 papers.

  1. 2026AIBuildAI: An AI Agent for Automatically Building AI Modelscited by 13
  2. 2025LoRAFusion: Efficient LoRA Fine-Tuning for LLMscited by 3
  3. 2025OctoPipe: Reducing Pipeline Bubbles for Heterogeneous Models via Co-Optimizing Partitioning, Placement, and Schedulingcited by 2
  4. 2025Scalable Synthesis of distributed LLM workloads through Symbolic Tensor Graphscited by 1
  5. 2026NEST: Network- and Memory-Aware Device Placement For Distributed Deep Learningcited by 0
  6. 2026Piper: A Programmable Distributed Training Systemcited by 0
  7. 2026ATTUNE: Agentic Throughput Tuning for LLM Trainingcited by 0