fix(megatron): sync fused shared-outer LoRA main grads

[zhoutong-hai]

Summary

• identify the replicated factor from each Megatron Bridge SharedOuterGroupedExpertAdapter during LoRA model construction
• when fused WGrad is enabled, SUM-all-reduce its optimizer-facing main_grad over the tensor+data group for the current pipeline stage
• derive the replica group and world size from Megatron rather than assuming a particular TP, DP, CP, EP, or PP layout
• fail closed when the ordered factor inventory or bounded post-reduction fc1/fc2 SHA-256 evidence differs across replicas
• retain the existing Bridge Parameter-hook path when WGrad fusion is disabled, avoiding a double reduction

Problem

Megatron Bridge registers a Parameter hook to synchronize the shared outer factor of routed-expert LoRA adapters. With fused weight-gradient accumulation, however, the real optimizer gradient is written directly into param.main_grad; only a dummy tensor reaches the Parameter hook. The optimizer therefore consumes rank-local shared-factor gradients.

The failure is delayed by LoRA initialization: LoRA-B starts at zero, so shared expert-fc1 LoRA-A first receives a useful gradient only after LoRA-B changes. In a GLM-5.2 TP8/PP4/EP8 run, the first update looked healthy, while the second update split every replicated expert-fc1 LoRA-A tensor into eight TP hashes before publication.

Disabling fusion on the adapter module alone did not repair the optimizer-facing buffer. Moving the SUM collective to final gradient processing did: eight consecutive full-scope updates produced replica-equal fc1/fc2 gradient evidence on all four pipeline stages, nonzero bitwise-equal early/middle/terminal native deltas, and maximum train/rollout absolute difference 0.0474493, KL 0.0162019, and TIS clip fraction 0.0125902. The full scope included routed-expert down-projection LoRA.

Implementation

Miles already finalizes explicitly tagged LoRA gradients before optimizer preparation. This change extends that mechanism with an intra-pipeline-stage domain matching the Bridge shared-factor ownership contract. Adapter class and factor dimensionality identify the target; model-specific parameter suffixes are not used. Stable per-chunk/module keys order collectives consistently across replicas.

For each present fc1/fc2 family, only one fixed-size SHA-256 digest is gathered after reduction. An inventory digest is gathered first so incompatible layouts fail before the gradient collective sequence.

Validation

• 66 LoRA utility tests passed
• added TP4/PP2 group-selection coverage that rejects crossing the pipeline axis
• added the second-update LoRA-A case with a dummy .grad and the real nonzero gradient in .main_grad
• added fc1/fc2 shared-factor discovery, non-fused no-double-reduce, and fail-closed digest tests
• Ruff checks passed
• Python byte-compilation passed
• rebased onto current main (cc92260df)

Non-goal

Long sparse-DSA prefill scoring exposed a separate SGLang indexer top-k determinism issue. It is intentionally excluded from this gradient-synchronization PR.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.