ShardTensor is not torch.compile-able: graph break in from_local resumes mid-construction before _spec is attached

Summary

Attempting to torch.compile a function that constructs a ShardTensor via from_local and calls redistribute fails on current main (post-#1556). Two distinct symptoms, one root cause: dynamo graph-breaks inside the ShardTensor construction path, and the resume machinery can't recover.

Repro (4 GPUs, torchrun --nproc_per_node=4)

import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.placement_types import Shard
from physicsnemo.domain_parallel import ShardTensor

dist.init_process_group("nccl")
rank, world = dist.get_rank(), dist.get_world_size()
torch.cuda.set_device(rank)
mesh = init_device_mesh("cuda", (world,), mesh_dim_names=["cp"])
S1, S2 = [Shard(1)], [Shard(2)]

def st_flat(local, g1):
    a = ShardTensor.from_local(local, mesh, S1, sharding_shapes="chunk", global_shape=g1)
    return a.redistribute(placements=S2).to_local()

local = torch.randn(1, 2, 12288, 1536, device=f"cuda:{rank}", dtype=torch.bfloat16)
torch.compile(st_flat)(local, (1, 8, 12288, 1536))

(The global_shape kwarg is from #1791; the same failure occurs with default sharding_shapes="infer" on unmodified main -- the graph break just happens slightly earlier in spec inference.)

Failure mode 1 -- flat call: InternalTorchDynamoError: AttributeError: 'ShardTensor' object has no attribute '_spec'

Traceback sequence (PyTorch 2.12 nightly, e53ae1a6):

1. Dynamo traces from_local -> _FromTorchTensor.apply -> forward, and graph-breaks at the _infer_shard_tensor_spec_from_local_chunks call (Python-heavy spec construction).
2. It synthesizes a resume frame (torch_dynamo_resume_in_forward_at_562) and continues at shard_tensor = ShardTensor(...).
3. Tracing resumes inside ShardTensor.__new__, mid-construction: torch.Tensor._make_wrapper_subclass(...) has returned, but _spec has not been assigned yet.
4. Dynamo introspects the half-constructed subclass; __tensor_flatten__ (added in #1556 for compile support) reads self._spec -> AttributeError -> InternalTorchDynamoError.

So the compile-support machinery itself (__tensor_flatten__ needing _spec) collides with dynamo's graph-break-resume landing in the window between wrapper allocation and _spec attachment.

Failure mode 2 -- nested calls: NameError: name 'Shard' is not defined

When the same chain is wrapped in nested helper functions (as in the HealDA reshard code, #1758), the graph break instead produces a resume frame that mis-resolves enclosing-scope names. This is the exact bug already documented in physicsnemo/experimental/models/healda/sharding.py's comment on @torch._dynamo.disable: "dynamo cannot usefully trace it (and tracing the resumed frame mis-resolves names)". Reproduced independently.

Why it matters

The eager-mode gap between ShardTensor resharding and manual all_to_all_single is now mostly closed by removing redundant collectives (#1779, #1791). The residual (~1.6 ms fwd+bwd per reshard round-trip at dit-5B shape, vs compiled manual) is autograd/dispatch overhead -- exactly what compile is supposed to remove. Measured on 4xGB200 with both fixes applied:

Until this works, @torch._dynamo.disable around ShardTensor reshard call sites (HealDA's current approach) is the correct workaround, and any model compiling around such a region pays a graph break there.

Possible directions

• Make _FromTorchTensor.forward fully traceable (hoist/allow-list the spec inference so no graph break lands inside construction), or
• Mark the construction path explicitly opaque (torch._dynamo.disable on _FromTorchTensor.forward / ShardTensor.__new__) so dynamo treats from_local as a black box instead of half-tracing it, then re-enter compiled code afterwards.

Related

• #1556 (merged) -- added __tensor_flatten__/__tensor_unflatten__ groundwork this collides with
• #1779, #1791 -- eager-mode collective fixes; this issue is the remaining piece of the ShardTensor-vs-manual gap from #1758

Start with the provided 4-GPU torch.compile reproduction, then inspect _FromTorchTensor.forward, ShardTensor.new, and tensor_flatten around the graph break. Compare the flat and nested-call failures, including the @torch._dynamo.disable workaround noted in physicsnemo/experimental/models/healda/sharding.py. Done means both from_local/redistribute paths compile without either reported failure.