feat(qat): add MXFP4 fake-QAT STE on TEGroupedLinear expert weights

[tonic-scitix]

What does this PR do ?

Adds an opt-in MXFP4 fake-QAT path for TEGroupedLinear expert weights. The forward simulates the exact quantization the FlashInfer MXFP4 rollout kernel applies (per-block (1 x 32) absmax → E8M0 power-of-2 scale → round to E2M1 grid → dequant multiply); backward is straight-through, so BF16 master weights and the optimizer state are unchanged. The goal is to let the learner see the same expert-weight quantization noise that an MXFP4-serving rollout engine sees, as preparation for BF16-train + MXFP4-rollout RL runs (e.g. DSV4 / DeepSeek-V4-Flash).

Bit-exact equivalence to flashinfer.fp4_quantize + mxfp4_dequantize_host has been verified on real DSV4-Flash expert tensors (128 experts × 2 grouped linears, ~200M elements, zero diff).

How to enable

Runtime-gated by env; no CLI changes, no default behavior change.

Env var	Effect
OPEN_TRAINING_MXFP4_FAKE_QAT_FLAG=1	Enable the STE on TEGroupedLinear expert weights. Default 0 (disabled).
OPEN_TRAINING_MXFP4_BLOCK_SIZE=32	Override the quantization block size. Default 32 (MXFP4 spec). Ablation only.

Mutually exclusive with the existing OPEN_TRAINING_INT4_FAKE_QAT_FLAG;
setting both raises at the enable site.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Loading

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

1. Attach the Expert Review label when your PR is ready for review.
2. GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

1. Add Final Review label
2. GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.