AutoParallel: Composable, PyTorch-Native Automatic Parallelism for Production Training

AutoParallel: Composable, PyTorch-Native Automatic Parallelism for Production Training Kaijian Wang1,2,∗ , Sanket Purandare1 , Aditya Venkataraman1 , Will Constable1 , Ivan Kobzarev1 , Elias Ellison1 , Simon Fan1 , Edward Z. Yang1 , Davide Italiano1 , Patrick Labatut1 , Yuke Wang2 , Francisco Massa1,† 1

arXiv:2611.08660v1 [cs.DC] 5 Nov 2026

†

Meta Superintelligence Labs, 2 Rice University Project lead, ∗ Work done during an internship at Meta

Training large models requires multiple interacting forms of parallelism. As models and clusters scale, finding and maintaining efficient parallelization plans becomes essential for training performance. Diverse model architectures, workloads, and hardware topologies demand flexible strategies, yet production PyTorch systems typically rely on hand-written recipes that are costly to adapt and tend to reuse familiar patterns, which may be suboptimal for a given workload. We present AutoParallel, a PyTorch-native planner for distributed model training. AutoParallel searches parallelization choices jointly across the training computation, allowing established techniques to be combined and unconventional strategies to emerge. It expresses its decisions using native PyTorch abstractions and produces ordinary PyTorch programs, avoiding a separate execution runtime while composing directly with the PyTorch compilation and training stack. AutoParallel also composes automatic decisions with model-specific distributed code, allowing dynamic or specialized regions to remain under user control. Plans are inspectable and replayable artifacts. Across dense, hybrid, and mixture-of-experts models on meshes of up to four dimensions and 128 NVIDIA H100 GPUs, AutoParallel matches or exceeds expert-written TorchTitan recipes on most configurations, with throughput gains of up to 18.3%, and discovers plans that uniform expert recipes do not express. Date: October 30, 2026 Code: https://github.com/meta-pytorch/autoparallel

1

Introduction

Model training has grown at an extraordinary pace. The Hugging Face Hub alone hosts over three million models (Tazi et al., 2025); inside organizations, teams routinely train and fine-tune models for tasks spanning natural language, vision, multimodal generation, and more. Widely adopted frameworks such as PyTorch (Paszke et al., 2019), combined with open model architectures and community libraries, have lowered the barrier to training models. Achieving efficient training, however, is far from straightforward. Scaling model training beyond a single accelerator requires carefully combining multiple forms of parallelism, including Data Parallelism (Rajbhandari et al., 2020; Zhao et al., 2023), Tensor Parallelism (Shoeybi et al., 2019; Korthikanti et al., 2023), Context Parallelism (Liu et al., 2023a), Pipeline Parallelism (Huang et al., 2019; Narayanan et al., 2019), and Expert Parallelism (Fedus et al., 2022), together with system-level optimizations including activation recomputation (Chen et al., 2016; Korthikanti et al., 2023), mixed precision (Micikevicius et al., 2018), and deep learning compilers (Ansel et al., 2024). Moreover, identifying an effective training configuration is highly context-specific: it depends jointly on model architecture, workload shapes, and hardware topology (Narayanan et al., 2021). Hand-tuning a distributed training recipe requires dedicated infrastructure expertise and iterative experimentation, and this cost is amortized only when a single configuration is reused across a long production run. In practice, this makes hand-optimization economical for a narrow regime: a small number of flagship models, each trained by a dedicated team over weeks or months. It does not scale to the broader landscape of model training, where researchers iterate on new architectures faster than recipes can be tuned, product 1

teams train many models for different tasks without a dedicated infrastructure effort for each, and open-source practitioners need efficient training out of the box. A recent benchmarking study by Hugging Face (Tazi et al., 2025) illustrates the scale of this challenge: of 3,306 parallelism configurations for LLaMA models on up to 512 H100 GPUs, nearly half failed with out-of-memory errors, and the sweep cost approximately $469K while covering only a narrow slice of the configuration space. A substantial body of prior work (Jia et al., 2019; Shazeer et al., 2018; Zheng et al., 2022; Xu et al., 2021; Miao et al., 2023; Jia et al., 2022; Unger et al., 2022; Lin et al., 2024, 2025) has demonstrated that automatic parallelization is feasible. Despite these advances, automatic parallelism has not replaced manual recipe development in practice. The key barriers are threefold: (1) many prior systems rely on non-PyTorch runtimes or specialized compiler stacks, placing them outside the frameworks most practitioners actually use; (2) they often produce opaque solutions that are difficult to inspect, modify, or debug, offering no guarantee of matching the quality of hand-tuned production recipes; and (3) they compose poorly with the optimizations that production training requires, including compilation, activation checkpointing, and mixed precision. Moreover, modern training recipes increasingly involve dynamic parallelism patterns, such as context parallelism and data-dependent expert routing in mixture-of-experts layers. These patterns require specialized communication primitives beyond static operator placement, and existing systems that operate on static computation graphs cannot readily support them. As a result, although PyTorch is widely used for both research and production training (Paszke et al., 2019; Ansel et al., 2024), PyTorch practitioners still lack an automatic path from model definition to efficient distributed training (§2.3, §7).

Figure 1 AutoParallel system overview. Colored boxes schematically illustrate tensor sharding, and bold outlines mark

the strategy selected for each operator. The orange local_map region keeps its user-declared placements, and user constraints feed back into the search.

In this paper, we present AutoParallel, a PyTorch-native single-program, multiple-data (SPMD) planner for distributed model training that closes this gap. As Figure 1 shows, we add a planning layer between graph capture and plan materialization in the PyTorch stack: it enumerates valid DTensor placements for every operator in the joint forward-backward graph, selects one per operator under a hardware-aware cost model, and emits an ordinary PyTorch module that the existing compiler stack executes without a separate runtime (§5). For dynamic parallelism patterns such as context parallelism and mixture-of-experts, we provide composition boundaries that let manually parallelized regions interoperate with auto-planned computation (§5.3). AutoParallel is open source1 (Meta PyTorch, 2026), has been presented at the PyTorch Conference (Massa et al., 2025; Purandare and Huang, 2026), and is integrated into TorchTitan’s GraphTrainer for production training (PyTorch Team, 2026). We evaluate AutoParallel on dense, hybrid, mixture-of-experts, and long-context workloads across multidimensional meshes and up to 128 NVIDIA H100 GPUs, comparing against expert-written TorchTitan (Liang et al., 2024) recipes under matched meshes and runtime optimizations (§6). AutoParallel matches or exceeds the expert recipes on most configurations without per-model tuning, and when the workload changes, replanning recovers the performance that a stale plan leaves behind. This paper makes the following contributions: • Graph-wide placement planning. We formulate SPMD parallelization as a single optimization over the joint forward-backward graph under a hardware-aware cost model. We introduce a scalable approximate 1 https://github.com/meta-pytorch/autoparallel

2

solver that produces near-optimal plans on instances small enough to verify against an exact solver and scales to meshes where exact solving is infeasible. • Native recipe generation. We materialize decisions as an ordinary PyTorch program with no separate runtime and integrate it with TorchTitan’s GraphTrainer (Purandare and Huang, 2026; PyTorch Team, 2026) for full-graph Inductor compilation with communication-computation overlap and collective bucketing. • Practical control and composition. We make plans inspectable and replayable, allow users to constrain placements, and compose them with user-controlled dynamic parallelism regions (context parallelism, mixture-of-experts) via explicit placement contracts (§3, §5). • Production evaluation. We report end-to-end performance against expert recipes, replanning across sequence lengths, planner scalability, and cost-model fidelity (§6 and Appendix B).

2

Background and Challenges

2.1

Parallelism as Per-Operator Placement

