[AOTAutograd] [DTensor] subclass output SymInt arity mismatch causes partitioner assertion
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
`torch.compile` fails when compiling a DTensor-returning module where AOTAutograd metadata records symbolic DTensor wrapper sizes, but joint graph capture later unwraps concrete DTensor instances and emits fewer forward outputs.
The failure is:
```text
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
AssertionError: Node view_3 was invalid, but is output
```
This appears to be an AOTAutograd traceable-wrapper-subclass output arity issue. During metadata collection, `subclass_fw_graph_out_meta` can count SymInt size/stride outputs for DTensor outputs. During joint graph capture, `unwrap_tensor_subclasses(..., append_symints=True)` appears to decide which SymInts to emit from the current DTensor instance. If that instance reports the same size as a concrete `int`, the traced graph emits fewer forward outputs than metadata counted. Then `num_inner_fwd_outputs` is too large, and the partitioner misclassifies a backward node as a forward output.
### Minimal repro
```bash
CUDA_VISIBLE_DEVICES=0,1 \
TORCHINDUCTOR_COMPILE_THREADS=1 \
torchrun --standalone --nproc_per_node=2 repro.py
```
```python
#!/usr/bin/env python3
import os
from functools import partial
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import DTensor, Partial, Replicate, distribute_module
from torch.distributed.tensor.parallel import (
ParallelStyle,
SequenceParallel,
parallelize_module,
)
TP_DEGREE = 2
BATCH_SIZE = 1
LOCAL_SEQ_LEN = 1
HIDDEN_SIZE = 8
class ReplicateParallel(ParallelStyle):
@staticmethod
def _prepare_input_fn(placement, mod, inputs, device_mesh):
x = inputs[0]
if not isinstance(x, DTensor):
x = DTensor.from_local(x, device_mesh, (placement,), run_check=False)
return (x, *inputs[1:])
@staticmethod
def _prepare_output_fn(placement, mod, outputs, device_mesh):
def replicate(x):
if isinstance(x, DTensor) and x.placements != (placement,):
x = x.redistribute(placements=(placement,), async_op=True)
return x
if isinstance(outputs, tuple):
return tuple(replicate(x) for x in outputs)
return replicate(outputs)
def _apply(self, module, device_mesh):
return distribute_module(
module,
device_mesh,
None,
partial(self._prepare_input_fn, Replicate()),
partial(self._prepare_output_fn, Replicate()),
)
class IdentityFront(nn.Module):
def forward(self, x):
return x
class Router(nn.Module):
def forward(self, x):
x = x.reshape(-1, HIDDEN_SIZE)
y = x * 2
return y, y + 1
def main():
local_rank = int(os.environ["LOCAL_RANK"])
rank = int(os.environ["RANK"])
torch.cuda.set_device(local_rank)
dist.init_process_group("nccl")
mesh = init_device_mesh("cuda", (TP_DEGREE,), mesh_dim_names=("tp",))
front = IdentityFront().cuda()
router = Router().cuda()
parallelize_module(front, mesh, SequenceParallel())
parallelize_module(router, mesh, ReplicateParallel())
front = torch.compile(front, backend="inductor", fullgraph=True)
router = torch.compile(router, backend="inductor", fullgraph=True)
x = torch.randn(
BATCH_SIZE,
LOCAL_SEQ_LEN,
HIDDEN_SIZE,
device="cuda",
requires_grad=True,
)
x = front(x)
assert isinstance(x, DTensor)
x = x.redistribute(placements=(Replicate(),), async_op=False)
x = x.to_local(grad_placements=(Partial(),))
y0, y1 = router(x)
out = y0 + y1
dist.barrier()
if rank == 0:
print(f"unexpected_success out_norm={out.norm().item():.6f}", flush=True)
dist.destroy_process_group()
if __name__ == "__main__":
main()
```
### Expected behavior
The script should compile and run successfully.
### Actual behavior
On an unfixed build, both ranks fail during Inductor/AOTAutograd partitioning with:
```text
AssertionError: Node view_3 was invalid, but is output
```
### Suspected root cause
AOTAutograd metadata and graph capture disagree on the flattened output arity for DTensor wrapper subclass outputs.
Specifically, metadata collection records extra SymInt outputs for DTensor outer sizes. Later graph capture unwraps the current DTensor instance and omits those outputs if the corresponding size appears concrete. This makes `num_inner_fwd_outputs` larger than the actual number of forward outputs in the joint graph, causing the partitioner to treat backward-only nodes as forward outputs.
A possible fix is to make output unwrapping during graph capture use `subclass_fw_graph_out_meta` as the source of truth for which size/stride SymInt outputs must be emitted, while taking the actual values from the current wrapper instance.
cc @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @weifengpy @ezyang @albanD @chauhang @penguinwu @tianyu-l @XilunWu @SherlockNoMad @ppwwyyxx @bdhirsh @bobrenjc93 @aorenste
Contributor guide
Assessment
This issue has not been assessed yet.