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

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.

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