diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.6_35b_a3b_lora_int4_qat.sh b/examples/train/megatron/run_megatron_dapo_qwen3.6_35b_a3b_lora_int4_qat.sh new file mode 100644 index 0000000000..1b8c3e2c3c --- /dev/null +++ b/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}" \ + $@ diff --git a/skyrl/backends/skyrl_train/workers/megatron/fake_int4_qat.py b/skyrl/backends/skyrl_train/workers/megatron/fake_int4_qat.py new file mode 100644 index 0000000000..9312e65dcd --- /dev/null +++ b/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)." + ) diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 9f6b52c795..36622840cc 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -325,6 +325,51 @@ def extract_weights(self, dtype: torch.dtype): class MegatronWorker: + def _maybe_setup_fake_int4_qat(self): + """Wire up INT4-served training and return the BF16 bridge-weights path. + + Reads the *policy's* ``model.fake_int4_qat`` (single source of truth for the + shared base model; the ref worker mirrors it). Two independent knobs: + + - ``bf16_base_path``: when set, the trainer loads its BF16 master weights + from here instead of ``model.path``. Needed whenever ``model.path`` is a + compressed-tensors INT4 checkpoint (which Megatron-Bridge cannot load) + served by the inference engine. This redirect happens regardless of + ``enabled`` -- so an INT4-served *baseline* WITHOUT fake-quant (to show + the uncorrected train/infer mismatch) still loads. + - ``enabled``: additionally install the ``TEGroupedLinear`` fake-quant STE + so the trainer's MoE experts match the INT4 grid the sampler serves. + + Returns ``bf16_base_path`` (or ``None`` for a plain BF16 ``model.path``). + """ + fq = getattr(self.cfg.policy.model, "fake_int4_qat", None) + if fq is None: + return None + + rank0 = getattr(self, "_rank", 0) == 0 + if fq.enabled: + from skyrl.backends.skyrl_train.workers.megatron.fake_int4_qat import ( + install_fake_int4_qat, + ) + + install_fake_int4_qat( + group_size=fq.group_size, + scale_divisor=fq.scale_divisor, + q_min=fq.q_min, + ) + if rank0: + logger.info( + f"fake-INT4 QAT enabled (group_size={fq.group_size}, scale_divisor={fq.scale_divisor}); " + f"trainer BF16 masters from {fq.bf16_base_path or 'model.path'}, " + "MoE experts fake-quantized to INT4 in forward (STE backward)." + ) + elif fq.bf16_base_path and rank0: + logger.info( + f"fake-INT4 QAT disabled; trainer loads BF16 masters from {fq.bf16_base_path} " + "while the inference engine serves INT4 model.path (uncorrected train/infer mismatch)." + ) + return fq.bf16_base_path or None + def init_configs( self, model_path, @@ -335,9 +380,17 @@ def init_configs( flash_attn=False, lora_config=None, language_model_only=False, + bridge_weights_path=None, ): """ Initialize the Megatron-Bridge bridge and provider objects + hf_config and tokenizer + + ``bridge_weights_path`` (fake-INT4 QAT): when set, the Megatron-Bridge loads + its BF16 master weights from this path instead of ``model_path``. Used when + ``model_path`` is a compressed-tensors INT4 checkpoint (which the bridge + cannot load) served by the inference engine, while the trainer keeps BF16 + masters and fake-quantizes them in the forward pass. Tokenizer + HF config + (the logical model identity) still come from ``model_path``. """ tokenizer = get_tokenizer(model_path, trust_remote_code=True) hf_config_original = AutoConfig.from_pretrained(model_path, trust_remote_code=True) @@ -368,7 +421,13 @@ def init_configs( for key in ("recompute_granularity", "recompute_method", "recompute_num_layers"): transformer_config_kwargs[key] = None - bridge = AutoBridge.from_hf_pretrained(model_path, trust_remote_code=True) + bridge_source = bridge_weights_path or model_path + if bridge_weights_path: + logger.info( + f"fake-INT4 QAT: loading BF16 master weights from {bridge_source} " + f"(logical model / inference checkpoint: {model_path})" + ) + bridge = AutoBridge.from_hf_pretrained(bridge_source, trust_remote_code=True) # For Qwen3.5, language_model_only routes to the native GPTModel + GDN # path (which supports sample packing) instead of the VL Qwen3VLModel @@ -436,6 +495,10 @@ def init_configs( self.provider = provider self.bridge = bridge + # Logical model identity (what the inference engine serves). Differs from + # the bridge weights path only under fake-INT4 QAT (INT4 model.path, BF16 + # bridge weights); used so saved LoRA adapters reference the INT4 base. + self._logical_model_path = model_path # strategy.hf_config is the on-disk source-of-truth used by # save_hf_configs and must NOT carry runtime overrides like @@ -789,6 +852,10 @@ def init_model(self, model_path, num_training_steps: int = 1e9): """ Initialize the model, optimizer, and scheduler for the policy worker. """ + # Fake-INT4 QAT: install the MoE expert fake-quant hook and (when the + # served checkpoint is INT4) redirect the trainer's BF16 master weights. + bridge_weights_path = self._maybe_setup_fake_int4_qat() + # initialize the bridge and provider objects self.init_configs( model_path, @@ -798,6 +865,7 @@ def init_model(self, model_path, num_training_steps: int = 1e9): bf16=self.cfg.bf16, flash_attn=self.cfg.flash_attn, language_model_only=self.cfg.policy.language_model_only, + bridge_weights_path=bridge_weights_path, ) if self.enable_router_replay: @@ -1294,7 +1362,8 @@ async def _save_lora_adapters_and_sync( set(infer_target_modules_from_adapter_weights(adapter_state.keys())) - {"base_layer"} ) base_model_name_or_path = str( - getattr(self.bridge.hf_pretrained, "model_name_or_path", "") + getattr(self, "_logical_model_path", "") + or getattr(self.bridge.hf_pretrained, "model_name_or_path", "") or getattr(self.bridge.hf_pretrained, "name_or_path", "") ) adapter_config = build_adapter_config_dict( @@ -1506,6 +1575,11 @@ def init_model(self, model_path, num_training_steps: int = 1e9): """ Initialize the model for the ref worker. """ + # Fake-INT4 QAT: the ref shares the policy's base model. Mirror the + # BF16-master redirect so it can load an INT4-served checkpoint, and the + # (global) fake-quant hook keeps the KL anchor in the same weight space. + bridge_weights_path = self._maybe_setup_fake_int4_qat() + # initialize the bridge and provider objects self.init_configs( model_path, @@ -1515,6 +1589,7 @@ def init_model(self, model_path, num_training_steps: int = 1e9): bf16=self.cfg.bf16, flash_attn=self.cfg.flash_attn, language_model_only=self.cfg.ref.language_model_only, + bridge_weights_path=bridge_weights_path, ) self.actor_module = self.make_megatron_module( diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 5bdc74214b..771be488e5 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -100,10 +100,45 @@ class SkyRLLoraConfig(BaseConfig): >= ``max_loras`` if explicitly set.""" +@dataclass +class FakeInt4QatConfig(BaseConfig): + """Fake-INT4 quantization-aware training for MoE experts (Megatron only). + + When the inference engine serves the MoE experts as real ``compressed-tensors`` + INT4 (e.g. ``casperhansen/Qwen3.6-35B-A3B-INT4-RTN``) but the trainer holds + BF16 masters, enabling this fake-quantizes the frozen expert GEMMs onto the + same INT4 grid in the forward pass (straight-through backward), removing the + train/infer weight mismatch. See + ``skyrl.backends.skyrl_train.workers.megatron.fake_int4_qat``. + """ + + enabled: bool = False + group_size: int = 32 + """Group size along the input dim; must match the served checkpoint (32).""" + symmetric: bool = True + scale_divisor: float = 7.5 + """Symmetric-INT4 scale divisor ``scale = amax / scale_divisor``: + ``7.5`` = llm-compressor / compressed-tensors RTN (``[-8, 7]``; matches + ``casperhansen/Qwen3.6-35B-A3B-INT4-RTN``); ``7.0`` = Kimi K2-Thinking / K2.6 / + Miles (``[-7, 7]``). Set ``q_min`` consistently.""" + q_min: float = -8.0 + """Lower clamp of the INT4 code range: ``-8`` for llm-compressor RTN + (``scale_divisor=7.5``), ``-7`` for Kimi/Miles (``scale_divisor=7.0``, whose + QAT never emits ``-8``).""" + bf16_base_path: Optional[str] = None + """Megatron-Bridge cannot load a compressed-tensors INT4 checkpoint, so when + ``model.path`` points at the INT4 model the trainer loads its BF16 master + weights from this path instead (the Miles ``--ref-load`` pattern). The INT4 + ``model.path`` remains what the inference engine serves and the logical name. + When None, the trainer loads weights from ``model.path`` directly (only valid + if that path is already a BF16 checkpoint).""" + + @dataclass class ModelConfig(BaseConfig): path: Optional[str] = None lora: SkyRLLoraConfig = field(default_factory=SkyRLLoraConfig) + fake_int4_qat: FakeInt4QatConfig = field(default_factory=FakeInt4QatConfig) # --------------------------------------------------------------------------- diff --git a/tests/backends/skyrl_train/_fake_int4_qat_golden.py b/tests/backends/skyrl_train/_fake_int4_qat_golden.py new file mode 100644 index 0000000000..fd3ae70dd7 --- /dev/null +++ b/tests/backends/skyrl_train/_fake_int4_qat_golden.py @@ -0,0 +1,36 @@ +"""Golden fixtures for fake-INT4 QAT bit-exactness tests. AUTO-GENERATED. + +Slices of real quantization artifacts, stored as bit patterns (uint16 for bf16, +int32 for compressed-tensors packed codes) so no float-repr rounding is involved. + + QWEN_RTN: llm-compressor RTN convention (scale_divisor=7.5, q_min=-8), from + Qwen3.6-35B-A3B (BF16 masters, true originals) and Qwen3.6-35B-A3B-INT4-RTN + (weight_packed / weight_scale), layer 0 expert 0: down_proj rows 0-3 + (normal groups, includes stored -8 codes) + gate_proj rows 0-3 + (near-denormal groups, amax ~ 7.8e-38 -- pins the no-eps-clamp behavior). + KIMI_QAT: Kimi K2.6 convention (scale_divisor=7.0, q_min=-7), from + moonshotai/Kimi-K2.6 layer 1 expert 0 gate_proj rows 0-7; masters are the + compressed-tensors dequant (the only BF16 form that exists for this model). + +Each entry: rows, cols, group_size, scale_divisor, q_min, and bit patterns for +masters (bf16), weight_packed (int32, 8 codes per int32), weight_scale (bf16). +Regenerate with verification/gen_golden_fixtures.py (needs the checkpoints). +""" + +# fmt: off +QWEN_RTN = dict( + rows=8, cols=64, group_size=32, scale_divisor=7.5, q_min=-8.0, + masters_u16=[[15493, 469, 469, 33237, 469, 15430, 33237, 469, 15219, 33237, 33237, 48238, 48179, 469, 469, 33237, 469, 48043, 469, 33237, 48152, 48154, 33237, 469, 48166, 48263, 469, 33237, 15301, 15251, 33237, 469, 469, 48278, 469, 33237, 469, 469, 469, 33237, 469, 48044, 469, 33237, 33237, 33237, 469, 469, 15258, 33237, 47970, 48157, 33237, 469, 33237, 469, 469, 47789, 48243, 469, 469, 33237, 469, 48271], [15190, 469, 33237, 33237, 469, 15530, 469, 33237, 48249, 33237, 469, 15268, 15250, 469, 469, 33237, 33237, 48220, 469, 469, 15509, 47746, 33237, 469, 48251, 15265, 469, 33237, 48041, 47592, 33237, 469, 33237, 15129, 469, 469, 469, 33237, 33237, 469, 33237, 15444, 33237, 33237, 469, 33237, 469, 33237, 48218, 33237, 47992, 48286, 33237, 469, 469, 469, 469, 48090, 48028, 469, 33237, 469, 33237, 15314], [15029, 469, 33237, 33237, 469, 15431, 469, 469, 15363, 33237, 469, 48136, 48199, 469, 469, 33237, 33237, 48221, 469, 469, 15244, 47878, 33237, 469, 14976, 48210, 469, 33237, 15155, 47374, 33237, 469, 33237, 15230, 469, 469, 469, 33237, 33237, 469, 469, 48076, 33237, 33237, 469, 33237, 469, 33237, 47987, 33237, 15434, 15433, 33237, 469, 469, 469, 33237, 47756, 15227, 33237, 33237, 469, 33237, 15076], [47880, 33237, 469, 33237, 33237, 48252, 33237, 33237, 47975, 33237, 33237, 15113, 48135, 33237, 33237, 469, 469, 48040, 33237, 33237, 15443, 48145, 33237, 33237, 48111, 48150, 33237, 469, 48052, 15350, 469, 33237, 469, 48263, 33237, 33237, 33237, 469, 469, 33237, 33237, 48083, 469, 33237, 33237, 33237, 469, 469, 48128, 33237, 48040, 48212, 469, 33237, 33237, 469, 469, 15067, 48105, 469, 469, 33237, 469, 15205], [15236, 15421, 15584, 15363, 48056, 15300, 48152, 15305, 48032, 15187, 15200, 48133, 15349, 15374, 14966, 47873, 15370, 47686, 15233, 15623, 15422, 48339, 48281, 48119, 15255, 15551, 48304, 48150, 15622, 48103, 48304, 48261, 15461, 48069, 48263, 15478, 47787, 15335, 48187, 15401, 47468, 15517, 48004, 15028, 15384, 48128, 48170, 48016, 48259, 14996, 48302, 48050, 48026, 48356, 15169, 15319, 15560, 15430, 15187, 47994, 47958, 15401, 47686, 48157], [33237, 33237, 33237, 469, 469, 33237, 33237, 469, 33237, 469, 469, 33237, 469, 469, 33237, 469, 469, 33237, 33237, 469, 469, 33237, 469, 469, 469, 469, 469, 469, 33237, 33237, 469, 33237, 33237, 469, 469, 469, 469, 33237, 33237, 33237, 469, 469, 469, 469, 469, 33237, 33237, 469, 33237, 469, 33237, 33237, 469, 469, 469, 33237, 33237, 469, 33237, 469, 33237, 33237, 469, 33237], [469, 33237, 469, 469, 33237, 469, 469, 469, 469, 469, 33237, 33237, 33237, 33237, 469, 33237, 33237, 33237, 469, 33237, 33237, 469, 469, 33237, 469, 469, 469, 33237, 33237, 469, 33237, 469, 33237, 33237, 33237, 33237, 33237, 469, 469, 469, 469, 469, 33237, 469, 469, 469, 469, 33237, 469, 33237, 469, 33237, 33237, 33237, 33237, 469, 469, 33237, 469, 33237, 33237, 469, 33237, 469], [33237, 469, 33237, 469, 469, 33237, 33237, 33237, 33237, 469, 469, 33237, 469, 469, 33237, 469, 469, 33237, 33237, 469, 469, 33237, 469, 469, 33237, 469, 469, 469, 469, 33237, 469, 33237, 33237, 469, 469, 469, 469, 33237, 33237, 33237, 469, 33237, 469, 33237, 33237, 33237, 33237, 469, 33237, 469, 33237, 33237, 469, 469, 469, 33237, 33237, 469, 33237, 33237, 469, 33237, 469, 33237]], + packed_i32=[[-1998026609, -2004674422, -2008774552, -2002024445, -2004318200, -2004318104, -2004334710, 411599480], [-1996978039, -2004178814, -2003859400, -2004449118, -2004318056, -2004317992, -2004347005, -1467447720], [-1996978039, -2004797300, -2005235704, -2004187111, -2004318040, -2004318136, -2004287610, -1735882120], [-2012706681, -2004576122, -2008119192, -2000320444, -2004318200, -2004318120, -2004343164, -1467447912], [-1768444231, -2002097769, 1680603530, 1131373529, -1247230868, 1969915864, -1459129724, 1488419263], [-249564911, -234938383, -14683873, 521273343, 287309809, -249561089, 536809969, 521269745], [-917729, 521212415, 535895825, -235855873, -978671, 536867327, -250536161, -235856097], [287306225, -234938383, -14683873, 522190833, 287309809, -250536161, 536809969, 522129905]], + scale_u16=[[15120, 15136], [15157, 15145], [15084, 15063], [15110, 15120], [15248, 15219], [114, 114], [114, 114], [114, 114]], +) +# fmt: on + +# fmt: off +KIMI_QAT = dict( + rows=8, cols=64, group_size=32, scale_divisor=7.0, q_min=-7.0, + masters_u16=[[0, 48452, 15528, 48452, 48096, 48096, 48096, 48296, 48224, 15456, 15456, 48096, 48296, 48224, 48224, 15328, 15328, 48296, 15528, 48296, 0, 15656, 15456, 48296, 48296, 15328, 15656, 48096, 48296, 48424, 15584, 15528, 0, 15578, 48218, 15450, 48218, 15652, 15450, 0, 0, 0, 48218, 0, 15450, 48346, 0, 48218, 48346, 0, 48474, 48346, 48218, 15578, 0, 15450, 15450, 0, 15450, 15652, 15450, 48420, 15807, 48218], [15377, 48273, 48145, 15505, 48273, 15742, 15377, 48145, 0, 48346, 48273, 15578, 15377, 15578, 48401, 15578, 15706, 15377, 48145, 15669, 0, 15377, 15377, 48437, 48145, 48401, 15578, 15633, 15633, 15377, 15578, 48273, 0, 15542, 15542, 48310, 15414, 48310, 15670, 48392, 48182, 48310, 15542, 15775, 15414, 15752, 48182, 48182, 0, 48182, 0, 48182, 0, 48182, 48438, 48182, 15542, 48310, 15414, 0, 15542, 15624, 15624, 15624], [0, 15366, 48262, 15366, 48262, 15494, 0, 15622, 48329, 48134, 48329, 48134, 48134, 15494, 48424, 15622, 15366, 48390, 15722, 15494, 0, 15656, 48134, 48262, 0, 15366, 48134, 48390, 48262, 48134, 48262, 15494, 15441, 15619, 48337, 48337, 15671, 48081, 48285, 48285, 48285, 0, 0, 0, 15517, 48209, 0, 48209, 0, 48209, 48209, 48209, 15313, 48081, 0, 48209, 15517, 15517, 15517, 48209, 15313, 15313, 0, 0], [15315, 15518, 15315, 15443, 48211, 48286, 15443, 15315, 15443, 15315, 48339, 48441, 48388, 15315, 15315, 0, 48339, 48211, 48339, 15315, 48286, 15673, 48083, 15518, 15620, 15518, 48286, 15443, 48339, 0, 48083, 15315, 48388, 48260, 15620, 48326, 15492, 15364, 0, 48260, 15620, 48132, 48132, 48132, 48260, 48260, 0, 48132, 15492, 48260, 48326, 15558, 0, 48132, 48487, 0, 48421, 48132, 48260, 48388, 15364, 48326, 48326, 15558], [48221, 15526, 48449, 48093, 48093, 48294, 48221, 15325, 15325, 15526, 48294, 48294, 15325, 15526, 15325, 15654, 15325, 48349, 15453, 15325, 15453, 0, 48093, 15626, 15453, 48093, 48349, 48093, 0, 48221, 0, 48294, 48375, 15545, 15642, 15704, 0, 48119, 0, 15351, 0, 48119, 48119, 48375, 48247, 48410, 48247, 15479, 48247, 48119, 15351, 48119, 48247, 48119, 48375, 48247, 15351, 15479, 15351, 0, 15351, 48247, 48375, 48119], [0, 48341, 0, 15573, 48270, 48142, 15374, 15502, 48142, 15701, 0, 15374, 48142, 15374, 15736, 48398, 15502, 48341, 48341, 48142, 48142, 48341, 15630, 48398, 48142, 0, 0, 15630, 48341, 15573, 0, 15502, 15458, 48397, 15330, 48298, 48226, 15330, 15530, 15458, 48226, 15530, 15458, 15458, 0, 48226, 15530, 0, 0, 15586, 48454, 48098, 48098, 15458, 48226, 48226, 15330, 0, 15586, 48298, 48098, 15658, 15458, 48397], [48159, 48287, 48159, 48455, 15391, 48159, 15519, 48159, 15391, 48287, 48287, 15598, 48287, 15519, 48159, 15519, 15519, 15391, 15687, 48366, 48159, 15519, 48159, 15391, 15755, 15598, 48287, 15519, 48159, 0, 15519, 15647, 15475, 15475, 15347, 15475, 48243, 48115, 0, 15640, 48243, 0, 15475, 15475, 15603, 0, 48310, 15475, 48371, 15603, 48438, 15475, 0, 0, 15347, 15475, 15542, 15701, 48115, 48115, 48371, 48243, 15542, 48371], [48373, 0, 15544, 15544, 48117, 48245, 48117, 15477, 48373, 0, 15702, 15349, 48117, 48440, 15349, 48245, 15349, 15477, 15702, 48373, 48245, 48245, 15605, 0, 15544, 15544, 15477, 48312, 48312, 48312, 15349, 15349, 15788, 48325, 15429, 15429, 48404, 48197, 15429, 48197, 15429, 0, 15429, 48197, 15429, 48197, 15557, 15429, 48404, 15429, 15429, 15557, 0, 15429, 0, 0, 0, 0, 15429, 0, 48197, 15429, 15429, 15557]], + packed_i32=[[1467423512, -1771734362, 1525177177, -1138393451, -1984456792, 2020181896, -1733860218, 2136586633], [2046207849, -1262897576, 966317982, 1805437767, 1550412456, 2011822695, 1954052216, -1145403030], [-928606568, -1012435595, 1742253897, -1502197864, 1434404058, 1751877765, 1752786536, -2003211333], [-1705596487, -1986849638, -1208642460, -1752914499, 1754946660, 2019981180, -2122795670, -1252440461], [-1772654154, -373729863, -678782391, 1483240570, -1736901196, -1506392200, 1685485942, 1953073577], [-1451837352, 1335335143, 1280800090, -1464481657, -1416210118, -1956074826, 1722249672, 988241033], [2054764391, -1482246551, -1750639206, -897079617, -663311958, -1517507962, -1450663228, 1264875515], [-1486374012, 1764204420, -1939451991, -1722459461, 2037750127, -1703315063, -2003261035, -1449686648]], + scale_u16=[[15328, 15450], [15377, 15414], [15366, 15313], [15315, 15364], [15325, 15351], [15374, 15330], [15391, 15347], [15349, 15429]], +) +# fmt: on diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fake_int4_qat.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fake_int4_qat.py new file mode 100644 index 0000000000..d2b7d7121e --- /dev/null +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fake_int4_qat.py @@ -0,0 +1,273 @@ +"""GPU integration tests for the fake-INT4 QAT TEGroupedLinear hook. + +Run with: +uv run --isolated --extra dev --extra megatron -- pytest -s tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fake_int4_qat.py + +These protect the feature against megatron-core / TransformerEngine version bumps. +The dangerous failure mode is SILENT: if a TE refactor stops routing +``_get_weight_tensors()`` into the grouped GEMM (or starts caching weights), the +hook keeps "installing" fine but no longer affects compute, and the train/infer +gap quietly returns. The canary test therefore asserts the hook's *effect on the +GEMM output*, not just its installation. +""" + +import inspect + +import pytest +import ray +import torch + +from skyrl.backends.skyrl_train.distributed.dispatch import ( + WorkerOutput, + loss_fn_outputs_to_tensor, +) +from skyrl.backends.skyrl_train.workers.megatron import fake_int4_qat as fq_mod +from skyrl.backends.skyrl_train.workers.megatron.megatron_worker import MegatronWorker +from skyrl.train.config import SkyRLTrainConfig +from skyrl.train.utils.utils import validate_cfg +from tests.backends.skyrl_train.gpu.utils import ( + init_worker_with_type, + ray_init_for_tests, +) + +MOE_MODEL_NAME = "Qwen/Qwen3-30B-A3B" +GS = 32 + + +# --------------------------------------------------------------------------- +# Tripwires: fail loudly (and readably) when megatron/TE refactor the surfaces +# the monkeypatch depends on. +# --------------------------------------------------------------------------- + + +@pytest.mark.megatron +def test_te_grouped_linear_still_exposes_get_weight_tensors(): + import transformer_engine.pytorch.module.grouped_linear as te_grouped_linear + from megatron.core.extensions.transformer_engine import TEGroupedLinear + + assert hasattr(TEGroupedLinear, "_get_weight_tensors"), ( + "TEGroupedLinear._get_weight_tensors is gone: the fake-INT4 QAT monkeypatch " + "(skyrl.backends.skyrl_train.workers.megatron.fake_int4_qat) no longer has an attach point. " + "Find where the new TE version fetches weights in GroupedLinear.forward and re-target the patch." + ) + fwd_src = inspect.getsource(te_grouped_linear.GroupedLinear.forward) + assert "_get_weight_tensors" in fwd_src, ( + "transformer_engine GroupedLinear.forward no longer calls _get_weight_tensors(): " + "the fake-INT4 QAT patch would install but silently stop affecting the expert GEMMs. " + "Re-target the patch to the new weight-fetch path before bumping TE." + ) + + +@pytest.mark.megatron +def test_maybe_setup_fake_int4_qat_wiring(): + """The worker helper: BF16 redirect is independent of the STE install.""" + + class _Stub: + _rank = 0 + cfg = SkyRLTrainConfig() + + stub = _Stub() + fq_cfg = stub.cfg.trainer.policy.model.fake_int4_qat + # SkyRLTrainConfig nests trainer.*; the worker reads self.cfg.policy (=cfg.trainer.policy) + stub.cfg = stub.cfg.trainer + + # disabled + no redirect + assert MegatronWorker._maybe_setup_fake_int4_qat(stub) is None + assert not fq_mod._installed + + # disabled + redirect: path returned, STE NOT installed (INT4-served baseline) + fq_cfg.bf16_base_path = "/data/bf16-masters" + assert MegatronWorker._maybe_setup_fake_int4_qat(stub) == "/data/bf16-masters" + assert not fq_mod._installed + + # enabled: STE installed once, path still returned + from megatron.core.extensions.transformer_engine import TEGroupedLinear + + orig = TEGroupedLinear._get_weight_tensors + try: + fq_cfg.enabled = True + assert MegatronWorker._maybe_setup_fake_int4_qat(stub) == "/data/bf16-masters" + assert fq_mod._installed + assert TEGroupedLinear._get_weight_tensors is not orig + finally: + TEGroupedLinear._get_weight_tensors = orig + fq_mod._installed = False + + +# --------------------------------------------------------------------------- +# Canary: the patched weights must actually flow into the grouped GEMM. +# Runs in a ray task so the process-global patch and torch.distributed state +# cannot leak into other tests. +# --------------------------------------------------------------------------- + + +@ray.remote(num_gpus=1) +def _run_hook_effect_canary(): + import os + + import torch + from megatron.core import parallel_state as mpu + from megatron.core.extensions.transformer_engine import ( + TEColumnParallelGroupedLinear, + TEGroupedLinear, + ) + from megatron.core.transformer import TransformerConfig + + from skyrl.backends.skyrl_train.workers.megatron.fake_int4_qat import ( + fake_int4_quantize_ste, + install_fake_int4_qat, + ) + + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29511") + torch.distributed.init_process_group(backend="nccl", world_size=1, rank=0) + mpu.initialize_model_parallel(tensor_model_parallel_size=1, expert_model_parallel_size=1) + torch.cuda.set_device(0) + + num_gemms, in_f, out_f, tokens = 4, 128, 64, 32 + cfg = TransformerConfig( + num_layers=1, + hidden_size=in_f, + num_attention_heads=4, + num_moe_experts=num_gemms, + add_bias_linear=False, + bf16=True, + params_dtype=torch.bfloat16, + use_cpu_initialization=False, + gradient_accumulation_fusion=False, + sequence_parallel=False, + ) + + def make_module(seed): + torch.manual_seed(seed) + return TEColumnParallelGroupedLinear( + num_gemms=num_gemms, + input_size=in_f, + output_size=out_f, + config=cfg, + init_method=cfg.init_method, + bias=False, + skip_bias_add=True, + is_expert=True, + ).cuda() + + module = make_module(0) + torch.manual_seed(1) + x = torch.randn(tokens, in_f, dtype=torch.bfloat16, device="cuda") + m_splits = [tokens // num_gemms] * num_gemms + masters = [getattr(module, f"weight{i}").data.clone() for i in range(num_gemms)] + + with torch.no_grad(): + out_unpatched, _ = module(x, m_splits) + + # reference: the same UNpatched GEMM, with weights replaced by their fake-quant + ref_module = make_module(0) + for i in range(num_gemms): + getattr(ref_module, f"weight{i}").data.copy_(fake_int4_quantize_ste(masters[i], GS, 7.5, -8.0)) + out_reference, _ = ref_module(x, m_splits) + + orig = TEGroupedLinear._get_weight_tensors + try: + install_fake_int4_qat(group_size=GS, scale_divisor=7.5, q_min=-8.0) + out_patched, _ = module(x, m_splits) + + # 1) the hook has an effect at all (catches silent no-op) + effect = not torch.equal(out_patched, out_unpatched) + # 2) and the effect is exactly quantization: bitwise-equal to the + # same TE kernel run on pre-quantized weights + exact = torch.equal(out_patched, out_reference) + + # 3) weights are re-read every forward (catches future TE caching): + # mutate a master in place and expect the output to track it + getattr(module, f"weight{0}").data.mul_(2.0) + out_mutated, _ = module(x, m_splits) + tracks_mutation = not torch.equal(out_mutated, out_patched) + finally: + TEGroupedLinear._get_weight_tensors = orig + + torch.distributed.destroy_process_group() + return effect, exact, tracks_mutation + + +@pytest.mark.megatron +def test_fake_int4_hook_affects_grouped_gemm(ray_init_fixture): + effect, exact, tracks_mutation = ray.get(_run_hook_effect_canary.remote()) + assert effect, "patched forward == unpatched forward: the fake-quant hook is a silent no-op" + assert exact, ( + "patched forward != TE GEMM on pre-quantized weights: the hook is quantizing " + "something other than what the GEMM consumes (weight-fetch path changed?)" + ) + assert tracks_mutation, ( + "output ignored an in-place master-weight update: TE appears to cache weights " + "across forwards, so QAT would train against stale quantized experts" + ) + + +# --------------------------------------------------------------------------- +# Worker-level end-to-end: enabling the flag must change the policy logprobs +# through SkyRL's real init path on a (truncated) MoE model. +# --------------------------------------------------------------------------- + + +def _moe_forward_cfg() -> SkyRLTrainConfig: + cfg = SkyRLTrainConfig() + cfg.trainer.policy.model.path = MOE_MODEL_NAME + cfg.trainer.micro_forward_batch_size_per_gpu = 2 + cfg.trainer.micro_train_batch_size_per_gpu = 2 + cfg.trainer.remove_microbatch_padding = True + cfg.trainer.logger = "console" + cfg.trainer.gradient_checkpointing_use_reentrant = True + cfg.trainer.strategy = "megatron" + cfg.trainer.placement.colocate_all = False # forward-only test, no inference engine + cfg.trainer.placement.policy_num_gpus_per_node = 4 + cfg.trainer.placement.ref_num_gpus_per_node = 4 + mcfg = cfg.trainer.policy.megatron_config + mcfg.tensor_model_parallel_size = 4 + mcfg.expert_model_parallel_size = 4 + mcfg.expert_tensor_parallel_size = 1 + if mcfg.transformer_config_kwargs is None: + mcfg.transformer_config_kwargs = dict() + mcfg.transformer_config_kwargs["num_layers"] = 2 # truncate the 30B MoE for CI + validate_cfg(cfg) + return cfg + + +async def _run_moe_forward(qat_enabled: bool): + from tests.backends.skyrl_train.gpu.gpu_ci.megatron.test_megatron_worker import ( + get_test_training_batch, + ) + + cfg = _moe_forward_cfg() + cfg.trainer.policy.model.fake_int4_qat.enabled = qat_enabled + batch = get_test_training_batch(4) + actor_group = init_worker_with_type( + "policy", + shared_pg=None, + colocate_all=False, + num_gpus_per_node=cfg.trainer.placement.policy_num_gpus_per_node, + cfg=cfg, + ) + refs = actor_group.async_run_ray_method("mesh", "forward", data=batch) + output = WorkerOutput.cat(actor_group.actor_infos, ray.get(refs)) + return loss_fn_outputs_to_tensor(output.loss_fn_outputs, key="logprobs") + + +@pytest.mark.asyncio +@pytest.mark.megatron +async def test_megatron_forward_qat_on_off(ray_init_fixture): + """QAT must change expert compute end-to-end, and only moderately (STE, not corruption).""" + try: + logprobs_off = await _run_moe_forward(qat_enabled=False) + ray.shutdown() + ray_init_for_tests() + logprobs_on = await _run_moe_forward(qat_enabled=True) + finally: + ray.shutdown() + + assert logprobs_off.shape == logprobs_on.shape + assert not torch.equal(logprobs_off, logprobs_on), ( + "logprobs identical with fake-INT4 QAT on/off: the hook did not reach the expert GEMMs " + "through the worker init path (megatron module structure or install wiring changed?)" + ) + mean_abs = (logprobs_on - logprobs_off).abs().mean().item() + assert mean_abs < 1.0, f"QAT moved logprobs implausibly far (mean |diff| = {mean_abs:.3f})" diff --git a/tests/backends/skyrl_train/test_fake_int4_qat.py b/tests/backends/skyrl_train/test_fake_int4_qat.py new file mode 100644 index 0000000000..df25a932b5 --- /dev/null +++ b/tests/backends/skyrl_train/test_fake_int4_qat.py @@ -0,0 +1,161 @@ +"""CPU unit tests for the fake-INT4 QAT straight-through estimator. + +Run with: +uv run --isolated --extra skyrl-train --extra dev pytest -s tests/backends/skyrl_train/test_fake_int4_qat.py + +The golden tests pin the STE to slices of REAL quantization artifacts (see +``_fake_int4_qat_golden.py``): any change to the grid arithmetic (fp32 vs bf16 +rounding, eps clamps, divisor/clamp-range handling) that diverges from what the +inference engine serves fails these tests bit-exactly, without needing +compressed-tensors, megatron, or a GPU. +""" + +import pytest +import torch + +from skyrl.backends.skyrl_train.workers.megatron.fake_int4_qat import ( + _FakeInt4QuantizeSTE, + fake_int4_quantize_ste, +) +from tests.backends.skyrl_train._fake_int4_qat_golden import KIMI_QAT, QWEN_RTN + +GS = 32 + +_DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + + +def _unpack_int4(packed: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """compressed-tensors unpack_from_int32 semantics (offset encoding: nibble = q + 8).""" + out = torch.zeros(rows, packed.shape[1] * 8, dtype=torch.int32) + for j in range(8): + out[:, j::8] = (packed >> (4 * j)) & 0xF + return (out - 8).to(torch.int8)[:, :cols] + + +def _load_golden(case: dict, device: str): + rows, cols = case["rows"], case["cols"] + masters = torch.tensor(case["masters_u16"], dtype=torch.uint16).view(torch.bfloat16).reshape(rows, cols) + packed = torch.tensor(case["packed_i32"], dtype=torch.int32) + scale = torch.tensor(case["scale_u16"], dtype=torch.uint16).view(torch.bfloat16).reshape(rows, cols // GS) + q_ref = _unpack_int4(packed, rows, cols) + # the served grid, exactly as compressed-tensors dequantize() produces it + served = (q_ref.view(rows, cols // GS, GS).to(torch.bfloat16) * scale.unsqueeze(-1)).reshape(rows, cols) + return masters.to(device), q_ref.to(device), scale.to(device), served.to(device) + + +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize("case", [QWEN_RTN, KIMI_QAT], ids=["qwen_rtn_7p5", "kimi_qat_7p0"]) +def test_golden_bit_exact_vs_served_artifact(case, device): + """STE output must equal the served INT4 grid bit-for-bit on real checkpoint slices.""" + masters, q_ref, scale_ref, served = _load_golden(case, device) + deq = fake_int4_quantize_ste(masters, case["group_size"], case["scale_divisor"], case["q_min"]) + assert deq.dtype == torch.bfloat16 + assert torch.equal(deq, served), "dequantized weights diverge from the served artifact grid" + + # the recomputed scale and codes must match the stored artifact too + rows, cols = case["rows"], case["cols"] + g = masters.view(rows, cols // GS, GS) + amax = g.abs().amax(dim=-1, keepdim=True).to(torch.float32) + scale = (amax / case["scale_divisor"]).to(masters.dtype) + assert torch.equal(scale.squeeze(-1), scale_ref), "recomputed scale != stored weight_scale" + safe = torch.where(scale == 0, torch.ones_like(scale), scale) + q = torch.clamp(torch.round(g / safe), case["q_min"], 7.0) + assert torch.equal(q.to(torch.int8).view(rows, cols), q_ref), "recomputed codes != stored codes" + + +def test_golden_covers_edge_cases(): + """The fixtures must keep covering the regressions they were built to catch.""" + _, q_qwen, _, _ = _load_golden(QWEN_RTN, "cpu") + assert (q_qwen == -8).any(), "RTN slice lost its stored -8 code coverage" + masters, _, _, _ = _load_golden(QWEN_RTN, "cpu") + amax = masters.view(QWEN_RTN["rows"], -1, GS).float().abs().amax(-1) + assert (amax < 1e-30).any(), "RTN slice lost its near-denormal group coverage (eps-clamp regression guard)" + _, q_kimi, _, _ = _load_golden(KIMI_QAT, "cpu") + mx = q_kimi.view(KIMI_QAT["rows"], -1, GS).abs().amax(-1) + assert mx.eq(7).all(), "Kimi slice must have max|q| == 7 in every group" + + +@pytest.mark.parametrize("device", _DEVICES) +def test_kimi_identity_on_dequant_masters(device): + """div=7 property that makes the Kimi flow exact: fake-quant of a dequant is the identity.""" + masters, _, _, served = _load_golden(KIMI_QAT, device) + assert torch.equal(masters, served) # masters ARE the dequant for this fixture + deq = fake_int4_quantize_ste(masters, GS, 7.0, -7.0) + assert torch.equal(deq, masters) + # and it is idempotent in general for freshly quantized tensors + w = torch.randn(16, 128, dtype=torch.bfloat16, device=device) + once = fake_int4_quantize_ste(w, GS, 7.0, -7.0) + twice = fake_int4_quantize_ste(once, GS, 7.0, -7.0) + assert torch.equal(once, twice) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_ste_backward_is_identity(device): + w = torch.randn(8, 64, dtype=torch.bfloat16, device=device, requires_grad=True) + out = fake_int4_quantize_ste(w, GS, 7.5, -8.0) + grad = torch.randn_like(out) + out.backward(grad) + assert torch.equal(w.grad, grad) + + +def test_main_grad_attribute_propagates(): + w = torch.randn(8, 64, dtype=torch.bfloat16) + w.main_grad = torch.zeros(8, 64) + out = fake_int4_quantize_ste(w, GS, 7.5, -8.0) + assert out.main_grad is w.main_grad + # and absent when the master has none + w2 = torch.randn(8, 64, dtype=torch.bfloat16) + assert not hasattr(fake_int4_quantize_ste(w2, GS, 7.5, -8.0), "main_grad") + + +@pytest.mark.parametrize("device", _DEVICES) +def test_padding_path_matches_manual_padding(device): + """in_features % group_size != 0: pad with zeros, quantize, slice back.""" + w = torch.randn(8, 40, dtype=torch.bfloat16, device=device) + out = fake_int4_quantize_ste(w, GS, 7.5, -8.0) + assert out.shape == w.shape and torch.isfinite(out).all() + padded = torch.zeros(8, 64, dtype=torch.bfloat16, device=device) + padded[:, :40] = w + ref = fake_int4_quantize_ste(padded, GS, 7.5, -8.0)[:, :40] + assert torch.equal(out, ref) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_all_zero_group_quantizes_to_zero(device): + w = torch.zeros(4, 64, dtype=torch.bfloat16, device=device) + w[:, GS:] = torch.randn(4, GS, dtype=torch.bfloat16, device=device) + out = fake_int4_quantize_ste(w, GS, 7.5, -8.0) + assert torch.isfinite(out).all(), "all-zero group must not produce NaN (0/0 guard)" + assert out[:, :GS].abs().sum() == 0 + + +@pytest.mark.parametrize(("div", "qmin"), [(7.5, -8.0), (7.0, -7.0)]) +def test_codes_stay_in_convention_range(div, qmin): + w = torch.randn(32, 256, dtype=torch.bfloat16) + g = w.view(32, -1, GS) + scale = (g.abs().amax(-1, keepdim=True).to(torch.float32) / div).to(w.dtype) + safe = torch.where(scale == 0, torch.ones_like(scale), scale) + codes = torch.round(g / safe).clamp(qmin, 7.0) + deq = fake_int4_quantize_ste(w, GS, div, qmin) + assert torch.equal(deq, (codes * scale).reshape(32, 256)) + assert codes.min() >= qmin and codes.max() <= 7.0 + if div == 7.0: + # /7 pins the max-|w| element of each group to +-7 and never emits -8 + assert codes.abs().amax(-1).eq(7).all() + + +def test_non_contiguous_input_matches_contiguous(): + base = torch.randn(16, 256, dtype=torch.bfloat16) + nc = base[:, ::2] + assert not nc.is_contiguous() + out = fake_int4_quantize_ste(nc, GS, 7.5, -8.0) + ref = fake_int4_quantize_ste(nc.contiguous(), GS, 7.5, -8.0) + assert torch.equal(out, ref) + + +def test_autograd_function_arity(): + """backward must return one gradient per forward input (x, group_size, scale_div, q_min).""" + w = torch.randn(4, 32, dtype=torch.bfloat16, requires_grad=True) + out = _FakeInt4QuantizeSTE.apply(w, GS, 7.5, -8.0) + out.sum().backward() + assert w.grad is not None diff --git a/tests/train/test_config.py b/tests/train/test_config.py index 2c6853c79b..37de30d458 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -337,6 +337,35 @@ def test_cross_field_defaults(): ) # same as `generator.sampling_params.max_generate_length` +def test_fake_int4_qat_defaults(): + """Defaults must stay pinned to the llm-compressor RTN convention, disabled.""" + cfg = SkyRLTrainConfig.from_cli_overrides([]) + fq = cfg.trainer.policy.model.fake_int4_qat + assert fq.enabled is False + assert fq.group_size == 32 + assert fq.scale_divisor == 7.5 + assert fq.q_min == -8.0 + assert fq.bf16_base_path is None + + +def test_fake_int4_qat_cli_overrides(): + """All convention knobs must be settable from the CLI (the Kimi K2.x convention).""" + cfg = SkyRLTrainConfig.from_cli_overrides( + [ + "trainer.policy.model.fake_int4_qat.enabled=true", + "trainer.policy.model.fake_int4_qat.group_size=32", + "trainer.policy.model.fake_int4_qat.scale_divisor=7.0", + "trainer.policy.model.fake_int4_qat.q_min=-7", + "trainer.policy.model.fake_int4_qat.bf16_base_path=/data/bf16-dump", + ] + ) + fq = cfg.trainer.policy.model.fake_int4_qat + assert fq.enabled is True + assert fq.scale_divisor == 7.0 + assert fq.q_min == -7.0 + assert fq.bf16_base_path == "/data/bf16-dump" + + class TestTrainerUseSamplePackingAlias: """`trainer.use_sample_packing` is a deprecated alias for `trainer.remove_microbatch_padding` on the RL entrypoint config (mirrors the ``fsdp2``->``fsdp`` alias)."""