ITADN

torch.compile fails with MoE + EP: PendingUnbackedSymbolNotFound on DTensor .view(-1, dim)

#3409OpenacisseJZhong 创建于 2026-05-20
A
acisseJZhongcommented
### 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 条评论