A training computation can be represented as a dataflow graph of tensor operators spanning both the forward and backward passes. A device mesh is a multi-dimensional array of accelerators, where each axis corresponds to a parallelism dimension with its own communication bandwidth. AutoParallel targets single-program, multiple-data (SPMD) execution (Xu et al., 2021; Lepikhin et al., 2020): every device in the mesh runs the same program on its local shard of each tensor, and devices synchronize through collective communication. A placement assigns each tensor one of three types along each mesh axis: Replicate (the full tensor is present on every device along that axis), Shard(d) (the tensor is partitioned along dimension d), or Partial (each device holds a partial result awaiting reduction). On a D-dimensional mesh, a tensor’s placement is a D-tuple drawn from these types. Under this view, named parallelism strategies such as FSDP, tensor parallelism, and sequence parallelism are not first-class concepts but emergent consequences of per-operator placement decisions within a single SPMD program; §4.1 develops this correspondence with concrete examples. Data-parallel axes always shard the input batch; on a DDP axis the weights are replicated, and on a DPS axis they are sharded as in FSDP. HSDP combines a DDP axis with a DPS axis.

2.2

PyTorch Compilation and Distribution Stack

AutoParallel builds on three layers of existing PyTorch infrastructure rather than introducing a separate intermediate representation or runtime. TorchDynamo and AOTAutograd (Ansel et al., 2024) capture the model’s forward computation and automatically derive the backward pass, producing a single joint computation graph at the ATen operator level. DTensor and DeviceMesh (Liang et al., 2024) provide the placement vocabulary and redistribution semantics: each distributed tensor carries its mesh and placement metadata, and transitions between placements are realized through collectives. SimpleFSDP (Zhang et al., 2024) showed that FSDP communication can be captured this way: storing parameters as sharded DTensors and triggering all-gather via parametrization makes the collectives ordinary graph nodes visible to the compiler; we generalize this from a fixed sharding rule to per-operator placement search (§4.1). TorchInductor (Ansel et al., 2024) compiles the resulting graph into fused kernels and schedules communication-computation overlap. Because AutoParallel produces an ordinary PyTorch module with DTensor parameters, the output flows through this stack unchanged: Inductor compilation, distributed checkpointing, mixed precision, and activation checkpointing all compose without adaptation. GraphTrainer (Purandare and Huang, 2026) is TorchTitan’s full-graph training backend. AutoParallel uses it as backend and an entry point to standard Inductor compilation.

2.3

Challenges

The per-operator placement view (§2.1) makes the plan space expressive, but it raises three challenges for automatic parallelization in production PyTorch training.

3

C1: Placement decisions are fine-grained and workload-specific. Even with the mesh and parallelism degrees

fixed, operator placements alone matter: constraining one or two projections of LLaMA 3 8B on a 4×8 mesh and re-solving yields 11 feasible plans whose measured step times span 2.2×, from 365 to 790 ms (Appendix B). The best plan also shifts with the workload; at 32K tokens, a fresh plan runs 51% faster than one tuned for 2K-token sequences (§6.4). A planner must therefore search per-operator placements over the whole training graph and be cheap enough to rerun. C2: The search space grows exponentially with mesh dimensionality. Production training stacks HSDP, tensor,

context, and expert parallelism into high-dimensional meshes whose axes differ in bandwidth by an order of magnitude. Because a placement is a D-tuple (§2.1), the candidate strategies per operator grow exponentially with the mesh dimensionality D. For the LLaMA 3 8B model, options grow from 12.9K decision variables on a 1D mesh to 12.2M on 3D, which takes 30 minutes to solve, and to 523M on 4D, where the problem cannot be built within an hour (§6.5). High-dimensional meshes therefore call for approximate search. C3: Plans should compose with the production stack. A plan is most useful in production when it runs on

the same compiler stack as hand-written recipes, composes with activation checkpointing, mixed precision, and distributed checkpointing, and leaves specialized regions, such as custom context-parallel kernels and data-dependent mixture-of-experts routing, under user control. Engineers also need to inspect, constrain, and replay plans; our own evaluation relies on all three (§6.3, §6.4, Appendix B). Systems with a separate runtime make these requirements hard to meet in PyTorch. We address C1 with graph-wide placement search (§4.1), C2 with a seeded MRF solver (§4.3), and C3 with native DTensor output and local_map boundaries (§5).

3

System Overview

AutoParallel takes a single-GPU model definition and a device mesh as input and produces a parallelized PyTorch module as output. Two principles guide its design. First, every component builds on the PyTorch compilation and distribution stack described in §2.2, eliminating the need for a separate intermediate representation or custom runtime. Second, we draw an explicit boundary between computation optimized automatically and computation the user controls directly, composing the two through placement contracts rather than attempting to absorb all distributed patterns into the search. Both principles address C3 in §2.3.

3.1

End-to-End Workflow

1 Figure 1 summarizes the workflow. Given a model, a device mesh, and representative inputs, AutoParallel ○ traces the model once with TorchDynamo and AOTAutograd into a joint forward-backward ATen graph, which 2 enumerates candidate DTensor strategies for every operator, each defining how requires no GPU. It then ○ the operator executes within a single SPMD program; because candidates come from DTensor’s propagation 3 the solver selects rules, every plan the solver can select is semantically valid by construction (§4.1). Next, ○ one strategy per operator to minimize estimated computation and communication cost under placement 4 the plan is materialized as DTensor parameters plus explicit redistribution constraints (§4.2, §4.3). Finally, ○ collectives, producing an ordinary PyTorch module that runs eagerly or through the PyTorch compiler stack (§5).

3.2

Optimization Boundary

Standard tensor operators, including matrix multiplications, normalizations, attention projections, and element-wise functions, participate in placement search. Regions with dynamic or specialized distributed behavior, such as mixture-of-experts routing and context-parallel attention, remain user-written through local_map placement contracts, which the solver treats as fixed constraints (§5.3).

3.3

User Interface

At its simplest, AutoParallel takes a model, a mesh, representative inputs, and the desired output placements:

4

mesh = init_device_mesh("cuda", (4, 8)) with torch.device("meta"): model = Transformer(config) par_model = auto_parallel(model, mesh, sample_inputs=(x,), out_shardings=(Shard(0), Replicate())) par_model.to_empty(device="cuda") par_model.init_weights()

Input placements come from representative DTensors; ordinary tensors are replicated. After weight initialization, the module supports a standard training loop. Users can constrain input/output placements and parameter storage, configure the solver, inspect operator costs, and serialize plans for replay. For auto-planned operators, mesh names are descriptive labels and placements are chosen freely across axes.

4

Placement Planning

We formulate parallelization as a graph-wide placement optimization over the joint forward-backward graph. Restricting the search to SPMD placements keeps the space tractable while still expressing data parallelism, FSDP, tensor and sequence parallelism, and their combinations; pipeline parallelism, which partitions the graph into stages rather than sharding operators, is orthogonal and discussed in §7.

4.1

Strategy Space

The search space is built from PyTorch DTensor’s placement vocabulary. On a D-dimensional device mesh, the placement of a tensor is a D-tuple, one element per mesh dimension, drawn from {Replicate, Shard(d), Partial} (§2.1). A strategy for an operator specifies the placement of every input and output tensor; the set of valid strategies for an operator is determined by DTensor’s propagation rules, which encode the semantics of each ATen operation. Each rule admits only input and output placements under which the per-device local computations, combined with the implied collectives, reproduce the result of the unsharded operator. Every candidate plan is therefore semantically equivalent to single-device execution up to floating-point reassociation in reductions, the same property that holds for hand-written data- and tensor-parallel recipes. Every enumerated strategy is valid in this sense, and so is any combination of them: plan application inserts the required collective wherever adjacent placements disagree (§5.1), so the selected plan computes correct results and the planner can optimize purely for performance. Training with AutoParallel plans reproduces the expert recipe’s loss curve within the variation caused by changing the mesh alone (Appendix D). Table 1 Strategies for Y = XW on a two-dimensional mesh. Each cell is a placement tuple with one entry per mesh

axis. S(d) = Shard(d), R = Replicate, P = Partial, AG = AllGather, AR = AllReduce; (0): on mesh axis 0 only.

DP FSDP Col-TP Row-TP FSDP+TP

X

W

Y

Comm.

S(0), S(0) S(0), S(0) R, R R, S(1) S(0), R

