Repository navigation
ShardTensor: fixes for torch.compile with domain parallelism (#1986) - #2210
Draft
coreyjadams wants to merge 8 commits into
Draft
coreyjadams wants to merge 8 commits into
coreyjadams wants to merge 8 commits into
Conversation
4 of 6 tasks
Contributor
CODEOWNERS review mapCurrent for commit ⏳ @coreyjadams — 9 file(s)
⏳ @negin513 — 9 file(s)
⏳ @peterdsharpe — 9 file(s)
⏳ @pzharrington — 1 file(s)
No CODEOWNER
Comment |
Dynamo folds has_torch_function to False for traceable wrapper subclasses, so a compiled na2d called the raw kernel on the ShardTensor and skipped the halo exchange (wrong results, or rank-divergent collectives). Under compile, route subclass calls through an allow_in_graph helper: Dynamo records one graph node and fake propagation runs the handler, as for torch ops. Test uses a pure-torch reference kernel, so it runs without NATTEN.
The dispatch handlers compared shard dims against dim as written: dim=-1 (stack's backward, e.g. RoPE rotate-half under compile) relabeled Shard(2) as Shard(1) on the gradient, and a negative dim on the sharded axis escaped the guard.
With dynamic shapes, tensor_meta holds SymInts, which are unhashable: the spec hash (DTensor sharding propagation) and validate_aligned_sharding's layout key both raised. The hash falls back to (mesh, placements); the alignment check skips symbolic layouts.
…compile
Under tracing, unpatched ops went straight to dispatch; a plain DTensor
operand (e.g. a Replicate parameter like the DiT mask token) then took the
op and rejected the ShardTensor ('got mixed torch.Tensor and DTensor').
Such ops now use the DTensor fallback, as in eager.
Under tracing, ops with a grad-requiring plain operand (exact Tensor, Parameter, FakeTensor or FunctionalTensor; grad mode on) take the DTensor fallback, whose differentiable promotion reduces the Partial grad to a plain tensor as in eager. Other tensor subclasses keep their dispatch.
Symbolic specs hash by (mesh, placements) only, so specs from different compiles collide in the spec cache. Comparing their SymInts raised 'SymInt ... not tracked' (1-GPU UNet force_sharding compile). Compare ints by value and SymInts by identity; anything else is a cache miss.
Dynamo treats ShardTensor as a plain tensor and never calls its __torch_function__, so a Python-level has_torch_function check (or allow_in_graph) can't reach the ShardTensor handler cleanly. na1d/2d/3d now call torch.library custom ops when an operand is a tensor subclass; ShardTensor registers the halo-exchange handler on them with register_named_function_handler, which runs in eager and under AOT. Plain tensors still call NATTEN directly.
coreyjadams
force-pushed
the
dp-06-compile-1986
branch
from
October 30, 2026 21:46
9626e5b to
93dc28f
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PhysicsNeMo Pull Request
Description
Fixes the
torch.compile+ ShardTensor failures from #1986 that remained afterthe domain-parallel stack (#1979 to #1983). Each fix has a test that fails without it.
tensor and never calls its
__torch_function__, so the wrapper'shas_torch_functioncheck was folded to False and a compiledna2dcalledthe raw kernel on the ShardTensor, skipping the halo exchange: wrong results,
or rank-divergent collectives. With a tensor-subclass operand,
na1d/2d/3dnow call
torch.libraryopsphysicsnemo::na{1,2,3}d, and ShardTensorregisters its halo-exchange handler on them with
register_named_function_handler. That handler runs in eager and under AOT,above autograd, so gradients flow through the halo exchange as in eager.
Plain tensors still call NATTEN directly. On the subclass path only
is_causal,scaleandbackendare passed through; other NATTEN optionsraise. NATTEN's own ops can't be hooked instead: its Hopper/Blackwell
backends permute tokens into tiles before calling them. The test uses a
pure-torch reference kernel, so it runs without NATTEN.
select/select_backwardwith a negative dim. The handlers comparedshard dims against
dimas written;dim=-1(stack's backward, e.g. RoPErotate-half under compile) relabeled
Shard(2)asShard(1)on the gradient.dynamic=True, the spec holdsSymInts, which are unhashable: the spec hash and the alignment check both
raised. The hash falls back to
(mesh, placements); the alignment checkskips symbolic layouts. Because symbolic specs now share a hash, specs from
different compiles meet in the spec cache;
__eq__compares their SymIntsby identity only, instead of raising "SymInt ... not tracked" (seen on the
1-GPU UNet
force_shardingcompile).plain DTensor operand (e.g. a
Replicateparameter such as the DiT masktoken) took the op and rejected the ShardTensor. Such ops now use the DTensor
fallback, as in eager.
ShardTensor got an unreduced
ShardTensor(Partial)gradient under compile,which FSDP2 rejects. Under tracing, ops with a grad-requiring plain operand
now take the differentiable promotion, so the gradient is reduced as in
eager. Limited to exact plain types (
Tensor,Parameter, and theirFakeTensor/FunctionalTensorstand-ins) with grad mode on, so othertensor subclasses keep their own dispatch.
Not addressed here:
from_localoutput across a graph break, e.g. around an apexLayerNorm). It does not affect the recipe without apex; it needs a test that
reproduces it first.
Validated on 4 GPUs:
test/domain_parallel(natten, natten compile, DiT,spec cache, compile) green; the regional weather diffusion recipe trains
compiled with domain parallelism on 1, 2 and 4 GPUs (next PR).
Checklist
Dependencies
None.