[megatron] 1/n towards Kimi K2.6: fake QAT training for Qwen3.6-35B-INT4#1862
Open
casper-hansen wants to merge 5 commits into
Open
[megatron] 1/n towards Kimi K2.6: fake QAT training for Qwen3.6-35B-INT4#1862casper-hansen wants to merge 5 commits into
casper-hansen wants to merge 5 commits into
Conversation
Fake-quantize the frozen fused MoE expert GEMMs (TEGroupedLinear) onto the served compressed-tensors INT4 grid inside the forward pass, with a straight-through-estimator backward so gradients still reach the BF16 masters (LoRA adapters stay BF16, matching "INT4 base + BF16 adapter" at inference). This removes the train/infer log-prob gap when vLLM serves the experts as INT4 (e.g. Qwen3.6-35B-A3B-INT4-RTN) while the Megatron trainer holds BF16 masters. - fake_int4_qat.py: group-symmetric INT4 STE; scale_divisor=7.5 matches llm-compressor/compressed-tensors RTN (bit-exact), 7.0 matches Kimi/Miles. - config: trainer.policy.model.fake_int4_qat (enabled, group_size, scale_divisor, bf16_base_path). - megatron_worker: install the STE hook and, since Megatron-Bridge cannot load a compressed-tensors INT4 checkpoint, redirect the trainer's BF16 master weights via bf16_base_path while keeping the INT4 model.path as the logical identity (so saved LoRA adapters reference the INT4 base). Ref worker mirrors the policy. The BF16 redirect is independent of the STE so an INT4-served baseline (fake-quant off) still loads. - example: new run_megatron_dapo_qwen3.6_35b_a3b_lora_int4_qat.sh -- a 1-node H200 INT4-QAT ablation (QAT_MODE off/on) with the fused LM head, leaving the original run_megatron_dapo_qwen3.6_35b_a3b_lora.sh untouched.
…ors grid The STE previously computed the quantization grid in fp32 with an eps floor on the scale, which diverged from the artifact the inference engine actually serves: llm-compressor/compressed-tensors produces the grid in the weight dtype (scale = rn_bf16(amax/divisor), divide+round in bf16, bf16-multiply dequant), and near-denormal groups (which real checkpoints contain) were zeroed by the eps clamp. Measured against Qwen3.6-35B-A3B-INT4-RTN, only ~89% of expert weight elements landed on the served grid (residual RMS 4.3e-4 of the 1.03e-3 no-QAT gap). Rework the STE forward to mirror the compressed-tensors quantize->dequantize arithmetic exactly: fp32 amax with a single rounding into the weight dtype for the scale, grid math in the weight dtype, and a zero-guard instead of the eps clamp. Verified bit-exact (100.000000% of elements, RMS 0) against both supported artifact conventions: - Qwen3.6-35B-A3B-INT4-RTN (scale_divisor=7.5, q_min=-8), 75.5M elements - moonshotai/Kimi-K2.6 (scale_divisor=7.0, q_min=-7), 50.7B elements across layers 1/30/60 with BF16 masters dequantized from the INT4 release Adds fake_int4_qat.q_min to the config for the Kimi convention, whose QAT pins codes to [-7, 7] and never emits -8. A/B on the INT4-QAT example (QAT_MODE=on, 5 steps, 8xH200): policy/rollout_train_logprobs_abs_diff_mean drops from 0.0156-0.0189 (old STE) to 0.0118-0.0141 (fixed), std 0.053 -> 0.040, max spikes 4.4 -> 2.2, TIS is_ratio_max 17.4 -> 3.6, tis_clip_frac 1e-4 -> 0.
The (scale_divisor, q_min) convention pair is fully specified by trainer.policy.model.fake_int4_qat, so the module-level SCALE_DIV_* / Q_MIN_* constants were redundant defaults that could silently drift from the config's. Remove them and make the STE / installer parameters required, so every caller states the convention explicitly (the worker already passes all values straight from the config).
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces fake-INT4 quantization-aware training (QAT) for Megatron MoE experts, addressing the train/infer log-prob mismatch when serving with INT4 but training with BF16. It adds a straight-through estimator (STE) for group-symmetric INT4 fake-quantization, integrates it into the Megatron worker, and updates the configuration schema and training scripts. The review feedback suggests replacing ".view()" with ".reshape()" in the fake-quantization forward pass to robustly handle potentially non-contiguous tensors and prevent runtime errors.
Black (pre-commit, 24.10.0) wants the conditional expression in the _get_weight_tensors patch parenthesized; the code-quality job runs it over all files so the branch was failing CI. Also switch the two grid .view() calls to .reshape() (Gemini review): the TE expert weights the hook sees are always contiguous (reshape is then exactly view, no copy), but fake_int4_quantize_ste is a public helper and a sliced/transposed 2D input would make .view raise; reshape handles it by copying. Verified bit-exactness is unchanged and non-contiguous inputs now round-trip correctly.
CPU tier (runs in the cheap CI job, no megatron/TE needed): - golden bit-exactness tests pinning the STE to slices of the real served artifacts (Qwen3.6 INT4-RTN with true BF16 masters incl. near-denormal groups and stored -8 codes; Kimi K2.6 with dequant masters), stored as bit patterns so any grid-arithmetic drift fails exactly - STE properties: identity backward, main_grad propagation, padding, all-zero groups, convention code ranges, div-7 idempotence/identity, non-contiguous inputs - config defaults + CLI overrides for fake_int4_qat GPU tier (gpu_ci/megatron, guards megatron-core/TE bumps): - tripwire that TEGroupedLinear._get_weight_tensors exists and is still consumed by TE GroupedLinear.forward, with actionable failure messages - hook-effect canary: patched forward must equal the same TE grouped GEMM run on pre-quantized weights bitwise, differ from the unpatched forward, and track in-place master mutations (catches silent no-op and future weight caching) - _maybe_setup_fake_int4_qat wiring (redirect independent of STE install) - worker-level forward on truncated Qwen3-30B-A3B (2 layers, TP4/EP4): QAT on/off must change logprobs moderately through the real init path
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
INT4 serving now adds essentially nothing over BF16 serving. The gap with fake-quant on equals the same-weights floor. The example for Megatron uses https://huggingface.co/casperhansen/Qwen3.6-35B-A3B-INT4-RTN which includes the code to use vLLM project's llm-compressor to produce these weights.
The difference between the Qwen3.6 INT4 weights and Kimi is a slight nuance in the SCALE_DIV and Q_MIN variable which is intended to be forward compatible when Kimi K2.6 is supported at a later time.
Results
STE removes ~55% of the uncorrected train/infer gap, halves the spread, and cuts the tail hardest.
Can we get closer to zero?
The remaining 0.0118 is the cross-engine numerics floor, and we're sitting exactly on it. The cleanest evidence is STE step 1: at that point the LoRA adapters are still zero and the expert grid is bit-exact, so trainer and vLLM score bit-identical weights and the gap is 0.0118.
Inspiration
This was inspired by the work done in miles https://github.com/radixark/miles/ on the Megatron side.