🐛[BUG]: torch.compile + ShardTensor failures with regional model recipe

Version

Newest from main

On which installation method(s) does this occur?

Source

Describe the issue

After torch compile support was added to ShardTensor, I tried to enable compiling the diffusion loss in a domain-parallel (ShardTensor) training run in the regional diffusion model recipe, but it fails with a sequence of independent errors. The same recipe trains fine in eager mode with domain_parallel_size > 1.

The errors go beyond my understanding of ShardTensor details. I included a (human sanity checked) AI-generated summary of the problems encountered when trying to debug this; in the "Reproduced" section below I include a config file and a startup command to replicate the issue.

Reproducer

First, remove the ShardTensor+compilation guard in the recipe: in utils/trainer.py, change

compile_loss = self.use_torch_compile and not self.use_shard_tensor

to

compile_loss = self.use_torch_compile

Then save the following config file as config/repro_shardtensor_compile.yaml:

defaults:
  - dataset/mock
  - model/stormscope
  - training/default
  - sampler/edm_deterministic
  - hydra/default
  - _self_

dataset:
  image_size: [512, 640]
  num_samples: 64
  num_state_channels: 3
  num_background_channels: 4
  num_invariant_channels: 2
  num_scalar_cond_channels: 2
  model_type: "hybrid"
  use_mask: true  # required by use_nan_mask_tokens

model:
  diffusion_conditions: ["background", "invariant"]
  hyperparameters:
    depth: 4
    hidden_size: 384
    num_heads: 6
    patch_size: 4
    attn_kernel_size: 31
    attention_backend: "natten2d_rope"
    layernorm_backend: "apex"
    use_nan_mask_tokens: true
    attn_kwargs:
      qk_norm: true

training:
  batch_size: 1  # domain parallelism requires a local batch size of 1
  domain_parallel_size: 2
  total_train_steps: 3
  print_progress_freq: 1
  checkpoint_freq: 1000    # above total_train_steps: never checkpoints
  validation_freq: 1000    # above total_train_steps: never validates
  validation_steps: 0
  num_data_workers: 1
  log_to_wandb: False
  log_to_tensorboard: False
  seed: 1
  clip_grad_norm: 1.0
  validation_plot_variables: ["state_0"]
  perf:
    fp_optimizations: "amp-bf16"
    torch_compile: True
  optimizer:
    name: "adamw"
    lr: 1.0e-3
    betas: [0.9, 0.99]
  scheduler:
    name: null
    lr_rampup_steps: 1
  loss:
    type: 'edm'
    sigma_distribution: loguniform
    sigma_min: 0.001
    sigma_max: 400.0
    sigma_data: 1.0

sampler:
  args:
    sigma_min: 0.001
    sigma_max: 400.0

Finally, reproduce by running on a node with at least 2 GPUs with:

torchrun --nproc_per_node=2 train.py --config-name repro_shardtensor_compile

(For an eager-mode baseline, you can run it with ++training.perf.torch_compile=false)

Issues (AI-generated)

Common theme: ShardTensor compiles when its ops reach Dynamo as graph nodes evaluated
under fake propagation, and breaks when Dynamo traces ShardTensor's own machinery as
bytecode — which drags to_local, from_local, their autograd.Functions, the
thread-local conversion guard and ShardTensorSpec into the trace. Two things force that
here: na2d calls handle_torch_function from plain Python
(nn/functional/natten.py:116), so the registered handler body is traced instruction by
instruction; and ShardTensor.from_local lacks the in-graph trace rule that
DTensor.from_local has (torch/_dynamo/trace_rules.py:178).

1. DTensor RNG is not traceable before torch 2.12. torch.randn_like(x0) in
diffusion/noise_schedulers/domain_parallel.py:301 gives a fake-tensor error on
aten.to.dtype_layout(tensor(size=(16,), dtype=torch.uint8), ..., device=cuda:N):
OffsetBasedRNGTracker._distribute_region touches the real Philox state under
FakeTensorMode. The traceable run_dtensor_rng_op path first lands in release/2.12.
Fix: torch ≥ 2.12. Worth asserting/documenting a minimum version, since the error gives
no hint.

