论文检索

Optimization · Large Scale, Parallel and Distributed

Zaifeng Pan, Yipeng Shen, Zhengding Hu, Zhuang Wang, Aninda Manocha, Zheng Wang, zhongkai yu, Yue Guan, Yufei Ding

LLM-based multi-agent simulations are increasingly adopted across application domains, but remain difficult to scale due to GPU memory pressure. Each agent maintains private GPU-resident states, including models, prefix caches, and adapters, which quickly exhaust device memory as the agent count grows. We identify two key properties of these workloads: sparse agent activation and an estimable agent invocation order. Based on an analysis of representative workload classes, we introduce invocation distance, a unified abstraction that estimates the relative order in which agents will issue future LLM requests. Leveraging this abstraction, we present ScaleSim, a memory-efficient LLM serving system for large-scale multi-agent simulations. ScaleSim enables proactive prefetching and priority-based eviction, supports diverse agent-specific memory through a modular interface, and achieves up to 1.74$\times$ speedup over SGLang on simulation benchmarks.

Optimization · Large Scale, Parallel and Distributed

Nozomi Hata, Kenta Niwa

For communication-efficient decentralized learning, advanced network (NW) topologies, such as exponential and 1-peer exponential graphs, have been studied under homogeneous communication delays. However, real-world NWs exhibit heterogeneous communication delays, making node assignment optimization crucial for minimizing the Bottleneck Communication Delay (BCD). We propose BTSP-MSR, an approximate method for minimizing BCD on circulant digraphs, including exponential and 1-peer exponential graphs. Leveraging the fact that circulant digraphs can be viewed as a union of (directed) ring graphs, we derive an upper bound on the BCD by combining the ring-graph BCD (BTSP) with a deviation term (MSR). We then construct a solver that sequentially minimizes these two terms. Numerical experiments show that BTSP-MSR consistently reduces BCD across several circulant digraphs with large numbers of nodes. Notably, incorporating the exponential or 1-peer exponential graph enables communication-efficient decentralized learning under heterogeneous delay settings.

Deep Learning · Foundation Models

Yulu Gan, Phillip Isola

Pretraining produces a learned parameter vector that is typically treated as a starting point for further iterative adaptation. In this work, we instead view the outcome of pretraining as a distribution over parameter vectors, whose support already contains task-specific experts. We show that in smaller or insufficiently trained models such expert solutions occupy a negligible fraction of the volume of this distribution, making their discovery reliant on structured optimization methods such as gradient descent. In contrast, in large, well-pretrained models the density of task-experts increases dramatically, so that diverse specialists populate a substantial fraction of the neighborhood around the pretrained weights. Motivated by this perspective, we explore a simple, fully parallel post-training method that samples $N$ parameter vectors at random, selects the top $K$, and ensembles them via majority vote to combine complementary expertise. Despite its simplicity, this approach is competitive with standard post-training methods such as PPO, GRPO, and ES for contemporary large-scale models.

Optimization · Large Scale, Parallel and Distributed

Hyunji Jung, Sungbin Shin, Namhoon Lee

Asynchronous pipeline parallelism maximizes hardware utilization by eliminating the pipeline bubbles inherent in synchronous execution, offering a path toward efficient large-scale distributed training. However, this efficiency gain can be compromised by gradient staleness, where the immediate model updates with delayed gradients introduce noise into the optimization process. Crucially, we identify a critical, yet often overlooked, pathology: this delay scales linearly with pipeline depth, fundamentally undermining the very scalability that the method originally intends to provide. In this work, we investigate this inconsistency and bridge the gap by rectifying delayed gradients through basis rotation, restoring scalable asynchronous training while maintaining performance. Specifically, we observe that the deleterious effects of delayed gradients are exacerbated when the Hessian eigenbasis is misaligned with the standard coordinate basis. We demonstrate that this misalignment prevents coordinate-wise adaptive schemes, such as Adam, from effectively leveraging curvature-aware adaptivity. This failure leads to significant oscillations in the optimization trajectory and, consequently, slower convergence. We substantiate these findings through both rigorous theoretical analysis and empirical evaluation. To address this challenge, we propose the use of basis rotation, demonstrating that it effectively mitigates the alignment issue and significantly accelerates convergence in asynchronous settings. For example, our training of a 1B-parameter LLM with basis rotation achieves the same training loss in 76.8% fewer iterations compared to the best-performing asynchronous pipeline parallel training baseline.

Optimization · Large Scale, Parallel and Distributed

Yihan Zhang, Xinwen Zhang, My T. Thai, Jie Wu, Hongchang Gao

