NeurIPS 2026 论文列表

Distributed training is essential for scaling LLM training across thousands of GPUs. However, as distributed training requires complex implementations, they are prone to silent bugs, which do not produce explicit error signals but lead to in correct training outcomes. Common debugging practices based on monitoring training loss or gradient norm curves are slow or unable to detect bugs and also do not help localize bugs. We design and implement nnTrace, the first systematic differential testing system for detecting and localizing silent bugs in distributed training. nnTrace aligns intermediate tensors from distributed training with those from a trusted reference implementation. To properly compare the floating-point values in the corresponding tensors, we propose a novel mathematical analysis that provides a guideline for setting tolerances, enabling nnTrace to distinguish bug-induced errors from numerical errors. Experimental results demonstrate that nnTrace effectively detects 11 existing bugs and 3 new bugs in the widely used Megatron-LM framework. nnTrace is effective in various training recipes, including low-precision recipes involving BF16 and FP8. Notably, Megatron-LM has already adopted the method proposed by nnTrace in its development workflow. Our code is available at https://anonymous.4open.science/r/NECK-3C61/.

Large-scale AI training runs span days within fixed HPC allocations, making checkpointing essential for bounded recovery after failure. Asynchronous checkpointing hides checkpoint latency from the training loop, but does not guarantee that the most recent state has reached a durable target: data may still be in flight to the shared parallel filesystem (PFS) when a failure occurs, forcing rollback to an older checkpoint. We formalize this durability gap through a distributional model that combines hardware crossover analysis with statistical models of PFS completion time under scale and contention. Rank-local completion times follow a lognormal body distribution, while wall-clock completion is governed by an extreme-value process determined by the slowest writer. The resulting analysis yields a contention-dependent threshold,~$C^\star_p(n)$, beyond which the PFS is no longer the fastest path to durable persistence. Guided by this analysis, we design CacheX, a transparent staging layer that composes with PyTorch Distributed Checkpoint~(DCP). CacheX stages bulk data to local DRAM and replicates it to a neighboring node's NVMe, providing allocation-local durability while preserving eventual persistence to the PFS. On Frontier, across IOR runs up to 256~nodes and end-to-end FSDP+DCP training up to 64~nodes, CacheX reduces write-time variability by up to~$43\times$, achieves a replica-durable effective throughput of~${\sim}\,2.5$ to $2.8$\,GiB/s per node decoupled from PFS contention, and reduces worst-case rollback exposure. Our results show that staging is not a universal accelerator; rather, it is a conditional policy for selecting the first durable checkpoint boundary under measured scale, load, and tail behavior.

In centralized, distributed, and federated learning with stochastic gradients and $n$ workers, it was recently shown that it is infeasible to find an $\varepsilon$-stationary point faster than $\tilde{\Omega}(\min[{\frac{d \kappa L \Delta}{\varepsilon}} + \frac{h L \Delta}{\varepsilon} + \frac{h \sigma^2 L \Delta}{n \varepsilon^2}, {\frac{h \sigma^2 L \Delta}{\varepsilon^2}} + \frac{h L \Delta}{\varepsilon}])$ seconds in both homogeneous and heterogeneous settings under standard assumptions: $L$-smoothness, $\sigma^2$-bounded unbiased stochastic gradients, and lower boundedness of the function, i.e., $f(x) \geq f^\star$ for all $x \in \mathbb{R}^d$, where $\Delta = f(x^0) - f^*$, $h$ is the computation time, $\kappa$ is the communication speed between the workers and the server, and $d$ is the dimension of the iterates and gradients. This result is pessimistic since it does not allow a complexity in which both ${\frac{d \kappa L \Delta}{\varepsilon}}$ and ${\frac{h \sigma^2 L \Delta}{\varepsilon^2}}$ improve with $n$, even when using random sparsification techniques; moreover, this lower bound can be matched by either non-distributed SGD or vanilla Synchronous SGD, which reduces the impact of recent progress in the design of compression-based methods. In this work, we challenge this limitation and propose new compressed methods, Inkheart SGD and M4, and show that under an additional structural assumption, which is necessary due to the lower bound and which does not restrict the class of considered problems, we achieve new state-of-the-art time complexities that break this pessimistic barrier and allow scaling with the number of workers $n$.

Existing benchmarks for multimodal memory reasoning largely evaluate systems within pre-assembled contexts, but under-evaluate whether agents can use evidence distributed across independently originated sources. We argue that source-distributed memory composition is an important and under-examined bottleneck in multimodal agent memory, especially when relevant evidence is fragmented across heterogeneous artifacts such as conversations, profiles, screenshots, tables, images, and documents. To address this gap, we introduce Source-distributed Multimodal Memory Benchmark, which measures whether agents can retrieve, align, and compose multimodal evidence scattered across multiple sources rather than reason within a single curated context. SMMBench evaluates four core capabilities: (1) cross-source multimodal reasoning; (2) conflict resolution; (3) preference reasoning; (4) memory-grounded action prediction. The benchmark contains $1,877$ samples grounded in $264$ sources. Experiments on representative memory-style and retrieval-based baselines show that current systems still struggle on these capabilities, positioning source-distributed multimodal memory as an important and still under-evaluated challenge for multimodal agents. Our data and code are available at https://anonymous.4open.science/r/NeurIPSBenchmarksW8kP6gV/.

Asynchronous stochastic gradient descent (ASGD) is a standard way to exploit heterogeneous compute resources in distributed learning: instead of forcing fast workers to wait for slow ones, the server updates the model whenever a gradient arrives. Vanilla ASGD applies each arriving gradient with the same weight. When local data distributions are heterogeneous, this becomes problematic: faster workers contribute more updates, and we show theoretically that the method is biased toward a frequency-weighted average of the local objectives rather than the desired global objective. Existing remedies typically move away from the simple ASGD template by introducing gathering phases, buffering, or extra memory. We show that this is unnecessary. Keeping the standard ASGD mechanism, we recover the correct objective by rescaling worker-specific stepsizes in proportion to their computation times, so that each worker contributes the same aggregate learning rate over a cycle. In the non-convex setting, under smoothness and bounded heterogeneity assumptions, we prove that the resulting method, Rescaled ASGD, converges to stationary points of the correct global objective in the fixed-computation model. Its time complexity matches the known lower bound in the leading term, while the effects of staleness and data heterogeneity appear only in lower-order terms. Experiments confirm that the method converges to the correct objective and is competitive with state-of-the-art baselines.

Pre-training large language models on massive GPU clusters has made hardware faults routine rather than rare, driving the need for resilient training systems. Yet existing frameworks either focus on specific parallelism schemes or risk drifting away from a failure-free training trajectory. We propose ReCoVer, a resilient LLM pre-training system that upholds a single invariant: each iteration keeps the number of microbatches constant, ensuring per-iteration gradients remain stochastically equivalent to a failure-free run. The framework is organized as three decoupled protocol layers: (1) Fault-tolerant collectives that isolate faults from propagating across replicas; (2) in-step fine-grained recovery that preserves intra-iteration progress and prevents gradient corruption; (3) versatile-workload policy that dynamically redistributes microbatch quotas across the survivors. The design is parallelism-agnostic, integrating directly with both 3D parallelism and Hybrid Sharded Data Parallel (HSDP) as a drop-in substrate. We evaluate our implementation on end-to-end pre-training tasks for up to 512 GPUs, ReCoVer successfully preserves the training trajectory from a failure-free reference despite of 256 GPUs lost spread across the run. For comparison with checkpoint-and-restart baselines, ReCoVer demonstrates $2.23\times$ higher effective throughput after successive failures. This advantage results in ReCoVer processing 74.9% more tokens at 234 GPU-hours, with the gap widening as the training prolongs.

