SPMD types in TorchTitan

TL;DR

TorchTitan now uses spmd_types as its default backend for distributed model computation. Authors specify sharding contracts and collectives explicitly, with optional typechecking to catch distributed-correctness errors during development. At runtime, the typechecking machinery can be erased so forward and backward execute on plain tensors. Coupled FWD-BWD typing and support for both global and local SPMD also make distributed behavior easier to express and reason about. Across our repeated debug-model benchmarks, spmd_types improved eager throughput by up to 46% and Inductor throughput by 5-9% over partial_dtensor, without a meaningful peak-memory increase. FSDP2 continues to use DTensor as the persistent storage representation outside model computation.

TorchTitan aims for a unified model definition across single-GPU execution and any 1D-to-ND composition of parallelisms, for both training and inference. A model’s forward() should describe its computation as if running on a single device, while composing well with any combination of parallelisms: FSDP, CP, TP, PP, or EP.

Historically, this meant keeping the forward definition free of distributed code and collectives, while DTensor provided the global SPMD (GSPMD) programming model. Model code read like single-device code. At compute time, DTensors carried a DeviceMesh and Replicate/Shard/Partial placements describing their distribution across GPUs, while sharding propagation (shard prop) computed and implicitly inserted any required collectives to make the program correct. Previous designs also relied on DTensor-based context managers to locally handle context parallelism, loss parallelism, and vocab-parallel embedding.

GSPMD abstractions kept model code clean, but as Titan model support and use cases expanded, we encountered more problems with DTensor in compute regions, especially concerning runtime overhead, BWD expressivity, and local region bugs.

DTensor pains in FWD/BWD

Runtime overhead. As a tensor subclass, DTensor intercepts every compute operation at the __torch_dispatch__ level, where it runs sharding propagation to check input placements and implicitly reshard or insert collectives as needed, executes the operation on inner tensors, and then rewraps the result with output placements and other metadata. Shard prop caching makes repeated operations cheaper, improving steady-state performance, while compilation or CUDA graphs move this work out of runtime, although overhead is still paid during warmup. In eager mode, shard prop can be prohibitively expensive, especially in inference scenarios where varying input token counts frequently cause cache misses, or for placements such as _StridedShard, whose redistribution planning on large meshes can take 10-20 seconds per uncached operation.

Implicit redistributions. When an operator’s input placements don’t match existing sharding rules, DTensor calculates the cheapest way to redistribute inputs to get a globally correct output, implicitly adding collectives if needed. This can be desirable during early authoring, when distributed correctness takes priority over explicit control, but it is also often a source of surprising sharding decisions and performance regressions, especially when rules are missing on the framework side. In upstream PyTorch, we have since added ExplicitRedistributionContext to error out when detected, and have more than doubled operator coverage this past year, but historically this has been a source of pain.

Forward and backward placements are independent. By design, sharding propagation is executed separately for the FWD and BWD of the same operation. Gradient placements therefore aren’t constrained by activation or parameter placements: the autograd graph enforces global shape consistency, but sharding and local shapes are opaque to it. Occasionally this means gradient reductions can occur in surprising spots, and in extreme cases FWD/BWD sharding can be on different tensor dimensions.

As __torch_dispatch__ intercepts after the autograd graph has been constructed, coupled FWD-BWD behavior has been difficult to retroactively enforce on the framework side, and difficult to author on the user side without swapping DTensor out in favor of custom autograd functions.

Local-region boundaries. DTensor APIs like local_map or to_local let authors drop into local regions when compute shapes are inexpressible with standard sharding (e.g. MoE EP, uneven multimodal shapes), or when computation is not GSPMD (e.g. TP vocab-parallel CE or embedding). Placements are completely untracked inside local regions, making re-annotation on exit bug-prone, and incoming gradient placements difficult to reason about.

spmd_types: erasure mode, explicit-only, coupled FWD-BWD, local+global SPMD

Motivated by these pain points, we developed spmd_types, a standalone companion library that preserves DTensor-style GSPMD placement-based authoring with different distributed characteristics:

Explicit collectives. As of August 2026, spmd_types does not insert communication implicitly. Authors write the required collective directly—for example, spmd.redistribute(x, "tp", src=spmd.S(i), dst=spmd.R) for an all-gather—and the typechecker verifies that its source and destination types are valid.

Erasure mode. Shard-propagation rules are used for validation rather than runtime planning. During development, authors annotate model inputs and parameters on mesh axes like DP, CP, TP, and EP, then run the program with typechecking enabled to validate operator inputs, outputs, and collectives. Fake process groups allow multi-GPU setups to be typechecked on a single device. In production training, the typechecker is disabled, leaving the validated plain-tensor program and its explicit collectives with no added overhead.

Coupled forward and backward behavior. Differentiable operations such as redistribute, convert, and all_gather are implemented with custom autograd functions. The transition selected in forward therefore fixes the corresponding backward collective, rather than asking shard propagation to make a new placement decision during backward. Typechecking observes the forward call at the __torch_function__ level and validates this contract before the autograd function executes. Forward and backward types are paired as follows:

On each mesh axis, R means replicated data whose independently computed gradients are P; P means pending contributions whose gradient is R. I means identical data and identical computation on every rank, so its gradient remains I. V only promises that values may vary by rank and therefore remains V in backward.

S(i) (the global SPMD refinement of V) is mirrored in BWD. Collective types are similarly tied:

A Megatron-TP MLP example with just the forward():

# x: sequence-sharded if SP is on; invariant otherwise
# w1: colwise TP; w2: rowwise TP

# annotate types (optional at runtime)
x_tp = spmd.S(1) if enable_sp else spmd.I
x_BLD = spmd.assert_type(x_BLD, {tp_axis: x_tp})
w1_DF = spmd.assert_type(w1_DF, {tp_axis: spmd.S(1)})
w2_FD = spmd.assert_type(w2_FD, {tp_axis: spmd.S(0)})

# compute:
# enter TP region -> colwise -> silu -> rowwise -> all-reduce or reduce-scatter back
x_BLD = spmd.redistribute(x_BLD, tp_axis, src=x_tp, dst=spmd.R)
x1_BLF = x_BLD @ w1_DF
hidden_BLF = F.silu(x1_BLF)
out_BLD = hidden_BLF @ w2_FD
out_BLD = spmd.redistribute(out_BLD, tp_axis, src=spmd.P, dst=x_tp)

With sequence parallelism on, inter-block activations are S(1). Forward executes an all-gather, colwise matrix multiply, SiLU, rowwise matrix multiply, and reduce-scatter. Backward is the mirror image: the forward S(1) -> R all-gather fixes a P -> S(1) reduce-scatter in backward, while the forward P -> S(1) reduce-scatter fixes an S(1) -> R all-gather in backward. Weights and intermediate activations retain the same shard dimension in their gradients.

With sequence parallelism off, inter-block activations are I. The forward I -> R conversion is a local no-op, but the corresponding P -> I is a gradient all-reduce in BWD. The MLP compute itself is unchanged.

Global and local SPMD. DTensor-style GSPMD works best when placements can be expressed using standard sharding, with APIs such as to_local and local_map providing escape hatches when they cannot. Common examples include MoE expert parallelism and multimodal encoder computation with uneven sharding, and TP vocab-parallel cross-entropy or embedding, where per-rank computation doesn’t match the global view.