Decentralized stochastic bi-level optimization has been actively studied in recent years. However, existing studies assume that the lower-level loss function is strongly convex, which limits their applicability to many machine learning models. To address this limitation, in this paper, we propose a novel decentralized stochastic first-order optimization algorithm, which does not require second-order Hessian or Jacobian matrices, for the setting where the lower-level loss function is nonconvex but satisfies the Polyak–Łojasiewicz (PL) condition. Additionally, unlike existing single-agent methods that introduce a regularization term to the lower-level loss function to artificially enforce strong convexity, our algorithm does not require such modification. Moreover, our algorithm employs a constant single-timescale learning rate for updating variables, which is different from the time-dependent and two-timescale learning rate schedules used in prior work. To establish the convergence rate, we develop a new convergence analysis framework for the pure PL condition, rather than relying on the artificial strong convexity introduced through regularization in existing single-agent methods. To the best of our knowledge, this is the first algorithm for nonconvex decentralized bi-level optimization that offers theoretical convergence guarantees under mild conditions. Finally, our extensive experimental results on hyperparameter optimization and model pruning applications validate the efficacy of the proposed algorithm.

Optimization · Large Scale, Parallel and Distributed

Xinwen Zhang, Richard Souvenir, Hongchang Gao

Federated stochastic multi-level compositional minimax optimization supports a growing number of machine learning applications. However, the interplay of multi-level compositional structure, minimax formulation, and federated setting poses significant optimization challenges, resulting in slow convergence rates for existing algorithms. In this paper, we propose a novel federated algorithm, LS$^2$MC-GDA, that leverages smoothing techniques and variance reduced stochastic compositional gradients. To support our theoretical analysis, we introduce a stage-wise extension of LS$^2$MC-GDA, which serves to bridge the gap between different stationarity measures. We establish that our algorithm achieves a sample complexity of $O(\kappa^{3/2}/N\epsilon^3)$ and a communication complexity of $O(\kappa/\epsilon^2)$, substantially improving existing theoretical results in terms of the condition number $\kappa$ and the solution accuracy $\epsilon$ and achieving a linear speedup with respect to the number of workers $N$. Finally, experimental results validate the effectiveness of our approach.

Optimization · Large Scale, Parallel and Distributed

Pierre Aguié, Mathieu Even, Laurent Massoulié

We analyze the Accelerated Noisy Power Method, an algorithm for Principal Component Analysis in the setting where only inexact matrix-vector products are available, which can arise for instance in decentralized PCA. While previous works have established that acceleration can improve convergence rates compared to the standard Noisy Power Method, these guarantees require overly restrictive upper bounds on the magnitude of the perturbations, limiting their practical applicability. We provide an improved analysis of this algorithm, which preserves the accelerated convergence rate under much milder conditions on the perturbations. We show that our new analysis is worst-case optimal, in the sense that the convergence rate cannot be improved, and that the noise conditions we derive cannot be relaxed without sacrificing convergence guarantees. We demonstrate the practical relevance of our results by deriving an accelerated algorithm for decentralized PCA, which has similar communication costs to non-accelerated methods. To our knowledge, this is the first decentralized algorithm for PCA with provably accelerated convergence.

Optimization · Large Scale, Parallel and Distributed

Yancheng Wu, Huikang Liu, Wenzhi Gao, Yuexin Su, Tongyang Li, Dongdong Ge, Yinyu Ye

Quantum computation offers the potential for a significant constant-factor speedup for the Ordered Search Problem (OSP). A classical construction is the $k$-query quantum ordered search algorithm, which can exactly search an $N$-element ordered list and achieves a query complexity improvement of a factor of $\frac{k}{\log_2 N}$. For larger $k$, stronger constant-factor improvements could be obtained by finding the largest admissible list size $N^\star$, a task that can be formulated as a structured semidefinite program (SDP). However, solving this SDP becomes computationally intractable beyond $k=6$, as existing CPU and GPU solvers rely on explicit construction of prohibitively large constraint matrices. In this paper, we introduce a matrix-free GPU SDP framework that evaluates the highly structured constraints in OSP on-the-fly using custom CUDA kernels, reducing memory complexity from quadratic to linear and shifting the bottleneck from memory to computation. Using this approach, we tightly bracket the optimal list size for $k=6$ as $90,000 \le N^\star < 94,000$, improving the best known upper bound on the query coefficient from $0.390$ to $0.365$. We further certify these results by constructing rigorous dual infeasibility certificates via matrix-free minimum-eigenvalue estimation.

Optimization · Large Scale, Parallel and Distributed

