[megatron] 2/n towards Kimi K2.6: INT4 -> BF16 masters dequantization tool

[avigyabb]

What does this PR do?

Follow-up to #1862 (fake-INT4 QAT). Adds a streaming converter that turns a compressed-tensors pack-quantized INT4 checkpoint into a plain BF16 HF checkpoint — producing the master weights that trainer.policy.model.fake_int4_qat.bf16_base_path loads.

For models whose BF16 originals were never released (Kimi K2-Thinking / K2.6), the dequantized INT4 weights are the only valid masters: under the (scale_divisor=7.0, q_min=-7) convention, fake-quantize of a dequant is the identity bit-for-bit, so training from the dequant reproduces the served grid exactly.

Implementation

• skyrl/backends/skyrl_train/workers/megatron/dequantize_int4_checkpoint.py:
  • Streams shard-by-shard with lazy cross-shard companion lookup (weight_scale can live in a different shard than its weight_packed); buffers ~5GB output shards and rewrites the safetensors index, so the 1T-scale Kimi conversion runs in bounded memory on a CPU node.
  • Unpack/dequant match compressed-tensors dequantize() bit-for-bit (nibble j of each int32 = code j + 8; group multiply in the scale dtype).
  • Everything not pack-quantized (attention, shared experts, dense MLPs, vision tower, lm_head — Kimi K2.6's entire ignore list) is copied through unchanged; quantization_config is stripped from config.json; sidecar files are copied.
  • --verify off|sample|full: each sampled module's dequant must be a fixed point of the QAT grid (recomputed weight_scale equals the stored one, fake-quantize is the identity). Fails by design on llm-compressor RTN (/7.5) checkpoints, whose dequant is not valid as QAT masters — those need the original BF16 release.
  • Rejects weight_zero_point / weight_g_idx (symmetric group quantization only).

Usage:

uv run --isolated -m skyrl.backends.skyrl_train.workers.megatron.dequantize_int4_checkpoint \
    <int4_checkpoint_dir> <bf16_masters_dir> --verify sample

Tests

CPU-only, no compressed-tensors dependency:

uv run --isolated --extra skyrl-train --extra dev pytest tests/backends/skyrl_train/test_dequantize_int4_checkpoint.py

• Unpack and dequant pinned bit-exactly against the real moonshotai/Kimi-K2.6 and Qwen3.6-35B-A3B-INT4-RTN golden slices from _fake_int4_qat_golden.py (the Kimi slice converts end-to-end with --verify full and reproduces the served grid).
• End-to-end directory conversion (sharded + single-file input, forced output sharding + index, config stripping, sidecar copy).
• Verify-rejection on the RTN convention; zero-point rejection.
• Existing test_fake_int4_qat.py suite still passes.

Remaining work (next steps in this series)

• LM-only Megatron-Bridge mapping for KimiK25ForConditionalGeneration (vision tower strip, analogous to the Qwen3.5 LM-only bridge).
• End-to-end W4A16 + LoRA validation (tracker SkyRL quantization (FP8, MXFP8, NVFP4) tracker #1943) and a multi-node K2.6 recipe.

🤖 Generated with Claude Code