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

• URL: http://arxiv.org/abs/2503.19050v1
• Date: Mon, 24 Mar 2025 18:21:08 GMT
• Title: Mist: Efficient Distributed Training of Large Language Models via Memory-Parallelism Co-Optimization
• Authors: Zhanda Zhu, Christina Giannoula, Muralidhar Andoorveedu, Qidong Su, Karttikeya Mangalam, Bojian Zheng, Gennady Pekhimenko,
• Abstract summary: We propose Mist, a memory, overlap, and imbalance-aware automatic distributed training system.<n>We show that Mist achieves an average of 1.28$times$ (up to 1.73$times$) speedup compared to state-of-the-art manual system Megatron-LM and state-of-the-art automatic system Aceso.
• Score: 11.223375172715722
• License: http://creativecommons.org/licenses/by-nc-sa/4.0/
• Abstract: 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$\times$ (up to 1.73$\times$) and 1.27$\times$ (up to 2.04$\times$) speedup compared to state-of-the-art manual system Megatron-LM and state-of-the-art automatic system Aceso, respectively.

Related papers

• AutoQRA: Joint Optimization of Mixed-Precision Quantization and Low-rank Adapters for Efficient LLM Fine-Tuning [23.59600455731982]
  We propose a joint optimization framework that simultaneously optimize the bit-width and LoRA rank configuration for each layer during the mixed quantized fine-tuning process.<n>Experiments show that AutoQRA achieves performance close to full-precision fine-tuning with a memory footprint comparable to uniform 4-bit methods.
  arXiv  Detail & Related papers  (2026-02-25T07:18:08Z)
• RaBiT: Residual-Aware Binarization Training for Accurate and Efficient LLMs [5.782015253162346]
  Residual binarization enables matmul-free inference by stacking binary layers.<n>We propose RaBiT, a novel quantization framework that resolves coadaptation by algorithmically enforcing a residual hierarchy.<n>RaBiT achieves state-of-the-art performance, rivals even hardware-intensive Vector Quantization (VQ) methods, and delivers a $4.49times$ inference speed-up over full-precision models.
  arXiv  Detail & Related papers  (2026-02-05T06:41:11Z)
• Canzona: A Unified, Asynchronous, and Load-Balanced Framework for Distributed Matrix-based Optimizers [36.650880799066215]
  Asynchronous approaches suffer from computational redundancy, while layer-wise partitioning fails to reconcile this conflict.<n>For Data Parallelism, we introduce an alpha-Balanced Static Partitioning strategy that respects atomicity while neutralizing the load imbalance.<n>Our approach achieves a 1.57x speedup in end-to-end time and reducing step latency by 5.8x compared to the baseline.
  arXiv  Detail & Related papers  (2026-02-04T07:38:24Z)
• Repulsor: Accelerating Generative Modeling with a Contrastive Memory Bank [65.00301565190824]
  mname is a plug-and-play training framework that requires no external encoders.<n>mname achieves a state-of-the-art FID of textbf2.40 within 400k steps, significantly outperforming comparable methods.
  arXiv  Detail & Related papers  (2025-12-09T14:39:26Z)
• ThreadWeaver: Adaptive Threading for Efficient Parallel Reasoning in Language Models [99.6720868215076]
  We introduce ThreadWeaver, a framework for adaptive parallel reasoning.<n> ThreadWeaver achieves accuracy on par with popular sequential reasoning models of comparable size.<n>We show that ThreadWeaver delivers up to 1.53x average speedup in token latency.
  arXiv  Detail & Related papers  (2025-11-24T18:55:59Z)
• Memory- and Latency-Constrained Inference of Large Language Models via Adaptive Split Computing [8.705453442427585]
  Large language models (LLMs) have achieved near-human performance across diverse reasoning tasks.<n>Their deployment on resource-constrained Internet-of-Things (IoT) devices remains impractical due to massive parameter footprints and memory-intensive autoregressive decoding.<n>This work introduces the first autoregressive-aware split computing framework designed explicitly for LLM deployment on edge devices.
  arXiv  Detail & Related papers  (2025-11-06T02:55:07Z)
• PT$^2$-LLM: Post-Training Ternarization for Large Language Models [52.4629647715623]
  Large Language Models (LLMs) have shown impressive capabilities across diverse tasks, but their large memory and compute demands hinder deployment.<n>We propose PT$2$-LLM, a post-training ternarization framework tailored for LLMs.<n>At its core is an Asymmetric Ternary Quantizer equipped with a two-stage refinement pipeline.
  arXiv  Detail & Related papers  (2025-09-27T03:01:48Z)
• AutoLayout: Closed-Loop Layout Synthesis via Slow-Fast Collaborative Reasoning [102.71841660031065]
  Auto is a fully automated method that integrates a closed-loop self-validation process within a dual-system framework.<n>The effectiveness of Auto was validated across 8 distinct scenarios, where it demonstrated a significant 10.1% improvement over SOTA methods.
  arXiv  Detail & Related papers  (2025-07-06T08:35:22Z)
• Pangu Embedded: An Efficient Dual-system LLM Reasoner with Metacognition [95.54406667705999]
  Pangu Embedded is an efficient Large Language Model (LLM) reasoner developed on Ascend Neural Processing Units (NPUs)<n>It addresses the significant computational costs and inference latency challenges prevalent in existing reasoning-optimized LLMs.<n>It delivers rapid responses and state-of-the-art reasoning quality within a single, unified model architecture.
  arXiv  Detail & Related papers  (2025-05-28T14:03:02Z)