Wenhao He, Youhe Jiang, Penghao Zhao, Quanqing Xu, Eiko Yoneki, Bin Cui, Fangcheng Fu

With the rapid evolution of Large Language Models (LLMs), multi-round workflows, such as autonomous agents and iterative retrieval, have become increasingly prevalent. However, this raises hurdles for serving LLMs under prefill-decode (PD) disaggregation, a widely adopted paradigm that separates the compute-bound prefill phase and memory-bound decode phase onto individual resources. Specifically, existing systems overlook the interleaved prefill-decode workload pattern in multi-round inference, leading to sub-optimal handling of the incremental prefill workloads and model deployment for the two phases. In this work, we present AMPD, a brand new disaggregated serving framework for multi-round LLM inference. The core of AMPD is to coordinate the prefill workloads based on real-time workloads by adaptively determining *where* to carry out these workloads and *how* they are scheduled, in order to maximize service level objective (SLO) attainment. In addition, we tailor a planning algorithm for our scenario, facilitating the deduction of optimal resource allocation and parallel strategies for the two phases. Empirical results demonstrate that AMPD substantially improves SLO attainment compared to state-of-the-art baselines.

Optimization · Large Scale, Parallel and Distributed

Andrej Jovanović, Alex Iacob, Mher Safaryan, Ionut-Vlad Modoranu, Lorenzo Sani, Shen, Xinchi Qiu, Dan Alistarh, Nicholas Lane

Distributed training of foundation models via $\texttt{DDP}$ is limited by interconnect bandwidth. While infrequent communication strategies reduce synchronization frequency, they remain bottlenecked by the memory and communication requirements of optimizer states. Low-rank optimizers can alleviate these constraints; however, in the local-update regime, workers lack access to the full-batch gradients required to compute low-rank projections, which degrades performance. We propose $\texttt{LoRDO}$, a principled framework unifying low-rank optimization with infrequent synchronization. We first demonstrate that, while global projections based on pseudo-gradients are theoretically superior, they permanently restrict the optimization trajectory to a low-rank subspace. To restore subspace exploration, we introduce a full-rank quasi-hyperbolic update. $\texttt{LoRDO}$ achieves near-parity with low-rank $\texttt{DDP}$ in language modeling and downstream tasks at model scales of $125$M--$720$M, while reducing communication by $\approx10\times$. Finally, we show that $\texttt{LoRDO}$ improves performance even more in very low-memory settings with small rank/batch size.

Optimization · Large Scale, Parallel and Distributed

Ling Chen, Houming Wu, Wenjie Yu

Pipeline parallelism is essential for large-scale model training, but existing asynchronous approaches often degrade convergence due to parameter mismatch between forward and backward passes. We propose Asynchronous Multi-Directional Pipeline parallelism (AMDP) to mitigate this issue while sustaining high utilization. AMDP limits the first stage of each pipeline to process at most two minibatches before backpropagation, bounding the number of parameter updates between forward and backward passes. To alleviate the resulting pipeline bubbles, AMDP launches multiple concurrent pipelines and adapts their number according to pipeline depth. In addition, AMDP accumulates gradients across minibatches and applies them in a single update, ensuring that only a bounded number of minibatches experience parameter mismatch, limited to within one optimization step. Experiments on GPT- and BERT-style models demonstrate that AMDP significantly accelerates training while preserving convergence. The source code based on Megatron-LM is available at https://anonymous.4open.science/r/Megatron-AMDP-59A7.

Optimization · Large Scale, Parallel and Distributed

Chamin Hewa Koneputugodage, Thalaiyasingam Ajanthan, Sameera Ramasinghe, Hadi Mohaghegh Dolatabadi, Shamane Siriwardhana, Gil Avraham, Violetta Shevchenko, Karol Pajak, James Snewin, Alexander Long

To make large-scale distributed training practical outside high-bandwidth datacenters, we must reduce blocking, high-volume synchronization. While DiLoCo communicates infrequently, its outer synchronization remains bandwidth-heavy and brittle to stragglers and transient failures. We relax exact synchronization to approximate synchronization via mixing/gossip, which degrades gracefully under delays and communication failures. This allows us to factorize DiLoCo synchronization into a non-blocking mixing step that overlaps computation with no staleness, and a blocking mixing step that tightens worker agreement, yielding a tunable trade-off between compute utilization and optimization stability. On up to billion-parameter language models in low-bandwidth settings, our method substantially improves compute utilization while matching DiLoCo’s training progress, and is more robust to failures.

Optimization · Large Scale, Parallel and Distributed

Minghao Yan, Zhuang Wang, Zhen Jia, Shivaram Venkataraman, Yida Wang

