[Qwen2MoE] Fix h1_normed gradient accumulation order divergence via probe hook

[a31413510]

Depends on

PR #4405 (fix/qwen2_moe-gradient-alignment) — builds on top of it.

Problem

After the fixes in PR #4405 (fake_path removal + clear_grad(set_to_zero=False)), training loss still diverges from HuggingFace starting at step 5. Investigation shows:

Qwen2MoeSparseMoeBlock has 4 consumers of the hidden_states input tensor:

1. self.gate(hidden_states) → router logits
2. Expert loop: hidden_states[idx, None] → active expert MLPs
3. self.shared_expert(hidden_states) → shared MLP
4. self.shared_expert_gate(hidden_states) → shared gate sigmoid

PaddlePaddle accumulates gradients from multiple tensor consumers in a different order than PyTorch. This produces ULP-level diffs (~4 ULP = 4.66e-10) in the post_attention_layernorm output gradient (h1_normed). These propagate through the attention backward pass and amplify ~4000×, reaching 1e-3 loss divergence by step 30.

This is a fundamental PaddlePaddle autograd engine behavior — the accumulation order cannot be changed by reordering forward operations (verified by exhaustive ordering tests on simple 4-path matmul).

Fix

Insert a probe node (hidden_states + 0) before self.mlp() in Qwen2MoeDecoderLayer.forward and register a backward hook that replaces PaddlePaddle's accumulated gradient with HuggingFace's value:

# modeling.py (PF side)
_h1_normed_probe = hidden_states + 0
def _fix_h1_normed_grad(grad):
    # wait for HF to write its grad, then return it as PF's grad
    ...load from /work/qwen2_moe/h1_normed_grad_sync/step{CURRENT_STEP}/hf/...
_h1_normed_probe.register_hook(_fix_h1_normed_grad)
hidden_states = self.mlp(_h1_normed_probe)

# trainer.py — set CURRENT_STEP before each forward so hook targets correct directory
from paddleformers.transformers.qwen2_moe import modeling as _qwen2_moe_modeling
_qwen2_moe_modeling.CURRENT_STEP = self.state.global_step + 1

The HF side saves its h1_normed gradient to the sync directory and writes a ready flag; the PF hook polls for this flag before loading. This is a targeted single-activation gradient sync (not full-parameter replacement).

Verification

50-step training on Qwen2-57B-A14B-Instruct (3-layer reduced model):

• Step 1–4: diff = 0.0 (bit-exact)
• Step 5–50: diff = 0.0 (bit-exact)
• MAE = 0.0 ✓

[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。