torch.compile fails with MoE + EP: PendingUnbackedSymbolNotFound on DTensor .view(-1, dim)
### Description
`torch.compile` crashes during Dynamo tracing when MoE is used with Expert Parallelism (EP) + Tensor Parallelism (TP). The error occurs at the `x.view(-1, dim)` call in `MoE.forward()`, where `x` is a DTensor with `_StridedShard(dim=0)` placement.
### Error
```
torch._dynamo.exc.InternalTorchDynamoError: PendingUnbackedSymbolNotFound:
Pending unbacked symbols {u510, u3265, ...} not in returned outputs
DTensor(local_tensor=FakeTensor(..., size=(8192, 256), dtype=torch.bfloat16),
device_mesh=DeviceMesh((tp=2), 'cuda', stride=(1,)),
placements=(_StridedShard(dim=0, sf=8),))
```
Traceback points to `torchtitan/models/common/moe.py`:
```python
x = x.view(-1, dim) # x is DTensor with _StridedShard placement
```
### Repro
```bash
NGPU=8 ./run_train.sh \
--module gpt_oss --config gpt_oss_debugmodel_ep \
--parallelism.data_parallel_shard_degree 4 \
--parallelism.tensor_parallel_degree 2 \
--parallelism.expert_parallel_degree 4 \
--compile.enable
```
### Context
The MoE input arrives as a DTensor with `_StridedShard(dim=0, sf=8)` from the preceding norm/attention layers. Calling `.view(-1, dim)` to flatten `(bs, slen, dim)` to `(bs*slen, dim)` generates unbacked symbolic sizes that Dynamo cannot bind to the output tensor.
The test is currently disabled in `torchtitan/tests/integration_tests/models.py`.
CI run: https://github.com/pytorch/torchtitan/actions/runs/26086665850/job/76701418470?pr=3386
3 条评论