[Torchtitan][Pytorch][ROCm 7.14] Qwen3 MoE TP+EP training produces extremely large(nan) gradient norms
module: rocm
## Problem
The ROCm 7.14 PyTorch wheel produces extremely large gradient
norms during Qwen3 MoE training with FSDP+TP+EP. The equivalent ROCm 7.2 wheel
produces stable gradient norms.
**In some runs this can lead to NaN values or loss-parity failures.**
## Failing test
ROCm 8-GPU feature test:
**Qwen3 MoE FSDP+TP+EP loss parity against reference**
Configuration:
- Model: `qwen3`
- Config: `qwen3_moe_debug`
- FSDP degree: 4
- Tensor parallel degree: 2
- Expert parallel degree: 4
- GPUs: 8
- Deterministic mode with seed 42
- GPU: AMD Instinct MI355X (`gfx950`)
## Reproduction command
```bash
MODULE=qwen3 CONFIG=qwen3_moe_debug NGPU=8 LOG_RANK=0 \
./run_train.sh \
--parallelism.tensor_parallel_degree 2 \
--parallelism.expert_parallel_degree 4 \
--debug.deterministic \
--debug.seed=42 \
--metrics.enable_tensorboard \
--metrics.log_freq=1 \
--training.steps=100 \
--metrics.save_tb_folder=tb_baseline \
--checkpoint.enable \
--checkpoint.export_dtype=bfloat16 \
--checkpoint.load_only
ROCm 7.14 environment — failing
torch: 2.14.0.dev20260805+rocm7.14
torch git: f012e05e101851dce4c495559849318cbb7ed119
HIP: 7.14.60850
spmd_types: 0.2.1
triton-rocm: 3.8.0+git10f6be36
Example output:
step: 1 loss: 8.10514 grad_norm: 21485.1250
step: 5 loss: 5.08495 grad_norm: 66.5681
step: 6 loss: 4.94442 grad_norm: 62844.9609
step: 26 loss: 3.65015 grad_norm: 284010.5938
step: 41 loss: 3.43462 grad_norm: 49713.1133
step: 90 loss: 2.96572 grad_norm: 101188.0547
ROCm 7.2 environment — working
torch: 2.14.0.dev20260802+rocm7.2
torch git: 944416f66ed4dd9502c3ff18a94881857ae6bab3
HIP: 7.2.53211
spmd_types: 0.2.1
triton-rocm: 3.8.0+git10f6be36
Example output:
step: 1 loss: 8.10514 grad_norm: 2.9300
step: 5 loss: 4.82658 grad_norm: 2.5541
step: 26 loss: 3.15735 grad_norm: 0.7451
step: 41 loss: 3.01872 grad_norm: 0.3288
step: 90 loss: 2.82196 grad_norm: 0.2069
Important observation
Both environments produce the same initial loss:
step 1 loss: 8.10514
However, the step-1 gradient norm differs significantly:
ROCm 7.2: 2.9300
ROCm 7.14: 21485.1250
The losses diverge after the first optimizer update:
ROCm 7.2 step 2 loss: 7.34490
ROCm 7.14 step 2 loss: 7.76990
This suggests the initialization, input data, and initial forward pass are equivalent, while the backward pass or distributed gradient-norm calculation differs.
**Expected behavior**
ROCm 7.14 should produce finite, stable gradient norms comparable to ROCm 7.2 and pass the Qwen3 MoE loss-parity test.
0 条评论