Trilevel learning has emerged as a powerful paradigm across various domains in machine learning and networking. In practical scenarios, data is frequently generated and distributed across multiple decentralized nodes. While existing distributed trilevel learning methods circumvent the transmission of raw data, they remain susceptible to privacy leakage. Furthermore, cutting plane methods, which are widely adopted in bilevel and trilevel learning, exhibit analogous intrinsic privacy risks. To address these critical issues, we propose a Differentially Private Distributed Trilevel Optimization (DPDTO) framework. In DPDTO, we first develop a value-function-embedded privacy-preserving cutting plane method for trilevel optimization problems, comprising an inner-layer value function reformulation and an outer-layer cutting plane relaxation. Building upon this, a distributed optimization algorithm is introduced to effectively address trilevel optimization while preserving privacy. Theoretically, we provide a comprehensive analysis for the proposed framework, including asymptotic convergent relaxation, non-asymptotic convergence rate, privacy guarantees, and communication complexity, uncovering a tripartite trade-off among privacy, utility, and complexity in trilevel optimization. Extensive experimental results on two distributed trilevel optimization problems further demonstrate the effectiveness and superiority of the proposed DPDTO.

Random Forests (RF) are among the most powerful and widely used predictive models for centralized tabular data, yet few methods exist to adapt them to the federated learning setting. Unlike most federated learning approaches, the piecewise-constant nature of RF prevents exact gradient-based optimization. As a result, existing federated RF implementations rely on unprincipled heuristics: by aggregating decision trees trained independently on clients, they fail to optimize the centralized decision-tree impurity criterion, which constitutes the ground-truth in federated settings, even under simple covariate shifts. We propose FedForest, a new federated RF algorithm for horizontally partitioned data that naturally accommodates diverse forms of client data heterogeneity, from covariate shift to more complex concept shift mechanisms. We prove that our splitting procedure, based on aggregated client statistics, closely approximates the split selected by a centralized algorithm. Moreover, FedForest allows splits on client indicators when beneficial, enabling a non-parametric form of personalization that is absent from prior federated random forest methods. Empirically, we demonstrate that the resulting federated forests closely match centralized performance across heterogeneous benchmarks while remaining communication-efficient.

Reasoning LLMs produce longer outputs, requiring speculative decoding drafters trained on extended sequences. Parallel drafting—predicting multiple tokens per forward pass—offers latency benefits over sequential generation, but training complexity scales quadratically with the product of sequence length and parallel positions, rendering long-context training impractical. We present P(arallel)-EAGLE, which transforms EAGLE from autoregressive to parallel multi-token prediction via a learnable shared hidden state. To scale training to long contexts, we develop a framework featuring attention mask pre-computation and sequence partitioning techniques, enabling gradient accumulation \textit{within} individual sequences for parallel-prediction training. We implement P-EAGLE in vLLM and demonstrate speedups of 1.10×–1.36× over autoregressive EAGLE-3 across GPT-OSS 120B, 20B, and Qwen3-Coder 30B.

A classic result in learning and optimization theory is that exponentiated gradient descent (EG) or the multiplicative weight update method computes an $\\epsilon$-approximate minimizer of an $\\ell_1$-Lipschitz function over the $d$-dimensional probability simplex $\\Delta_d :=\\{x \in \\mathbb{R}^d_{\\geq0}:\\sum_{i\in[d]}x_i=1\\}$ with $\\tilde{O}(\epsilon^{-2})$ queries to a first-order oracle. We investigate whether there are improved parallel algorithms for this fundamental problem. We show that, up to logarithmic factors, the answer is negative for $\\epsilon=\\tilde \\Omega(d^{-1/6})$ - any randomized algorithm that makes $O(\\mathrm{poly}(d))$ subgradient queries per round requires $\\tilde{\\Omega}(1/\\epsilon^2)$ rounds to output an $\epsilon$-optimal point with constant success probability. Moreover, to obtain this result we provide an analogous characterization of the parallel complexity of minimizing a $\ell_1$-Lipschitz convex function over the unit $\\ell_1$ ball. Previous lower bounds for non-constant $\\epsilon$ for this $\\ell_1$-Lipschitz convex optimization problem either made additional assumptions on the domain, or instead attained a bound of $\\Omega(\\epsilon^{-2/3})$ [DG19].

