PyTorch DevLog

Making PT2 Symbolic Tracing Reliable for Distributed Workloads

Sanket Purandare (@sanketpurandare) · July 10, 2026 · 11 min read
dynamic_shapesunbackeddistributeddtensorflex_attentioninductortracing

TL;DR – To capture a whole distributed training step as one FX graph, PT2 has to trace models whose batch and sequence dimensions are unbacked SymInts. When we tried this on the TorchTitan DeepSeek-V3 MoE trainer, every layer of the stack either silently specialized those dims to concrete ints or blew up with a data-dependent error (DDE). This post walks through the 11 fixes — spanning ATen meta kernels, ProxyTensor, the ShapeEnv, Inductor, collective bucketing, DTensor, and FlexAttention — that make each layer keep tensor semantics symbolic while allowing hints for policy decisions only. That contract is what unblocked end-to-end symbolic tracing for the graph trainer and expert-parallel (EP) compute/communication overlap.

Background / Motivation

The graph trainer captures an entire training step — forward, backward, optimizer, and collectives — as a single FX graph, then rewrites that graph to go faster: bucketing collectives, chunking work across the token grid, and overlapping expert-parallel (EP) communication with expert compute.

The workload driving this effort is EP-overlap for a large MoE model (DeepSeek-V3). To overlap the all-to-all / all-gather / reduce-scatter traffic of expert parallelism with compute, the graph pass chunks the token grid into slices. Those slices are not nice round numbers — they are derived symbolic expressions over the unbacked batch/sequence symbols, things like u0 // 2, 2 * (u0 // 2), and (u1 + 1) // 2.

For the captured graph to be valid across variable-length batches, the batch and sequence dimensions must stay symbolic and unbacked through the entire trace. The moment any layer specializes them to a concrete integer, two things break:

  1. The graph is only correct for the one shape it happened to be traced with.
  2. Rewrites that depend on derived symbolic extents (EP-overlap chunking) can no longer be expressed at all.

But PT2’s tracing path was full of spots that quietly assumed a concrete integer. They failed in two recurring ways:

Every subsystem the trainer touches — matmul/SDPA meta kernels, FlexAttention, DTensor layout propagation, collective bucketing, Inductor codegen — had at least one of these. This is the story of removing them, layer by layer, and of the single contract that ties the fixes together.

Design / Approach

The one contract

One principle runs through all 11 fixes and is the main reusable takeaway:

Semantics stay symbolic; hints are for policy only. A value that describes runtime tensor behavior — a size, a stride, a split, a reshape — must remain symbolic. An optimization hint may be consumed only where a Python integer is genuinely required to choose an optimization (bucket byte sizing, foreach grouping, layout scoring). On those policy paths, an unhinted unbacked symbol fails fast with a clear error rather than guessing.

A hint is never a promise that a dimension really is that value. Everything else falls out of this:

The rest of this section walks the stack bottom-up, because if a lower layer specializes, nothing above it can stay symbolic.

1. ATen / FakeTensor meta kernels — keep the math symbolic

The meta and decomposition kernels are the foundation. Two of them were still forcing concrete sizes on the matmul/attention path.

The folded matmul path (ND @ 1D/2D) forced concrete sizes twice — an empty-tensor numel() check and a DimVector fold/view built from sizes(). Reworked to the symbolic equivalents (sym_numel(), sym_sizes(), sym_strides(), reshape_symint(), _unsafe_view_symint(), symbolic output resize), with contiguity guarded through TORCH_GUARD_OR_FALSE. It also lets fake view metadata recover from hinted unbacked contiguity relations — e.g. a shape like 2 * (u0 // 2) whose stride is expressed through u0 — by emitting symbolic equality checks that hold at the traced hint instead of specializing u0. (#183397, fixes torchtitan#3322)

SDPA lowers through batched matmul, and the broadcasted bmm path still converted batch dims to int64_t before expand/reshape/view — breaking rank-4 inputs with an unbacked batch symbol, including math SDPA. Carried symbolic sizes through with sym_sizes(), SymDimVector, infer_size_symdimvector(), expand_symint(), reshape_symint(), _unsafe_view_symint(). Separately, aten.size.default has a non-symbolic int[] schema, so running it on a fake tensor with unbacked sizes forces the symbol; during torch-function metadata tracing we redirect it to aten.sym_size.default so downstream consumers get symbolic size proxies (this covers cuDNN SDPA fake outputs whose batch dim stays symbolic). (#183398, fixes torchtitan#3324)

Finally, trace tooling serializes FakeTensor storage metadata to JSON. We keep the size field symbolic and add a separate size_hint field only when every free symbol in the storage-size expression has an explicit hint override — giving diagnostic/policy tooling a concrete expected extent without ever specializing the shape. (#183839, fixes #183835)

2. Inductor codegen — bind and order unbacked extents without guarding

Several Inductor codegen paths needed to reason about unbacked extents without turning policy into semantic guards (#183840):

(fixes #183834, #185341)

3. Collective bucketing — isolate and tolerate symbols

Bucketing rewrites collectives with nested make_fx traces, and it turned out to be a hotspot for both failure modes.

First, isolation. Bucketing reuses the surrounding FakeTensorMode / ShapeEnv. When that env already has pending fresh unbacked symbols from an outer dynamic trace, compute_unbacked_bindings tries to account for symbols that have nothing to do with the bucket’s inputs and outputs, and errors. The fix snapshots and clears pending/ignorable fresh unbacked symbols around the nested bucketing trace, then restores the ambient state afterward — so bucketing is responsible only for the symbols it produces. (#183495, fixes #183679)

Then, tolerance — the diff that crystallized the whole project’s contract. Bucketing uses tensor metadata for two distinct purposes:

(#183544, fixes #183676)

4. Dynamic Shapes / ShapeEnv core — the substrate

rebind_unbacked already recorded equivalences when a retraced binding mapped an unbacked symbol to another symbol or to a constant, but it asserted that any non-symbol replacement with free symbols was invalid. That is too strong for legitimate derived unbacked shapes like (u1 + 1) // 2: the old symbol still has a concrete binding relationship and should be eliminated in favor of that expression. We now record the replacement via _eliminate_unbacked (the existing path for replacing an unbacked symbol by a non-symbol expression). Notably, this fixed the FlexAttention/HOP reproducer without restoring the previously-reverted broad HOP fake-trace suppression that had caused cond / AOTInductor / Executorch / FlexAttention regressions. (#183837, fixes #183677)

Non-strict tracing can also receive raw SymInt inputs from an outer fake-tensor trace — this is the FlexAttention BlockMask path used by graph chunking, where tensor sizes and raw SymInt captures refer to the same outer unbacked sequence length. We taught ShapeEnv to transfer a foreign unbacked SymInt expression into the local ShapeEnv: it first rebuilds the expression from any already-transferred foreign symbols, and for any unresolved remainder mints one opaque local unbacked symbol for the whole foreign expression, recording its source, range, and hint. A shared cache means tensor dimensions and scalar captures that originate from the same foreign symbols preserve their sharing. Guard printing now emits torch.sym_max / torch.sym_min for SymPy Max/Min so guards evaluate symbolically instead of via Python truthiness. Raw foreign wrapping stays gated to non-strict tracing, and data-dependent branching on the transferred symbol still fails through the normal DDE path rather than specializing on a hint. (#187273, fixes #187272)

5. ProxyTensor / make_fx / Regional Inductor — preserve provenance at the boundary

Metadata should be preserved at the boundary where a value becomes an FX proxy, not rediscovered later in FlexAttention or some other downstream subsystem (#187231):

(fixes #187230)

6. DTensor — recognize equivalent symbolic layouts

Compiled DTensor paths can see symbolic local layout metadata that is semantically valid but not syntactically identical to what was saved during forward propagation. For to_local() backward, AOTAutograd can produce a local gradient stride like (Max(1, u3), 1) while the saved DTensor spec uses (u1, 1); recomputing the global gradient stride forced compute_global_tensor_info() to evaluate symbolic stride relations and raised a DDE for the default same-placement backward.

We now reuse the original spec for that default backward when the strides are compatible: provable equality is accepted directly, contiguous symbolic forms like Max(1, u*) go through check_contiguous_sizes_strides(..., false_if_dde=True), and otherwise we emit torch._check assertions for the required stride equalities before taking the shortcut. Placement changes and uneven channels-last shards keep the recomputation path so autograd can repair the physical layout. The same class of issue in aten.t sharding propagation is fixed by registering transpose with allow_unbacked_sharding=False, so unproven candidates (8 < 2*u0) are pruned instead of being evaluated as a Python bool. (#187026, fixes #187025)

7. FlexAttention HOP — the top of the stack

Finally, the FlexAttention BlockMask path itself — handling unbacked symbolic predicates and scalar shape captures. This is the direct consumer of everything below it, and the workload that motivated the whole stack. (#183838, fixes #183833)

Results / Benchmarks

With the full stack landed:

Each fix landed with targeted regression coverage: fake-tensor unit tests (rank-3/4 unbacked matmul, math + cuDNN SDPA, symbolic aten.size.default, hinted-view contiguity), dynamic-shapes tests (foreign-ShapeEnv transfer, rebind-to-expression, symbolic sym_max/sym_min guard printing), Inductor unbacked-symint tests (stride ordering, fallback binding reuse), collective bucketing trace tests, and DTensor compile tests. On top of the unit coverage, the TorchTitan graph_trainer H100 integration run (aot_fx_trace_deepseek_v3_sdpa_full_inductor_ep_overlap_moe_seq) exercises the full path end-to-end.

Open questions / Future work

References