2. Constructing a ShardTensor in traced Python.
InternalTorchDynamoError: AttributeError: 'ShardTensor' object has no attribute '_spec',
from ret._spec = spec in torch_dynamo_resume_in___new__. Dynamo can't graph
Tensor._make_wrapper_subclass (trace_rules.py:253, and __new__ is in
disallowed_tensor_methods at :3110), so it breaks inside __new__; the resume frame's
STORE_ATTR then calls is_fake() → __tensor_flatten__ → self._spec, which doesn't
exist yet. Reordering __new__ can't help. Hit from the recipe's Apex FusedLayerNorm
monkeypatch and from natten_patches.py:209 (_partial_natten).
Fix: the exemption DTensor already has — torch._dynamo.allow_in_graph(ShardTensor.from_local)
at module scope. Keys on id(fn); from_local is a staticmethod so one marking at import
suffices and eager is unaffected.

3. Mixed DTensor/ShardTensor elementwise ops.
RuntimeError: aten.mul.Tensor got mixed torch.Tensor and DTensor at dit_layers.py:500.
mask_token is a Replicate() DTensor parameter, alpha a Shard(1) ShardTensor;
__torch_function__'s _is_tracing bypass sends the op to ATen, where DTensor rejects
the foreign subclass. Workaround: use_nan_mask_tokens=False — a real functionality loss
for us, the EMEA domain has invalid regions. Suggested fix: make the bypass route through
_dispatch_fallback_via_dtensor when a plain DTensor is present.

4. delattr on a threading.local (torch bug + workaround).
TypeError: UserDefinedObjectVariable.method_setattr_standard() missing 1 required positional argument: 'value'
from shard_tensor.py:249. In torch/_dynamo/variables/user_defined.py:2618,
is_standard_setattr(method) or isinstance(self.value, threading.local) swallows
__delattr__ before the delattr branch at :2621 that supplies DeletedVariable().
Present in 2.13 and main; I can file upstream separately. Fix: in _conversion_scope's
finally, assign _conversion_guard.depth = previous_depth unconditionally —
_conversion_active() is the only reader and already defaults to 0.

5. _ToTorchTensor split setup_context + requires_grad mutation.
AssertionError: False != True in meta_utils.assert_metadata_eq (real vs fake
requires_grad). res.requires_grad_(input.requires_grad) inside forward desyncs the real
tensor from the fake Dynamo recorded. Fix: merge into combined forward(ctx, ...) and
drop the requires_grad_, mirroring DTensor's _ToTorchTensor (the compile-tested
variant; the split style is only needed for torch.func/vmap). _FromTorchTensor
(setup_context at shard_tensor.py:682) likely needs the same.

6. ShardTensorSpec.__hash__ raises on symbolic shapes.
TypeError: unhashable type: non-nested SymInt at _shard_tensor_spec.py:102, via
meta_utils._backward_error → __torch_dispatch__ → DTensor fallback → sharding
propagator cache key. Under fakeification tensor_meta.shape/.stride hold unbacked
SymInts — the mechanism the ShardTensorSpec docstring already documents for
_sharding_shapes. (Upstream has the same latent hole: DTensorSpec.__hash__ notes it
"must be lazy so that Dynamo does not try to hash non-singleton SymInts for the stride".)
Fix: collisions are legal, so except TypeError: return hash((self.mesh, self.placements)).

7. qk_norm promotes q/k to fp32, so NATTEN sees mixed dtypes. This is a general bug that occurs also without ShardTensor when using torch.compile + bf16. The diff below includes a fix.

8. OPEN: fakeifying a real non-leaf ShardTensor can't work.
RuntimeError: Attempted to set the storage of a tensor on device "cuda:1" to a storage on different device "meta" at meta_utils.py:2124. Appeared after I additionally marked
_partial_natten with @torch._dynamo.disable, which leaves its output — a real non-leaf
requires_grad ShardTensor — to be fakeified at the resume point. meta_tensor rebuilds
the subclass, finds its stride/storage_offset don't match (meta_utils.py:2063), and falls
back to forcing the layout with r.set_(r_s, ...). That works for plain FakeTensors
because in_kernel_invocation_manager makes them report meta; it can't work for a
wrapper subclass, whose device (shard_tensor.py:951) is real metadata saying cuda:1.
DTensor never reaches that branch.