Decoupled forward and backward behavior caused problems for these local-region APIs: the incoming grad_output placement and the reductions required when leaving the region were not always clear, although explicit gradient-placement APIs (PyTorch #155181, #173454, and #175867) and SPMD typechecking inside local_map have recently improved this behavior for DTensor.

SPMD types supports transitioning into local SPMD typechecking under a spmd.local() context or spmd.local_map wrapper. Within the region, spmd.S(i) becomes spmd.V (data that varies across ranks), while spmd.R/I/V/P continue to track whether values are identical or varying and whether forward or backward reductions are required. Only the tensor’s sharding relationship to the logical global tensor is discarded.

It is also useful to adopt only local SPMD: the typechecker can validate existing distributed computation and collectives without requiring the framework to be rewritten in GSPMD.

For more detail, see ezyang’s detailed writeup in the spmd_types local SPMD types documentation.

# Dense transformer compute on the DP/CP/TP mesh.
x = norm(attention(x))
x, ... = router_and_permute(x)

# Enter the EFSDP/EP mesh and use local SPMD semantics.
with spmd.set_current_mesh(sparse_mesh), spmd.local():
    x = spmd.all_to_all(...)  # dispatch
    x = grouped_experts(x, ...)
    x = spmd.all_to_all(...)  # combine

# Return to dense computation and continue the transformer block.
x = token_combine(x, ...)
x = next_layer(...)

Mesh reinterpretation

Unlike a DTensor, a plain tensor in SPMD types is not permanently associated with one DeviceMesh. spmd.set_current_mesh() pushes an ambient mesh onto a thread-local stack. Typechecking and collectives resolve mesh axis names against it, allowing operations like spmd.redistribute(tensor, "tp", src=spmd.S(0), dst=...) to be written without threading process groups through every module.

This is useful because a training step naturally uses several meshes:

1. FSDP/HSDP stores parameters on a storage mesh, treating FSDP as a unified axis separate from HSDP’s replicate axis.
2. At compute time, the same tensors are interpreted on a DP/CP/TP mesh for DP batch sharding, CP all-gathers, and TP collectives.
3. MoE expert computation reinterprets them on a sparse EFSDP/EP view of the same global ranks, repurposing the dense mesh for EP.
4. After expert computation, tensors return to the dense mesh for the rest of the transformer block.
5. After forward and backward, they can be interpreted on a loss or global mesh to reduce losses, gradient statistics, and logged scalars across ranks.

TorchTitan integration

TorchTitan has enabled spmd_types as the default backend for eager pretraining and RL, using spmd.* collectives and an optional typechecking mode across all current models, parallelisms, and runtime features. While the previous default backend used DTensor compute for TP, spmd_types covers every mesh axis except PP (DP/CP/TP/EFSDP/EP), building on the Config-Based Sharding work. The model computation and distributed behavior remain unchanged: both backends emit the same underlying collective patterns—which is important for GraphTrainer’s compilation passes—and produce bitwise-identical numerics.

With config-based sharding, most spmd.redistribute collectives are currently injected at module boundaries. This design is still evolving, and the collectives may eventually move explicitly into forward():

feed_forward_cfg.sharding_config = ShardingConfig(
    in_src_shardings={"x": attn_x_layout},
    in_dst_shardings={
        "x": dense_activation_placement(tp=spmd.R, cp=spmd.S(0)),
    },
)
feed_forward_cfg.w1.sharding_config = colwise_config()
feed_forward_cfg.w3.sharding_config = colwise_config()
feed_forward_cfg.w2.sharding_config = rowwise_config(output_sp=enable_sp)

The feed-forward config owns the input redistribution, the colwise configs own the w1/w3 parameter shardings, and the rowwise config declares a P output followed by either a sequence-parallel reduce-scatter or a non-SP all-reduce. During model initialization, TorchTitan reads these contracts, shards the model state, and installs module-boundary wrappers while leaving the model’s forward() parallelism-invariant.

TorchTitan has also moved toward explicit collective implementations for TP vocab-parallel embedding and cross-entropy, replacing DTensor’s MaskPartial and the implicit loss_parallel() context.

Performance comparison

Erasure mode is intended to remove DTensor’s runtime dispatch and shard-propagation overhead without requiring compilation or cudagraphs. The measurements below use TorchTitan’s debug models on 8 H100 GPUs, LBS 8, and context length 2048. Llama 3 uses selective AC; the MoE configurations use full AC.

The repeated runs show large eager gains of 42-46%. Inductor gains are smaller at 5-9%, and CUDA graphs largely erase the runtime-dispatch difference. Llama 3 retains a small gain, while the Qwen3 and DeepSeek V3 confidence intervals include zero. Peak memory is effectively unchanged.

Inductor compile perf has a clearer gap because TorchTitan’s compile usage is not full-graph but regional, and DTensor subclass-related overhead can still be obvious, when outside of compiled regions. TorchTitan’s chunked loss compiles the inner loss function, but leaves the chunking wrapper, applied on DTensor activations in eager. This leaves obvious tensor subclass wrap/unwrap overhead, particularly for the chunking local_map() call:

def _chunk_local(t):
    chunk_len = t.shape[0] // num_chunks
    return tuple(
        chunk.contiguous()
        for chunk in torch.split(t, [chunk_len] * num_chunks, dim=0)
    )

local_map(
    _chunk_local,
    out_placements=(t.placements,) * num_chunks,
    in_placements=(t.placements,),
    device_mesh=t.device_mesh,
)(t)

We see an overall >20% slowdown for loss computation, and this demonstrates how local-global SPMD transitions aren’t so seamless with DTensor, both in terms of authoring (refactor and wrap a simple chunk call because it’s not GSPMD; chunk then shard != shard then chunk) as well as overhead (8 added _FromTorchTensor calls from the profile, 1 per chunk):

Shape-varying inference decoding is another likely beneficiary because DTensor’s shard-propagation cache keys on shape. Preliminary eager TitanRL measurements showed substantial gains, but a controlled public benchmark has not landed yet; compiled or padded CUDA graph execution can also amortize much of this overhead.

DTensor / FSDP2 status

TorchTitan is deprecating the previously default partial_dtensor backend in favor of spmd_types, removing DTensor from forward and backward computation. In the near term, however, FSDP2 will continue to use DTensor for persistent state. A useful mental model is: DTensor at rest, spmd_types in compute.

FSDP2 continues to use DTensor as the rest-time/storage representation for sharded parameters and gradients. This remains useful for preserving global shape and sharding metadata, integrating with Distributed Checkpointing (DCP), and existing implementations like grad-norm compute. The spmd_types + FSDP2 upstream integration bridges the two representations: FSDP2 reads parameter annotations and translates them into DTensor placements for storage, then restores the spmd_types annotations when parameters are materialized for computation if typechecking is enabled. The corresponding backward types also tell FSDP2 how incoming gradients should be interpreted and which reductions it owns.

For SimpleFSDP in GraphTrainer, we have explored replacing its internal DTensor representation with plain tensors and spmd_types. The compiler passes only need to recognize the resulting collective patterns, so they do not fundamentally depend on DTensor. Even there, however, a DTensor representation at rest remains useful today for DCP.

Longer term, FlexShard may replace FSDP2 for these use cases. FlexShard explicitly separates persistent storage layouts from temporary compute layouts and does not require DTensor as the model-compute representation. Until that path matures, TorchTitan keeps a deliberate storage/compute separation.

Acknowledgments

Many thanks to Chien-Chin Huang (@fegin), Edward Z. Yang (@ezyang), and Tianyu Liu (@tianyu-l) for their design work, implementation, reviews, and discussions throughout this project, and Pei Zhang (@zpcore) and Vishal Nandavanam (@vishal9-team) for being core contributors to the design and implementation of spmd_types.