Data heterogeneity and low client participation are two key challenges in federated learning (FL). Client-reshuffling-based FL methods were recently introduced to improve participation efficiency by visiting each client once per meta-epoch; however, the resulting \textit{without-replacement} sampling induces inter-round dependence and conditional bias. As a consequence, existing client-reshuffling methods can still suffer from the data heterogeneity challenge due to this dependence. To bridge this gap, we propose \textbf{FedCDR}, a client-reshuffling FL algorithm built on Douglas–Rachford splitting. FedCDR supports \textit{inexact} local proximal updates via iterative solvers, enabling a practical communication–computation trade-off. For smooth nonconvex objectives, FedCDR with inexact local solvers attains a state-of-the-art $O(\epsilon^{-1})$ communication complexity to reach an $\epsilon$-approximate stationary point (i.e., $\mathbb{E}\|\nabla f(\tilde{x})\|^2 \le \epsilon)$, with the leading constant that is \textbf{independent of data heterogeneity} (i.e., it does not scale with common measures of heterogeneity). Technically, our analysis operates at the meta-epoch level: we control the deviation between reshuffled and full-client updates, construct a tailored potential function with provable descent, and sum over each meta-epoch to eliminate reshuffling-induced dependence. Experiments on synthetic tasks and benchmark datasets under heterogeneous partitions, including a 10,000-client setting, demonstrate consistent improvements over strong baselines and their client reshuffling variants.

Modern large language models keep the feed-forward network (FFN) intermediate width identical across every layer, a convention inherited from the original Transformer. Yet mechanistic work suggests layers at different depths do qualitatively different work, raising a natural question: are their capacity demands also non-uniform, and if so, can a fixed parameter budget be allocated more wisely than uniformly? To probe this, we apply principal component analysis (PCA) to the input of each FFN down-projection and measure, for every layer, the smallest number of principal components needed to reconstruct the activation up to a fixed cosine similarity --- the layer's intrinsic dimensionality. Scanning 58 open-weight LLMs spanning 14 families and 70M to 72B parameters, we find a recurring middle-heavy profile in which middle layers carry higher intrinsic dimensionality than edge layers; the profile is robust to the calibration corpus, even across English and Chinese, and is largely preserved through post-training fine-tuning. As a direct test that this signal is structurally meaningful, we use it to guide width pruning in a knowledge-distilled student under a matched parameter budget: the PCA-guided allocation outperforms the uniform-width convention by $+2.7$ points on Qwen2.5-3B and $+2.3$ on Qwen3-8B, while a same-scale teacher (LLaMA-3-8B) shows only $+0.9$, and a budget-matched inverted control consistently does worst. Across these three teachers, the size of the gap scales with how unevenly the teacher's intrinsic dimensionality is distributed across layers --- a single-number summary we call $\sigma_{99}$ --- and the advantage persists through downstream supervised fine-tuning.

Communication is a major bottleneck in distributed learning, especially in large-scale settings and in federated learning environments with slow links. Three standard ways to reduce this cost are communication compression, local training, and communication-computation overlap. Methods that combine these ingredients are used in practice and have been found to be effective for large-scale training, but there is little theory for methods that combine all three. We study a heterogeneous-compute setting in which different workers may take different numbers of local steps, and we propose LOSCAR-SGD, a Local SGD method that communicates only a sparse subset of model coordinates and continues optimizing while communication is in flight. A key ingredient is a delay-corrected merge rule that incorporates delayed synchronized information without discarding the progress made during the overlap phase. We give convergence guarantees for smooth non-convex objectives and show how sparsity, overlap, and worker heterogeneity affect the rate. To the best of our knowledge, this is the first theory for this combination of ingredients. Experiments further show that communication-computation overlap reduces training time and that the delay-corrected merge outperforms naive overwriting.