R, R S(0), S(0) R, S(1) R, S(0) S(0), S(1)

S(0), S(0) S(0), S(0) R, S(1) R, P S(0), S(1)

AR ∇W AG W — AR Y AG W (0)

To ground this abstraction, consider a matrix multiplication Y = XW on a two-dimensional mesh, illustrated in Table 1. Sharding X along the batch dimension across the entire mesh while replicating W produces data parallelism; additionally sharding W across the mesh produces FSDP (the table shows W in its stored placement; FSDP all-gathers it before the multiplication). These patterns are not declared; they are the operational consequences of particular placement tuples. Combining FSDP on one axis with column-parallel TP on another produces the FSDP + TP strategy: the planner assigns each axis its role based on cost, not naming convention. The optimizer selects among all valid candidates based on their estimated cost on the target hardware. In general, the candidate set for each operator is the Cartesian product of per-dimension placements, filtered 5

by operator semantics. On a D-dimensional mesh, a tensor with n dimensions admits up to (1 + n)D candidate placements (Replicate plus Shard along each of n dimensions, raised to the Dth power across mesh dimensions). After semantic filtering, which removes placements that shard a dimension smaller than the mesh degree or that violate operator constraints, the typical count per node is 10–30 on a two-dimensional mesh. On meshes with three or more dimensions, the combinatorial growth makes even enumeration expensive, motivating the seeding technique described in §4.3. Operators wrapped in local_map (§5.3) receive a single-strategy candidate set determined by the user’s declared boundary placements, removing them from the search. Given the per-node candidate sets, the placement problem is to select exactly one strategy sv ∈ Sv per node v to minimize the sum of per-operator computation cost and per-edge redistribution cost: X X min ccomp (v, sv ) + credist (su , sv ) (1) {sv }

v∈V

(u,v)∈E

where V is the set of nodes and E the set of edges in the computation graph, ccomp (v, sv ) is the computation cost of executing node v under strategy sv , and credist (su , sv ) is the communication cost of the redistribution collective required when the output placement of u disagrees with the input placement expected by v (zero when they agree). The optimization is subject to: (i) uniqueness: each node selects exactly one strategy; (ii) boundary constraints: placements of model inputs, outputs, and local_map boundaries are fixed; (iii) parameter-storage budget; and (iv) forward-backward consistency: the forward and backward graphs must agree on the placements of tensors that cross the boundary between them, specifically model inputs and outputs and their corresponding gradients; intermediate activations within the forward or backward graph are not individually constrained. Both cost terms are estimated by the hardware-aware model described next.

4.2

Hardware-Aware Cost Model

The cost model only needs to rank strategies in the same order as actual execution, so predicting relative costs matters more than predicting absolute time (evaluated in Appendix B). It is also pluggable: users can recalibrate or replace it without changing the solver. Each strategy implies a per-device tensor shape, to which the model applies a roofline estimate (Williams et al., 2009): the larger of FLOPs over peak throughput and memory traffic over memory bandwidth, plus a minimum kernel-launch overhead so that many small operators are not underestimated. Computation.

Communication. Each redistribution maps to a collective: Shard → Replicate to all-gather, Partial → Shard

to reduce-scatter, Partial → Replicate to all-reduce, and Shard → Shard across dimensions to all-to-all; Replicate → Partial is disallowed. A plain latency-bandwidth model misranks these collectives, because NCCL selects among algorithms (ring, tree, NVLS) and protocols by message size and topology, which can change effective bandwidth by an order of magnitude. To stay close to real execution, the model mirrors NCCL’s selection logic, adds corrections calibrated against NCCL benchmarks on the target cluster, and uses per-axis bandwidth and latency, so that a large collective costs more on a slow inter-node axis than on a fast intra-node axis.

4.3

Scalable Solver

The placement problem (Equation 1) is an integer program whose difficulty grows with mesh dimensionality D: the per-node search space grows as O((1+n)D ) (§4.1), making strategy enumeration and pairwise cost computation expensive even before the solver runs, and the combinatorial optimization over the full graph compounds this cost. As §2.3 shows, an exact ILP formulation with CBC (Forrest and Lougee-Heimer, 2005) becomes impractical beyond two dimensions; we keep it as a reference for solution quality (§6.5). From ILP to MRF. The ILP formulation constrains parameter storage through the sum of the parameters’

local-to-global size ratios. Parameter placement choices couple this budget across the graph. In production model training, however, parameter memory redundancy across devices is not acceptable: the standard practice is to fully shard all parameters across the non-DDP mesh dimensions, and this is the tightest packing any feasible plan can achieve. When we enforce parameter sharding in the candidate strategies, restricting weight nodes to Shard placements on all non-DDP mesh axes, persistent parameter storage is 6

fixed: each device stores exactly its 1/N share of every parameter, where N is the product of non-DDP mesh degrees. The parameter-storage constraint then no longer depends on the remaining decisions. Appendix C examines placement choices under different parameter-storage budgets. The ILP comparison in §6.5 uses the same parameter-sharding restriction, so both solvers optimize the same objective. With inputs sharded and redundant computation charged by the cost model, the selected plans in our experiments keep activations sharded, and peak activation memory is governed by selective activation checkpointing. With parameter storage fixed, the objective (Equation 1) decomposes into per-node terms (computation cost) and per-edge terms (redistribution cost), each depending on at most two adjacent nodes’ strategy choices. This structure is a pairwise Markov Random Field (MRF) (Kolmogorov, 2006), enabling efficient approximate inference via message passing rather than integer programming. Making the MRF tractable. Two reductions shrink the problem before solving.

Graph clustering. Models with repeated structure, such as the L identical layers of a transformer, produce L copies of the same subproblem in the computation graph. We identify copies via content-based hashing of each node’s operator, arguments, and strategy set, and share decision variables across identical groups. Split-dimension seeding. Even after clustering, candidate enumeration and pairwise cost computation grow exponentially with D (Table 3). We therefore decompose the D-dimensional problem into D independent one-dimensional subproblems, one per mesh axis, each using a projected cost model reflecting that axis’s bandwidth. Each 1D problem is small enough to solve exactly. The per-dimension solutions are assembled into a D-tuple seed per node, and the full problem is rebuilt with enumeration restricted to a Hamming ball of radius r (default r=2) around each seed, retaining only strategies that differ from the seed in at most r dimensions. This reduces the effective candidate count from O((1+n)D ) to O(Dr · nr ), making enumeration and cost computation tractable at high D. Solving the MRF. Nodes that must choose jointly, including clustered copies, forward-backward pairs, and

flow-coupled producers, are merged into groups via union-find, each group selecting one joint assignment. Flow-coupled producers feed the same input position of clustered operator copies and therefore share one placement decision. The unary potential for a group aggregates computation and self-redistribution costs across its members, weighted by the number of cluster copies each represents. The pairwise potential between adjacent groups is the redistribution cost between their selected strategies. We solve the resulting MRF with Tree-Reweighted Sequential message passing (TRW-S), which alternates forward and backward sweeps over monotonic chains induced by a topological ordering and keeps the best decoded assignment across sweeps. Two rounds of local search then refine the solution: coordinate descent over single groups, and star-block search over high-degree groups together with their neighbors.

5

Execution and Composition

5.1

Plan Materialization

We materialize the placement plan through a two-pass lowering of the joint forward-backward graph. In the first pass, an interpreter walks every node and inserts redistribution collectives (all-gather, reduce-scatter, all-to-all) wherever the output placement of one operator disagrees with the input placement expected by the next. View and factory operations have their shape arguments adjusted from global to per-device local dimensions. In the second pass, standard Inductor decompositions lower any remaining compound operations into base ATen operators. Model parameters are distributed according to the plan: each parameter is wrapped as a DTensor with the selected placement, sharded across the device mesh. The lowered graph and sharded parameters are packaged into an ordinary nn.Module that inherits the original model’s class, preserving isinstance checks and module hierarchy. The output is a standard PyTorch module with no custom runtime state, ready for training loops, eager debugging, or compilation.

7

5.2

Compilation and Communication Overlap

