diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 03bdb027296..d40ba8de1ce 100755 --- a/CHANGELOG.rst +++ b/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 `_) 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/``. diff --git a/examples/llm_qat/ARGUMENTS.md b/examples/llm_qat/ARGUMENTS.md index 0a3e2b7a12d..35f788318c9 100644 --- a/examples/llm_qat/ARGUMENTS.md +++ b/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. | diff --git a/examples/llm_qat/configs/train/lr/lr_config_amax.yaml b/examples/llm_qat/configs/train/lr/lr_config_amax.yaml new file mode 100644 index 00000000000..6b2b5f303f9 --- /dev/null +++ b/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 diff --git a/examples/llm_qat/configs/train/lr_config_example.yaml b/examples/llm_qat/configs/train/lr/lr_config_example.yaml similarity index 95% rename from examples/llm_qat/configs/train/lr_config_example.yaml rename to examples/llm_qat/configs/train/lr/lr_config_example.yaml index 844e5199e8b..ce3c6a5a5a2 100644 --- a/examples/llm_qat/configs/train/lr_config_example.yaml +++ b/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. diff --git a/examples/llm_qat/configs/train/qad_scale_only.yaml b/examples/llm_qat/configs/train/qad_scale_only.yaml new file mode 100644 index 00000000000..39c999d02e1 --- /dev/null +++ b/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 diff --git a/examples/llm_qat/configs/train/qad_with_learnt_amax.yaml b/examples/llm_qat/configs/train/qad_with_learnt_amax.yaml new file mode 100644 index 00000000000..1923849f588 --- /dev/null +++ b/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 diff --git a/modelopt/torch/kernels/quantization/gemm/fp4_kernel.py b/modelopt/torch/kernels/quantization/gemm/fp4_kernel.py index 2b655f4bbcb..c9499b612f6 100644 --- a/modelopt/torch/kernels/quantization/gemm/fp4_kernel.py +++ b/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) diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py index 6853abcf85c..f72715d410c 100644 --- a/modelopt/torch/opt/plugins/transformers.py +++ b/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." ), }, ) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 0ca30d18448..8973ec28ed0 100644 --- a/modelopt/torch/quantization/config.py +++ b/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 `_. + 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 diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index fa7a8a0a128..00187d291c0 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -37,10 +37,10 @@ normalize_quant_cfg_list, ) from .nn import ( - NVFP4StaticQuantizer, QuantModule, QuantModuleRegistry, SequentialQuantizer, + StaticBlockScaleQuantizer, SVDQuantLinear, TensorQuantizer, ) @@ -100,10 +100,11 @@ def restore_quantized_model( def maybe_promote_nvfp4_static_quantizer(module: nn.Module, quantizer_state: dict) -> None: - if quantizer_state.get("_is_nvfp4_static_quantizer") and not isinstance( - module, NVFP4StaticQuantizer - ): - NVFP4StaticQuantizer.from_tensor_quantizer(module) + if ( + quantizer_state.get("_is_static_block_scale_quantizer") + or quantizer_state.get("_is_nvfp4_static_quantizer") + ) and not isinstance(module, StaticBlockScaleQuantizer): + StaticBlockScaleQuantizer.from_tensor_quantizer(module) def _restore_shared_quant_state_aliases( @@ -153,6 +154,7 @@ def restore_quantizer_state(model: nn.Module, config: QuantizeConfig, metadata: if isinstance(module, TensorQuantizer): name = get_unwrapped_name(name, model) state = quantizer_state_dict[name] + # TODO: Add a registry for TensorQuantizers and avoid this manual conversion. maybe_promote_nvfp4_static_quantizer(module, state) module.set_from_modelopt_state(state) diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py index 2db966ddbed..80be73daa2a 100644 --- a/modelopt/torch/quantization/mode.py +++ b/modelopt/torch/quantization/mode.py @@ -39,6 +39,7 @@ CompressConfig, GPTQCalibConfig, LocalHessianCalibConfig, + LSQConfig, MaxCalibConfig, MseCalibConfig, QuantizeAlgoCfgType, @@ -62,6 +63,7 @@ gptq, layerwise_calibrate, local_hessian_calibrate, + lsq, max_calibrate, mse_calibrate, smoothquant, @@ -531,3 +533,15 @@ def config_class(self) -> type[QuantizeAlgorithmConfig]: return GPTQCalibConfig _calib_func = gptq + + +@CalibrateModeRegistry.register_mode +class LSQModeDescriptor(BaseCalibrateModeDescriptor): + """Mode for LSQ (Learned Scale Quantization) algorithm.""" + + @property + def config_class(self) -> type[QuantizeAlgorithmConfig]: + """Specifies the config class for the mode.""" + return LSQConfig + + _calib_func = lsq diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 7e5bb85c09b..2a8f9dba1bb 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -29,6 +29,7 @@ import torch.nn.functional as F from tqdm import tqdm +from modelopt.torch.opt.config import ModeloptBaseConfig from modelopt.torch.opt.searcher import ForwardLoop from modelopt.torch.quantization.utils.layerwise_calib import ( LayerActivationCollector, @@ -42,7 +43,7 @@ from .calib import MseCalibrator, NVFP4MSECalibrator, _Calibrator from .conversion import create_and_replace_svdquant_linear_on_the_fly, set_quantizer_by_cfg_context -from .nn import NVFP4StaticQuantizer, QuantModule, SequentialQuantizer, TensorQuantizer +from .nn import QuantModule, SequentialQuantizer, StaticBlockScaleQuantizer, TensorQuantizer from .utils import ( SHARED_PATTERNS, SharedWeightGlobalAmaxState, @@ -54,7 +55,7 @@ is_quantized_linear, is_quantized_row_parallel_linear, persistent_materialization, - promote_nvfp4_static_quantizers, + promote_static_block_weight_quantizers, ) from .utils.calib_utils import _GPTQ_HELPER_REGISTRY, GPTQHelper @@ -63,6 +64,7 @@ "awq", "layerwise_calibrate", "local_hessian_calibrate", + "lsq", "max_calibrate", "smoothquant", "svdquant", @@ -76,7 +78,7 @@ def _collect_weight_stats(quantizer: nn.Module, weight: torch.Tensor) -> None: def _is_calibrated_nvfp4_static(q) -> bool: """True iff ``q`` is an enabled NVFP4-static weight quantizer with ``_amax`` set.""" return ( - isinstance(q, NVFP4StaticQuantizer) + isinstance(q, StaticBlockScaleQuantizer) and not q._disabled and q.is_nvfp4_static and getattr(q, "_amax", None) is not None @@ -131,15 +133,16 @@ def _check_grouped_weight_global_amax_synced(model: nn.Module) -> None: def _finalize_with_shared_state(model: nn.Module, weight_patterns: list[str]) -> None: - """Finalize quantization from the attached shared state: aggregate, promote, verify. + """Finalize calibrated static quantizers and attached shared state. Aggregates each fusible group's shared weight ``global_amax`` and promotes it onto the member NVFP4-static quantizers, so siblings read the unified value instead of their own - ``_amax``; under the default patterns, verifies the name groups were actually synced. - Call once ``_amax`` is final: single-process, or after the distributed amax sync. + ``_amax``. Promotes static-block weight quantizers after their ``_amax`` is final. Under + the default patterns, verifies the name groups were actually synced. Call once ``_amax`` + is final: single-process, or after the distributed amax sync. """ SharedWeightGlobalAmaxState.populate(model) - promote_nvfp4_static_quantizers(model) + promote_static_block_weight_quantizers(model) # Under the default patterns, verify the fusible name groups were actually synced. if weight_patterns == list(SHARED_PATTERNS): _check_grouped_weight_global_amax_synced(model) @@ -314,7 +317,7 @@ def max_calibrate( for name, module in model.named_modules(): if isinstance(module, QuantModule) and _has_expert_parallelism(module): for child in module.children(): - if isinstance(child, (TensorQuantizer, SequentialQuantizer)): + if isinstance(child, TensorQuantizer | SequentialQuantizer): _check_moe_calibration_complete(child, module.parallel_state) def sync_quantizer_amax_across_dp_ep(quantizer, parallel_state, parent_name, child_name): @@ -333,7 +336,7 @@ def sync_quantizer_amax_across_dp_ep(quantizer, parallel_state, parent_name, chi for name, module in model.named_modules(): if isinstance(module, QuantModule): for child_name, child in module.named_children(): - if isinstance(child, (TensorQuantizer, SequentialQuantizer)): + if isinstance(child, TensorQuantizer | SequentialQuantizer): sync_quantizer_amax_across_dp_ep(child, module.parallel_state, name, child_name) # Step 3: TP sync # Objective: the quantization parameters when TP = 8 then changed to TP=4 then back to TP=8 should be the same @@ -1985,7 +1988,7 @@ def gptq( Per-module steps: 1. ``max_calibrate`` to set amax values from the current activations. - 2. Promote eligible quantizers to ``NVFP4StaticQuantizer`` (two-level scaling). + 2. Promote eligible quantizers to ``StaticBlockScaleQuantizer`` (two-level scaling). 3. Collect per-linear-layer Hessian matrices via forward hooks. 4. Blockwise weight updates using the inverse Hessian to compensate for rounding error (the core GPTQ column-wise update). @@ -2045,3 +2048,73 @@ def _make_gptq_handle(name, m): if torch.cuda.is_available(): torch.cuda.empty_cache() print_rank_0(f"GPTQ time: {time.time() - total_start:.2f}s") + + +def _run_scale_calibration(model, forward_loop, scale_algorithm): + """Run scale calibration.""" + if scale_algorithm is None: + scale_algorithm = {"method": "mse"} + + if isinstance(scale_algorithm, ModeloptBaseConfig): + scale_algorithm = scale_algorithm.model_dump(exclude_unset=True) + + method = scale_algorithm.get("method") + algo_kwargs = {k: v for k, v in scale_algorithm.items() if k != "method"} + calib_funcs = { + "mse": mse_calibrate, + "local_hessian": local_hessian_calibrate, + "max": max_calibrate, + } + calib_funcs[method](model, forward_loop=forward_loop, **algo_kwargs) + + +@torch.no_grad() +def lsq( + model: nn.Module, + forward_loop: ForwardLoop | None = None, + scale_algorithm: dict | None = None, + learnable_amax: list | str = ("post",), + tied_amax: bool = False, + quantize_pre_scale: bool = True, + **kwargs, +): + """Run scale calibration then convert to LSQ mode. + + Uses separate pre (quant) and post (dequant) amax values. + Forward: ``w_q = Q_STE(w / s_pre) * s_post`` where ``s = amax / Q_max``. + + Args: + model: Quantized model. + forward_loop: Calibration data forward loop. + scale_algorithm: Calibration algorithm config to run first. + Dict with 'method' key: 'mse', 'local_hessian', or 'max'. + Defaults to {'method': 'mse'} if None. + learnable_amax: Which amax params are learnable: 'pre', 'post', + ['pre', 'post'], or []. + tied_amax: If True, pre and post share a single tensor. + quantize_pre_scale: If False, skip FP8 quantization for the LSQ pre scale. + """ + _run_scale_calibration(model, forward_loop, scale_algorithm) + + name_to_module = dict(model.named_modules()) + seen_modules: set[int] = set() + seen_quantizers: set[int] = set() + for module in name_to_module.values(): + if id(module) in seen_modules or not isinstance(module, QuantModule): + continue + seen_modules.add(id(module)) + with enable_weight_access_and_writeback(module, model, name_to_module): + for weight, quantizer in module.iter_weights_for_calibration(): + if id(quantizer) in seen_quantizers: + continue + seen_quantizers.add(id(quantizer)) + if not isinstance(quantizer, StaticBlockScaleQuantizer) or not hasattr( + quantizer, "_amax" + ): + continue + quantizer.enable_lsq( + learnable_amax=learnable_amax, + tied_amax=tied_amax, + quantize_pre_scale=quantize_pre_scale, + dtype=weight.dtype, + ) diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 98d9e0dcb1e..ddc7ac2b045 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -57,6 +57,8 @@ from ...tensor_quant import ( dynamic_block_quant, fake_tensor_quant, + fp4_cast_ste, + int_cast_ste, scaled_e4m3, static_blockwise_fp4_fake_quant, ) @@ -64,10 +66,15 @@ from ...utils.numeric_utils import fp8_max_for_normalization from ..functional import normalized_hadamard_transform +# torch.finfo(...).tiny gives the smallest normal E4M3 value; scale clamping needs +# the smallest positive subnormal value representable by the 3-bit mantissa. +_FP8_E4M3_MIN_POSITIVE = torch.finfo(torch.float8_e4m3fn).smallest_normal / (2**3) + __all__ = [ "HardDisabledTensorQuantizer", "NVFP4StaticQuantizer", "SequentialQuantizer", + "StaticBlockScaleQuantizer", "TensorQuantizer", "TensorQuantizerCache", "is_registered_quant_backend", @@ -1436,19 +1443,57 @@ def set_from_attribute_config(self, attribute_cfg): self._disabled = True -class NVFP4StaticQuantizer(TensorQuantizer): - """TensorQuantizer for NVFP4 static block quantization with two-level scaling. +def _clamp_scale(scale: torch.Tensor, min_value: float | torch.Tensor = 1e-8) -> torch.Tensor: + """Clamp per-block scale to guard against small/zero values.""" + return torch.where(scale <= min_value, min_value, scale) + + +def _amax_to_scale( + amax: torch.Tensor, max_bound: float, min_value: float | torch.Tensor = 1e-8 +) -> torch.Tensor: + """Convert amax to per-block scale, guarding against small/zero values.""" + return _clamp_scale(amax.float() / max_bound, min_value) + + +def _to_local(t: torch.Tensor) -> torch.Tensor: + """Convert DTensor to local tensor (no-op for regular tensors). + + Under FSDP2, learnable parameters are DTensors but the quantizer forward + operates on local tensors (see TensorQuantizer.forward DTensor handling). + to_local() preserves autograd so gradients flow back to the DTensor parameter. + """ + if DTensor is not None and isinstance(t, DTensor): + return t.to_local() + return t + + +class StaticBlockScaleQuantizer(TensorQuantizer): + """TensorQuantizer for static block quantization with two-level scaling. + Supports both FP4 (E2M1) and INT block quantization formats with configurable + block_size and optional FP8 scale quantization. Uses _global_amax and inherited _amax for per-block amax values. - Preserves both amax states in fp32. + Preserves static amax states in fp32. """ + _lsq: bool = False + _learnable_amax: list = [] + _tied_amax: bool = False + # FP4 default; overwritten on promotion with the format-specific bound, including INT. + _quant_max_bound: float = 6.0 + _quantize_scales: bool = True + _quantize_pre_scale: bool = True + def _preserve_amax_in_fp32(self): amax = getattr(self, "_amax", None) - if amax is not None: + if amax is not None and not isinstance(amax, nn.Parameter): self._amax = amax.to(dtype=torch.float32) global_amax = getattr(self, "_global_amax", None) - if global_amax is not None and global_amax.dtype != torch.float32: + if ( + global_amax is not None + and not isinstance(global_amax, nn.Parameter) + and global_amax.dtype != torch.float32 + ): if "_global_amax" in self.__dict__.get("_shared_quant_tied_attrs", set()): global_amax.data = global_amax.to(dtype=torch.float32) else: @@ -1461,8 +1506,8 @@ def _amax_setter_helper(self, value): @classmethod def from_tensor_quantizer( cls, tq: TensorQuantizer, global_amax: torch.Tensor | None = None - ) -> "NVFP4StaticQuantizer": - """Convert a TensorQuantizer to NVFP4StaticQuantizer in-place. + ) -> "StaticBlockScaleQuantizer": + """Convert a TensorQuantizer to StaticBlockScaleQuantizer in-place. Args: tq: The TensorQuantizer to convert. @@ -1477,11 +1522,48 @@ def _preserve_and_set_global_amax(tq): if isinstance(tq, cls): _preserve_and_set_global_amax(tq) return tq + is_nvfp4_static = getattr(tq, "is_nvfp4_static", False) tq.__class__ = cls - tq._is_nvfp4_static_quantizer = True + tq._is_static_block_scale_quantizer = True + if is_nvfp4_static: + tq._is_nvfp4_static_quantizer = True + tq._quant_max_bound = float(tq.maxbound) _preserve_and_set_global_amax(tq) return tq + @property + def amax_pre(self): + """Pre (quantization) amax. Returns _amax_post when tied.""" + if self._tied_amax: + return self._amax_post + return self._amax_pre + + @property + def amax_post(self): + """Post (dequantization) amax.""" + return self._amax_post + + @property + def amax(self): + """Return amax, derived from learnable amax parameters if in LSQ mode.""" + if self._lsq and not self._tied_amax: + raise RuntimeError( + "LSQ with untied amaxes has separate pre and post parameters. " + "Access them via amax_pre / amax_post." + ) + if self._lsq: + return self._amax_post + if not hasattr(self, "_amax"): + return None + return self._amax + + @amax.setter + def amax(self, value): + assert value is not None, "amax cannot be set to None." + if not isinstance(value, torch.Tensor): + value = torch.tensor(value) + self._amax_setter_helper(value) + @property def global_amax(self): """Return global_amax for quantization.""" @@ -1506,6 +1588,11 @@ def global_amax(self, value): global_amax.data.copy_(value.clone().detach().to(global_amax.device)) self._preserve_amax_in_fp32() + @property + def has_quantized_block_scale(self): + """True when per-block scales are FP8 (E4M3) quantized (format-only check).""" + return self._block_sizes is not None and self._block_sizes.get("scale_bits") == (4, 3) + def _apply(self, fn, recurse=True): """Apply module transforms without rounding static scale state.""" amax = getattr(self, "_amax", None) @@ -1513,20 +1600,116 @@ def _apply(self, fn, recurse=True): module = super()._apply(fn, recurse=recurse) self._preserve_amax_in_fp32() - if amax is not None: + if amax is not None and amax.device.type != "meta": self.amax = amax - if global_amax is not None: + if global_amax is not None and global_amax.device.type != "meta": self.global_amax = global_amax return module + def _short_amax(self, fmt=".4f"): + """Short description of amax, accounting for LSQ mode.""" + if not self._lsq: + return super()._short_amax(fmt) + learn = self._learnable_amax + learn_str = "frozen" if not learn else f"learn=[{','.join(learn)}]" + if self._tied_amax: + return f"LSQ(tied={self._short_tensor(self._amax_post.data, fmt)}, {learn_str})" + return ( + f"LSQ(pre={self._short_tensor(self._amax_pre.data, fmt)}, " + f"post={self._short_tensor(self._amax_post.data, fmt)}, {learn_str})" + ) + + def enable_lsq( + self, + quantize_scales: bool | None = None, + learnable_amax: list | str = ("post",), + tied_amax: bool = False, + quantize_pre_scale: bool = True, + dtype: torch.dtype | None = None, + ): + """LSQ mode with configurable learnable/frozen amax tensors. + + The per-block amax params are initialized from the calibrated ``_amax``. The + per-tensor scale is derived from ``global_amax`` at runtime so shared-group + updates are always reflected. + + Args: + quantize_scales: Whether to FP8-quantize per-block scales (NVFP4). When None, + defaults to ``has_quantized_block_scale``. + learnable_amax: Which amax params are learnable: 'pre', 'post', + ['pre', 'post'], or []. + tied_amax: If True, pre and post share a single tensor. + quantize_pre_scale: Whether to FP8-quantize the LSQ pre scale. + dtype: Optional dtype for the amax params. Kept at weight dtype for FSDP2 + mixed-precision support (see TODO below). + """ + assert hasattr(self, "_amax"), "enable_lsq requires a calibrated _amax." + if quantize_scales is None: + quantize_scales = self.has_quantized_block_scale + if quantize_scales: + assert self.global_amax is not None, ( + "enable_lsq(quantize_scales=True) requires global_amax to be set." + ) + + # TODO: Support fp32 learnable amax values once a stable PyTorch release + # includes FSDP2 mixed-precision parameter dtype support. + amax = self._amax.float() + if dtype is not None: + amax = amax.to(dtype) + delattr(self, "_amax") + learn = {learnable_amax} if isinstance(learnable_amax, str) else set(learnable_amax) + + if "post" in learn: + self._amax_post = nn.Parameter(amax.clone(), requires_grad=True) + else: + self.register_buffer("_amax_post", amax.clone()) + + if not tied_amax: + if "pre" in learn: + self._amax_pre = nn.Parameter(amax.clone(), requires_grad=True) + else: + self.register_buffer("_amax_pre", amax.clone()) + + self._quantize_scales = quantize_scales + self._quantize_pre_scale = quantize_pre_scale + self._lsq = True + self._learnable_amax = sorted(learn) + self._tied_amax = tied_amax + + def _cast_ste(self, inputs): + """Cast inputs to quantized representable values (no scaling).""" + if isinstance(self._num_bits, tuple): + return fp4_cast_ste(inputs) + return int_cast_ste(inputs, self._num_bits, self._unsigned, self._narrow_range) + + def _block_scale_from_amax(self, amax: torch.Tensor, quantize: bool) -> torch.Tensor: + """Compute the per-block scale from a per-block amax, optionally FP8-quantizing it.""" + if quantize: + per_tensor_scale = _amax_to_scale(self.global_amax, self._quant_max_bound) + min_value = _FP8_E4M3_MIN_POSITIVE * per_tensor_scale.view(-1) + scale = _amax_to_scale(amax, self._quant_max_bound, min_value=min_value) + return scaled_e4m3(scale, per_tensor_scale, None, 4, 3) + return _amax_to_scale(amax, self._quant_max_bound, min_value=1e-8) + def _fake_quantize(self, inputs): """Fake quantization using two-level scaling with _amax and _global_amax.""" - if self.amax is not None: + if self._lsq: + scale_post = self._block_scale_from_amax( + _to_local(self.amax_post), self._quantize_scales + ) + scale_pre = self._block_scale_from_amax( + _to_local(self.amax_pre), self._quantize_scales and self._quantize_pre_scale + ) + quant_input = inputs.float() / scale_pre.float().view(-1, 1) + w_cast = self._cast_ste(quant_input) + return (w_cast * scale_post.view(-1, 1).to(w_cast.dtype)).to(inputs.dtype) + + if self.amax is not None and self.is_nvfp4_static: return static_blockwise_fp4_fake_quant( inputs, self.amax, - self.global_amax, # Can be None, will be computed internally - True, # quantize_block_scales + self.global_amax, + True, fp8_max_for_normalization(self), inputs.dtype, self._pass_through_bwd, @@ -1534,6 +1717,9 @@ def _fake_quantize(self, inputs): return super()._fake_quantize(inputs) +NVFP4StaticQuantizer = StaticBlockScaleQuantizer + + class SequentialQuantizer(nn.Sequential): """A sequential container for :class:`TensorQuantizer` modules. diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index cb48b2bf304..20e083491aa 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -645,7 +645,63 @@ def _tensor_quant(inputs, amax, num_bits=8, unsigned=False, narrow_range=True): return outputs +class FP4CastSTEFunction(Function): + """FP4 cast with STE backward -- no scale/descale, just rounding.""" + + @staticmethod + def forward(ctx, x, out_dtype=None): + """Forward pass: cast to FP4 using triton kernel. + + Args: + x: Input tensor of shape [NUM_BLOCKS, BLOCK_SIZE]. + out_dtype: Output dtype. Defaults to x.dtype. + """ + if not triton_kernel.IS_AVAILABLE: + raise RuntimeError("FP4CastSTEFunction requires triton.") + ctx.save_for_backward(x) + return triton_kernel.static_blockwise_fp4_cast(x, out_dtype) + + @staticmethod + def backward(ctx, grad_outputs): + """Backward pass: STE with clip mask at ``|x| <= 6.0``.""" + (x,) = ctx.saved_tensors + grad = torch.where(x.abs() <= 6.0, grad_outputs, torch.zeros_like(grad_outputs)) + return grad, None + + +class IntCastSTEFunction(Function): + """Integer quantization cast with STE backward, analogous to FP4CastSTEFunction.""" + + @staticmethod + def forward(ctx, x, num_bits, unsigned=False, narrow_range=True): + """Forward pass: clamp-round to integer range.""" + max_bound = (2.0 ** (num_bits - 1 + int(unsigned))) - 1.0 + if unsigned: + min_bound = 0 + elif narrow_range: + min_bound = -max_bound + else: + min_bound = -max_bound - 1 + ctx.save_for_backward(x) + ctx.min_bound = min_bound + ctx.max_bound = max_bound + return torch.clamp(x.round(), min_bound, max_bound) + + @staticmethod + def backward(ctx, grad_outputs): + """Backward pass: STE with clip mask.""" + (x,) = ctx.saved_tensors + grad = torch.where( + (x >= ctx.min_bound) & (x <= ctx.max_bound), + grad_outputs, + torch.zeros_like(grad_outputs), + ) + return grad, None, None, None + + fake_tensor_quant = FakeTensorQuantFunction.apply scaled_e4m3 = ScaledE4M3Function.apply dynamic_block_quant = DynamicBlockQuantizationFunction.apply static_blockwise_fp4_fake_quant = StaticBlockwiseFP4FakeQuantFunction.apply +fp4_cast_ste = FP4CastSTEFunction.apply +int_cast_ste = IntCastSTEFunction.apply diff --git a/modelopt/torch/quantization/utils/__init__.py b/modelopt/torch/quantization/utils/__init__.py index 69969b2554d..9fb7eacffaf 100644 --- a/modelopt/torch/quantization/utils/__init__.py +++ b/modelopt/torch/quantization/utils/__init__.py @@ -34,6 +34,8 @@ "is_quantized_linear", "is_quantized_row_parallel_linear", "iter_shared_quant_states", + "promote_nvfp4_static_quantizers", + "promote_static_block_weight_quantizers", "reduce_amax", "reduce_sum", "replace_function", diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 478788c4f1c..5c839f050ad 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -962,8 +962,8 @@ def update_quant_cfg_with_kv_cache_quant( return quant_cfg -def promote_nvfp4_static_quantizers(model: nn.Module) -> int: - """Convert eligible TensorQuantizers to NVFP4StaticQuantizer in-place. +def promote_static_block_weight_quantizers(model: nn.Module) -> int: + """Convert eligible static-block weight TensorQuantizers in-place. After max calibration sets per-block amax values, NVFP4 static quantizers need to be promoted so they use the two-level scaling path (global amax + @@ -973,9 +973,14 @@ def promote_nvfp4_static_quantizers(model: nn.Module) -> int: ``model``, the promoted quantizer's ``_global_amax`` buffer is tied to that canonical state buffer instead of receiving an independent copy. - Returns the number of quantizers converted. + Returns the number of NVFP4 quantizers converted. """ - from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, TensorQuantizer + from modelopt.torch.quantization.nn import ( + QuantModule, + SequentialQuantizer, + StaticBlockScaleQuantizer, + TensorQuantizer, + ) from modelopt.torch.quantization.utils.shared_input import ( SharedWeightGlobalAmaxState, iter_shared_quant_states, @@ -988,33 +993,51 @@ def promote_nvfp4_static_quantizers(model: nn.Module) -> int: for state in iter_shared_quant_states(model, SharedWeightGlobalAmaxState) for quantizer in state._member_quantizers() } - converted = 0 for _name, module in list(model.named_modules()): - if not isinstance(module, TensorQuantizer) or not module.is_enabled: - continue - if not module.is_nvfp4_static: + if not isinstance(module, QuantModule): continue - amax = module.amax - if amax is None: - continue - - # Grouped siblings share one canonical global_amax (common FP8 grid); otherwise - # fall back to this quantizer's own per-block amax. - already_promoted = isinstance(module, NVFP4StaticQuantizer) - shared = shared_by_quantizer.get(id(module)) - if shared is not None and shared.global_amax is not None: - NVFP4StaticQuantizer.from_tensor_quantizer(module) - shared.tie_member_quantizer(module) - else: - if shared is not None and not amax.is_meta: - raise RuntimeError( - f"{_name}: weight quantizer is in a shared group whose global_amax was not " - "populated before promotion; run populate after calibration so siblings " - "share one scale instead of falling back to their own." - ) - global_amax = reduce_amax(amax.clone().detach(), axis=None) - NVFP4StaticQuantizer.from_tensor_quantizer(module, global_amax=global_amax) - if not already_promoted: - converted += 1 + for _, quantizer in module.iter_weights_for_calibration(): + if isinstance(quantizer, SequentialQuantizer): + if len(quantizer) == 0: + continue + quantizer = quantizer[0] + if not isinstance(quantizer, TensorQuantizer): + continue + quantizer_id = id(quantizer) + if not quantizer.is_enabled or not quantizer.is_static_block_quant: + continue + amax = quantizer.amax + if amax is None: + continue + if quantizer.is_nvfp4_static: + # Grouped siblings share one canonical global_amax (common FP8 grid); otherwise + # fall back to this quantizer's own per-block amax. + already_promoted = isinstance(quantizer, StaticBlockScaleQuantizer) + shared = shared_by_quantizer.get(quantizer_id) + if shared is not None and shared.global_amax is not None: + StaticBlockScaleQuantizer.from_tensor_quantizer(quantizer) + shared.tie_member_quantizer(quantizer) + else: + if shared is not None and not amax.is_meta: + raise RuntimeError( + f"{_name}: weight quantizer is in a shared group whose global_amax was " + "not populated before promotion; run populate after calibration so " + "siblings share one scale instead of falling back to their own." + ) + global_amax = reduce_amax(amax.clone().detach(), axis=None) + StaticBlockScaleQuantizer.from_tensor_quantizer( + quantizer, global_amax=global_amax + ) + if not already_promoted: + converted += 1 + elif isinstance(quantizer._num_bits, int): + # Integer static-block weights are promoted so LSQ can use + # StaticBlockScaleQuantizer. + StaticBlockScaleQuantizer.from_tensor_quantizer(quantizer) return converted + + +def promote_nvfp4_static_quantizers(model: nn.Module) -> int: + """Compatibility wrapper for static-block weight quantizer promotion.""" + return promote_static_block_weight_quantizers(model) diff --git a/modelopt_recipes/general/ptq/nvfp4_default-kv_fp8.yaml b/modelopt_recipes/general/ptq/nvfp4_default-kv_fp8.yaml index 6a65efef57a..9be27b7bae9 100644 --- a/modelopt_recipes/general/ptq/nvfp4_default-kv_fp8.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_default-kv_fp8.yaml @@ -24,8 +24,8 @@ imports: metadata: recipe_type: ptq description: >- - Composes dynamic NVFP4 W4A4 model quantization with FP8 KV-cache quantization; uses max - calibration. + Composes dynamic NVFP4 W4A4 model quantization with FP8 KV-cache quantization for PTQ, + QAT, and QAD; uses max calibration. quantize: algorithm: max quant_cfg: diff --git a/modelopt_recipes/general/qad/nvfp4_default-kv_fp8.yaml b/modelopt_recipes/general/qad/nvfp4_default-kv_fp8.yaml new file mode 120000 index 00000000000..97cca50092a --- /dev/null +++ b/modelopt_recipes/general/qad/nvfp4_default-kv_fp8.yaml @@ -0,0 +1 @@ +../ptq/nvfp4_default-kv_fp8.yaml \ No newline at end of file diff --git a/modelopt_recipes/general/qad/nvfp4_dual_lsq-mse_init-fp8_kv.yaml b/modelopt_recipes/general/qad/nvfp4_dual_lsq-mse_init-fp8_kv.yaml new file mode 100644 index 00000000000..53a9651ed24 --- /dev/null +++ b/modelopt_recipes/general/qad/nvfp4_dual_lsq-mse_init-fp8_kv.yaml @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +metadata: + recipe_type: ptq + description: >- + Learns separate pre-quantization and post-quantization NVFP4 weight scales with MSE + initialization and FP8 scale sweep; uses dynamic NVFP4 activations and FP8 KV cache. + QAT/QAD-only: the LSQ scales are learned during training, so this recipe is not + suitable for calibration-only PTQ. +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + kv_fp8: configs/ptq/units/kv_fp8 + nvfp4: configs/numerics/nvfp4 + nvfp4_static: configs/numerics/nvfp4_static +quantize: + algorithm: + method: lsq + learnable_amax: + - pre + - post + tied_amax: false + quantize_pre_scale: false + scale_algorithm: + method: mse + fp8_scale_sweep: true + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: nvfp4_static + - quantizer_name: '*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8 + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/general/qad/nvfp4_lsq-mse_init-fp8_kv.yaml b/modelopt_recipes/general/qad/nvfp4_lsq-mse_init-fp8_kv.yaml new file mode 100644 index 00000000000..dba0f7d98fe --- /dev/null +++ b/modelopt_recipes/general/qad/nvfp4_lsq-mse_init-fp8_kv.yaml @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +metadata: + recipe_type: ptq + description: >- + Learns one shared pre-quantization and post-quantization NVFP4 weight scale with MSE + initialization and FP8 scale sweep; uses dynamic NVFP4 activations and FP8 KV cache. + QAT/QAD-only: the LSQ scales are learned during training, so this recipe is not + suitable for calibration-only PTQ. +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + kv_fp8: configs/ptq/units/kv_fp8 + nvfp4: configs/numerics/nvfp4 + nvfp4_static: configs/numerics/nvfp4_static +quantize: + algorithm: + method: lsq + learnable_amax: + - pre + - post + tied_amax: true + quantize_pre_scale: true + scale_algorithm: + method: mse + fp8_scale_sweep: true + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: nvfp4_static + - quantizer_name: '*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8 + - $import: default_disabled_quantizers diff --git a/tests/gpu/torch/quantization/test_fsdp2.py b/tests/gpu/torch/quantization/test_fsdp2.py index 55648ac26e2..6d47e1620ab 100644 --- a/tests/gpu/torch/quantization/test_fsdp2.py +++ b/tests/gpu/torch/quantization/test_fsdp2.py @@ -27,6 +27,7 @@ import modelopt.torch.quantization as mtq from modelopt.torch.opt.dynamic import _pytorch_managed +from modelopt.torch.quantization.nn import StaticBlockScaleQuantizer, TensorQuantizer from modelopt.torch.quantization.utils import ( enable_weight_access_and_writeback, persistent_materialization, @@ -136,6 +137,50 @@ def test_nested_fsdp2_backward(quant_cfg, dist_workers): dist_workers.run(partial(_test_nested_fsdp2_backward, quant_cfg=quant_cfg)) +class _LSQBf16Linear(nn.Module): + """Minimal bf16 module with LSQ learnable amax parameters.""" + + def __init__(self, dim=16): + super().__init__() + self.weight = nn.Parameter(torch.randn(dim, dim, dtype=torch.bfloat16)) + + tq = TensorQuantizer() + tq._num_bits = 4 + tq._unsigned = False + tq._narrow_range = True + tq._disabled = False + tq._block_sizes = {-1: dim} + tq._pass_through_bwd = True + tq.register_buffer("_amax", torch.ones(dim, dtype=torch.bfloat16)) + self.weight_quantizer = StaticBlockScaleQuantizer.from_tensor_quantizer(tq) + self.weight_quantizer.enable_lsq( + quantize_scales=False, + learnable_amax=["pre", "post"], + dtype=torch.bfloat16, + ) + + def forward(self, inputs): + weight = self.weight_quantizer._fake_quantize(self.weight) + return torch.nn.functional.linear(inputs, weight) + + +def _test_lsq_bf16_learnable_amax_fsdp2(rank, size): + torch.manual_seed(1) + model = _LSQBf16Linear().cuda(rank) + inputs = torch.randn(2, 16, device=rank, dtype=torch.bfloat16) + synchronize_state_dict(model) + + assert {p.dtype for p in model.parameters()} == {torch.bfloat16} + + model = fully_shard(model) + output = model(inputs) + output.float().sum().backward() + + +def test_lsq_bf16_learnable_amax_fsdp2(dist_workers): + dist_workers.run(_test_lsq_bf16_learnable_amax_fsdp2) + + class _DecoderBlock(nn.Module): """Minimal decoder block for FSDP2 sequential tests.""" diff --git a/tests/gpu/torch/quantization/test_lsq_cuda.py b/tests/gpu/torch/quantization/test_lsq_cuda.py new file mode 100644 index 00000000000..e2d802bc6b7 --- /dev/null +++ b/tests/gpu/torch/quantization/test_lsq_cuda.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GPU unit tests for the LSQ algorithm using FP4 (NVFP4) quantization.""" + +import pytest +import torch +from torch import nn + +import modelopt.torch.quantization as mtq + +NVFP4_LSQ_POST_MSE_CFG = { + "quant_cfg": { + "*weight_quantizer": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "static", "scale_bits": (4, 3)}, + "axis": None, + "enable": True, + }, + "*input_quantizer": { + "enable": False, + }, + }, + "algorithm": { + "method": "lsq", + "learnable_amax": ["post"], + "scale_algorithm": {"method": "mse", "fp8_scale_sweep": True}, + }, +} + +NVFP4_LSQ_PRE_POST_MSE_CFG = { + "quant_cfg": { + "*weight_quantizer": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "static", "scale_bits": (4, 3)}, + "axis": None, + "enable": True, + }, + "*input_quantizer": { + "enable": False, + }, + }, + "algorithm": { + "method": "lsq", + "learnable_amax": ["pre", "post"], + "scale_algorithm": {"method": "mse", "fp8_scale_sweep": True}, + }, +} + +NVFP4_LSQ_TIED_MSE_CFG = { + "quant_cfg": { + "*weight_quantizer": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "static", "scale_bits": (4, 3)}, + "axis": None, + "enable": True, + }, + "*input_quantizer": { + "enable": False, + }, + }, + "algorithm": { + "method": "lsq", + "learnable_amax": ["pre", "post"], + "tied_amax": True, + "scale_algorithm": {"method": "mse", "fp8_scale_sweep": True}, + }, +} + +NVFP4_LSQ_SKIP_PRE_SCALE_MSE_CFG = { + "quant_cfg": { + "*weight_quantizer": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "static", "scale_bits": (4, 3)}, + "axis": None, + "enable": True, + }, + "*input_quantizer": { + "enable": False, + }, + }, + "algorithm": { + "method": "lsq", + "learnable_amax": ["post"], + "quantize_pre_scale": False, + "scale_algorithm": {"method": "mse", "fp8_scale_sweep": True}, + }, +} + + +class SimpleModel(nn.Module): + """Minimal model for LSQ testing.""" + + def __init__(self): + super().__init__() + self.linear = nn.Linear(64, 64, bias=False) + + def forward(self, x): + return self.linear(x) + + +def _make_forward_loop(model, device): + x = torch.randn(2, 64, device=device) + + def forward_loop(m): + m(x) + + return forward_loop + + +@pytest.mark.parametrize( + "config", + [ + NVFP4_LSQ_POST_MSE_CFG, + NVFP4_LSQ_PRE_POST_MSE_CFG, + NVFP4_LSQ_TIED_MSE_CFG, + NVFP4_LSQ_SKIP_PRE_SCALE_MSE_CFG, + ], + ids=["post_only", "pre_and_post", "tied", "skip_pre_scale"], +) +def test_lsq_quantize_e2e(config): + """End-to-end: quantize a small model with LSQ + NVFP4 on GPU.""" + device = torch.device("cuda") + model = SimpleModel().to(device) + forward_loop = _make_forward_loop(model, device) + + model = mtq.quantize(model, config, forward_loop=forward_loop) + assert model.linear.weight_quantizer._quantize_pre_scale is config["algorithm"].get( + "quantize_pre_scale", True + ) + + # Verify the model still produces output of the correct shape + x = torch.randn(2, 64, device=device) + out = model(x) + assert out.shape == (2, 64) + + +def test_lsq_fp4_fake_quantize_differentiable(): + """Test that _fake_quantize in FP4 LSQ mode is differentiable.""" + from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( + StaticBlockScaleQuantizer, + TensorQuantizer, + ) + + device = torch.device("cuda") + tq = TensorQuantizer() + tq._num_bits = (2, 1) + tq._unsigned = False + tq._narrow_range = True + tq._disabled = False + tq._block_sizes = {-1: 16, "type": "static", "scale_bits": (4, 3)} + tq._pass_through_bwd = True + tq.register_buffer("_amax", torch.ones(4, device=device)) + tq.to(device) + sbsq = StaticBlockScaleQuantizer.from_tensor_quantizer( + tq, global_amax=torch.tensor(1.0, device=device) + ) + + # global_amax=1.0 with NVFP4 _quant_max_bound=6.0 yields per_tensor_scale = 1/6. + sbsq.amax = torch.ones(4, device=device) * 3.0 + sbsq.enable_lsq( + quantize_scales=True, + learnable_amax=["post"], + ) + + x = torch.randn(4, 16, device=device) + out = sbsq._fake_quantize(x) + assert out.shape == x.shape + out.sum().backward() + assert sbsq._amax_post.grad is not None + + +def test_lsq_fp4_cast_ste(): + """Test fp4_cast_ste on GPU.""" + from modelopt.torch.quantization.tensor_quant import fp4_cast_ste + + device = torch.device("cuda") + x = torch.tensor([[-3.0, 1.5, 0.0, 6.0, -6.0, 0.5, -0.5, 2.0]], device=device) + x.requires_grad_(True) + # fp4_cast_ste expects [NUM_BLOCKS, BLOCK_SIZE] -- pad to block size 16 + x_padded = torch.zeros(1, 16, device=device, requires_grad=True) + with torch.no_grad(): + x_padded[:, : x.shape[1]] = x.detach() + x_padded = x_padded.clone().detach().requires_grad_(True) + y = fp4_cast_ste(x_padded) + assert y.shape == x_padded.shape + y.sum().backward() + assert x_padded.grad is not None diff --git a/tests/unit/recipe/test_lsq_recipes.py b/tests/unit/recipe/test_lsq_recipes.py new file mode 100644 index 00000000000..dfc864a4009 --- /dev/null +++ b/tests/unit/recipe/test_lsq_recipes.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for LSQ QAD recipes.""" + +from pathlib import Path + +import pytest + +from modelopt.recipe.loader import load_recipe + +CONFIGS_DIR = Path(__file__).resolve().parents[3] / "modelopt_recipes" / "general" / "qad" + +# filename: (tied_amax, quantize_pre_scale) +_LSQ_RECIPES = { + "nvfp4_dual_lsq-mse_init-fp8_kv.yaml": (False, False), + "nvfp4_lsq-mse_init-fp8_kv.yaml": (True, True), +} + + +def _load_lsq_recipe(filename): + return load_recipe(CONFIGS_DIR / filename).quantize + + +def test_expected_lsq_recipe_files(): + assert {path.name for path in CONFIGS_DIR.glob("*lsq*")} == set(_LSQ_RECIPES) + + +def test_qad_default_nvfp4_recipe_reuses_ptq_recipe(): + recipe = CONFIGS_DIR / "nvfp4_default-kv_fp8.yaml" + + assert recipe.is_symlink() + assert recipe.resolve() == CONFIGS_DIR.parent / "ptq" / recipe.name + assert load_recipe(recipe).recipe_type.value == "ptq" + + +@pytest.mark.parametrize( + ("filename", "expected_tied", "expected_quantize_pre_scale"), + [(filename, *settings) for filename, settings in _LSQ_RECIPES.items()], +) +def test_lsq_recipe_loads_with_expected_algorithm( + filename, expected_tied, expected_quantize_pre_scale +): + algorithm = _load_lsq_recipe(filename).algorithm + + assert algorithm["method"] == "lsq" + assert algorithm["learnable_amax"] == ["pre", "post"] + assert algorithm["tied_amax"] is expected_tied + assert algorithm["quantize_pre_scale"] is expected_quantize_pre_scale + assert algorithm["scale_algorithm"] == {"method": "mse", "fp8_scale_sweep": True} + + +@pytest.mark.parametrize("filename", _LSQ_RECIPES) +def test_lsq_recipe_resolves_modular_quant_cfg(filename): + quantize = _load_lsq_recipe(filename) + entries = {entry.quantizer_name: entry for entry in quantize.quant_cfg} + + weight_cfg = entries["*weight_quantizer"].cfg.model_dump(exclude_unset=True) + input_cfg = entries["*input_quantizer"].cfg.model_dump(exclude_unset=True) + kv_cfg = entries["*[kv]_bmm_quantizer"].cfg.model_dump(exclude_unset=True) + + assert weight_cfg["block_sizes"]["type"] == "static" + assert weight_cfg["num_bits"] == (2, 1) + assert input_cfg["block_sizes"]["type"] == "dynamic" + assert input_cfg["num_bits"] == (2, 1) + assert kv_cfg["num_bits"] == (4, 3) diff --git a/tests/unit/torch/quantization/test_lsq.py b/tests/unit/torch/quantization/test_lsq.py new file mode 100644 index 00000000000..50fd1ffc870 --- /dev/null +++ b/tests/unit/torch/quantization/test_lsq.py @@ -0,0 +1,505 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU unit tests for the LSQ algorithm using INT4 quantization.""" + +import types +from unittest.mock import Mock, create_autospec + +import pytest +import torch +from torch import nn + +import modelopt.torch.quantization.model_calib as model_calib_module +import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module +from modelopt.torch.quantization.config import ( + LocalHessianCalibConfig, + LSQConfig, + MaxCalibConfig, + MseCalibConfig, + QuantizerAttributeConfig, +) +from modelopt.torch.quantization.model_calib import lsq, max_calibrate +from modelopt.torch.quantization.nn import QuantLinear +from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( + _FP8_E4M3_MIN_POSITIVE, + StaticBlockScaleQuantizer, + TensorQuantizer, + _amax_to_scale, +) +from modelopt.torch.quantization.tensor_quant import int_cast_ste +from modelopt.torch.quantization.utils.shared_input import SharedWeightGlobalAmaxState +from modelopt.torch.utils import to_empty_if_meta_device + + +def _make_int4_static_quantizer(): + tq = TensorQuantizer() + tq._num_bits = 4 + tq._unsigned = False + tq._narrow_range = True + tq._disabled = False + tq._block_sizes = {-1: 16} + tq._pass_through_bwd = True + tq.register_buffer("_amax", torch.ones(8)) + return StaticBlockScaleQuantizer.from_tensor_quantizer(tq) + + +def _skip_scale_calibration(monkeypatch): + monkeypatch.setattr( + "modelopt.torch.quantization.model_calib._run_scale_calibration", + lambda *args, **kwargs: None, + ) + + +@pytest.mark.parametrize( + ("num_bits", "expected_dispatch"), + [pytest.param((2, 1), "nvfp4", id="nvfp4"), pytest.param((4, 3), "generic", id="fp8")], +) +def test_non_lsq_static_float_dispatches_only_nvfp4_to_fp4_kernel( + monkeypatch, num_bits, expected_dispatch +): + tq = TensorQuantizer() + tq._num_bits = num_bits + tq._block_sizes = {-1: 16, "type": "static", "scale_bits": (4, 3)} + tq.register_buffer("_amax", torch.ones(4)) + quantizer = StaticBlockScaleQuantizer.from_tensor_quantizer(tq, global_amax=torch.tensor(1.0)) + dispatches = [] + + def fake_nvfp4(inputs, *_args): + dispatches.append("nvfp4") + return inputs + + def fake_generic(_self, inputs): + dispatches.append("generic") + return inputs + + monkeypatch.setattr(tensor_quantizer_module, "static_blockwise_fp4_fake_quant", fake_nvfp4) + monkeypatch.setattr(TensorQuantizer, "_fake_quantize", fake_generic) + + quantizer._fake_quantize(torch.ones(4, 16)) + + assert dispatches == [expected_dispatch] + + +class TestLSQConfig: + """Tests for LSQConfig validation.""" + + def test_default_config(self): + cfg = LSQConfig() + assert cfg.method == "lsq" + assert cfg.learnable_amax == ["post"] + assert cfg.tied_amax is False + assert cfg.quantize_pre_scale is True + assert cfg.scale_algorithm is None + + @pytest.mark.parametrize( + ("method", "config_type"), + [ + ("max", MaxCalibConfig), + ("mse", MseCalibConfig), + ("local_hessian", LocalHessianCalibConfig), + ], + ) + def test_scale_algorithm(self, method, config_type): + cfg = LSQConfig(scale_algorithm={"method": method}) + assert isinstance(cfg.scale_algorithm, config_type) + + def test_unsupported_scale_algorithm(self): + with pytest.raises(ValueError): + LSQConfig(scale_algorithm={"method": "smoothquant"}) + + def test_scale_algorithm_preserves_sparse_dict(self, monkeypatch): + cfg = LSQConfig(scale_algorithm={"method": "mse", "fp8_scale_sweep": True}) + assert cfg.model_dump()["scale_algorithm"] == { + "method": "mse", + "fp8_scale_sweep": True, + } + + calibrate = create_autospec(model_calib_module.mse_calibrate) + monkeypatch.setattr(model_calib_module, "mse_calibrate", calibrate) + model = Mock() + model_calib_module._run_scale_calibration(model, None, cfg.scale_algorithm) + calibrate.assert_called_once_with(model, forward_loop=None, fp8_scale_sweep=True) + + @pytest.mark.parametrize( + ("learnable_amax", "tied_amax"), + [ + (["post"], False), + (["pre"], False), + (["pre", "post"], False), + (["pre", "post"], True), + ([], False), + ([], True), + ("post", False), + ("pre", False), + ], + ) + def test_valid_combinations(self, learnable_amax, tied_amax): + cfg = LSQConfig(learnable_amax=learnable_amax, tied_amax=tied_amax) + assert cfg.tied_amax is tied_amax + + @pytest.mark.parametrize( + "learnable_amax", + [["post"], ["pre"], "post", "pre"], + ) + def test_invalid_tied_with_single_learnable(self, learnable_amax): + with pytest.raises(ValueError, match="tied_amax=True requires"): + LSQConfig(learnable_amax=learnable_amax, tied_amax=True) + + +class TestEnableLSQ: + """Tests for StaticBlockScaleQuantizer.enable_lsq() with INT4 format.""" + + def _make_quantizer(self): + """Create a StaticBlockScaleQuantizer configured for INT4.""" + sbsq = _make_int4_static_quantizer() + assert sbsq._quant_max_bound == 7.0 + return sbsq + + def test_post_only_learnable(self): + q = self._make_quantizer() + q.enable_lsq(quantize_scales=False, learnable_amax=["post"], tied_amax=False) + assert q._lsq is True + assert isinstance(q._amax_post, nn.Parameter) + assert q._amax_post.requires_grad is True + assert not isinstance(q._amax_pre, nn.Parameter) + assert not q._amax_pre.requires_grad + + def test_pre_only_learnable(self): + q = self._make_quantizer() + q.enable_lsq(quantize_scales=False, learnable_amax=["pre"], tied_amax=False) + assert isinstance(q._amax_pre, nn.Parameter) + assert q._amax_pre.requires_grad is True + assert not isinstance(q._amax_post, nn.Parameter) + + def test_both_learnable(self): + q = self._make_quantizer() + q.enable_lsq(quantize_scales=False, learnable_amax=["pre", "post"], tied_amax=False) + assert isinstance(q._amax_pre, nn.Parameter) + assert isinstance(q._amax_post, nn.Parameter) + + def test_tied_both_learnable(self): + q = self._make_quantizer() + q.enable_lsq(quantize_scales=False, learnable_amax=["pre", "post"], tied_amax=True) + assert q._tied_amax is True + assert isinstance(q._amax_post, nn.Parameter) + assert not hasattr(q, "_amax_pre") + assert q.amax_pre is q._amax_post + + def test_frozen(self): + q = self._make_quantizer() + q.enable_lsq(quantize_scales=False, learnable_amax=[], tied_amax=False) + assert not isinstance(q._amax_post, nn.Parameter) + assert not isinstance(q._amax_pre, nn.Parameter) + + def test_old_amax_deleted(self): + q = self._make_quantizer() + assert hasattr(q, "_amax") + q.enable_lsq(quantize_scales=False) + assert not hasattr(q, "_amax") + + def test_can_skip_pre_scale_quantization(self): + q = self._make_quantizer() + q.enable_lsq( + quantize_scales=False, + quantize_pre_scale=False, + ) + assert q._quantize_pre_scale is False + + def test_quantize_scales_without_global_amax_raises(self): + q = self._make_quantizer() + assert q.global_amax is None + with pytest.raises(AssertionError, match="global_amax"): + q.enable_lsq(quantize_scales=True) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_learnable_amax_uses_input_dtype(self, dtype): + q = self._make_quantizer() + q.enable_lsq( + quantize_scales=False, + learnable_amax=["pre", "post"], + dtype=dtype, + ) + + assert q._amax_pre.dtype == dtype + assert q._amax_post.dtype == dtype + + def test_dtype_cast_updates_learnable_amax_dtype(self): + q = self._make_quantizer() + q.enable_lsq( + quantize_scales=False, + learnable_amax=["pre", "post"], + ) + + q.to(dtype=torch.bfloat16) + + assert q._amax_pre.dtype == torch.bfloat16 + assert q._amax_post.dtype == torch.bfloat16 + + def test_to_empty_if_meta_device_materializes_static_amax(self): + q = self._make_quantizer() + q._amax = q._amax.to("meta") + q.global_amax = torch.tensor(1.0, device="meta") + + to_empty_if_meta_device(q, device=torch.device("cpu")) + + assert q._amax.device.type == "cpu" + assert q.global_amax.device.type == "cpu" + + +class TestLSQWeightIteration: + """Tests LSQ conversion for each weight exposed by QuantModule's iterator contract.""" + + def test_multiple_singular_weight_quantizers_use_their_weight_dtypes(self, monkeypatch): + _skip_scale_calibration(monkeypatch) + module = QuantLinear(16, 8, bias=False, dtype=torch.bfloat16) + module.weight_quantizer = _make_int4_static_quantizer() + module.proj = nn.Parameter(torch.ones(8, 16, dtype=torch.float16)) + module.proj_weight_quantizer = _make_int4_static_quantizer() + + lsq(module) + + assert module.weight_quantizer._lsq + assert module.proj_weight_quantizer._lsq + assert module.weight_quantizer._amax_post.dtype == torch.bfloat16 + assert module.proj_weight_quantizer._amax_post.dtype == torch.float16 + + def test_plural_expert_weight_quantizers_enter_lsq(self, monkeypatch): + _skip_scale_calibration(monkeypatch) + module = QuantLinear(16, 8, bias=False) + module.expert_weight = nn.Parameter(torch.ones(2, 8, 16)) + module.expert_weight_quantizers = nn.ModuleList( + [_make_int4_static_quantizer(), _make_int4_static_quantizer()] + ) + + def iter_expert_weights(self): + yield from zip(self.expert_weight, self.expert_weight_quantizers) + + module.iter_weights_for_calibration = types.MethodType(iter_expert_weights, module) + + lsq(module) + + assert all(quantizer._lsq for quantizer in module.expert_weight_quantizers) + + def test_shared_weight_quantizer_enters_lsq_once(self, monkeypatch): + _skip_scale_calibration(monkeypatch) + module = QuantLinear(16, 8, bias=False) + shared_quantizer = _make_int4_static_quantizer() + + def mark_lsq_enabled(*_args, **_kwargs): + shared_quantizer._lsq = True + + shared_quantizer.enable_lsq = Mock(side_effect=mark_lsq_enabled) + module.weight_quantizer = shared_quantizer + module.proj = nn.Parameter(torch.ones(8, 16)) + module.proj_weight_quantizer = shared_quantizer + + lsq(module) + + assert shared_quantizer._lsq + assert shared_quantizer.enable_lsq.call_count == 1 + assert module.weight_quantizer is module.proj_weight_quantizer + + @pytest.mark.parametrize("distributed_sync", [False, True]) + def test_max_calibrate_promotes_static_int_quantizer(self, distributed_sync): + module = QuantLinear(16, 8, bias=False) + config = QuantizerAttributeConfig(num_bits=4, block_sizes={-1: 16, "type": "static"}) + module.weight_quantizer.set_from_attribute_config(config) + module.input_quantizer.set_from_attribute_config(config) + + max_calibrate( + module, + forward_loop=lambda model: model(torch.randn(2, 16)), + distributed_sync=distributed_sync, + ) + + assert isinstance(module.weight_quantizer, StaticBlockScaleQuantizer) + assert module.weight_quantizer.export_amax() is not None + assert not isinstance(module.input_quantizer, StaticBlockScaleQuantizer) + assert module.input_quantizer.amax is not None + + +class TestIntCastSTE: + """Tests for int_cast_ste (INT4 STE function).""" + + def test_round_trip(self): + x = torch.tensor([[-3.2, 1.8, 0.0, 6.5, -7.1]], requires_grad=True) + y = int_cast_ste(x, 4) + assert y.shape == x.shape + max_bound = 7.0 + assert y.min() >= -max_bound + assert y.max() <= max_bound + y.sum().backward() + assert x.grad is not None + + def test_ste_gradient(self): + x = torch.tensor([[2.3, -2.3]], requires_grad=True) + y = int_cast_ste(x, 4) + y.sum().backward() + assert torch.all(x.grad == 1.0) + + +class TestFakeQuantizeLSQ: + """Tests for _fake_quantize() LSQ path with INT4.""" + + def _make_lsq_quantizer(self, learnable_amax=("post",), tied_amax=False): + tq = TensorQuantizer() + tq._num_bits = 4 + tq._unsigned = False + tq._narrow_range = True + tq._disabled = False + tq._block_sizes = {-1: 16} + tq._pass_through_bwd = True + tq.register_buffer("_amax", torch.ones(4) * 3.5) + sbsq = StaticBlockScaleQuantizer.from_tensor_quantizer(tq) + sbsq.enable_lsq(quantize_scales=False, learnable_amax=learnable_amax, tied_amax=tied_amax) + return sbsq + + def test_output_shape(self): + q = self._make_lsq_quantizer() + x = torch.randn(4, 16) + out = q._fake_quantize(x) + assert out.shape == x.shape + + def test_differentiable_post(self): + q = self._make_lsq_quantizer(learnable_amax=["post"]) + x = torch.randn(4, 16) + out = q._fake_quantize(x) + out.sum().backward() + assert q._amax_post.grad is not None + assert q._amax_pre.grad is None + + def test_differentiable_pre(self): + q = self._make_lsq_quantizer(learnable_amax=["pre"]) + x = torch.randn(4, 16) + out = q._fake_quantize(x) + out.sum().backward() + assert q._amax_pre.grad is not None + assert q._amax_post.grad is None + + def test_differentiable_both(self): + q = self._make_lsq_quantizer(learnable_amax=["pre", "post"]) + x = torch.randn(4, 16) + out = q._fake_quantize(x) + out.sum().backward() + assert q._amax_pre.grad is not None + assert q._amax_post.grad is not None + + def test_tied_shares_tensor(self): + q = self._make_lsq_quantizer(learnable_amax=["pre", "post"], tied_amax=True) + x = torch.randn(4, 16) + out = q._fake_quantize(x) + out.sum().backward() + assert q._amax_post.grad is not None + + def test_skip_pre_scale_quantization_still_quantizes_post(self, monkeypatch): + q = self._make_lsq_quantizer() + q._quantize_scales = True + q._quantize_pre_scale = False + # per_tensor_scale of 1.0: INT4 _quant_max_bound is 7.0, so scale = global_amax / 7. + q.global_amax = torch.tensor(float(q._quant_max_bound)) + quantize_flags = [] + orig_block_scale = q._block_scale_from_amax + + def spy_block_scale(amax, quantize): + quantize_flags.append(quantize) + return orig_block_scale(amax, quantize) + + monkeypatch.setattr(q, "_block_scale_from_amax", spy_block_scale) + + out = q._fake_quantize(torch.randn(4, 16)) + + assert out.shape == (4, 16) + # post scale is FP8-quantized, pre scale is not (quantize_pre_scale=False). + assert quantize_flags == [True, False] + + def test_skip_pre_scale_quantization_uses_raw_scale_floor(self, monkeypatch): + q = self._make_lsq_quantizer() + q._quantize_scales = True + q._quantize_pre_scale = False + q.global_amax = torch.tensor(float(q._quant_max_bound)) + min_values = [] + + def fake_amax_to_scale(amax, maxbound, min_value=1e-8): + # Only record the per-block (shape-4) scale calls, not global scale derivation. + if amax.numel() == 4: + min_values.append(min_value) + return torch.ones_like(amax) + + monkeypatch.setattr( + "modelopt.torch.quantization.nn.modules.tensor_quantizer._amax_to_scale", + fake_amax_to_scale, + ) + + out = q._fake_quantize(torch.randn(4, 16)) + + assert out.shape == (4, 16) + assert torch.equal(min_values[0], torch.tensor([_FP8_E4M3_MIN_POSITIVE])) + assert min_values[1] == 1e-8 + + +class TestLSQSharedGlobalAmax: + """Regression: LSQ must honor the shared/tied weight global_amax invariant. + + A q/k/v-style fusible group ties ``_global_amax`` to a single shared buffer object. + Since LSQ derives the per-tensor scale from ``global_amax`` at runtime (no snapshot), + an in-place update of the shared buffer (e.g. export unification) must propagate to + every member. Uses INT4 (FP8-quantized scales) members so the forward runs on CPU; + the shared-buffer mechanism under test is format-agnostic. + """ + + def _make_member(self, amax_value=2.0): + tq = TensorQuantizer() + tq._num_bits = 4 + tq._unsigned = False + tq._narrow_range = True + tq._disabled = False + tq._block_sizes = {-1: 16, "type": "static", "scale_bits": (4, 3)} + tq._pass_through_bwd = True + tq.register_buffer("_amax", torch.ones(4) * amax_value) + return StaticBlockScaleQuantizer.from_tensor_quantizer(tq) + + def _make_tied_lsq_group(self, global_amax=3.0, n_members=3): + members = [self._make_member(amax_value=2.0) for _ in range(n_members)] + state = SharedWeightGlobalAmaxState() + state.global_amax = torch.tensor(float(global_amax)) + for member in members: + assert state.tie_member_quantizer(member) + # All members must alias the single shared buffer object. + assert all(m._global_amax is members[0]._global_amax for m in members) + for member in members: + member.enable_lsq(quantize_scales=True) + return members + + def test_block_scale_tracks_shared_update(self): + members = self._make_tied_lsq_group(global_amax=3.0) + new_value = 5.0 + # Mutate the shared buffer in place, mimicking export unification. + members[0]._global_amax.data.fill_(new_value) + + for member in members: + expected = _amax_to_scale(torch.tensor(new_value), member._quant_max_bound) + scale = member._block_scale_from_amax(member.amax_post, quantize=True) + assert scale.shape == member.amax_post.shape + assert torch.all(scale >= _FP8_E4M3_MIN_POSITIVE * expected) + + def test_members_produce_identical_output_after_shared_update(self): + members = self._make_tied_lsq_group(global_amax=3.0) + members[0]._global_amax.data.fill_(5.0) + + x = torch.randn(4, 16) + outputs = [member._fake_quantize(x) for member in members] + for out in outputs[1:]: + assert torch.equal(out, outputs[0]) diff --git a/tests/unit/torch/quantization/test_mse_calibrator.py b/tests/unit/torch/quantization/test_mse_calibrator.py index df63e51de20..b472d379000 100644 --- a/tests/unit/torch/quantization/test_mse_calibrator.py +++ b/tests/unit/torch/quantization/test_mse_calibrator.py @@ -26,7 +26,7 @@ _register_fp8_sweep_calibrator, mse_calibrate, ) -from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, TensorQuantizer +from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, QuantLinear, TensorQuantizer from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( _QUANT_FUNCTIONAL_BACKENDS, register_quant_backend, @@ -645,7 +645,7 @@ def test_modelopt_static_nvfp4_uses_fp8_scale_sweep(self): ), amax=torch.tensor([1.0, 2.0]), ) - model = torch.nn.Module() + model = QuantLinear(16, 1, bias=False) model.weight_quantizer = q promote_nvfp4_static_quantizers(model) @@ -737,10 +737,9 @@ def forward(self, x): class TestStaticNVFP4Promotion: - class _LinearLike(torch.nn.Module): + class _LinearLike(QuantLinear): def __init__(self, amax): - super().__init__() - self.weight = torch.nn.Parameter(torch.empty(1, 16)) + super().__init__(16, 1, bias=False) cfg = QuantizerAttributeConfig( num_bits=(2, 1), block_sizes={-1: 16, "type": "static", "scale_bits": (4, 3)},