Add LSQ (Learned Scale Quantization) support and recipes

CHANGELOG.rst

@@ -19,6 +19,7 @@ Changelog
 
 **New Features**
 
+- Add Learned Scale Quantization (LSQ) and Dual-LSQ support for quantization-aware distillation, including learnable ``amax`` parameters, tied-scale and pre-scale options, focused NVFP4 recipes, and scale-only training.
 - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 <https://arxiv.org/abs/2605.18810>`_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change).
 - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred.
 - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``.

examples/llm_qat/ARGUMENTS.md

@@ -64,7 +64,7 @@ Extends [HuggingFace TrainingArguments](https://huggingface.co/docs/transformers
 |----------|------|---------|-------------|
 | `--trainable_params` | `list[str]` | `None` | Glob patterns (fnmatch) for parameters that should be trainable. All other parameters will be frozen. Mutually exclusive with frozen_params. |
 | `--frozen_params` | `list[str]` | `None` | Glob patterns (fnmatch) for parameters that should be frozen. Mutually exclusive with trainable_params. |
-| `--lr_config` | `str` | `None` | Path to a YAML file mapping fnmatch patterns to optimizer kwargs (e.g. lr, weight_decay). First matching pattern wins per parameter. See examples/llm_qat/configs/train/lr_config_example.yaml. |
+| `--lr_config` | `str` | `None` | Path to a YAML file mapping fnmatch patterns to optimizer kwargs (e.g. lr, weight_decay). First matching pattern wins per parameter. See examples/llm_qat/configs/train/lr/lr_config_example.yaml. |
 | `--manual_gc` | `bool` | `False` | Run `gc.collect()` before each training/prediction step to work around GPU memory leaks during QAT/distillation. |
 | `--liger_ce_label_smoothing` | `float` | `0.0` | Label smoothing for Liger fused CE loss. Only used when --use_liger_kernel is enabled. |
 | `--lora` | `bool` | `False` | Whether to add LoRA (Low-Rank Adaptation) adapter before training. When using real quantization, the LoRA adapter must be set, as quantized weights will be frozen during training. |

examples/llm_qat/configs/train/lr/lr_config_amax.yaml

@@ -0,0 +1,5 @@
+# Override the learning rate for LSQ's learnable amax parameters to 1e-4.
+"*weight_quantizer._amax_pre":
+  lr: 1e-4
+"*weight_quantizer._amax_post":
+  lr: 1e-4

examples/llm_qat/configs/train/lr_config_example.yaml → examples/llm_qat/configs/train/lr/lr_config_example.yaml

@@ -12,7 +12,7 @@
 #   eps          - term added to denominator for numerical stability
 #
 # Usage:
-#   --lr_config configs/train/lr_config_example.yaml
+#   --lr_config configs/train/lr/lr_config_example.yaml
 #
 # Tip: use `model.named_parameters()` to find the exact parameter names
 # for your model.

examples/llm_qat/configs/train/qad_scale_only.yaml

@@ -0,0 +1,51 @@
+# Scale-only QAD for LSQ-quantized checkpoints
+
+# Model
+model_name_or_path:  # e.g., qwen3-8b-lsq-quantized
+output_dir:  # e.g., qwen3-8b-lsq-scale-qad
+attn_implementation: flash_attention_2
+
+# Distillation
+distill: true
+teacher_model:  # e.g., Qwen/Qwen3-8B
+
+# Dataset
+dataset_config: configs/dataset/blend.yaml
+train_samples: 20000
+eval_samples: 2000
+
+# Train only LSQ amax scale parameters. Tied LSQ exposes only _amax_post.
+trainable_params:
+  - "*weight_quantizer._amax_pre"
+  - "*weight_quantizer._amax_post"
+
+# Hyperparameters
+num_train_epochs: 1.0
+# LSQ amax parameter requires higher learning rate than quantized weights
+learning_rate: 1e-4
+weight_decay: 0.0
+per_device_train_batch_size: 2
+per_device_eval_batch_size: 2
+gradient_accumulation_steps: 2
+model_max_length: 8192
+warmup_ratio: 0.05
+lr_scheduler_type: cosine
+use_liger_kernel: true
+manual_gc: true
+seed: 42
+do_train: true
+do_eval: true
+
+# Checkpointing
+load_best_model_at_end: true
+save_total_limit: 2
+
+# Evaluation
+eval_on_start: true
+eval_strategy: steps
+eval_steps: 50
+
+# Logging
+logging_steps: 1
+report_to:
+  - tensorboard

examples/llm_qat/configs/train/qad_with_learnt_amax.yaml

@@ -0,0 +1,46 @@
+# Full-parameter QAD for LSQ-quantized checkpoints
+
+# Model
+model_name_or_path:  # e.g., qwen3-8b-lsq-quantized
+output_dir:  # e.g., qwen3-8b-lsq-full-qad
+attn_implementation: flash_attention_2
+
+# Distillation
+distill: true
+teacher_model:  # e.g., Qwen/Qwen3-8B
+
+# Dataset
+dataset_config: configs/dataset/blend.yaml
+train_samples: 20000
+eval_samples: 2000
+
+# Hyperparameters
+num_train_epochs: 1.0
+learning_rate: 1e-5
+# Learnable LSQ amax may need a higher learning rate than quantized weights.
+lr_config: configs/train/lr/lr_config_amax.yaml
+per_device_train_batch_size: 2
+per_device_eval_batch_size: 2
+gradient_accumulation_steps: 2
+model_max_length: 8192
+warmup_ratio: 0.05
+lr_scheduler_type: cosine
+use_liger_kernel: true
+manual_gc: true
+seed: 42
+do_train: true
+do_eval: true
+
+# Checkpointing
+load_best_model_at_end: true
+save_total_limit: 2
+
+# Evaluation
+eval_on_start: true
+eval_strategy: steps
+eval_steps: 50
+
+# Logging
+logging_steps: 1
+report_to:
+  - tensorboard

modelopt/torch/kernels/quantization/gemm/fp4_kernel.py

@@ -28,7 +28,12 @@
 
 from ..common.nvfp4_quant import nvfp4_scalar_quant
 
-__all__ = ["compute_fp4_scales", "fp4_dequantize", "static_blockwise_fp4_fake_quant"]
+__all__ = [
+    "compute_fp4_scales",
+    "fp4_dequantize",
+    "static_blockwise_fp4_cast",
+    "static_blockwise_fp4_fake_quant",
+]
 
 
 _TORCH_TO_TL_DTYPE = {
@@ -309,3 +314,87 @@ def static_blockwise_fp4_fake_quant(
         )
 
     return y_flat.view(original_shape)
+
+
+@triton.jit
+def static_blockwise_fp4_cast_kernel(
+    x_ptr,  # [NUM_ELEMENTS] flattened pre-scaled input
+    y_ptr,  # [NUM_ELEMENTS] flattened output
+    NUM_ELEMENTS,
+    TILE_SIZE: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,
+):
+    """Round pre-scaled values to nearest FP4 representable value (no scale)."""
+    pid = tl.program_id(axis=0)
+    offset = pid * TILE_SIZE + tl.arange(0, TILE_SIZE)
+    mask = offset < NUM_ELEMENTS
+
+    x = tl.load(x_ptr + offset, mask=mask).to(tl.float32)
+    x_abs = tl.abs(x)
+
+    # FP4 E2M1 representable values: 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0
+    q_val = tl.where(
+        x_abs <= 0.25,
+        0.0,
+        tl.where(
+            x_abs < 0.75,
+            0.5,
+            tl.where(
+                x_abs <= 1.25,
+                1.0,
+                tl.where(
+                    x_abs < 1.75,
+                    1.5,
+                    tl.where(
+                        x_abs <= 2.5,
+                        2.0,
+                        tl.where(
+                            x_abs < 3.5,
+                            3.0,
+                            tl.where(x_abs <= 5.0, 4.0, 6.0),
+                        ),
+                    ),
+                ),
+            ),
+        ),
+    )
+
+    y = tl.where(x >= 0, q_val, -q_val)
+    tl.store(y_ptr + offset, y.to(OUT_DTYPE), mask=mask)
+
+
+def static_blockwise_fp4_cast(
+    x: torch.Tensor,
+    out_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Round pre-scaled values to nearest FP4 E2M1 representable value.
+
+    Unlike ``static_blockwise_fp4_fake_quant``, this does **not** apply any
+    scale -- the caller is responsible for pre-dividing by scale_pre and
+    post-multiplying by scale_post (as in LSQ).
+
+    Args:
+        x: Input tensor (any shape) on CUDA.
+        out_dtype: Output dtype. Defaults to x.dtype.
+    """
+    if out_dtype is None:
+        out_dtype = x.dtype
+
+    x_flat = x.contiguous().view(-1)
+    y_flat = torch.empty_like(x_flat, dtype=out_dtype)
+    NUM_ELEMENTS = x_flat.numel()
+    TILE_SIZE = 1024
+
+    tl_out_dtype = _torch_dtype_to_tl(out_dtype)
+    grid = ((NUM_ELEMENTS + TILE_SIZE - 1) // TILE_SIZE,)
+
+    with torch.cuda.device(x.device):
+        static_blockwise_fp4_cast_kernel[grid](
+            x_flat,
+            y_flat,
+            NUM_ELEMENTS,
+            TILE_SIZE=TILE_SIZE,
+            OUT_DTYPE=tl_out_dtype,
+        )
+
+    return y_flat.view_as(x)

modelopt/torch/opt/plugins/transformers.py

@@ -231,7 +231,7 @@ class ModelOptTrainerArguments(ModelOptHFArguments):
             "help": (
                 "Path to a YAML file mapping fnmatch patterns to optimizer kwargs "
                 "(e.g. lr, weight_decay). First matching pattern wins per parameter. "
-                "See examples/llm_qat/configs/train/lr_config_example.yaml."
+                "See examples/llm_qat/configs/train/lr/lr_config_example.yaml."
             ),
         },
     )

modelopt/torch/quantization/config.py

@@ -153,9 +153,16 @@
 import re
 import warnings
 from collections.abc import Mapping, Sequence
-from typing import Any, Literal
-
-from pydantic import AliasChoices, Field, ValidationInfo, field_validator, model_validator
+from typing import Any, ClassVar, Literal, TypeAlias
+
+from pydantic import (
+    AliasChoices,
+    Field,
+    ValidationInfo,
+    field_serializer,
+    field_validator,
+    model_validator,
+)
 
 from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField
 from modelopt.torch.opt.config_loader import load_config
@@ -1204,6 +1211,104 @@ def _gptq_qdq_default(self):
         return self
 
 
+_ScaleCalibConfig: TypeAlias = MaxCalibConfig | MseCalibConfig | LocalHessianCalibConfig
+
+
+class LSQConfig(QuantizeAlgorithmConfig):
+    """Config for LSQ (Learnt Scale Quantization) and Dual-LSQ algorithms.
+
+    In LSQ, the scale used for quantization is learnt. ModelOpt's LSQ is similar to the
+    original `Learned Step Size Quantization paper <https://arxiv.org/pdf/1902.08153>`_.
+    Its forward pass is ``w_q = Q_STE(w / s) * s``, where ``s`` is learnt.
+
+    Dual-LSQ learns separate pre-quantization and post-quantization scales. Its forward
+    pass is ``w_q = Q_STE(w / s_pre) * s_post``, where ``s_pre`` and ``s_post`` are
+    learnt. Dual-LSQ generally performs better than LSQ for learning NVFP4 per-block
+    weight scales.
+
+    Currently, only NVFP4 per-block weight-scale learning is supported. Both LSQ and
+    Dual-LSQ use a reparameterization that learns ``amax`` instead of scale directly,
+    where ``scale = amax / max_bound``.
+
+    ``learnable_amax`` controls which amax parameters are learnable vs frozen:
+        - ``["pre", "post"]``: both learnable
+        - ``"post"`` or ``["post"]``: only post learnable, pre frozen
+        - ``"pre"`` or ``["pre"]``: only pre learnable, post frozen
+        - ``[]``: both frozen (static scales)
+
+    ``tied_amax`` makes pre and post share a single tensor (requires both to
+    have the same learnable state, i.e. ``learnable_amax`` must be
+    ``["pre", "post"]`` or ``[]``).
+
+    ``quantize_pre_scale=False`` leaves the pre-quantization scale unquantized
+    while preserving the existing post-scale quantization behavior.
+    """
+
+    ScaleCalibConfig: ClassVar[Any] = _ScaleCalibConfig
+
+    method: Literal["lsq"] = ModeloptField("lsq")
+
+    learnable_amax: list[Literal["pre", "post"]] | Literal["pre", "post"] = ModeloptField(
+        default=["post"],
+        title="Which amax parameters are learnable.",
+        description=(
+            "Which amax params are learnable. "
+            "'pre', 'post', ['pre', 'post'], or []. "
+            "Defaults to ['post'] (post-only learnable)."
+        ),
+    )
+
+    tied_amax: bool = ModeloptField(
+        default=False,
+        title="Tie pre and post amax into a single tensor.",
+        description=(
+            "If True, pre and post share one underlying tensor. "
+            "Requires both to have the same learnable state."
+        ),
+    )
+
+    quantize_pre_scale: bool = ModeloptField(
+        default=True,
+        title="FP8-quantize the LSQ pre-quantization scale.",
+        description=(
+            "If False, LSQ uses the raw pre-quantization scale while keeping post-scale "
+            "quantization controlled by the quantizer's block-scale settings."
+        ),
+    )
+
+    scale_algorithm: _ScaleCalibConfig | None = ModeloptField(
+        default=None,
+        title="Scale calibration algorithm to run first.",
+        description=(
+            "Dict with 'method' key: 'mse', 'local_hessian', or 'max'. "
+            "Optional keys include 'fp8_scale_sweep' for FP4 formats. "
+            "Defaults to {'method': 'mse'} if None."
+        ),
+    )
+
+    @field_serializer("scale_algorithm")
+    def _serialize_scale_algorithm(self, value: _ScaleCalibConfig | None):
+        """Preserve the sparse public dict shape accepted by this field."""
+        if value is None:
+            return None
+        return {"method": value.method, **value.model_dump(exclude={"method"}, exclude_unset=True)}
+
+    @model_validator(mode="after")
+    def _validate_tied_amax(self):
+        """Validate tied_amax is compatible with learnable_amax."""
+        learn = self.learnable_amax
+        if isinstance(learn, str):
+            learn = [learn]
+        learn_set = set(learn)
+        if self.tied_amax:
+            if learn_set not in (set(), {"pre", "post"}):
+                raise ValueError(
+                    f"tied_amax=True requires learnable_amax to be [] or ['pre', 'post'], "
+                    f"got {self.learnable_amax}"
+                )
+        return self
+
+
 QuantizeQuantCfgType = list[QuantizerCfgEntry]
 QuantizerCfgListConfig = QuantizeQuantCfgType