For production training, the complete forward, loss, and backward computation of the materialized module, including the redistribution collectives inserted during materialization, is traced into a single FX graph and compiled with Inductor through TorchTitan’s GraphTrainer (Purandare and Huang, 2026). We rely on generic communication reordering, overlap scheduling, and collective bucketing passes already present in Inductor rather than reimplementing scheduling logic in the planner.

5.3

Composing with Dynamic Parallelism

Not all distributed computation can or should be folded into static placement search: some patterns involve dynamic collectives, and others benefit from user-controlled implementations that have been tuned for specific hardware or workloads. We handle both cases through local_map, a PyTorch higher-order operator that defines a placement contract at the composition boundary: the user declares how inputs should be distributed before entering the region and how outputs are distributed when exiting, then implements the internal computation and collectives using standard PyTorch distributed primitives. The optimizer treats these declared placements as fixed single-strategy constraints, optimizing all surrounding computation to be compatible with the boundary. Context parallelism. Context parallelism illustrates a region kept user-written because it relies on a custom

implementation. For long-context training, context parallelism distributes the sequence dimension across devices so that each device processes a chunk of the sequence. We wrap the attention computation in local_map with input placements that shard query, key, and value tensors along the sequence dimension on the CP mesh axis, allowing users to plug in their preferred implementation (Figure 2a). Mixture-of-experts. Mixture-of-experts illustrates the other reason: its routing produces data-dependent

shapes that static placement search cannot express. Expert parallelism distributes experts across devices with data-dependent token routing and all-to-all dispatch. We construct a device mesh where the expertparallel axis is folded into the existing mesh dimensions and wrap the token-dispatch, expert-computation, and token-combine phases in local_map with placements declaring which axis carries expert parallelism (Figure 2b).

Figure 2 Composition through local_map. Dashed boxes are user-written local_map regions; AutoParallel plans the

operators outside them (gray).

In both cases the composition boundary is explicit and inspectable: the user controls the internal implementation and placement declarations, and the optimizer handles everything outside.

6

Evaluation

We compare AutoParallel with expert-written recipes (§6.2), examine the placement decisions behind the differences (§6.3), measure the benefit of replanning when the workload changes (§6.4), and evaluate solver scalability on high-dimensional meshes (§6.5). Appendix B evaluates the cost model’s ranking fidelity.

6.1

Setup

All experiments run on nodes each with eight NVIDIA H100 80 GB SXM GPUs connected via NVLink within a node and RoCE across nodes. The software stack is PyTorch 2.14.0a0, CUDA 13.0, NCCL 2.30.7, and Python 3.12. All configurations train in bfloat16 mixed precision with float32 reductions; torch.compile, selective activation checkpointing policy, and communication settings are matched across arms at every mesh. 8

We evaluate LLaMA 3 8B and 70B (Llama Team, 2024), Qwen3 32B (Qwen Team, 2025), and MuseGlimmer 30B (Meta, 2026), whose hybrid attention combines sliding-window and full attention. Our 16B MoE configuration uses the DeepSeek-V3 architecture (DeepSeek-AI, 2024) with fewer layers and heads. The baseline is TorchTitan MainTrainer (Liang et al., 2024): expert-designed parallel strategies executed with the standard PyTorch training loop. AutoParallel receives the logical mesh and boundary placement contracts for local_map regions (context parallelism, mixture-of-experts) but no operator-level placement hints; it discovers the parallel strategy automatically and executes through GraphTrainer with full-graph Inductor compilation, communication-computation overlap, and collective bucketing. Both arms use the full set of Inductor optimizations: MainTrainer through JIT torch.compile and AutoParallel through a single full-graph trace. In GraphTrainer, AutoParallel keeps the default graph cleanup and selective activation checkpointing passes but disables GraphTrainer’s own collective bucketing, overlap-scheduling; communication is instead scheduled by Inductor’s overlap scheduling and collective bucketing. We measure a complete training iteration, excluding data loading and one-time tracing, planning, and compilation. We run 5 warm-up iterations followed by 20 measured iterations and report average throughput in tokens per second.

6.2

End-to-End Performance

Table 2 End-to-end throughput (K tokens/s) on scaling ladders from 8 to 128 GPUs. ∆ is AutoParallel relative to

MainTrainer. All runs use local batch size 2 per data-parallel rank.

Mesh

GPUs

MT

AP

∆%

LLaMA 3 8B (seq 8 K; 4D uses 16 K with CP) 2×4 8 2×8 16 4×8 32 2×4×8 64 4×4×2×4 128

54.2 82.1 165.3 323.8 575.5

58.4 88.0 181.5 335.4 554.7

+7.9 +7.2 +9.8 +3.6 −3.6

DeepSeek-V3 16B MoE (seq 8 K) 1×8, EP 8 8 4×8, EP 8 32 8×8, EP 8 64 4×4×8, EP 32 128

44.2 153.5 280.6 349.8

52.3 162.6 285.1 373.5

+18.3 +5.9 +1.6 +6.8

MuseGlimmer 30B Hybrid (seq 4 K) 8×2 16 16×2 32 32×2 64

36.0 72.3 142.1

36.2 72.0 137.3

+0.5 −0.4 −3.4

Qwen3 32B Dense (seq 8 K) 16×8

128

199.7

195.1

−2.3

LLaMA 3 70B Dense (seq 8 K) 16×8 128

101.0

99.2

−1.8

Table 2 reports aggregate throughput on scaling ladders from 8 to 128 GPUs. Each ladder increases mesh dimensionality with GPU count, from two-dimensional dp×tp at small scale to three- and four-dimensional meshes at larger scales. Mesh shapes list axes outermost first: the LLaMA 3 8B 3D and 4D tuples are (DDP, DPS, TP) and (DDP, DPS, CP, TP), respectively. EP gives the expert-parallel degree folded into the mesh; in DeepSeek-V3, EP 8 uses the TP axis, while EP 32 spans the final two axes of the 4×4×8 mesh. Across the 14 configurations, AutoParallel outperforms MainTrainer on 9 and stays within 3.6% on the remaining 5. Graph-wide search discovers non-obvious placements. On LLaMA 3 8B with two-dimensional dp×tp meshes (8–32 GPUs), AutoParallel consistently outperforms the expert baseline, improving throughput by 7.2–9.8%. The graph-wide search discovers per-operator placement refinements that the uniform manual strategy does not explore: AutoParallel tailors placements to individual operators within each layer based on their computation

9

and communication costs on the target mesh. Because the search spans the joint forward-backward graph, the forward and backward computations of the same layer can also take different layouts, which module-level recipes do not express. §6.3 examines these plans in detail. On DeepSeek-V3 16B, where local_map controls expert routing and AutoParallel optimizes the surrounding attention and shared-expert layers, the improvement reaches 18.3% at 8 GPUs (1×8, EP 8). These gains likely come from per-operator placements in the attention and shared-expert layers surrounding the local_map block, which the uniform expert recipe keeps fixed. On MuseGlimmer 30B, LLaMA 3 70B, and Qwen3 32B, AutoParallel matches expert throughput within 3.4%, indicating that the expert recipes for these models are already close to the plans AutoParallel selects, which it reaches without per-model tuning. Notably, Qwen3 is an external model family with no AutoParallel-specific recipe engineering, demonstrating generalization beyond the models used during development. Plans scale to high-dimensional meshes. As GPU count grows, production training moves to higher-dimensional

meshes that combine HSDP, tensor, context, and expert parallelism. On LLaMA 3 8B at 64 GPUs with a three-dimensional HSDP+TP mesh (2×4×8), AutoParallel improves throughput by 3.6% over the expert baseline. At 128 GPUs with a four-dimensional HSDP+CP+TP mesh (4×4×2×4), which adds context parallelism via local_map for 16K sequences, AutoParallel is within 3.6% of the expert. Context parallelism places most of the attention computation inside the manually controlled region, limiting the scope of autooptimization to the surrounding projections and feed-forward layers; we attribute the bounded gap mainly to this smaller optimization scope, a trade-off we accept by design, although we have not isolated it from other factors. The same pattern holds on DeepSeek-V3, where the mesh transitions from two-dimensional dp×ep at 8–64 GPUs to three-dimensional (4×4×8, EP 32) at 128 GPUs, and AutoParallel achieves a 6.8% improvement at the highest dimensionality.

