ITADN

[Torchtitan][Pytorch][ROCm 7.14] Qwen3 MoE TP+EP training produces extremely large(nan) gradient norms

#4205Openvidushi8 创建于 1 天前
module: rocm
V
vidushi8commented
## 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 条评论