Checklist
Detailed Information
Describe the bug
When using FSDP2 engine with Tensor Parallelism (TP > 1) for Qwen family models (Qwen3, etc.), the forward pass crashes with a DTensor dispatch error.
The root cause is that Qwen models have intermediate operations (aten.alias, aten.slice) between the final LayerNorm and lm_head in their forward() method. These ops are not registered in PyTorch's DTensor dispatch table, so when the norm output is a DTensor with Shard(1) placement, the subsequent slice/alias ops fail to propagate the sharding metadata, causing a dispatch error before reaching lm_head.
Expected behavior
Qwen models should work correctly with FSDP2 + TP > 1, the same as Llama models. The forward pass through final_norm → lm_head should complete without DTensor dispatch errors.
Full logs
To Reproduce
Commit ID
Please provide your Git commit ID.
Environment
Please provide your software and hardware information if you're not using a
containerized environment.
Script
The bash script or YAML configuration to run:
Checklist
cause, not a secondary error caused by peer workers.
Detailed Information
Describe the bug
When using FSDP2 engine with Tensor Parallelism (TP > 1) for Qwen family models (Qwen3, etc.), the forward pass crashes with a DTensor dispatch error.
The root cause is that Qwen models have intermediate operations (aten.alias, aten.slice) between the final LayerNorm and lm_head in their forward() method. These ops are not registered in PyTorch's DTensor dispatch table, so when the norm output is a DTensor with Shard(1) placement, the subsequent slice/alias ops fail to propagate the sharding metadata, causing a dispatch error before reaching lm_head.
Expected behavior
Qwen models should work correctly with FSDP2 + TP > 1, the same as Llama models. The forward pass through final_norm → lm_head should complete without DTensor dispatch errors.
Full logs
To Reproduce
Commit ID
Please provide your Git commit ID.
Environment
Please provide your software and hardware information if you're not using a
containerized environment.
Script
The bash script or YAML configuration to run: