Add NVFP4 quantization-aware distillation (QAD)

Summary:
QAD recovers the accuracy lost when a model is quantized to NVFP4 by
distilling from the original bf16 model into an NVFP4 fake-quantized copy. The
student's MoE experts and dense linear layers are fake-quantized with a
straight-through estimator during training; a frozen bf16 teacher supervises
via KL divergence (plus an optional hard-label cross-entropy term).

Adds:
- NVFP4 fake-quant primitives: torchtitan/components/quantization/
  nvfp4_fake_quant.py -- self-contained two-level block-scale E2M1 fake-quant
  (STE), plus a single entry point apply_nvfp4_fake_quant(model) that quantizes
  both the MoE experts and the dense linears (excluding lm_head/router/gate).
- KD/QAD trainer + configs: torchtitan/experiments/kd/ -- KDTrainer (frozen
  bf16 teacher + fake-quant student), kd_loss (KL + CE), and qwen3 QAD configs
  (debug + 30B-A3B).
- MoE fake-quant hook: GroupedExperts.fake_quant_fn in models/common/moe.py,
  applied to expert weights and activations inside _experts_forward.
- Post-init hook: post_model_init_fn in trainer.py, used to apply the
  fake-quant to the student after model init/parallelization.
- Single-file HF checkpoint support: protocols/state_dict_adapter.py.

The NVFP4 fake-quant primitives in nvfp4_fake_quant.py are intentionally
self-contained (no torchao dependency), but could be upstreamed to torchao if a
shared home is preferred.

Test Plan:
Distill a qwen3-30B-A3B checkpoint (student + bf16 teacher both init from the
same HF checkpoint):

  MODULE=kd.qwen3 CONFIG=qwen3_30b_a3b_qad NGPU=8 ./run_train.sh \
      --training.steps 200 \
      --checkpoint.interval 50 \
      --checkpoint.initial_load_path /path/to/qwen3-30b-a3b-hf \
      --hf_assets_path /path/to/qwen3-30b-a3b-hf

Quick self-contained smoke test (tiny debug model, no external checkpoint):

  MODULE=kd.qwen3 CONFIG=qwen3_moe_debug_qad NGPU=2 ./run_train.sh

Results (MATH accuracy, bf16 -> nvfp4):

  qwen3-30B-A3B
    baseline (200-step RL ckpt):  88.6% -> 86.6%   (-2.0)
    + QAD (800 steps):            nvfp4 88.6%       (+2.0)
    -> 100% of the gap recovered

  gpt-oss-20B
    baseline (100-step RL ckpt):  62.6% -> 50.8%   (-11.8)
    + QAD (1400 steps):           nvfp4 55.2%       (+4.4)
    -> 37% of the gap recovered

QAD closes the full NVFP4 gap on qwen3-30B-A3B and recovers a third of the
(much larger) gap on gpt-oss-20B.

Author: andrewor14

torchtitan/components/quantization/nvfp4_fake_quant.py

@@ -0,0 +1,170 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""NVFP4 fake-quantization primitives.
+
+Note: these primitives are self-contained, can be moved to torchao if needed.
+
+A minimal, self-contained implementation of NVFP4 fake quantization applied to a
+model's MoE experts and dense linear layers during training (e.g. for
+quantization-aware distillation or QAT), using a straight-through estimator.
+
+NVFP4 numerics (two-level block-scale quantization):
+
+* **Global scale (per tensor):** ``(448 * 6) / amax(x)`` -- 448 is the max
+  FP8-E4M3 value, 6 is the max FP4-E2M1 value.
+* **Block scale (per 16 elements):** each block gets its own FP8-E4M3 scale
+  ``block_amax / 6``, so precision adapts to local magnitude.
+
+Every value is rounded to the nearest E2M1 grid point
+(0, 0.5, 1, 1.5, 2, 3, 4, 6) then dequantized. Training uses a straight-through
+estimator (STE): forward returns the dequantized value, backward is identity.
+
+The single public entry point is :func:`apply_nvfp4_fake_quant`, which
+fake-quantizes BOTH the MoE experts and the dense linear layers (weight + input
+activation) -- the full set of layers an NVFP4 eval quantizes (everything
+except ``lm_head`` / ``router`` / ``gate``).
+"""
+
+import logging
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+
+# Representable positive E2M1 values and the round-to-nearest midpoint
+# boundaries between consecutive values.
+_E2M1_VALUES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
+_E2M1_BOUNDARIES = [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0]
+_SF_BLOCK_SIZE = 16
+
+# Module-name substrings the NVFP4 eval does NOT quantize, so QAD skips them too
+# to match the eval scope exactly: the LM head and the MoE router gate.
+_LINEAR_FQ_EXCLUDE = ("lm_head", "router", "gate")
+
+
+def _nvfp4_quant_dequant(x: torch.Tensor) -> torch.Tensor:
+    """NVFP4 E2M1 block-scale quantize-dequantize roundtrip.
+
+    For 3D tensors (expert weights) each expert slice is quantized independently
+    to avoid OOM from materializing the full tensor in float32. Returns the
+    dequantized tensor in the same dtype as *x*.
+    """
+    if x.ndim == 3:
+        out = torch.empty_like(x)
+        for i in range(x.shape[0]):
+            out[i] = _nvfp4_quant_dequant(x[i])
+        return out
+
+    device = x.device
+    orig_dtype = x.dtype
+
+    # Level 1: global scale factor maps tensor range into FP8*FP4 range
+    gsf = (448 * 6) / x.float().abs().nan_to_num().max()
+    x = x.float() * gsf
+
+    # Level 2: per-block FP8 scale factor adapts to local magnitudes
+    orig_shape = x.shape
+    x_flat = x.reshape(-1, _SF_BLOCK_SIZE)
+    block_amax = x_flat.abs().amax(dim=-1, keepdim=True)
+    block_sf = (block_amax / 6.0).to(torch.float8_e4m3fn).float()
+
+    # Round to nearest E2M1 value
+    x_norm = x_flat / block_sf.clamp(min=torch.finfo(torch.float8_e4m3fn).tiny)
+    sign = x_norm.sign()
+    x_abs = x_norm.abs()
+    boundaries = torch.tensor(_E2M1_BOUNDARIES, device=device)
+    values = torch.tensor(_E2M1_VALUES, device=device)
+    indices = torch.bucketize(x_abs, boundaries)
+    x_quant = sign * values[indices]
+
+    # Dequantize: reverse block scale, reshape, reverse global scale
+    x_dq = (x_quant * block_sf).reshape(orig_shape) / gsf
+    return x_dq.to(orig_dtype)
+
+
+def _nvfp4_fake_quantize_ste(x: torch.Tensor) -> torch.Tensor:
+    """NVFP4 block-scale fake quantize with STE (straight-through estimator).
+
+    Forward returns the dequantized value; backward treats quantization as
+    identity. Used for the MoE ``fake_quant_fn`` hook and for linear weights.
+    """
+    dq = _nvfp4_quant_dequant(x.detach())
+    return x + (dq - x).detach()
+
+
+def _nvfp4_fake_quantize_act_ste(x: torch.Tensor) -> torch.Tensor:
+    """STE NVFP4 fake-quant for a linear layer's INPUT activation of any rank.
+
+    ``_nvfp4_quant_dequant`` treats a 3D tensor as independent per-expert slices
+    (looping dim 0), which is wrong for an activation shaped
+    ``(batch, seq, features)``. Flatten to 2D ``(tokens, features)`` first so the
+    per-tensor global scale spans all tokens and the FP8 block scales tile along
+    the feature dim -- matching how the eval quantizes a linear's input -- then
+    restore the original shape. STE preserved through the reshape.
+    """
+    if x.ndim <= 2:
+        return _nvfp4_fake_quantize_ste(x)
+    orig_shape = x.shape
+    flat = x.reshape(-1, orig_shape[-1])
+    return _nvfp4_fake_quantize_ste(flat).reshape(orig_shape)
+
+
+def _wrap_linear_nvfp4_fake_quant(module: nn.Linear) -> None:
+    """Override a linear's forward so its WEIGHT and INPUT ACTIVATION are both
+    NVFP4 fake-quantized (STE) before the matmul, matching the NVFP4 eval.
+    Instance-level forward override (eager only)."""
+
+    def fq_forward(x: torch.Tensor) -> torch.Tensor:
+        w = _nvfp4_fake_quantize_ste(module.weight)
+        x = _nvfp4_fake_quantize_act_ste(x)
+        return F.linear(x, w, module.bias)
+
+    module.forward = fq_forward
+    module._nvfp4_fake_quant_linear = True
+
+
+def apply_nvfp4_fake_quant(model: nn.Module) -> nn.Module:
+    """Enable NVFP4 fake quantization on BOTH the MoE experts AND the dense
+    linear layers -- the full set of layers quantized by the NVFP4 eval.
+
+    * **MoE experts:** sets ``fake_quant_fn`` on every ``GroupedExperts`` module
+      (identified by having both ``num_experts`` and ``fake_quant_fn``), which
+      fake-quantizes the expert weights and activations inside
+      ``_experts_forward`` (FSDP2/DTensor-compatible).
+    * **Dense linears:** overrides the forward of every ``nn.Linear`` whose name
+      does not match :data:`_LINEAR_FQ_EXCLUDE` (``lm_head``/``router``/``gate``)
+      so its weight and input activation are fake-quantized.
+
+    Applied in-place; the same model is returned. Eager-only (linear forwards are
+    overridden at the instance level).
+    """
+    # Apply fake quant to the MoE experts (via the GroupedExperts hook).
+    n_moe = 0
+    for module in model.modules():
+        if hasattr(module, "num_experts") and hasattr(module, "fake_quant_fn"):
+            module.fake_quant_fn = _nvfp4_fake_quantize_ste
+            n_moe += 1
+
+    # Apply fake quant to the dense linear layers (excl lm_head/router/gate).
+    n_lin = 0
+    for name, module in model.named_modules():
+        if isinstance(module, nn.Linear) and not any(
+            excl in name for excl in _LINEAR_FQ_EXCLUDE
+        ):
+            _wrap_linear_nvfp4_fake_quant(module)
+            n_lin += 1
+
+    logger.info(
+        "Applied NVFP4 fake quant to %d GroupedExperts modules and %d linear "
+        "layers (excluded lm_head/router/gate)",
+        n_moe,
+        n_lin,
+    )
+    return model

torchtitan/experiments/__init__.py

@@ -13,6 +13,7 @@
         "autoparallel.llama3",
         "autoparallel.local_map_deepseek_v3",
         "torchft.llama3",
+        "kd.qwen3",
         "rl",
         # RL examples own a per-example config_registry under rl/examples/<name>;
         # listed here so `--module <name>` resolves (see ConfigManager).

torchtitan/experiments/kd/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.

torchtitan/experiments/kd/loss.py

@@ -0,0 +1,56 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Knowledge distillation loss combining KL divergence with cross-entropy."""
+
+import torch
+import torch.nn.functional as F
+
+from torchtitan.components.loss import cross_entropy_loss, IGNORE_INDEX
+
+
+def kd_loss(
+    student_logits: torch.Tensor,
+    teacher_logits: torch.Tensor,
+    labels: torch.Tensor,
+    temperature: float,
+    alpha: float,
+) -> torch.Tensor:
+    """Compute knowledge distillation loss.
+
+    Args:
+        student_logits: [B, L, V] student model output logits.
+        teacher_logits: [B, L, V] teacher model output logits (detached).
+        labels: [B, L] target token ids.
+        temperature: Softmax temperature for soft targets.
+        alpha: Weight for distillation loss vs hard-label loss.
+            1.0 = pure distillation, 0.0 = pure cross-entropy.
+
+    Returns:
+        Scalar loss (sum reduction, to be divided by global_valid_tokens).
+    """
+    # Mask for valid (non-padding) tokens
+    valid_mask = labels != IGNORE_INDEX  # [B, L]
+    flat_mask = valid_mask.flatten()  # [B*L]
+
+    # Hard-label cross-entropy loss (sum reduction)
+    ce = cross_entropy_loss(student_logits, labels)
+
+    # Soft-target KL divergence loss
+    flat_student = student_logits.flatten(0, 1).float()  # [B*L, V]
+    flat_teacher = teacher_logits.flatten(0, 1).float()  # [B*L, V]
+
+    # Only compute KL on valid tokens
+    flat_student = flat_student[flat_mask]
+    flat_teacher = flat_teacher[flat_mask]
+
+    student_log_probs = F.log_softmax(flat_student / temperature, dim=-1)
+    teacher_probs = F.softmax(flat_teacher / temperature, dim=-1)
+
+    # KL(teacher || student) with sum reduction, scaled by T^2
+    kl = F.kl_div(student_log_probs, teacher_probs, reduction="sum") * (temperature**2)
+
+    return alpha * kl + (1 - alpha) * ce

torchtitan/experiments/kd/qwen3/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.

torchtitan/experiments/kd/qwen3/config_registry.py

@@ -0,0 +1,120 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Config registry for knowledge distillation (KD) / quantization-aware
+distillation (QAD) experiments with Qwen3 models.
+
+QAD distills from a bf16 teacher into an NVFP4 fake-quantized student. The
+student's MoE experts and dense linears are fake-quantized via
+:func:`~torchtitan.components.quantization.nvfp4_fake_quant.apply_nvfp4_fake_quant`
+(wired as the ``post_model_init_fn``); the teacher stays bf16.
+"""
+
+from torchtitan.components.checkpoint import CheckpointManager
+from torchtitan.components.loss import CrossEntropyLoss
+from torchtitan.components.lr_scheduler import LRSchedulersContainer
+from torchtitan.components.metrics import MetricsProcessor
+from torchtitan.components.optimizer import default_adamw
+from torchtitan.components.quantization.nvfp4_fake_quant import (
+    apply_nvfp4_fake_quant,
+)
+from torchtitan.config import (
+    ParallelismConfig,
+    TrainingConfig,
+)
+from torchtitan.distributed.activation_checkpoint import SelectiveAC
+from torchtitan.experiments.kd.trainer import KDTrainer
+from torchtitan.hf_datasets.text_datasets import HuggingFaceTextDataLoader
+from torchtitan.models.qwen3 import model_registry
+
+
+def qwen3_moe_debug() -> KDTrainer.Config:
+    """Debug KD config with a tiny Qwen3 MoE model (no fake quant)."""
+    return KDTrainer.Config(
+        loss=CrossEntropyLoss.Config(),
+        hf_assets_path="./tests/assets/tokenizer",
+        model_spec=model_registry("debugmodel_moe"),
+        optimizer=default_adamw(lr=3e-4),
+        lr_scheduler=LRSchedulersContainer.Config(warmup_steps=2),
+        training=TrainingConfig(
+            local_batch_size=4,
+            seq_len=4096,
+            steps=10,
+        ),
+        dataloader=HuggingFaceTextDataLoader.Config(
+            dataset="c4_test",
+        ),
+        metrics=MetricsProcessor.Config(log_freq=1),
+        checkpoint=CheckpointManager.Config(
+            interval=10,
+            last_save_model_only=False,
+        ),
+        activation_checkpoint=SelectiveAC.Config(),
+        parallelism=ParallelismConfig(
+            expert_parallel_degree=1,
+        ),
+        # KD-specific settings
+        temperature=2.0,
+        alpha=0.5,
+    )
+
+
+def qwen3_moe_debug_qad() -> KDTrainer.Config:
+    """Debug QAD (quantization-aware distillation) config with a tiny Qwen3 MoE
+    model.
+
+    Same as :func:`qwen3_moe_debug` but with NVFP4 fake quantization applied to
+    the student's MoE experts AND dense linears (w4a4). The teacher stays bf16.
+    Self-contained (tiny model, ``c4_test``, no external checkpoint) -- the
+    quickest way to exercise the full QAD path end to end.
+    """
+    config = qwen3_moe_debug()
+    config.post_model_init_fn = apply_nvfp4_fake_quant
+    return config
+
+
+def qwen3_30b_a3b_qad() -> KDTrainer.Config:
+    """QAD config for Qwen3-30B-A3B.
+
+    Distills from a bf16 teacher into an NVFP4 fake-quantized student. Both the
+    student and the (frozen) teacher are initialized from the same HF checkpoint;
+    the student's MoE experts and dense linears are NVFP4 fake-quantized.
+    """
+    return KDTrainer.Config(
+        loss=CrossEntropyLoss.Config(),
+        hf_assets_path="./assets/hf/Qwen3-30B-A3B",
+        model_spec=model_registry("30B-A3B", attn_backend="flex"),
+        optimizer=default_adamw(lr=5e-6, betas=(0.9, 0.999), weight_decay=0.0),
+        lr_scheduler=LRSchedulersContainer.Config(
+            warmup_steps=100,
+            decay_ratio=0.1,
+            decay_type="cosine",
+        ),
+        training=TrainingConfig(
+            local_batch_size=2,
+            seq_len=2048,
+            steps=100,
+        ),
+        dataloader=HuggingFaceTextDataLoader.Config(
+            dataset="c4",
+        ),
+        metrics=MetricsProcessor.Config(log_freq=10),
+        checkpoint=CheckpointManager.Config(
+            enable=True,
+            initial_load_in_hf=True,
+            last_save_in_hf=True,
+        ),
+        activation_checkpoint=SelectiveAC.Config(),
+        parallelism=ParallelismConfig(
+            data_parallel_shard_degree=-1,
+            expert_parallel_degree=1,
+        ),
+        # KD-specific settings
+        temperature=2.0,
+        alpha=0.5,
+        # QAD: fake-quantize the student (experts + linears); teacher stays bf16
+        post_model_init_fn=apply_nvfp4_fake_quant,
+    )