Running dense, open-weight large language models on commercially accessible workstations or single-GPU cloud setups is increasingly desirable due to cost and privacy constraints, particularly on tasks like coding and reasoning. Starting at around 24B parameters, however, they exceed the available GPU VRAM on these setups, forcing reliance on offloading methods that limit throughput, while leaving substantial CPU compute underutilized. Heterogeneous CPU–GPU speculative decoding is the dominant lossless-acceleration framework in this regime, but existing methods are limited by what we refer to as the **heterogeneity gap**: their CPU and GPU do not co-execute, their draft is a separate set of weights rather than a subnetwork of the verifier, and their VRAM footprint cannot scale to fill the available budget. We observe that channel-saliency methods induce an ordering on the transformer’s FFN channels so that the top-$k$ prefix closely approximates the full output, with graceful degradation as $k$ decreases. Building on this observation, we introduce **Speculative Tensor Parallelism (STeP)**, a training-free, self-speculative method that extracts a subnetwork of the verifier as a GPU-resident draft, places the remainder on the CPU, and verifies via concurrent CPU and GPU computation, provably preserving the verifier’s sampling distribution. Across dense models from 24B to 123B, GPU memory budgets from 32 to 192 GB, and five benchmarks, STeP outperforms SpecExec and SubSpec by up to $1.4 \times$ and $1.8 \times$ respectively, surpasses either tensor parallelism or speculative decoding alone, maximally utilizes the available VRAM, and closes the heterogeneity gap.

The rapid growth of large language models is driving organizations to expand their GPU clusters, often with GPUs from multiple vendors. However, current deep learning frameworks lack support for communication across heterogeneous vendor GPUs, leading to inefficiency and higher costs. We present HetCCL, a collective communication library that unifies vendor-specific backends and enables RDMA-based communication across GPUs without requiring driver modifications. HetCCL introduces two novel mechanisms that enable cross-vendor communication while leveraging optimized vendor libraries, NVIDIA NCCL and AMD RCCL. Evaluations on a multi-vendor GPU cluster show that HetCCL matches NCCL and RCCL performance in homogeneous setups while extending support to heterogeneous environments, enabling practical and efficient training across NVIDIA and AMD GPUs without changes to existing deep learning applications.

Synthesizing Mixed-Boolean Arithmetic (MBA) expressions from input-output examples is central to program deobfuscation and also useful for compiler optimization, reverse engineering, and cryptanalysis. Existing MBA synthesizers are typically CPU-based and scale poorly on large specifications or complex targets. Recent GPU-accelerated synthesis methods achieve large speedups in qualitative settings, but they depend on caching observationally equivalent candidates; this strategy breaks down for MBA because candidate outputs are quantitative bitvectors and the behavioral space is enormous. We present SIMBA (Synthesis of Mixed-Boolean Arithmetic), a GPU-accelerated MBA synthesizer built around cache-free bottom-up enumeration. SIMBA avoids language caches entirely and uses a GPU-oriented enumeration design that keeps work local and highly parallel. In experiments, SIMBA is substantially faster than prior MBA synthesis tools, handles larger specifications, and reaches expression sizes that existing methods fail to solve. These results establish cache-free GPU synthesis as a practical and scalable approach for quantitative domains, and identify it as a strong alternative to cache-centric designs.

