[megatron] 1/n towards Kimi K2.6: fake QAT training for Qwen3.6-35B-INT4

examples/train/megatron/run_megatron_dapo_qwen3.6_35b_a3b_lora_int4_qat.sh

@@ -0,0 +1,172 @@
+set -x
+
+# DAPO LoRA training for Qwen3.6-35B-A3B on ONE 8xH200 node, colocated, with the
+# INT4 fake-quant QAT stack. vLLM serves the model as compressed-tensors INT4
+# (W4A16); the Megatron trainer holds BF16 masters and (when QAT is on)
+# fake-quantizes the MoE expert GEMMs onto the same INT4 grid in the forward pass
+# (STE backward), with TIS correcting the residual train/infer mismatch.
+#
+# Two ablation modes (QAT_MODE), both serving INT4 in vLLM so the comparison
+# isolates the effect of fake-quant + TIS:
+#   QAT_MODE=off : fake-quant OFF + TIS OFF  -> uncorrected BF16(train)/INT4(infer) mismatch
+#   QAT_MODE=on  : fake-quant ON  + TIS ON   -> corrected (on-policy)
+#
+#   QAT_MODE=off bash examples/train/megatron/run_megatron_dapo_qwen3.6_35b_a3b_lora.sh
+#   QAT_MODE=on  bash examples/train/megatron/run_megatron_dapo_qwen3.6_35b_a3b_lora.sh
+
+# INT4 actor served by vLLM; BF16 masters loaded by the trainer (Megatron-Bridge
+# can't load compressed-tensors, so it reads BF16 from FAKE_QUANT_BF16_PATH).
+MODEL_NAME="${MODEL_NAME:-/data/qwen36-int4/Qwen3.6-35B-A3B-INT4-RTN}"
+FAKE_QUANT_BF16_PATH="${FAKE_QUANT_BF16_PATH:-/data/qwen36-int4/Qwen3.6-35B-A3B}"
+
+DATA_DIR="$HOME/data/dapo"
+TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet"
+TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet"
+
+# --- ONE 8xH200 node, colocated. num_policy_gpus (8) == num_rollout_gpus (1*8). ---
+NUM_NODES=1
+NUM_GPUS_PER_NODE=8
+NUM_INFERENCE_ENGINES=1
+INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=8
+LOGGER="wandb"
+
+# --- QAT / TIS ablation toggle ---
+QAT_MODE="${QAT_MODE:-on}"   # on | off
+if [ "$QAT_MODE" = "on" ]; then
+  FAKE_QUANT_ENABLED=true
+  TIS_TYPE=token
+  RUN_SUFFIX="int4qat_tis_ON"
+else
+  FAKE_QUANT_ENABLED=false
+  TIS_TYPE=null            # disables off_policy_correction TIS
+  RUN_SUFFIX="int4qat_tis_OFF"
+fi
+
+CLIP_RATIO_LOW=0.2
+CLIP_RATIO_HIGH=0.28
+LOSS_REDUCTION="token_mean"
+# Keep overlong (truncated) responses so the batch is never empty after filtering
+# (Qwen3.6 math CoT often exceeds the short response cap used for a quick run).
+APPLY_OVERLONG_FILTERING=false
+OVERLONG_BUFFER_LEN=$((1024 * 1))
+OVERLONG_BUFFER_PENALTY_FACTOR=1.0
+
+USE_KL_LOSS=false
+TEMPERATURE=1.0
+TOP_P=1.0
+EVAL_TOP_P=0.7
+CLIP_RATIO_C=10.0
+MAX_PROMPT_LENGTH=$((1024 * 2))
+MAX_RESPONSE_LENGTH=$((1024 * 4))
+
+# --- reduced scale so the OFF vs ON comparison is quick to eyeball in wandb ---
+TRAIN_BATCH_SIZE=16
+MINI_BATCH_SIZE=16
+N_SAMPLES_PER_PROMPT=8
+EVAL_N_SAMPLES_PER_PROMPT=8
+ENFORCE_EAGER=false
+LR=1e-5
+
+LORA_RANK=32
+LORA_ALPHA=32
+
+# megatron config (8 GPUs: TP=4, EP=8/ETP=1 -> DP=2)
+MEGATRON_TP=4
+MEGATRON_PP=1
+MEGATRON_CP=1
+MEGATRON_EP=8
+MEGATRON_ETP=1
+
+TIS_IMP_RATIO_CAP=2.0
+
+OPTIMIZER_OFFLOAD=true
+OPTIMIZER_OFFLOAD_FRACTION=1.0
+
+# Qwen3.6 flags
+LANGUAGE_MODEL_ONLY=True
+ENGINE_INIT_KWARGS='{"gdn_prefill_backend": "triton", "compilation_config": {"cudagraph_mode": "FULL_DECODE_ONLY"}}'
+DISTRIBUTED_EXECUTOR_BACKEND="mp"
+export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800
+
+uv run --no-sync --extra megatron -m examples.train.algorithms.dapo.main_dapo \
+  data.train_data="['$TRAIN_FILE']" \
+  data.val_data="['$TEST_FILE']" \
+  trainer.algorithm.advantage_estimator="grpo" \
+  trainer.algorithm.policy_loss_type="dual_clip" \
+  trainer.algorithm.overlong_buffer_len=$OVERLONG_BUFFER_LEN \
+  trainer.algorithm.overlong_buffer_penalty_factor=$OVERLONG_BUFFER_PENALTY_FACTOR \
+  trainer.algorithm.loss_reduction=$LOSS_REDUCTION \
+  generator.inference_engine.enforce_eager=$ENFORCE_EAGER \
+  generator.apply_overlong_filtering=$APPLY_OVERLONG_FILTERING \
+  generator.sampling_params.temperature=$TEMPERATURE \
+  generator.sampling_params.top_p=$TOP_P \
+  generator.eval_sampling_params.top_p=$EVAL_TOP_P \
+  generator.eval_sampling_params.temperature=$TEMPERATURE \
+  generator.eval_sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \
+  trainer.algorithm.use_kl_loss=$USE_KL_LOSS \
+  trainer.algorithm.clip_ratio_c=$CLIP_RATIO_C \
+  trainer.policy.model.path="$MODEL_NAME" \
+  trainer.policy.model.fake_int4_qat.enabled=$FAKE_QUANT_ENABLED \
+  trainer.policy.model.fake_int4_qat.group_size=32 \
+  trainer.policy.model.fake_int4_qat.scale_divisor=7.5 \
+  trainer.policy.model.fake_int4_qat.bf16_base_path="$FAKE_QUANT_BF16_PATH" \
+  trainer.policy.megatron_config.lora_config.merge_lora=false \
+  trainer.fused_lm_head_logprob=true \
+  trainer.policy.language_model_only=$LANGUAGE_MODEL_ONLY \
+  generator.inference_engine.language_model_only=$LANGUAGE_MODEL_ONLY \
+  trainer.placement.colocate_all=true \
+  trainer.strategy=megatron \
+  generator.inference_engine.distributed_executor_backend="mp" \
+  trainer.placement.policy_num_nodes=$NUM_NODES \
+  trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \
+  generator.inference_engine.engine_init_kwargs="$ENGINE_INIT_KWARGS" \
+  generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \
+  generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \
+  trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \
+  trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \
+  trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \
+  trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \
+  trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \
+  trainer.policy.model.lora.rank=$LORA_RANK \
+  trainer.policy.model.lora.alpha=$LORA_ALPHA \
+  trainer.policy.megatron_config.optimizer_config_kwargs.overlap_cpu_optimizer_d2h_h2d=$OPTIMIZER_OFFLOAD \
+  trainer.policy.megatron_config.optimizer_config_kwargs.use_precision_aware_optimizer=$OPTIMIZER_OFFLOAD \
+  trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_cpu_offload=$OPTIMIZER_OFFLOAD \
+  trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_offload_fraction=$OPTIMIZER_OFFLOAD_FRACTION \
+  trainer.algorithm.off_policy_correction.tis_ratio_type=$TIS_TYPE \
+  trainer.algorithm.off_policy_correction.token_tis_ratio_clip_high=$TIS_IMP_RATIO_CAP \
+  trainer.epochs=1 \
+  trainer.algorithm.eps_clip_low=$CLIP_RATIO_LOW \
+  trainer.algorithm.eps_clip_high=$CLIP_RATIO_HIGH \
+  trainer.eval_batch_size=64 \
+  trainer.eval_before_train=false \
+  trainer.eval_interval=0 \
+  trainer.update_epochs_per_batch=1 \
+  trainer.train_batch_size=$TRAIN_BATCH_SIZE \
+  trainer.policy_mini_batch_size=$MINI_BATCH_SIZE \
+  trainer.micro_forward_batch_size_per_gpu=1 \
+  trainer.micro_train_batch_size_per_gpu=1 \
+  trainer.ckpt_interval=0 \
+  trainer.max_prompt_length=$MAX_PROMPT_LENGTH \
+  generator.sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \
+  trainer.policy.optimizer_config.lr=$LR \
+  trainer.policy.optimizer_config.num_warmup_steps=0 \
+  trainer.policy.optimizer_config.weight_decay=0.1 \
+  trainer.policy.optimizer_config.max_grad_norm=1.0 \
+  generator.inference_engine.backend=vllm \
+  generator.inference_engine.run_engines_locally=true \
+  generator.inference_engine.weight_sync_backend=nccl \
+  generator.batched=true \
+  environment.env_class=aime \
+  generator.n_samples_per_prompt=$N_SAMPLES_PER_PROMPT \
+  generator.eval_n_samples_per_prompt=$EVAL_N_SAMPLES_PER_PROMPT \
+  generator.inference_engine.gpu_memory_utilization=0.6 \
+  trainer.logger="$LOGGER" \
+  trainer.project_name="qwen3_6_dapo_lora_int4qat" \
+  trainer.run_name="dapo_lora_r32_qwen3_6_35b_a3b_1node_${RUN_SUFFIX}" \
+  trainer.export_path="$HOME/exports/dapo_lora_qwen3_6_${RUN_SUFFIX}" \
+  trainer.hf_save_interval=0 \
+  trainer.resume_mode=none \
+  trainer.max_ckpts_to_keep=1 \
+  trainer.ckpt_path="$HOME/ckpts/dapo_lora_qwen3_6_${RUN_SUFFIX}" \
+  $@