6.3

Discovered Plans

We inspect two plans: one that reproduces the expert recipe and one whose departures from it are instructive. We highlight individual decisions rather than enumerate every difference. We write Wqkv for the query, key, and value projections collectively, Wo for the attention output projection, W1 /W3 for the gate/up projections of the gated MLP, and W2 for its output projection. SDPA denotes scaled dot-product attention. MuseGlimmer 8×2: recovering the expert plan. On every MuseGlimmer mesh, AutoParallel independently

arrives at the same strategy as MainTrainer: FSDP on the dp axis, Megatron-style column/row tensor parallelism with sequence parallelism on tp, and attention heads sharded across tp. There are zero role-level deviations across all operators. This indicates that on this model the expert recipe is already close to the cost-model optimum, and the solver reaches it without guidance.

Figure 3 Schematic attention-block placement of the expert plan and AutoParallel. Vertical and horizontal splits

denote head and sequence sharding. LLaMA 3 8B: mixing Megatron TP with Ulysses-style all-to-all. Figure 3 shows the attention block of the plan AutoParallel selects on the 128-GPU 16×8 mesh. In the forward pass, instead of applying row-parallel Wo and reduce-scattering the partial sum, AutoParallel converts the head-sharded attention output to sequence-sharded with an all-to-all, as in DeepSpeed-Ulysses (Jacobs et al., 2023), and applies Wo with the full, all-gathered weight. With A the activation size of one TP group and p the TP degree, the all-to-all sends about A(p−1)/p2 bytes per rank, a factor of p less than the reduce-scatter; the plan pays an extra weight gather, which Inductor can overlap with computation. The same Wo change appears on the 4×8 and 2×4×8 meshes, including the 4×8 configuration with the largest LLaMA gain (9.8%), although we have not isolated it from the rest of the plan. Each ingredient (Megatron TP, the Ulysses head-to-sequence all-to-all, and FSDP-style weight gathering) is individually well known, but the evaluated uniform TorchTitan recipe

10

does not mix them at a single operator within a single layer; the per-operator search finds this combination automatically. Different forward and backward layouts. The backward paths in Figure 3 show a second departure. Module-level

recipes assign each projection one TP role, which fixes both its forward layout and its backward communication; AutoParallel places forward and backward operators independently and constrains only tensors that cross the forward-backward boundary (§4.1). It keeps the Q/K/V outputs head-sharded in the forward pass, as SDPA requires, but runs their grad-input matmuls sequence-sharded, so two all-to-alls replace the expert’s backward all-gather and reduce-scatter. The cost model estimates a saving of about 0.34 ms of backward communication per layer, which the current implementation partly loses to weight gathers over the combined mesh. These two examples are representative, not exhaustive: other configurations exhibit further deviations of varying complexity, but the pattern is consistent. On MuseGlimmer, the solver recovers the expert strategy; on LLaMA, it combines operator placements that differ from the uniform recipe. We note that throughput differences also reflect execution-level factors (compilation, bucketing, communication-computation overlap) that are not isolated in these comparisons.

6.4

Adaptation Through Replanning

A key motivation for automatic planning is that workload changes invalidate stale recipes. We test this by generating a reference plan P0 for LLaMA 3 8B on a 4×8 mesh at sequence length 2,048 and evaluating two strategies as the sequence length increases: Replay reuses P0 ’s placements unchanged, and Fresh invokes AutoParallel with the new sequence length.

Figure 4 Throughput of fresh and replayed plans on LLaMA 3 8B (4×8 mesh) across sequence lengths.

Figure 4 shows the results; labels give the throughput gain of Fresh over Replay. At 4K, only 36 nodes change placement and the gain is modest (1.6%). As the sequence length grows further from the reference, the stale plan falls increasingly behind the fresh plan: at 8K, Fresh improves by 41.5% with 2,437 changed nodes; at 32K, the advantage reaches 51.3%. The widening gap reflects the coupled nature of parallelization decisions: as activation sizes grow, the memory-throughput balance shifts and placements optimized for 2K become inappropriate. Planning on this 2D mesh takes under 40 s (§6.5), negligible relative to production runs of tens of thousands of steps.

6.5

Solver Scalability

We evaluate the approximate solver (§4.3) against the exact ILP on LLaMA 3 8B across one- through four-dimensional meshes, all with 32 or fewer GPUs so that the ILP remains verifiable where possible. We report build time (strategy enumeration and pairwise cost computation) and solve time separately. Graph clustering is enabled in both paths. Table 3 reports three results. Quality. On these 1D–3D instances, the approximate solver reaches the same objective as the ILP. Scaling. Build and solve are separate bottlenecks: seeding cuts the 3D build time by 6.7×, and MRF inference keeps solve time under 21 s on every mesh, so the approximate solver finishes the 3D mesh 30× faster than the ILP; on the 4D mesh the ILP cannot finish building within an hour. On the 11

Table 3 Solver scalability on LLaMA 3 8B. ILP vars: decision variables of the unseeded ILP; T/O: ILP build exceeds

one hour; objectives in µs.

Build (s)

Solve (s)

Objective

Mesh

ILP vars

ILP

Approx

ILP

Approx

ILP

Approx

1×8 4×8 4×2×4 2×2×2×4

12.9K 341K 12.2M 522.7M

4.7 17.9 428.9 T/O

4.9 12.4 64.3 194.3

0.3 12.3 1,795 —

5.1 6.1 10.4 20.8

626,547 311,290 300,156 —

626,547 311,290 300,156 301,828

trivially small 1D mesh, the approximate solver’s fixed overhead makes it slightly slower. End-to-end cost. Including one-time graph tracing, planning takes 25–233 s from 1D to 4D, a one-time cost amortized over training runs that typically span hours to weeks.

7

Related Work

Operator-level automatic parallelism. FlexFlow (Jia et al., 2019) searches a per-operator parallelism space with MCMC simulation over the Legion runtime; Unity (Unger et al., 2022) extends this with joint algebraicparallel optimization. Alpa (Zheng et al., 2022) formulates intra-operator sharding as an ILP and inter-operator pipeline partitioning as dynamic programming, targeting JAX/XLA with a Ray runtime. PartIR (Alabed et al., 2025) decomposes parallelization strategies into composable tactics with optional MCTS-based search, using an MLIR-based partitioning stack; users can inspect the rewritten module after each tactic. On the propagation side, GSPMD (Xu et al., 2021; Lepikhin et al., 2020) completes a full-graph sharding from sparse user annotations via heuristic propagation without a cost model; it supports manually partitioned subgraphs in XLA and has been deployed at production scale, but the key placement decisions are made by the user rather than a solver. These systems use compiler stacks and runtimes outside native PyTorch execution, and their partitioned programs do not directly compose with the PyTorch compilation ecosystem. Layer-level planners. A second line of work plans at layer or stage granularity, choosing among named parallelism strategies such as DP, TP, PP, and ZeRO. Galvatron (Miao et al., 2023) uses decision-tree pruning and dynamic programming to assign a strategy per layer. Mist (Zhu et al., 2025) and UniAP (Lin et al., 2025) plan under memory constraints via MILP and MIQP, respectively. These systems operate at a granularity too coarse to express the per-operator mixed placements from which strategies like FSDP and TP emerge naturally in our formulation (§4.1), and they rely on Megatron or custom runtimes. Many of them also search pipeline parallelism, which AutoParallel does not. PP targets regimes where inter-node bandwidth is insufficient to overlap FSDP and TP communication, whereas AutoParallel targets the SPMD placement decisions within a pipeline stage. The two problems are orthogonal and can be composed hierarchically, as in Alpa’s two-level design (Zheng et al., 2022). PyTorch-based systems. In the PyTorch ecosystem, TorchTitan (Liang et al., 2024) and veScale (Li et al., 2025) build on DTensor but rely on hand-written parallelization plans. Colossal-Auto (Liu et al., 2023b) adapts Alpa’s ILP to torch.fx but analyzes only the forward graph, outputs to a custom sharded module rather than DTensor, and does not integrate with torch.compile. nnScaler (Lin et al., 2024) generates multi-device PyTorch source code, but extending its search space requires expert-written constraints on its primitives, and it depends on its own runtime, which adds integration work when composing with the PyTorch compilation stack. JAX’s shard_map (JAX, 2026) provides a similar boundary for manually partitioned computation. AutoParallel embeds local_map regions as fixed-placement constraints inside the global solver (§5.3). Table 4 summarizes the key distinctions. To our knowledge, AutoParallel is the first system that automatically plans whole-graph operator placements for unmodified PyTorch models and emits a native PyTorch program composable with the production compilation stack.

12

Table 4 Comparison with representative parallelization systems. Auto: searches for a plan automatically. Op:

per-operator decisions. PT in: takes an unmodified PyTorch nn.Module (local_map regions optional). No RT : no additional runtime. Stack : composes with the compiler, activation checkpointing, mixed precision, and distributed checkpointing. Comp.: composes with hand-written dynamic regions such as CP and MoE. † XLA-based. ‡ Expertdefined search space.

TorchTitan GSPMD † Alpa Galvatron nnScaler ‡

AutoParallel

8

Auto

Op

PT in

No RT

Stack

Comp.

✗ ✗ ✓ ✓ ✓ ✓

– ✓ ✓ ✗ ✓ ✓

✓ ✗ ✗ ✓ ✓ ✓

✓ ✓ ✗ ✗ ✗ ✓

✓ ✓ ✗ ✗ ✗ ✓

✓ ✓ ✗ ✗ ✗ ✓

Conclusion

We presented AutoParallel, a PyTorch-native automatic parallelism planner that jointly optimizes per-operator placements across the full training graph, produces ordinary PyTorch programs without a separate runtime, and composes with both the PyTorch compilation stack and user-controlled dynamic parallelism regions. On up to 128 H100 GPUs, AutoParallel matches or exceeds expert-written TorchTitan recipes on most configurations and discovers plans that uniform expert recipes do not express. Crafting efficient parallelization strategies for distributed training has historically required scarce systems expertise, and we believe AutoParallel represents a step toward making high-performance distributed training accessible to the broader PyTorch community.

13

Appendix A

Placement Solutions on Community Models

We apply AutoParallel to 17 Hugging Face models on a single-node 1×8 H100 mesh and compare the cost-model objective against SimpleFSDP Zhang et al. (2024) as a zero-configuration baseline. All numbers are cost-model objectives (µs, no overlap modeling), not measured step times; no GPU execution is performed. Setup. Global batch size is 16 (local 2 per rank); inputs are constrained to Shard(0). Mixed precision uses bfloat16 compute with float32 reductions. Sequence length is 4096 by default, capped at each model’s position-table limit where applicable. The parameter-storage constraint uses high = 1/world_size (0.125), yielding the minimum feasible storage ratio for each parameter. AutoParallel solves the ILP exactly with CBC (status Optimal on all 17 instances, solve time 0.3–1.2 s). Both arms share the same cost model, strategy space, and parameter-storage budget; the only difference is the placement assignment. Table 5 shows the results. AutoParallel reduces the cost-model objective on all 17 models (0.2–40.8%). The plans fall into three groups. FSDP with embedding/head optimization (8 models, 559M–7.1B). The repeated transformer layers use the same

all-gather/reduce-scatter pattern as SimpleFSDP; AutoParallel only changes the embedding and LM-head layers, storing the token embedding sharded along the vocabulary dimension and, in 6 of 8 models, applying column-parallel sharding to the LM head. The gains are 0.2–1.1%. Tensor parallelism (4 models, 67M–355M). AutoParallel replaces per-parameter collectives with activation

collectives throughout the model: column-parallel QKV and FFN projections, row-parallel output projections with all-reduce, replicated layer norms. Communication cost drops by 21–41 K µs at the expense of a small compute increase. The large gains (22–41%) reflect the per-collective latency floor on NVSwitch (∼160 µs per all-gather/reduce-scatter): SimpleFSDP issues 211–585 small per-parameter collectives whose cost is 80–84% latency overhead. Mixed (5 models, 61M–6.1B). AutoParallel applies different strategies to different operators within the same

layer. For example, on GPT-J 6B, attention runs data-parallel while the MLP uses column/row tensor parallelism; on T5 Base, forward weights are gathered for data-parallel compute but nearly half of the weight gradients are produced sharded without a reduce-scatter. Gains range from 0.9% to 9.3%. The cost model does not model overlap or collective bucketing; Inductor’s bucketing would reduce SimpleFSDP’s latency-floor overhead in practice. These results show that AutoParallel generalizes across a broad range of popular architectures, from 61M to 7.1B parameters, consistently finding plans at least as good as SimpleFSDP without any per-model tuning, and discovering tensor-parallel or mixed strategies that a uniform FSDP baseline cannot express.

14

Table 5 Cost-model objective change (%) for AutoParallel vs. SimpleFSDP on 17 Hugging Face models (1×8 H100

mesh); negative is better.

Model

B

Params

Seq

∆%

FSDP with embedding/head optimization BLOOM 560M 559M BLOOM 7B 7.1B StableLM 1.6B 1.6B Phi-1 1.4B GPT-Neo 2.7B 2.7B Phi-2 2.8B Falcon 7B 6.9B Pythia 6.9B 6.9B

4096 4096 4096 4096 2048 4096 4096 4096

−1.1 −0.9 −0.9 −0.6 −0.5 −0.2 −0.2 −0.2

Tensor parallelism DistilBERT RoBERTa BERT GPT-2 Medium

67M 125M 110M 355M

512 512 512 1024

−40.8 −40.4 −40.2 −22.3

Mixed strategies GPT-2 XL T5 Small GPT-J 6B Pythia 160M T5 Base

1.6B 61M 6.1B 162M 223M

1024 4096 2048 4096 4096

−9.3 −3.1 −3.0 −1.0 −0.9

Cost Model Ranking Fidelity

We validate the cost model’s ranking quality on LLaMA 3 8B with a 4×8 mesh (FSDP 4, TP 8; 32 H100 GPUs, sequence length 8192, bfloat16 mixed precision). We generate 11 feasible placement plans that share the same mesh but differ in per-operator placements across the 7,803-node computation graph. Starting from the ILP-optimal plan (p00), computed with only input, output, and memory constraints, we produce alternatives by fixing the output placement of one or two key projections (Wqkv , Wo , or W2 ) to a specific strategy and re-solving the full ILP; the solver re-optimizes all remaining nodes jointly. Table 6 lists the resulting projection-output placements and cost-model objectives. All 11 candidates run on the same four-node allocation with the same checkpoint, data, torch.compile configuration, and selective activation checkpointing; only the placement assignment differs.

15

Table 6 Placement plans and cost-model ranking on LLaMA 3 8B (4×8 mesh, 32 GPUs). Each plan assigns placements

to all 7,803 graph nodes; the table shows outputs of the Q/K/V (Wqkv ), attention output (Wo ), MLP gate/up (W1/3 ), and MLP output (W2 ) projections. Placement tuples are (dp, tp); Sd = Shard(d), R = Replicate, P = Partial. C: cost-model objective (K µs); T : measured step time (ms); rk: rank.

Output placement (dp, tp) p00 p01 p02 p03 p04 p05 p06 p07 p08 p09 p10

Pred.

Meas.

Wqkv

Wo

W1/3

W2

C

rk

T

rk

S0 ,S2 S0 ,S2 S0 ,S2 S0 ,S2 S0 ,S2 S0 ,S2 P,R P,P P,P P,P P,P

S0 ,S1 S0 ,S1 P,S1 S0 ,S1 P,S0 P,S2 S0 ,S1 S1 ,S2 R,S1 P,S0 S2 ,S2