• Automatic Operator-level Parallelism Planning for Distributed Deep Learning -- A Mixed-Integer Programming Approach [6.449961842220686]
  We propose a bi-level solution framework balancing optimality with computational efficiency. Our framework achieves comparable or superior performance, reducing computational bubbles by half under the same memory constraints. Such capabilities position our solution as both a valuable research tool for exploring optimal parallelization strategies and a practical industrial solution for large-scale AI deployment.
  arXiv  Detail & Related papers  (2025-03-12T13:00:29Z)
• Seesaw: High-throughput LLM Inference via Model Re-sharding [8.840996987380484]
  We present Seesaw, an inference engine optimized for throughput-oriented tasks.<n>Key idea behind Seesaw is dynamic model re-sharding, a technique that facilitates the dynamic reconfiguration of parallelization strategies.
  arXiv  Detail & Related papers  (2025-03-09T04:14:06Z)
• AutoHete: An Automatic and Efficient Heterogeneous Training System for LLMs [68.99086112477565]
  Transformer-based large language models (LLMs) have demonstrated exceptional capabilities in sequence modeling and text generation.<n>Existing heterogeneous training methods significantly expand the scale of trainable models but introduce substantial communication overheads and CPU workloads.<n>We propose AutoHete, an automatic and efficient heterogeneous training system compatible with both single- GPU and multi- GPU environments.
  arXiv  Detail & Related papers  (2025-02-27T14:46:22Z)
• Efficient Adaptive Optimization via Subset-Norm and Subspace-Momentum: Fast, Memory-Reduced Training with Convergence Guarantees [5.399838579600896]
  We introduce two complementary techniques for memory optimization. One technique, Subset-Norm, reduces the momentum state's memory footprint by a low-dimensional subspace. The other technique, Subspace-Momentum, reduces the momentum state's memory footprint by a low-dimensional subspace.
  arXiv  Detail & Related papers  (2024-11-11T16:48:07Z)
• FusionLLM: A Decentralized LLM Training System on Geo-distributed GPUs with Adaptive Compression [55.992528247880685]
  Decentralized training faces significant challenges regarding system design and efficiency. We present FusionLLM, a decentralized training system designed and implemented for training large deep neural networks (DNNs) We show that our system and method can achieve 1.45 - 9.39x speedup compared to baseline methods while ensuring convergence.
  arXiv  Detail & Related papers  (2024-10-16T16:13:19Z)
• A Convex-optimization-based Layer-wise Post-training Pruner for Large Language Models [24.185245582500876]
  We introduce FISTAPruner, the first post-training pruner based on convex optimization models and algorithms. FISTAPruner incorporates an intra-layer cumulative error correction mechanism and supports parallel pruning. We evaluate FISTAPruner on models such as OPT, LLaMA, LLaMA-2, and LLaMA-3 with 125M to 70B parameters under unstructured and 2:4 semi-structured sparsity.
  arXiv  Detail & Related papers  (2024-08-07T12:33:46Z)
• Bypass Back-propagation: Optimization-based Structural Pruning for Large Language Models via Policy Gradient [57.9629676017527]
  We propose an optimization-based structural pruning on Large-Language Models. We learn the pruning masks in a probabilistic space directly by optimizing the loss of the pruned model. Our method operates for 2.7 hours with around 35GB memory for the 13B models on a single A100 GPU.
  arXiv  Detail & Related papers  (2024-06-15T09:31:03Z)
• Efficient Methods for Non-stationary Online Learning [61.63338724659592]
  We present efficient methods for optimizing dynamic regret and adaptive regret, which reduce the number of projections per round from $mathcalO(log T)$ to $1$. We also study an even strengthened measure, namely the interval dynamic regret'', and reduce the number of projections per round from $mathcalO(log2 T)$ to $1$.
  arXiv  Detail & Related papers  (2023-09-16T07:30:12Z)
• Improving Automatic Parallel Training via Balanced Memory Workload Optimization [36.87527680184956]
  Transformer models have emerged as the leading approach for achieving state-of-the-art performance across various application domains. We present Galvatron-BMW, a novel system framework that integrates multiple parallelism prevalent dimensions and automatically identifies the most efficient hybrid parallelism strategy. Our evaluations on different Transformer models demonstrate the capabilities of Galvatron-BMW in automating distributed training under varying GPU memory constraints.
  arXiv  Detail & Related papers  (2023-07-05T05:28:38Z)
• Fine-Tuning Language Models with Just Forward Passes [92.04219196752007]
  Fine-tuning language models (LMs) has yielded success on diverse downstream tasks, but as LMs grow in size, backpropagation requires a large amount of memory. We propose a memory-efficient zerothorder (MeZO) to operate in-place, thereby fine-tuning LMs with the same memory footprint as inference.
  arXiv  Detail & Related papers  (2023-05-27T02:28:10Z)
• On the Convergence of Distributed Stochastic Bilevel Optimization Algorithms over a Network [55.56019538079826]
  Bilevel optimization has been applied to a wide variety of machine learning models. Most existing algorithms restrict their single-machine setting so that they are incapable of handling distributed data. We develop novel decentralized bilevel optimization algorithms based on a gradient tracking communication mechanism and two different gradients.
  arXiv  Detail & Related papers  (2022-06-30T05:29:52Z)
• Joint inference and input optimization in equilibrium networks [68.63726855991052]
  deep equilibrium model is a class of models that foregoes traditional network depth and instead computes the output of a network by finding the fixed point of a single nonlinear layer. We show that there is a natural synergy between these two settings. We demonstrate this strategy on various tasks such as training generative models while optimizing over latent codes, training models for inverse problems like denoising and inpainting, adversarial training and gradient based meta-learning.
  arXiv  Detail & Related papers  (2021-11-25T19:59:33Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.