We study price-adjustment dynamics for computing competitive equilibria (CE) in Fisher markets with chores. Unlike in classical goods markets, prices in chores markets are payments for taking on undesirable tasks, and natural excess-demand dynamics can fail; even the naïve analogue of Walrasian tâtonnement may diverge. Recent work of Chaudhury et al. (2025) overcomes this obstacle via \emph{relative tâtonnement}, which subtracts the average excess-demand signal from the excess demand vector. This recovers convergence, but at the cost of coupling the price updates across all chores. This leaves open whether such global coupling is inherent, or whether convergent tâtonnement can be recovered through a genuinely local update in which each chore reacts only to its own excess demand. We answer this question affirmatively through \emph{multiplicative tâtonnement}, a fully distributed dynamics in which each chore price is updated using only its current price and its own excess-demand signal. Although the update contains no explicit normalization term, Walras' law and the multiplicative form of the update implicitly preserve the relevant aggregate price geometry. We prove that multiplicative tâtonnement converges to a CE in any chores Fisher market with continuous, convex, and 1-homogeneous (CCH) disutilities. For convex CES disutilities, we further prove an approximate-CE convergence rate with the same O(1/\varepsilon^2) dependence as relative tâtonnement, but with improved dependence on problem constants. Experiments on real-world and simulated instances show that multiplicative tâtonnement is substantially faster in practice, often by an order of magnitude.

Data heterogeneity and multiple local update steps in federated learning can increase client drift and make global parameter aggregation less stable. Existing Sharpness Aware Minimization (SAM) based federated learning methods usually try to ease this problem by flattening local neighborhoods around client models. However, they do not explicitly optimize on the geometry region where the global parameters are actually aggregated. As a result, the aggregated global model often cannot fully benefit from local flattening. To address this problem, we propose a federated restricted SAM method called FedRSAM, which improves the stability of federated optimization by flattening the parameter aggregation region. First, we build a cone-shaped geometry region using the global update direction from the previous round and the variance of client parameter aggregation, in order to predict where the aggregated parameters may fall in the next round. Then, we restrict the search of SAM perturbation directions to the cone-manifold and only flatten the loss surface inside this region. Extensive experiments on multiple benchmarks and on both vision and natural language backbones verify the superiority of our method in terms of convergence stability and performance under heterogeneous settings.

Long-context training of large language and multimodal models couples expressive structured attention masks (document, prefix-LM, share-question, sliding-window, block-sparse) with multi-node context parallelism, exposing the mask as a dominant bottleneck. The column-wise sparse mask of FlashMask compresses such masks from $O(n^2)$ to $O(n)$ via four query-axis boundary vectors and covers a broad family of structured patterns. Implemented for Ampere GPUs, FlashMask does not exploit FlashAttention-3's TMA/WGMMA warp-specialized pipeline on Hopper GPUs, and offers no distributed support. We observe that the column-wise sparse mask is indexed by key positions but stores query-axis interval boundaries. Under any multi-chunk context-parallel assignment, this property reduces per-rank mask localization to a provably zero-communication \emph{clip--shift--max} on the boundary vectors, directly compatible with the single-GPU kernel. Building on this property, FlashMask-3 introduces a mask-aware Hopper kernel with a three-role warp specialization atop the FlashAttention-3 TMA/WGMMA pipeline, together with a mask-aware distributed extension: zero-communication mask localization, a sparsity-aware load balancer (HCC-IPO), and a mask-aware compute-communication overlap with mask-driven communication pruning, whose forward all-gather additionally adopts topology-aware hierarchical routing. On a single H100 GPU across 12 structured masks at context lengths from 4K to 128K, FlashMask-3 improves achieved TFLOPs over FlashMask, FlexAttention, and MagiAttention by 45\% to 141\%, 34\% to 71\%, and 2\% to 44\%, respectively. Under context-parallel degrees from 4 to 32 with 8K sequence per device, it attains 31\% to 75\% compute-communication overlap efficiency and improves throughput over Megatron-Core CP (ring), Megatron-Core CP (all-gather) and MagiAttention by $3.25\times$ to $5.11\times$, $1.97\times$ to $3.26\times$ and $1.06\times$ to $4.09\times$ respectively. The source code will be released on GitHub.