skyrl/backends/skyrl_train/workers/megatron/fake_int4_qat.py

@@ -0,0 +1,146 @@
+"""Fake-INT4 quantization-aware training for Megatron MoE experts.
+
+When vLLM serves the experts as real ``compressed-tensors`` INT4 but the trainer
+holds BF16 masters, the two disagree (a train/infer log-prob gap). This fake-
+quantizes the frozen fused expert GEMMs (``TEGroupedLinear``) onto the same
+group-symmetric 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).
+
+The grid is computed with the *same arithmetic* the serving artifact was produced
+with, so the fake-quantized weights are bit-exact to what the inference engine
+dequantizes (verified element-for-element against real checkpoints):
+
+    amax  = max|w| per group (exact)
+    scale = rn_dtype(amax / scale_divisor)      # single fp32->bf16 rounding,
+                                                # equals the stored ``weight_scale``
+    q     = clamp(round(w / scale), q_min, 7)   # divide+round in the weight dtype,
+                                                # matches compressed-tensors quantize()
+    deq   = q * scale                           # bf16 multiply, matches dequantize()
+
+All-zero groups quantize to zero (guarded division; no eps clamp -- an eps floor
+would distort near-denormal groups that real checkpoints do contain).
+
+Two conventions, selected by ``(scale_divisor, q_min)``:
+
+- ``(7.5, -8)``: llm-compressor / compressed-tensors RTN. Verified bit-exact
+  against ``Qwen3.6-35B-A3B-INT4-RTN`` (requires the *original* BF16 weights as
+  masters; a dequantized INT4 checkpoint does NOT reproduce a /7.5 grid).
+- ``(7.0, -7)``: Kimi K2-Thinking / K2.6 / Miles QAT. Verified bit-exact against
+  ``moonshotai/Kimi-K2.6`` with masters dequantized from the INT4 release (the
+  max-|w| element of every group codes to +-7, which makes the recomputed grid
+  reproduce the stored one exactly).
+
+Enabled and parameterised entirely by ``trainer.policy.model.fake_int4_qat``.
+"""
+
+from __future__ import annotations
+
+import torch
+from loguru import logger
+
+# Symmetric signed-INT4 upper bound; shared by both conventions. The convention
+# knobs (scale_divisor, q_min) come from ``trainer.policy.model.fake_int4_qat``.
+_Q_MAX = 7.0
+
+
+def _ceil_div(a: int, b: int) -> int:
+    return (a + b - 1) // b
+
+
+class _FakeInt4QuantizeSTE(torch.autograd.Function):
+    """Group-symmetric INT4 fake-quantize with a straight-through backward.
+
+    The forward reproduces the compressed-tensors quantize->dequantize pipeline
+    bit-exactly in the weight dtype (see module docstring); the backward is the
+    identity, so gradients pass straight through to the BF16 master weight.
+    """
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, group_size: int, scale_div: float, q_min: float) -> torch.Tensor:  # noqa: D401
+        out_features, in_features = x.shape
+
+        # Pad the input dim up to a whole number of groups. Real MoE dims divide
+        # evenly (2048 / 512 by 32), but stay defensive so odd shapes don't crash.
+        n_padded = _ceil_div(in_features, group_size) * group_size
+        if n_padded != in_features:
+            x_p = x.new_zeros((out_features, n_padded))
+            x_p[:, :in_features] = x
+        else:
+            x_p = x
+
+        # reshape (not view): free for the always-contiguous TE weights, and
+        # copy-on-noncontiguous keeps the public helper safe for sliced inputs.
+        g = x_p.reshape(out_features, n_padded // group_size, group_size)
+        # amax is exact in the weight dtype; the fp32 divide + cast back applies
+        # exactly one rounding, matching compressed-tensors' calculate_qparams
+        # (the stored ``weight_scale``). Grid arithmetic below stays in the
+        # weight dtype so q and deq match quantize()/dequantize() bit-for-bit.
+        amax = g.abs().amax(dim=-1, keepdim=True).to(torch.float32)
+        scale = (amax / scale_div).to(x.dtype)
+        safe_scale = torch.where(scale == 0, torch.ones_like(scale), scale)
+
+        q = torch.clamp(torch.round(g / safe_scale), q_min, _Q_MAX)
+        deq = (q * scale).reshape(out_features, n_padded)
+        out = deq[:, :in_features].contiguous()
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        # Straight-through estimator: identity gradient to the BF16 master weight.
+        return grad_output, None, None, None
+
+
+def fake_int4_quantize_ste(
+    x: torch.Tensor,
+    group_size: int,
+    scale_div: float,
+    q_min: float,
+) -> torch.Tensor:
+    """Apply the fake-INT4 STE to a 2D ``[out, in]`` weight, preserving Megatron's
+    ``main_grad`` bookkeeping so the fused optimizer still finds its grad buffer.
+
+    ``(scale_div, q_min)`` selects the convention: ``(7.5, -8)`` llm-compressor
+    RTN, ``(7.0, -7)`` Kimi/Miles."""
+    out = _FakeInt4QuantizeSTE.apply(x, group_size, scale_div, q_min)
+    if hasattr(x, "main_grad"):
+        out.main_grad = x.main_grad
+    return out
+
+
+_installed = False
+
+
+def install_fake_int4_qat(
+    group_size: int,
+    scale_divisor: float,
+    q_min: float,
+) -> None:
+    """Monkeypatch ``TEGroupedLinear._get_weight_tensors`` to fake-quantize the
+    fused MoE expert weights. Call once per worker when
+    ``trainer.policy.model.fake_int4_qat.enabled`` is set (the config also supplies
+    ``group_size``, ``scale_divisor`` and ``q_min``)."""
+    global _installed
+    if _installed:
+        return
+
+    from megatron.core.extensions.transformer_engine import TEGroupedLinear
+
+    original = TEGroupedLinear._get_weight_tensors
+
+    def _patched(self):
+        return [
+            (
+                fake_int4_quantize_ste(w, group_size, scale_divisor, q_min)
+                if isinstance(w, torch.Tensor) and w.dim() == 2
+                else w
+            )
+            for w in original(self)
+        ]
+
+    TEGroupedLinear._get_weight_tensors = _patched
+    _installed = True
+    logger.info(
+        f"fake-INT4 QAT: patched TEGroupedLinear MoE experts "
+        f"(group_size={group_size}, scale_divisor={scale_divisor}, q_min={q_min}, STE backward)."
+    )