Low-Rank Adaptation (LoRA) has gained popularity as a fine-tuning approach for Large Language Models (LLMs) due to its low resource requirements and good performance. While numerous studies have investigated improving LoRA serving efficiency by serving multiple LoRAs concurrently, existing methods assume that a wide range of LoRA adapters are available for serving. In our work, we conduct extensive empirical studies to show that current LoRA training paradigms do not efficiently utilize hardware resources and incur high overhead to obtain a performant LoRA adapter. Leveraging these insights, we propose PLoRA, which automatically orchestrates concurrent LoRA fine-tuning jobs under given hardware and model constraints and develops performant kernels to improve training efficiency. Across a range of LLMs and LoRA configurations, PLoRA improves training throughput by up to 12.8x and reduces the overall fine-tuning makespan by up to 7.52x compared to existing approaches.

Optimization · Large Scale, Parallel and Distributed

zhixin wang, Jiaming Xu, Tianyi Zhou, Mingjun Zhang, Liming Liu, JiaruiHu, Dian Yang, TongYu Wang, Ping Zhang, Jinlong Hou 等

Effectively scaling Reinforcement Learning (RL) is crucial for enhancing the reasoning and alignment of Large Language Models. The massive data and complex execution flows inherent in these tasks require a distributed architecture capable of efficient scaling. However, to simplify programming and dependency management, mainstream frameworks often rely on a centralized architecture where a single node dispatches both control and data. This inherent coupling creates significant communication bottlenecks, severely limiting system scalability and efficiency. We present DistFlow, a novel, fully distributed RL framework that adopts a multi-controller paradigm. By decoupling data transmission from control dispatch, DistFlow establishes a parallelism-aware, decentralized Data Coordinator that leverages local caching, load balancing, and asynchronous double buffer to minimize communication overhead and mitigate straggler effects. For control logic, it introduces a task scheduler built upon Directed Acyclic Graph (DAG) that facilitates fine-grained, independent execution. Experimental results demonstrate that DistFlow achieves near-linear scalability up to 512 GPUs and delivers up to a 2.63x throughput improvement over state-of-the-art (SOTA) frameworks.

Optimization · Large Scale, Parallel and Distributed

Feihu Huang, Yuning Luo, Songcan Chen

Large models recently are widely applied in machine learning, so efficient training of large models has received widespread attention. More recently, a useful Muon optimizer is specifically designed for matrix-structured parameters of large models. Although some works have begun to studying the Muon optimizer, the existing Muon and its variants still suffer from high sample complexity or high memory for large models. To fill this gap, we propose a light and fast Muon (LiMuon) optimizer for training large models, which builds on the momentum-based variance reduced technique and randomized Singular Value Decomposition (SVD). In particular, our LiMuon optimizer simultaneously has a lower memory and lower sample complexity than the Muon. Moreover, we prove that our LiMuon has a lower sample complexity of $O(\epsilon^{-3})$ for finding an $\epsilon$-stationary solution of non-convex stochastic optimization under the generalized smooth condition. Numerical experimental results on training Mamba-130M, Qwen2.5-0.5B and ViT models demonstrate effectiveness of our LiMuon optimizer.

Deep Learning · Large Language Models

Xuan-Phi Nguyen, Shrey Pandit, Austin Xu, Caiming Xiong, Shafiq Joty

Mixture-of-Experts (MoE) models are typically pre-trained with explicit load-balancing constraints to ensure statistically balanced expert routing. Despite this, we observe that even well-trained MoE models exhibit significantly imbalanced routing. This behavior is arguably natural—and even desirable—as imbalanced routing allows models to concentrate domain-specific knowledge within a subset of experts. Expert parallelism (EP) is designed to scale MoE models by distributing experts across multiple devices, but with a less-discussed assumption of balanced routing. Under extreme imbalance, EP can funnel a disproportionate number of tokens to a small number of experts, leading to compute- and memory-bound failures on overloaded devices during post-training or inference, where explicit load balancing is often inapplicable. We propose Least-Loaded Expert Parallelism (LLEP), a novel EP algorithm that dynamically reroutes excess tokens and associated expert parameters from overloaded devices to underutilized ones. This ensures that all devices complete their workloads within the minimum collective latency while respecting memory constraints. Across different model scales, LLEP achieves up to 5x speedup and 4x reduction in peak memory usage compared to standard EP. This enables faster and higher-throughput post-training and inference, with ~1.9x faster for gpt-oss-120b. We support our method with extensive theoretical analysis and comprehensive empirical evaluations, including ablation studies. These results illuminate key trade-offs and enable a principled framework for hardware-specific hyper-parameter tuning to achieve optimal performance.

Optimization · Large Scale, Parallel and Distributed

Yihan Zhang, Xinghua Shi, Meikang Qiu, Yu Wang, Hongchang Gao

The stochastic minimax optimization problem has widespread applications in machine learning. Recently, numerous distributed minimax optimization algorithms have been developed to handle distributed training data. However, most of these algorithms suffer from high communication costs. To address this issue, we develop a novel communication-efficient decentralized stochastic gradient descent ascent with momentum algorithm based on the error feedback mechanism. Importantly, our algorithm demonstrates how to balance the full-precision update and the compression residual with novel designs for coefficients regarding variables and gradients to guarantee convergence. However, compressing the primal and dual variables (and their gradients) of stochastic minimax optimization problems with the error feedback mechanism presents significant challenges for convergence analysis. In particular, it incurs the circle dependence among consensus errors and compression errors. To overcome this challenge, we propose novel strategies that enable the establishment of the convergence rate for our algorithm. Our theoretical results demonstrate how the compression operator influences the convergence rate. Finally, extensive experimental results confirm the efficacy of our proposed algorithm.

Optimization · Large Scale, Parallel and Distributed

Yuki Takezawa, Anastasiia Koloskova, Sebastian Stich

Decentralized SGD is a fundamental algorithm in decentralized learning, although the influence of an underlying network topology on its convergence behavior is not yet fully understood. Existing convergence analyses have shown that topologies with a small spectral gap significantly deteriorate the convergence rate of Decentralized SGD in both homogeneous and heterogeneous cases. However, many prior papers have reported that indeed the choice of the topology has a significant experimental impact in the heterogeneous case, but has little experimental impact on training behavior in the homogeneous case. In this paper, we present a tighter convergence analysis of Decentralized SGD, offering a more precise understanding of how topologies affect the convergence rate than the prior analysis. Specifically, unlike existing convergence analyses that used only the spectral gap as a property of the topology, our novel analysis shows that all eigenvalues of the mixing matrix affect the convergence rate. Throughout the experiments, we carefully evaluated the convergence behavior of Decentralized SGD and demonstrated that our novel convergence analysis can more accurately describe the effect of topology on the convergence rate.

Optimization · Stochastic

Hiroki Naganuma, Shagun Gupta, Youssef Briki, Ioannis Mitliagkas, Irina Rish, Parameswaran Raman, Hao-Jun Shi

To maximize hardware utilization, modern machine learning systems typically employ large constant or manually tuned batch size schedules, relying on heuristics that are brittle and costly to tune. Existing adaptive strategies based on gradient noise scale (GNS) offer a principled alternative. However, their assumption of SGD's Euclidean geometry creates a fundamental mismatch with popular optimizers based on generalized norms, such as signSGD / Signum ($\ell_\infty$) and stochastic spectral descent (specSGD) / Muon (Schatten-$\infty$). In this work, we derive non-Euclidean gradient noise scales for signSGD and spectral descent that naturally emerge from the geometry of its dual norm. We propose an efficient variance estimation procedure that leverages distributed data-parallel systems with negligible overhead. Our experiments demonstrate that adaptive batch size strategies using non-Euclidean GNS enable us to match the validation loss of constant-batch baselines while reducing training steps by up to 66\% for Signum and Muon.

Applications · Chemistry, Physics, and Earth Sciences

Yitian Wang, Zhao Yang, Angxiao Yue, Wentao Guo, Yaning Cui, Hongteng Xu

Modeling peptide cyclization is critical for the virtual screening of candidate peptides with desirable physical and pharmaceutical properties. This task is challenging because a cyclic peptide often exhibits diverse, ring-shaped conformations, which cannot be well captured by deterministic prediction models derived from linear peptide folding. In this study, we propose MuCO (Multi-stage Conformation Optimization), a generative peptide cyclization method that models the distribution of cyclic peptide conformations conditioned on the corresponding linear peptide. In principle, MuCO decouples the peptide cyclization task into three stages: topology-aware backbone design, generative side-chain packing, and physics-aware all-atom optimization, thereby generating and optimizing conformations of cyclic peptides in a coarse-to-fine manner. This multi-stage framework enables an efficient parallel sampling strategy for conformation generation and allows for rapid exploration of diverse, low-energy conformations. Experiments on the large-scale CPSea dataset demonstrate that MuCO significantly outperforms state-of-the-art methods in consistently in physical stability, structural diversity, secondary structure recovery, and computational efficiency, making it a promising computational tool for exploring and designing cyclic peptides.