GLM-4.5 model accuracy alignment

[zhanghonggeng]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Adds bit-exact numerical alignment patches between Megatron-LM and PaddleFleet for the GLM-4.5 MoE model, gated by USE_ACCURACY_COMPATIBLE=1 environment variable. When disabled, all original code paths remain unchanged.

Changes Overview

Router scoring (moe_utils.py, router.py)

• Sigmoid/softplus scoring keeps bf16 intermediate dtype (matching PaddleFleet) and uses fp64 denominator accumulation for deterministic probs normalization.
• Disables local_tokens_per_expert accumulation for expert bias updates during alignment mode.

Permute/Unpermute (moe_utils.py)

• New _PermuteAlignedAutogradFn: deterministic gather-reshape-sum backward (fp32) matching PaddleFleet's _PermuteAlignedPyLayer.
• Deterministic unpermute via embedding lookup + reshape-sum instead of non-deterministic scatter_add_.

SwiGLU activation (fused_bias_swiglu.py)

• Eager (non-fused) forward/backward paths for BiasSwiGLUFunction, SwiGLUFunction, and WeightedSwiGLUFunction to match PaddleFleet's unfused computation order.

Cross-entropy loss (language_module.py)

• Uses F.cross_entropy in fp32 (non-TP-split) when alignment mode is active, matching PaddleFleet's per-token loss calculation.

Embedding (layers.py)

• Forces weight[masked_input] indexing (deterministic backward) instead of F.embedding.

Expert forward (experts.py)

• Pads small expert batches (M < 17) to 32 tokens to match PaddleFleet's GEMM shape alignment.

Grad finalization (finalize_model_grads.py)

• Skips global num_tokens scaling when loss is already normalized in the autograd graph (PaddleFleet's fixed-loss path).
• Skips router expert bias update in alignment mode.

Optimizer (distrib_optimizer.py)

• Allows torch.optim.AdamW in DistributedOptimizer for alignment with PaddleFleet's optimizer choice.

MoE layer (moe_layer.py)

• Clones hidden_states into independent autograd nodes for router/dispatcher/shared-expert paths, preserving bf16 gradient summation order.

Code quality

• Added docstrings to resolve pylint C0116 warnings.
• Removed unused import os.
• Applied black formatting fixes.

Testing

• Verified 3-step training loss matches PaddleFleet bit-exactly on GLM-4.5-Air-tiny-2L (EP=8):
  • Step 1: 12.1779623
  • Step 2: 7.3173542
  • Step 3: 4.90155077
• All changes are gated by USE_ACCURACY_COMPATIBLE=1; default behavior is unchanged.