S0 ,S2 S2 ,S1 S0 ,S2 S2 ,S1 S0 ,S2 S0 ,S2 S0 ,S2 S1 ,S2 S0 ,S2 S1 ,S2 S1 ,S2

S0 ,P P,S1 S0 ,P P,S0 S0 ,P S0 ,P S0 ,P S1 ,P S0 ,P S1 ,P S1 ,P

642 694 703 707 710 722 773 835 860 868 921

1 2 3 4 5 6 7 8 9 10 11

364.9 494.8 421.9 507.0 430.8 453.3 600.5 638.7 790.4 618.7 784.0

1 5 2 6 3 4 7 9 11 8 10

Figure 5 Left: Cost-model objective vs. measured step time for 11 placement plans on LLaMA 3 8B, 4×8 mesh; right:

predicted vs. measured rank.

Figure 5 shows the results. The ILP-selected plan (p00) is both the lowest-cost and the measured-fastest candidate (364.9 ms), giving a selection regret of 0%. Selection regret is (Tselected / mini Ti − 1) × 100% over the 11 candidates. The Spearman rank correlation between cost-model objective and measured step time is ρ = 0.855 across the 11 candidates, which span a 2.2× range of measured times (365–790 ms). The cost model does not model communication-computation overlap or collective bucketing; both are handled at runtime by Inductor and are likely contributors to the local ranking inversions visible in the middle of the range. The cost model’s ranking quality is sufficient for plan selection among these candidates, even without modeling runtime optimizations.

C

Plan Stability Under Memory Budget

We sweep the normalized parameter-storage budget on LLaMA 3 70B with a 4×8 mesh (FSDP 4, TP 8; 32 GPUs, bfloat16 mixed precision) from fully sharded (high = 1/32, i.e. each parameter stored at 1/world_size) to unconstrained (high = 1, every parameter allowed to be fully replicated). For this model, high bounds the average local-to-global element-count ratio across parameter tensors. All results are ILP optimal; no GPU execution is performed. We additionally run three counterfactual solves at high = 1 that force all parameters to Replicate on the TP axis, the FSDP axis, or both.

16

Table 7 Memory budget sweep on LLaMA 3 70B (4×8 mesh). Placement tuples are (dp, tp).

Obj.

Comm

Norm placement

Wt/GPU

high

(K µs)

(K µs)

attn

ffn

(GB)

1/32 1/16 1/8 1/4 1/2 1

3,854 3,828 3,826 3,816 3,807 3,807

1,878 1,852 1,850 1,840 1,831 1,831

S0 ,S0 R,S0 S0 ,R S0 ,R R,R R,R

S0 ,S0 R,S0 S0 ,R R,R R,R R,R

8.819 8.820 8.820 8.822 8.824 8.824

Table 7 shows the results. Across the full budget range, every matrix projection (Wqkv , Wo , W1 /W3 , W2 , embedding, LM head) maintains identical placements, and compute cost is constant at 1,976 K µs. The entire 1.2% objective variation comes from the 161 norm-weight tensors (80 layers × 2 norms + final norm), which progressively transition from fully sharded to replicated as the budget increases. These 161 tensors account for 28.6% of parameter tensor count but only 0.0019% of parameter bytes (5.3 MB of 282 GB), so weights per GPU change by just 0.005 GB. Table 8 Counterfactual solves forcing all parameters to Replicate on given mesh axes (high = 1). ∆ is relative to the

unconstrained baseline.

Constraint

Obj. (K µs)

∆%

Wt/GPU (GB)

None R on tp R on fsdp R on both

3,807 3,911 4,053 4,856

— +2.7 +6.5 +27.6

8.8 70.6 35.3 282.2

Table 8 confirms why the solver does not replicate large matrices even when unconstrained. Forcing all parameters to Replicate on the TP axis, the FSDP axis, or both increases the objective by 2.7%, 6.5%, and 27.6% respectively. The mechanism is gradient synchronization cost: sharded storage requires a bfloat16 all-gather in the forward pass and a float32 reduce-scatter for gradients, while replicated storage eliminates the forward gather but requires a float32 all-reduce whose data volume is roughly double that of a reduce-scatter. At 70B-scale weight dimensions, this cost asymmetry dominates. Norm weights behave differently because they have only 8,192 elements, making their communication cost latency-dominated. Sharded storage requires two cross-axis collectives per norm per layer (forward all-gather + backward reduce-scatter); replicated storage requires only one (backward all-reduce), saving approximately 294 µs per norm. The total norm savings of 47,338 µs account for the entire communication reduction from high = 1/32 to 1. In summary, the solver’s placement decisions for the dominant parameters are driven by communication cost, not memory pressure, and are robust to the memory budget setting.

D

Numerical Validation

DTensor propagation rules restrict the search to plans that are semantically equivalent to single-device execution up to floating-point reassociation in reductions (§4.1). To validate the materialized programs, we train LLaMA 3 8B for 1,000 steps on a DPS4×CP2×TP4 mesh (32 GPUs, sequence length 16K, global batch size 8, AdamW with learning rate 3 × 10−4 , 200 warm-up steps, and linear decay) with MainTrainer and with AutoParallel, starting from the same seed checkpoint and replaying the same data and learning-rate schedule. As a reference for how much reduction order alone changes the loss, we also run MainTrainer on the equivalent DPS4×CP1×TP8 mesh.

17

Figure 6 Training loss of LLaMA 3 8B with MainTrainer and AutoParallel.

Figure 6 shows that the two loss curves overlap. Over the last 100 steps, AutoParallel’s loss differs from MainTrainer’s by 0.80% on average, the same as MainTrainer on the TP8 mesh (0.82%); both comparisons show larger transient deviations between steps 100 and 500. AutoParallel’s numerical behavior is therefore within the variation that any change of parallel layout introduces.

18

References Sami Alabed, Daniel Belov, Bart Chrzaszcz, Juliana Franco, Dominik Grewe, Dougal Maclaurin, James Molloy, Tom Natan, Tamara Norman, Xiaoyue Pan, Adam Paszke, Norman A. Rink, Michael Schaarschmidt, Timur Sitdikov, Agnieszka Swietlik, Dimitrios Vytiniotis, and Joel Wee. PartIR: Composing SPMD partitioning strategies for machine learning. In Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2025. Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, David Berard, Evgeni Burovski, Geeta Chauhan, Anjali Chourdia, Will Constable, Alban Desmaison, Zachary DeVito, Elias Ellison, Will Feng, Jiong Gong, Michael Gschwind, Brian Hirsh, Sherlock Huang, Kshiteej Kalambarkar, Laurent Kirsch, Michael Lazos, Mario Lezcano, Yanbo Liang, Jason Liang, Yinghai Lu, C. K. Luk, Bert Maher, Yunjie Pan, Christian Puhrsch, Matthias Reso, Mark Saroufim, Marcos Yukio Siraichi, Helen Suk, Michael Suo, Phil Tillet, Eikan Wang, Xiaodong Wang, William Wen, Shunting Zhang, Xu Zhao, Keren Zhou, Richard Zou, Ajit Mathews, Gregory Chanan, Peng Wu, and Soumith Chintala. PyTorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2024. Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174, 2016. DeepSeek-AI. DeepSeek-V3 technical report. arXiv preprint arXiv:2412.19437, 2024. William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1–39, 2022. John Forrest and Robin Lougee-Heimer. CBC user guide. INFORMS Tutorials in Operations Research, 2005. https://github.com/coin-or/Cbc. Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia Xu Chen, Dehao Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, and Zhifeng Chen. GPipe: Efficient training of giant neural networks using pipeline parallelism. In Advances in Neural Information Processing Systems (NeurIPS), 2019. Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. DeepSpeed Ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309.14509, 2023. JAX. Manual parallelism with shard_map. shard-map.html. Accessed October 2026.

JAX documentation, 2026.

https://docs.jax.dev/en/latest/201/

