[hychiang-git]

What does this PR do?

Type of change: New feature, new example, bug fix, new tests.

Enable the standard Megatron-Bridge PTQ → mergeable LoRA QAD → HF export workflow for DeepSeek-V4-Flash-0731 routed experts. Bridge imports the official MXFP4/FP8 checkpoint as floating-point training tensors; importance-weighted IQ2_XXS GPTQ calibrates the experts, and distill.py trains the adapters against a frozen online teacher.

• Preserve exact decoded fake-quant values with straight-through gradients and extend the existing frozen-input gradient wrapper to HybridStack for reentrant checkpointing.
• Add V4 CSA/HCA/native HF mappings, retain FP32 mHC controls and integer routing tables, and pack fused gate/up weights before splitting so exported IQ2 values match the training forward.
• Add the Flash recipe, standard-entry-point instructions, optional deterministic/mHC/router-bias controls, and focused regressions. Non-expert weights remain floating point; this PR does not add NVFP4 non-expert quantization.

Stacked draft: based on #2539 (kmorabia/mbridge-glm53-flash) to reuse its shared mHC exporter helpers. Retarget to main after that dependency merges. The PR adds no custom training driver.

Usage

See modelopt_recipes/models/deepseek-ai/DeepSeek-V4-Flash-0731/README.md for the three commands and input preparation requirements.

quantize.py --recipe models/deepseek-ai/DeepSeek-V4-Flash-0731/ptq/iq2_xxs_experts-gptq
            --lora_rank 8 --lora_alpha 16
    → distill.py --student_megatron_path <PTQ checkpoint>
    → export_quantized_megatron_to_hf.py --megatron_path <trained iteration>

The initial V4 export path supports a main-only HF view (MTP excluded), TP1/PP1 and SequentialMLP with expert parallelism. Omit --compress when creating the Megatron training checkpoint. Train only routed-expert rank-8/alpha-16 adapters; export packed IQ2 or a separate folded BF16 HF copy.

Testing

This PR branch: 131 focused CPU tests pass across Flash/GLM recipes, V4 export, GGML codecs/GPTQ and mergeable LoRA. Includes the final fused gate/up and FP32 mHC regressions ported into the normal test suite. New distillation CLI coverage and HybridStack GPU tests are included but were not executed on this CPU-only machine.

Recorded full-model v02 validation: two nodes × four GB300 GPUs; TP1/EP8/PP1/CP1, SequentialMLP. The run used the earlier ModelOpt base f1d68fd4dd50 plus #2539 at 6d6b91752677 and these Flash changes, with the Bridge/Core dependency fixes below.

• IQ2 GPTQ: 512 × 4096-token calibration windows across 43 main blocks; exact zero-adapter attachment and independent eight-rank PTQ reload.
• Four QAD updates at sequence length 512, microbatch 1/global batch 8. Fresh step-two resume matches uninterrupted step four for 156,293 model/adapter/optimizer/RNG entries.
• Exact pre/post-merge probe outputs on all eight ranks.
• Both HF exports checked across 34,223 tensors: all 33,024 packed expert tensors decode exactly to folded BF16, with matching remaining tensors.
• Fresh folded-HF import reproduces merged Bridge probe outputs on all eight ranks.

Output parity uses one calibration-derived 32-token input. These are integration checks, not held-out quality or production-training measurements. The rebased branch has not repeated the full-model GPU run, and optimized packed-IQ2 serving is not covered.

Ruff passes with the exporter's six existing PLR0917 findings excluded; the same six findings are reproduced on the unchanged #2539 base. Formatting, mypy, recipe validation, Bandit, license, RST, whitespace and file-size checks pass. Markdownlint 0.18.1 passes directly; its pre-commit installer is blocked by the local npm git-package restriction.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅ / ❌ / N/A
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: ✅ / ❌ / N/A
• Did you write any new necessary tests?: ✅ / ❌ / N/A
• Did you update Changelog?: ✅ / ❌ / N/A
• Did you get Claude approval on this PR?: ✅ / ❌ / N/A

Additional Information

Dependency pins for the recorded GPU run:

• Megatron-Bridge 5a876442eba8433336c1491d49a6261a0b20749d, plus a data-iterator setup fix preserving restored RNG.
• Megatron-Core/LM 6e7e6ed435e41af04fc41e7007ef862049153fdc, plus support for saving/reloading empty frozen optimizer branches.

Those two fixes belong in their respective repositories and are not included here. Before marking this draft ready, resolve their upstream disposition, complete the framework/GPU checks on the final combined stack, and complete the repository review checklist. No accuracy-retention or serving-performance claim is made.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}