[Qwen2MoE] Fix gradient alignment: remove fake_path and use clear_grad(set_to_zero=False)

[a31413510]

Problem

When training Qwen2MoE with PaddleFormers, post-optimizer weights diverge from HuggingFace (PyTorch) starting from step 2. Two root causes were identified, both related to MoE expert activation behavior.

Root Causes & Fixes

Fix 1: Remove fake_path from Qwen2MoeSparseMoeBlock (modeling.py)

The original code iterated over all 64 experts for every step, using a "fake path" for inactive ones:

fake_state = expert_layer(fake_current_state * 0)  # runs the full MLP with zero input

This produced grad=0 tensors for all inactive expert parameters. AdamW treats grad=0 differently from grad=None — it still applies weight_decay to parameters with grad=0, causing inactive expert weights to slowly drift away from the HuggingFace reference after each optimizer step.

Fix: Replace the all-experts loop + fake_path with nonzero()-based iteration that only processes experts with at least one assigned token — matching the HF implementation exactly and keeping inactive expert gradients as None.

# Before: loop over all 64 experts, fake_path for inactive ones
for expert_idx in range(self.num_experts):
    if tokens_per_expert[expert_idx] <= 0.1:
        fake_state = expert_layer(x * 0)  # produces grad=0, diverges from HF

# After: only process active experts, inactive stay grad=None
expert_hit = paddle.greater(...).nonzero()
for expert_idx in expert_hit:
    expert_idx = int(expert_idx[0])
    ...

Fix 2: clear_grad(set_to_zero=False) in trainer.py

The default clear_grad() (equivalent to set_to_zero=True) zeros out all gradients at the end of each step, which revives the grad=0 problem for the next step even after Fix 1. Using set_to_zero=False keeps grad=None for parameters that were not activated, so AdamW correctly skips them — matching PyTorch/HF behavior.

This fix is consistent with the same change already applied for Qwen3MoE in PR #4358.

Verification

Tested on Qwen2-57B-A14B-Instruct (reduced to 3 hidden layers for speed). Training loss aligned with Swift/HuggingFace reference across 50 steps (MAE = 0.0, bit-exact).

Related

• PR [Qwen3MoE] Fix gradient alignment between PaddlePaddle and PyTorch/HF #4358: same clear_grad(set_to_zero=False) fix for Qwen3MoE

[paddle-bot]

Thanks for your contribution!

[github-actions]

This Pull Request is stale because it has been open for 60 days with no activity. 当前Pull Request 60天内无活动，被标记为stale。