Xianyan Jia, Le Jiang, Ang Wang, Wencong Xiao, Ziji Shi, Jie Zhang, Xinyuan Li, Langshi Chen, Yong Li, Zhen Zheng, Xiaoyong Liu, and Wei Lin. Whale: Efficient giant model training over heterogeneous GPUs. In 2022 USENIX Annual Technical Conference (USENIX ATC), 2022. Zhihao Jia, Matei Zaharia, and Alex Aiken. Beyond data and model parallelism for deep neural networks. In Proceedings of Machine Learning and Systems (MLSys), 2019. Vladimir Kolmogorov. Convergent tree-reweighted message passing for energy minimization. IEEE Transactions on Pattern Analysis and Machine Intelligence, 28(10):1568–1583, 2006. Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. In Proceedings of Machine Learning and Systems (MLSys), 2023. Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006.16668, 2020. Youjie Li, Cheng Wan, Zhiqi Lin, Hongyu Zhu, Jiacheng Yang, Ziang Song, Xinyi Di, Jiawei Wu, Huiyao Shu, Wenlei Bao, Yanghua Peng, Haibin Lin, and Li-Wen Chang. veScale: Consistent and efficient tensor programming with eager-mode SPMD. arXiv preprint arXiv:2509.07003, 2025. Wanchao Liang, Tianyu Liu, Less Wright, Will Constable, Andrew Gu, Chien-Chin Huang, Iris Zhang, Wei Feng, Howard Huang, Junjie Wang, Sanket Purandare, Gokul Nadathur, and Stratos Idreos. TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training. arXiv preprint arXiv:2410.06511, 2024.

19

Hao Lin, Ke Wu, Jie Li, Jun Li, and Wu-Jun Li. UniAP: Unifying inter- and intra-layer automatic parallelism by mixed integer quadratic programming. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2025. Zhiqi Lin, Youshan Miao, Quanlu Zhang, Fan Yang, Yi Zhu, Cheng Li, Saeed Maleki, Xu Cao, Ning Shang, Yilei Yang, Weijiang Xu, Mao Yang, Lintao Zhang, and Lidong Zhou. nnScaler: Constraint-guided parallelization plan generation for deep learning training. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2024. Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2310.01889, 2023a. Yuliang Liu, Shenggui Li, Jiarui Fang, Yanjun Shao, Boyuan Yao, and Yang You. Colossal-auto: Unified automation of parallelization and activation checkpoint for large-scale models. arXiv preprint arXiv:2302.02599, 2023b. Llama Team. The Llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024. Francisco Massa, Will Constable, Brian Hirsh, Edward Yang, Ivan Kobzarev, Elias Ellison, Pei Zhang, Sanket Purandare, Simon Fan, Ruisi Zhang, Milad Mohammadi, Davide Italiano, Peng Wu, and Patrick Labatut. Automatic parallelization in PyTorch. Poster, PyTorch Conference 2025, San Francisco, 2025. https://static.sched.com/hosted_files/pytorchconference/c4/Automatic%20Parallelism%20in% 20PyTorch%20-%20Will%20Constable%20&%20Francisco%20Massa,%20Meta.pdf. Meta. Muse-Glimmer-30B. Hugging Face model card, 2026. https://huggingface.co/meta-models/Muse-Glimmer-30B. Accessed October 2026. Meta PyTorch. AutoParallel: Compiler-driven automatic sharding for PyTorch. https://github.com/meta-pytorch/ autoparallel, 2026. Xupeng Miao, Yujie Wang, Youhe Jiang, Chunan Shi, Xiaonan Nie, Hailin Zhang, and Bin Cui. Galvatron: Efficient transformer training over multiple GPUs using automatic parallelism. Proceedings of the VLDB Endowment, 16(3): 470–479, 2023. Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training. In International Conference on Learning Representations (ICLR), 2018. Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R. Devanur, Gregory R. Ganger, Phillip B. Gibbons, and Matei Zaharia. PipeDream: Generalized pipeline parallelism for DNN training. In Proceedings of the 27th ACM Symposium on Operating Systems Principles (SOSP), 2019. Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. Efficient large-scale language model training on GPU clusters using Megatron-LM. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC), 2021. Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems (NeurIPS), 2019. Sanket Jayant Purandare and Sherlock Huang. TorchTitan GraphTrainer: LLM distributed training with PyTorch’s compiler toolkit. Talk, PyTorch Conference North America 2026, San Jose, 2026. https://events.linuxfoundation. org/pytorch-conference-north-america/program/schedule/?id=1257898. PyTorch Team. AutoParallel integration in TorchTitan GraphTrainer. https://github.com/pytorch/torchtitan/tree/ main/torchtitan/experiments/graph_trainer, 2026. Qwen Team. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025. Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. ZeRO: Memory optimizations toward training trillion parameter models. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC), 2020.

20

Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, Ryan Sepassi, and Blake Hechtman. Mesh-TensorFlow: Deep learning for supercomputers. In Advances in Neural Information Processing Systems (NeurIPS), 2018. Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. MegatronLM: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019. Nouamane Tazi, Ferdinand Mom, Haojun Zhao, Phuc Nguyen, Mohamed Mekkouri, Leandro von Werra, and Thomas Wolf. The ultra-scale playbook: Training LLMs on GPU clusters. 2025. https://huggingface.co/spaces/nanotron/ ultrascale-playbook?section=high-level_overview. Hugging Face Technical Report. Colin Unger, Zhihao Jia, Wei Wu, Sina Lin, Mandeep Baines, Carlos Efrain Quintero Narvaez, Vinay Ramakrishnaiah, Nirmal Prajapati, Pat McCormick, Jamaludin Mohd-Yusof, Xi Luo, Dheevatsa Mudigere, Jongsoo Park, Misha Smelyanskiy, and Alex Aiken. Unity: Accelerating DNN training through joint optimization of algebraic transformations and parallelization. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2022. Samuel Williams, Andrew Waterman, and David Patterson. Roofline: An insightful visual performance model for multicore architectures. Communications of the ACM, 52(4):65–76, 2009. doi: 10.1145/1498765.1498785. Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Blake Hechtman, Yanping Huang, Rahul Joshi, Maxim Krikun, Dmitry Lepikhin, Andy Ly, Marcello Maggioni, Ruoming Pang, Noam Shazeer, Shibo Wang, Tao Wang, Yonghui Wu, and Zhifeng Chen. GSPMD: General and scalable parallelization for ML computation graphs. arXiv preprint arXiv:2105.04663, 2021. Ruisi Zhang, Tianyu Liu, Will Feng, Andrew Gu, Sanket Purandare, Wanchao Liang, and Francisco Massa. SimpleFSDP: Simpler fully sharded data parallel with torch.compile. arXiv preprint arXiv:2411.00284, 2024. Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li. PyTorch FSDP: Experiences on scaling fully sharded data parallel. Proceedings of the VLDB Endowment, 16(12):3848–3860, 2023. Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Eric P. Xing, Joseph E. Gonzalez, and Ion Stoica. Alpa: Automating inter- and intraoperator parallelism for distributed deep learning. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2022. Zhanda Zhu, Christina Giannoula, Muralidhar Andoorveedu, Qidong Su, Karttikeya Mangalam, Bojian Zheng, and Gennady Pekhimenko. Mist: Efficient distributed training of large language models via memory-parallelism co-optimization. In Proceedings of the Twentieth European Conference on Computer Systems (EuroSys), 2025.

AI Usage Disclosure We used AI assistants based on Claude to help with writing and language polishing, to launch and monitor experiment jobs, and to organize experimental data into tables and figures. All ideas, the system design and implementation, the experimental design, the structure and writing approach of the paper, and the analysis and interpretation of the data are our own. We reviewed all AI-assisted code and text, and we take full responsibility for the content of this paper.

21

Related documents

• FluxMoE: Decoupling Expert Residency for High-Performance MoE Serving#998172
• Terastal: Layer-Variant-based Scheduling for Real-Time Multi-DNN Workloads on Heterogeneous Accelerators#266148
• JET: Justification Evaluation in Transformer#1118947
• BatchParScale: Sparse Batch-Mixing for Sequence Models#1400086
• Explicit State and Resource Contracts for Low-Precision Pipeline Parallel Training under Captured Graphs#1049275