Likely contributing: _partial_natten hands from_local the output of unhalo_padding,
a narrowed view (nonzero storage offset, non-contiguous), while
_infer_shard_tensor_spec_from_local_chunks fabricates a C-contiguous global stride
(_shard_tensor_spec.py:634) — the hazard shard_tensor.py:222-227 already documents and
mitigates for _resolve_spec_for_dtensor but not for from_local.

What would help most

1. In-graph exemption for ShardTensor.from_local (#2) — the biggest blocker, with no
   config workaround because physicsnemo's own NATTEN handler trips it.
2. A compile test over the sharded NATTEN path (na2d on ShardTensors with halo
   exchange, inside torch.compile). #4, #5, #6 and #8 all live there.
3. Mixed DTensor/ShardTensor op support (#3), so we can keep masked tokens. Any
   recipe that shards activations but replicates parameters as DTensors will hit this.

--- a/physicsnemo/domain_parallel/_shard_tensor_spec.py
+++ b/physicsnemo/domain_parallel/_shard_tensor_spec.py
@@ -98,8 +98,13 @@ class ShardTensorSpec(DTensorSpec):
         if self._sharding_shapes is not None:
             hash_items.append(tuple(sorted(self._sharding_shapes.items())))
-        hash_tuple = tuple(hash_items)
-        return hash(hash_tuple)
+        try:
+            return hash(tuple(hash_items))
+        except TypeError:
+            # Fakeification turns tensor_meta's shape/stride into unbacked
+            # SymInts, which are unhashable. Hash collisions are permitted, so
+            # degrade to the symbol-free fields; __eq__ still discriminates.
+            return hash((self.mesh, self.placements))

--- a/physicsnemo/domain_parallel/shard_tensor.py
+++ b/physicsnemo/domain_parallel/shard_tensor.py
@@ -245,10 +245,7 @@ def _conversion_scope():
     finally:
-        if previous_depth == 0:
-            delattr(_conversion_guard, "depth")
-        else:
-            _conversion_guard.depth = previous_depth
+        _conversion_guard.depth = previous_depth
@@ class _ToTorchTensor(torch.autograd.Function):
-    @staticmethod
-    def forward(
-        input: "ShardTensor",
-        grad_placements: Sequence[Placement] | None = None,
-    ) -> torch.Tensor:
-        local_tensor = input._local_tensor
-        res = local_tensor.view_as(local_tensor)
-        res.requires_grad_(input.requires_grad)
-        return res
-
-    @staticmethod
-    def setup_context(ctx, inputs, output) -> None:
-        input, grad_placements = inputs
-        ctx.shard_tensor_spec = input._spec
-        ctx.grad_placements = grad_placements
+    @staticmethod
+    def forward(
+        ctx,
+        input: "ShardTensor",
+        grad_placements: Sequence[Placement] | None = None,
+    ) -> torch.Tensor:
+        ctx.shard_tensor_spec = input._spec
+        ctx.grad_placements = grad_placements
+        local_tensor = input._local_tensor
+        return local_tensor.view_as(local_tensor)
@@ -2012,3 +2030,5 @@ install_aot_plain_tangent_coercion()
+
+torch._dynamo.allow_in_graph(ShardTensor.from_local)

--- a/physicsnemo/nn/module/dit_layers.py
+++ b/physicsnemo/nn/module/dit_layers.py
@@ -639,7 +639,7 @@ class RopeNatten2DSelfAttention(Natten2DSelfAttention):
         q, k, v = qkv.unbind(0)
-        q, k = self.q_norm(q), self.k_norm(k)
+        q, k = self.q_norm(q).to(dtype=v.dtype), self.k_norm(k).to(dtype=v.dtype)

Minimum reproducible example

Relevant log output

Environment details

Start with the regional diffusion recipe and the provided config, running torchrun --nproc_per_node=2 train.py --config-name repro_shardtensor_compile. Read utils/trainer.py, physicsnemo/domain_parallel/shard_tensor.py, _shard_tensor_spec.py, natten_patches.py, and dit_layers.py alongside the listed PyTorch trace locations. Done means isolating the independent failures in compile-focused tests and covering the sharded NATTEN path, mixed tensor operations, and the documented version constraints.