diff --git a/.github/actions/megatron-from-source/action.yml b/.github/actions/megatron-from-source/action.yml new file mode 100644 index 00000000000..50eef77d41d --- /dev/null +++ b/.github/actions/megatron-from-source/action.yml @@ -0,0 +1,27 @@ +name: Megatron-Bridge from source +description: >- + Replace the NeMo container's Megatron-Bridge and Megatron-Core sources with a Megatron-Bridge commit and + its pinned Megatron-LM, so CI covers features newer than the container. Remove once the container has them. + +inputs: + commit: + description: Megatron-Bridge commit. Keep in sync with examples/megatron_bridge/README.md. + default: 1d9caf7324e608aca7181de40be26a6387976210 + +runs: + using: composite + steps: + - shell: bash + run: | + set -euxo pipefail + src=$(mktemp -d) + git -C "$src" init -q + git -C "$src" fetch -q --depth 1 https://github.com/NVIDIA-NeMo/Megatron-Bridge.git "${{ inputs.commit }}" + git -C "$src" checkout -q FETCH_HEAD + git -C "$src" submodule update -q --init --depth 1 3rdparty/Megatron-LM + # Both are editable installs in the NeMo container, so replacing the sources is enough. + for dir in src/megatron/bridge 3rdparty/Megatron-LM/megatron; do + rm -rf "/opt/Megatron-Bridge/$dir" + cp -r "$src/$dir" "/opt/Megatron-Bridge/$dir" + done + rm -rf "$src" diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml index 18979ce0f48..67f6eb7935e 100644 --- a/.github/workflows/_example_tests_runner.yml +++ b/.github/workflows/_example_tests_runner.yml @@ -64,6 +64,8 @@ jobs: run: | echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu:/usr/local/tensorrt/targets/x86_64-linux-gnu/lib" >> $GITHUB_ENV echo "PATH=${PATH}:/usr/local/tensorrt/targets/x86_64-linux-gnu/bin" >> $GITHUB_ENV + - if: contains(inputs.docker_image, 'nvidia/nemo:') + uses: ./.github/actions/megatron-from-source - name: Install dependencies shell: bash run: | diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml index 6340788026f..fd8b8278cbd 100644 --- a/.github/workflows/gpu_tests.yml +++ b/.github/workflows/gpu_tests.yml @@ -90,6 +90,8 @@ jobs: - name: Setup environment variables run: | echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu" >> $GITHUB_ENV + - if: contains(matrix.container_image, 'nvidia/nemo:') + uses: ./.github/actions/megatron-from-source - name: Run gpu tests env: COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1b4c6cbe2be..4200b39beee 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -23,7 +23,7 @@ Changelog - Add ``--sft_in_batch_packing``, ``--moe_aux_loss_coeff`` and ``--dsa_kernel_backend`` to ``examples/megatron_bridge/distill.py`` (``--dsa_kernel_backend`` also to ``quantize.py``): pack ``--sft_hf_dataset`` conversations per micro-batch instead of padding each to ``--seq_length``, override the MoE load-balancing loss weight (e.g. 0 for padded SFT data), and select fused DeepSeek sparse-attention kernels for long sequences. - Add ``--document_boundaries`` to ``examples/megatron_bridge/distill.py`` to attend within each document of a pre-tokenized window, and ``tokenizer`` / ``max_length`` / ``pack_max_tokens`` to ``chat_jsonl`` blends to drop conversations with nothing to train and pair the rest for ``--sft_in_batch_packing``. Packed distillation of Gated DeltaNet models needs ``flash-linear-attention`` >= 0.5.2 (earlier versions compute it wrongly); install ``examples/megatron_bridge/requirements.txt`` in ``nemo:26.08``. - Add ``--num_layers_in_first_pipeline_stage``, ``--num_layers_in_last_pipeline_stage`` and ``--recompute_num_layers_per_stage`` to ``examples/megatron_bridge/distill.py`` to balance pipeline-parallel distillation: give the last stage, which also runs the output layers and KD loss, fewer layers, and recompute fewer layers on later stages. -- Add ``--vllm_fakequant_export`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` to evaluate formats vLLM has no kernels for (e.g. GGML IQ) as BF16 fake-quant checkpoints, and allow ``--ep_size`` > 1 for checkpoints built with ``--no_moe_grouped_gemm``. +- Add ``--vllm_fakequant_export`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` to evaluate PTQ / QAD checkpoints in vLLM for formats it has no kernels for (e.g. NVFP4 on Hopper, GGML IQ): it folds the weight quantizers into BF16 weights and saves through Megatron-Bridge at any TP / EP / PP, LLMs and VLMs, with the activation and KV-cache quantizers in sidecars that ``examples/vllm_serve/vllm_serve_fakequant.py`` loads automatically. Also allow ``--ep_size`` > 1 for checkpoints built with ``--no_moe_grouped_gemm``. - GGML GPTQ checkpoints from ``examples/megatron_bridge`` keep their calibrated codes when the Megatron checkpoint is reloaded for QAD or export, and layerwise calibration now runs on Megatron vision-language models. - The Megatron unified HF export (``export_mcore_gpt_to_hf``) now copies the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged, from local disk only. The ``examples/megatron_bridge`` scripts make a Hugging Face Hub model local first with ``modelopt.torch.export.ensure_local_checkpoint``. @@ -45,6 +45,7 @@ Changelog - The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root. - Fix GPT-OSS draft training with pre-normalization teacher hidden states by enabling the target model's final normalization. - Fix vanilla DSpark exports to use the Qwen3 DSpark architecture, head names and configuration expected by vLLM. Re-export existing training checkpoints to use this format. +- ``mtq.fold_weight`` now folds sequential weight quantizers (e.g. W4A8 INT4 + FP8 weights) into the weights; it used to leave them unfolded and enabled. - Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``. - Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source. - Fix GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU. diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 66683526675..4d67602186c 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -46,6 +46,22 @@ docker run \ ${DOCKER_IMAGE} bash ``` +Some features (e.g. [`--vllm_fakequant_export`](#evaluating-a-qad-checkpoint-in-vllm-with-fake-quantization)) need a newer Megatron-Bridge than the one in `nvcr.io/nvidia/nemo:26.08`. Until a NeMo container includes it, check out Megatron-Bridge at the commit ModelOpt's CI tests against, together with its Megatron-LM submodule: + +```bash +export MBRIDGE_DIR=${PWD}/Megatron-Bridge +git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git ${MBRIDGE_DIR} +git -C ${MBRIDGE_DIR} checkout 1d9caf7324e608aca7181de40be26a6387976210 +git -C ${MBRIDGE_DIR} submodule update --init 3rdparty/Megatron-LM +``` + +and add these mounts to the `docker run` command above: + +```bash + -v ${MBRIDGE_DIR}/src/megatron/bridge:/opt/Megatron-Bridge/src/megatron/bridge \ + -v ${MBRIDGE_DIR}/3rdparty/Megatron-LM/megatron:/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron \ +``` + > [!WARNING] > Use `python -m pip` instead of `pip` to avoid conflicts with the system-wide installed packages in the NeMo containers. You may also refer to this [doc](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docker/common/README.md#installing-packages-inside-the-container) on how to correctly install packages in the NeMo containers without breaking existing torch installation. @@ -371,6 +387,27 @@ sequence length, and GPU memory: measure per-stage time and memory over a few it The distilled checkpoint retains the ModelOpt quantization state, so it can be converted to a deployable HuggingFace checkpoint with [export_quantized_megatron_to_hf.py](export_quantized_megatron_to_hf.py) (point `--megatron_path` at `/output/qwen3_8b_nvfp4_qad/checkpoints`), exactly like the PTQ checkpoint in [step 2 above](#post-training-quantization). +#### Evaluating a QAD checkpoint in vLLM with fake quantization + +To evaluate a format vLLM has no kernels for (e.g. NVFP4 on Hopper, GGML IQ2, or a custom research format), export with `--vllm_fakequant_export`. The weight quantizers are folded into BF16 weights and Megatron-Bridge saves the model at any TP / EP / PP, VLMs included. The enabled activation and KV-cache quantizers go to `quantizer_state.pth` (ranges) and `quant_recipe.yaml` (settings), both under HF names: + +```bash +torchrun --nproc_per_node 8 export_quantized_megatron_to_hf.py \ + --hf_model_name_or_path Qwen/Qwen3-8B \ + --megatron_path /output/qwen3_8b_nvfp4_qad/checkpoints \ + --tp_size 8 \ + --vllm_fakequant_export \ + --export_unified_hf_path /tmp/Qwen3-8B-NVFP4-QAD-vllm-fq +``` + +Serve the export with [vllm_serve_fakequant.py](../vllm_serve/vllm_serve_fakequant.py), which loads both sidecars automatically, then run your evaluations against the server: + +```bash +python ../vllm_serve/vllm_serve_fakequant.py /tmp/Qwen3-8B-NVFP4-QAD-vllm-fq -tp 8 --port 8000 +``` + +Weight-only formats leave no quantizer enabled after folding, so no sidecars are written and stock `vllm serve` loads the export as a plain BF16 checkpoint. This mode needs a Megatron-Bridge that writes the quantizer sidecars; see [Pre-Requisites](#pre-requisites). + ### LoRA QAD Add `--lora_rank` to PTQ to train low-rank weight updates with the usual QAD flow: diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py index 08a58102ce7..3b7871e5ebb 100644 --- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py +++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py @@ -33,12 +33,28 @@ --pp_size 2 \ --export_unified_hf_path /tmp/Qwen3-8B-FP8-hf +With --vllm_fakequant_export, the weight quantizers are instead folded into BF16 weights and the +model is saved through Megatron-Bridge (any TP / EP / PP, LLMs and VLMs). The enabled activation and +KV-cache quantizers go to quantizer_state.pth and quant_recipe.yaml sidecars, which +examples/vllm_serve/vllm_serve_fakequant.py loads to evaluate the checkpoint with vLLM. Formats +without them (e.g. weight-only GGML IQ2) need no sidecars and serve with stock vLLM: + + torchrun --nproc_per_node 2 export_quantized_megatron_to_hf.py \ + --hf_model_name_or_path Qwen/Qwen3-8B \ + --megatron_path /tmp/Qwen3-8B-NVFP4-QAD/checkpoints \ + --tp_size 2 \ + --vllm_fakequant_export \ + --export_unified_hf_path /tmp/Qwen3-8B-NVFP4-QAD-vllm-fq + See `README.md` in this directory for more details. """ import argparse +import json +from pathlib import Path import torch +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge from megatron.bridge.models.hf_pretrained.utils import is_safe_repo from megatron.core.utils import unwrap_model from mlflow_utils import ( @@ -53,10 +69,13 @@ import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as dist from modelopt.torch.export import ensure_local_checkpoint, export_mcore_gpt_to_hf -from modelopt.torch.export.plugins.vllm_fakequant_megatron import export_mcore_gpt_to_hf_vllm_fq +from modelopt.torch.export.plugins.vllm_fakequant_megatron import ( + check_mcore_vllm_fakequant_quantizers, +) from modelopt.torch.utils import print_args, print_rank_0 from modelopt.torch.utils.mlflow import Tool, masked_args from modelopt.torch.utils.plugins.mbridge import ( + is_vlm_config, load_mbridge_model_from_hf, load_modelopt_megatron_checkpoint, use_moe_grouped_gemm, @@ -111,9 +130,10 @@ def get_args() -> argparse.Namespace: "--vllm_fakequant_export", action="store_true", help=( - "Fold the weight quantizers into BF16 weights, for accuracy evaluation in a stock vLLM " - "of formats it has no kernels for (e.g. GGML IQ2). Activation quantizers go to a " - "quantizer_state.pth sidecar." + "Fold the weight quantizers into BF16 weights and save through Megatron-Bridge, for " + "accuracy evaluation in vLLM of formats it has no kernels for (e.g. NVFP4 QAD on " + "Hopper, GGML IQ2). Enabled activation and KV-cache quantizers go to " + "quantizer_state.pth and quant_recipe.yaml sidecars for vllm_serve_fakequant.py." ), ) parser.add_argument( @@ -122,7 +142,12 @@ def get_args() -> argparse.Namespace: help="Export extra modules such as Medusa heads, EAGLE, or MTP.", ) - # Only pipeline and expert parallelism are supported for export + parser.add_argument( + "--tp_size", + type=int, + default=1, + help="Tensor parallel size (--vllm_fakequant_export only; the unified exporter needs 1)", + ) parser.add_argument("--pp_size", type=int, default=1, help="Pipeline parallel size") parser.add_argument( "--ep_size", @@ -157,6 +182,63 @@ def get_args() -> argparse.Namespace: return args +def export_vllm_fakequant(bridge, model: list, export_dir: str) -> None: + """Fold the weight quantizers into BF16 weights and save them with Megatron-Bridge. + + Megatron-Bridge writes the enabled quantizers' ranges and settings under HF names to + ``quantizer_state.pth`` and ``quant_recipe.yaml``, the sidecars ``vllm_serve_fakequant.py`` + loads; it writes none when no quantizer is left enabled (e.g. weight-only formats). + """ + if not hasattr(MegatronModelBridge, "build_modelopt_quantizer_recipe"): + raise RuntimeError( + "--vllm_fakequant_export needs a Megatron-Bridge that writes the quantizer sidecars " + "(quantizer_state.pth and quant_recipe.yaml); see the Pre-Requisites in " + "examples/megatron_bridge/README.md." + ) + for chunk in model: + check_mcore_vllm_fakequant_quantizers(unwrap_model(chunk)) + mtq.fold_weight(unwrap_model(chunk)) + bridge.save_hf_pretrained(model, export_dir) + if dist.is_master(): + # The weights are BF16, so a pre-quantized source's quantization config would misdescribe them. + config_path = Path(export_dir) / "config.json" + config = json.loads(config_path.read_text()) + if config.pop("quantization_config", None) is not None: + config_path.write_text(json.dumps(config, indent=2)) + (Path(export_dir) / "hf_quant_config.json").unlink(missing_ok=True) + dist.barrier() + + +def export_unified_hf(unwrapped_model, args: argparse.Namespace, trust_remote_code: bool) -> None: + """Export the quantized model to a unified HF checkpoint with ModelOpt's Megatron exporter.""" + # Extra modules (Medusa / EAGLE / MTP) only exist on the last pipeline stage. Use an all-reduce + # MAX over all ranks (rather than a broadcast from a hard-coded source rank) so the decision is + # correct regardless of pipeline placement / global rank ordering. + has_extra_modules = hasattr(unwrapped_model, "eagle_module") or hasattr( + unwrapped_model, "medusa_heads" + ) + if torch.distributed.is_initialized(): + flag = torch.tensor( + [int(has_extra_modules)], dtype=torch.int, device=torch.cuda.current_device() + ) + torch.distributed.all_reduce(flag, op=torch.distributed.ReduceOp.MAX) + has_extra_modules = bool(flag.item()) + export_extra_modules = has_extra_modules and args.export_extra_modules + + print_rank_0( + f"Exporting to HuggingFace (unified) checkpoint at {args.export_unified_hf_path}..." + ) + # TODO: Gemma3-VL is not in export_mcore_gpt_to_hf's per-arch mappings yet. + export_mcore_gpt_to_hf( + unwrapped_model, + args.local_checkpoint_path, + export_extra_modules=export_extra_modules, + dtype=torch.bfloat16, + export_dir=args.export_unified_hf_path, + trust_remote_code=trust_remote_code, + ) + + def main(args: argparse.Namespace): # --hf_model_name_or_path stays as given. Later steps read the local copy of the whole # checkpoint in local_checkpoint_path, since exporters read the source's files from local disk @@ -167,9 +249,13 @@ def main(args: argparse.Namespace): trust_remote_code = is_safe_repo( trust_remote_code=args.trust_remote_code, hf_path=args.local_checkpoint_path ) + if not args.vllm_fakequant_export and args.tp_size != 1: + raise ValueError("--tp_size > 1 needs --vllm_fakequant_export.") + if args.vllm_fakequant_export and args.export_extra_modules: + raise ValueError("--export_extra_modules is not supported with --vllm_fakequant_export.") # Build the model structure from HF - _bridge, _provider, model, _unwrapped_model, _tokenizer = load_mbridge_model_from_hf( + bridge, _provider, model, _unwrapped_model, _tokenizer = load_mbridge_model_from_hf( hf_model_name_or_path=args.local_checkpoint_path, trust_remote_code=trust_remote_code, moe_grouped_gemm=use_moe_grouped_gemm( @@ -178,7 +264,7 @@ def main(args: argparse.Namespace): force_sequential=args.no_moe_grouped_gemm, ), provider_overrides={ - "tensor_model_parallel_size": 1, # Tensor parallelism is not supported + "tensor_model_parallel_size": args.tp_size, "pipeline_model_parallel_size": args.pp_size, "expert_model_parallel_size": args.ep_size, "expert_tensor_parallel_size": 1, # Expert tensor parallelism is not supported @@ -187,7 +273,11 @@ def main(args: argparse.Namespace): "pipeline_dtype": torch.bfloat16, }, init_model_parallel=True, - load_weights=False, # The weights come from the Megatron checkpoint, so HF weights are not loaded + # The weights come from the Megatron checkpoint. Megatron-Bridge saves the whole in-memory + # model, and a VLM's QAD checkpoint holds only the language model, so load the HF vision + # tower first. + load_weights=args.vllm_fakequant_export + and is_vlm_config(args.local_checkpoint_path, trust_remote_code=trust_remote_code), ) # Load the quantized checkpoint (with the correct layer spec) rather than reconstructing it from the checkpoint @@ -197,35 +287,11 @@ def main(args: argparse.Namespace): unwrapped_model = unwrap_model(model[0]) mtq.merge_quant_lora(unwrapped_model) - # Extra modules (Medusa / EAGLE / MTP) only exist on the last pipeline stage. Use an all-reduce - # MAX over all ranks (rather than a broadcast from a hard-coded source rank) so the decision is - # correct regardless of pipeline placement / global rank ordering. - has_extra_modules = hasattr(unwrapped_model, "eagle_module") or hasattr( - unwrapped_model, "medusa_heads" - ) - if torch.distributed.is_initialized(): - flag = torch.tensor( - [int(has_extra_modules)], dtype=torch.int, device=torch.cuda.current_device() - ) - torch.distributed.all_reduce(flag, op=torch.distributed.ReduceOp.MAX) - has_extra_modules = bool(flag.item()) - export_extra_modules = has_extra_modules and args.export_extra_modules - - print_rank_0( - f"Exporting to HuggingFace (unified) checkpoint at {args.export_unified_hf_path}..." - ) - # TODO: Gemma3-VL is not in export_mcore_gpt_to_hf's per-arch mappings yet. - export_fn = ( - export_mcore_gpt_to_hf_vllm_fq if args.vllm_fakequant_export else export_mcore_gpt_to_hf - ) - export_fn( - unwrapped_model, - args.local_checkpoint_path, - export_extra_modules=export_extra_modules, - dtype=torch.bfloat16, - export_dir=args.export_unified_hf_path, - trust_remote_code=trust_remote_code, - ) + if args.vllm_fakequant_export: + print_rank_0(f"Exporting vLLM fake-quant checkpoint to {args.export_unified_hf_path}...") + export_vllm_fakequant(bridge, model, args.export_unified_hf_path) + else: + export_unified_hf(unwrapped_model, args, trust_remote_code) args.checkpoint_exported = True print_rank_0(f"Exported HuggingFace checkpoint to {args.export_unified_hf_path}") diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index b958646a140..93a6aae13c9 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -258,6 +258,7 @@ python3 examples/hf_ptq/hf_ptq.py \ Note: `--pyt_ckpt_path` can point to either an HF checkpoint or a ModelOpt-saved checkpoint (e.g., a QAT/QAD checkpoint produced by `examples/llm_qat/train.py`). If the input checkpoint is already quantized, the script will **skip re-quantization** and only export artifacts for vLLM fakequant reload. +- For **Megatron-Bridge** models, export with `export_quantized_megatron_to_hf.py --vllm_fakequant_export` as described in the [Megatron-Bridge example README](../megatron_bridge/README.md#evaluating-a-qad-checkpoint-in-vllm-with-fake-quantization). This generates `quant_recipe.yaml` (versioned, one ModelOpt quantizer state per HF quantizer name) and `quantizer_state.pth` for vLLM fakequant reload. - For **MCore** models, export the model with flag `--export-vllm-fq` as described in [Megatron-LM README](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/post_training/modelopt#-nvfp4-quantization-qauntization-aware-training-and-model-export). This generates `quant_recipe.yaml` and `quantizer_state.pth` for vLLM fakequant reload. Step 2: use the exported artifacts when serving: diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index 30723624739..64c4deb4ed4 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py @@ -15,6 +15,7 @@ """Export Megatron Core Model to HuggingFace vLLM fakequant checkpoint.""" import os +import re import tempfile from collections.abc import Callable, Mapping from contextlib import contextmanager @@ -35,7 +36,7 @@ from modelopt.torch.utils import get_unwrapped_name from modelopt.torch.utils.distributed import DistributedProcessGroup, is_master -__all__ = ["export_mcore_gpt_to_hf_vllm_fq"] +__all__ = ["check_mcore_vllm_fakequant_quantizers", "export_mcore_gpt_to_hf_vllm_fq"] def _quantizer_configs(module: torch.nn.Module) -> tuple[dict[str, dict], str]: @@ -97,6 +98,39 @@ def _quantizer_configs(module: torch.nn.Module) -> tuple[dict[str, dict], str]: return configs, "" +def check_mcore_vllm_fakequant_quantizers(model: torch.nn.Module) -> None: + """Raise if vLLM fake-quant reload cannot reproduce a quantizer of a Megatron model. + + Checks the settings the vLLM fake-quant export and reload support, with weight quantizers + folded into the weights. Collective: a failure on any rank raises on every rank. + """ + failure = "" + tied = any(getattr(m, "share_embeddings_and_output_weights", False) for m in model.modules()) + # MTP and lm_head quantizers are rejected outright, before their individual settings. + for name, quantizer in model.named_modules(): + if not isinstance(quantizer, TensorQuantizer): + continue + active = quantizer.is_enabled and quantizer._if_quant + transform = quantizer.rotate_is_enabled or quantizer.pre_quant_scale is not None + if re.search(r"(^|\.)mtp\.", name) and (active or transform): + failure = f"MTP quantization is not supported by vLLM fakequant export/reload: {name}" + elif re.search(r"(^|\.)output_layer\.", name) and ( + (active and (tied or "weight_quantizer" not in name)) or (tied and transform) + ): + failure = f"Unsupported lm_head quantization for vLLM fakequant: {name}" + if failure: + break + if not failure: + _, failure = _quantizer_configs(model) + failure = DistributedProcessGroup.get_dist_syncd_obj( + failure, + DistributedProcessGroup(None), + lambda failures: next((message for message in failures if message), ""), + ) + if failure: + raise ValueError(failure) + + def _save_quantizer_state(path: Path, save: Callable[[Path], object]) -> None: """Save quantizer state or recipes and share write completion or failure across ranks.""" failure = "" @@ -226,24 +260,8 @@ def __init__(self, *args, **kwargs): self._grouped_fq_modules: set[torch.nn.Module] = set() def _gather_exclude_modules(self) -> list[str]: - """Validate settings before gathering excluded modules.""" - # Tied embeddings bypass the output-layer route, so validate the live head here too. - head = getattr(self.model, "output_layer", None) - if head is not None: - for name, quantizer in head.named_modules(): - if not isinstance(quantizer, TensorQuantizer): - continue - tied = self.model.share_embeddings_and_output_weights - active = quantizer.is_enabled and quantizer._if_quant - if (active and (tied or "weight_quantizer" not in name)) or ( - tied and (quantizer.rotate_is_enabled or quantizer.pre_quant_scale is not None) - ): - self._quantizer_validation_failure = ( - self._quantizer_validation_failure - or f"Unsupported lm_head quantization for vLLM fakequant: {name}" - ) - break - # All ranks reach this after MTP collection and before writing weights. + """Raise export-time failures (e.g. conflicting packed recipes) before gathering.""" + # All ranks reach this after collecting the shards and before writing weights. failure = DistributedProcessGroup.get_dist_syncd_obj( self._quantizer_validation_failure, DistributedProcessGroup(None), @@ -439,23 +457,10 @@ def _get_quantized_state( """ name_to_value = {} qformat: str = self._get_quantization_format(module) + # MTP quantizers are rejected up front (check_mcore_vllm_fakequant_quantizers). if prefix.startswith("mtp."): - for name, quantizer in module.named_modules(): - if isinstance(quantizer, TensorQuantizer) and ( - (quantizer.is_enabled and quantizer._if_quant) - or quantizer.rotate_is_enabled - or quantizer.pre_quant_scale is not None - ): - self._quantizer_validation_failure = self._quantizer_validation_failure or ( - f"MTP quantization is not supported by vLLM fakequant export/reload: {name}" - ) - break - return self._get_weight_bias(module, dtype), qformat, 0 - quantizer_configs, error = _quantizer_configs(module) - if error: - self._quantizer_validation_failure = self._quantizer_validation_failure or error - # Skip unsupported kernels while peers complete the export collectives. return self._get_weight_bias(module, dtype), qformat, 0 + quantizer_configs, _ = _quantizer_configs(module) source_prefix = prefix if not prefix or prefix.endswith(".") else prefix + "." for qname, qstate in quantizer_configs.items(): marker_id = len(self._quantizer_recipe_markers) @@ -540,6 +545,7 @@ def export_mcore_gpt_to_hf_vllm_fq( Raises: ValueError: If a quantizer uses settings that cannot be restored. """ + check_mcore_vllm_fakequant_quantizers(model) exporter = VllmFqGPTModelExporter( model, pretrained_model_name_or_path, diff --git a/modelopt/torch/quantization/nn/modules/quant_linear.py b/modelopt/torch/quantization/nn/modules/quant_linear.py index d0569ff5b71..a817b3025cd 100644 --- a/modelopt/torch/quantization/nn/modules/quant_linear.py +++ b/modelopt/torch/quantization/nn/modules/quant_linear.py @@ -171,26 +171,14 @@ def fold_weight(self, keep_attrs: bool = False): and isinstance(self.weight_quantizer, TensorQuantizer) and self.weight_quantizer.fake_quant ) + # Read the LoRA factors first: folding drops the weight quantizer's buffers. + lora_a = lora_b = None + if should_fold and not keep_attrs and self._not_sequential_quantizers(): + lora_a = self.weight_quantizer.svdquant_lora_a + lora_b = self.weight_quantizer.svdquant_lora_b super().fold_weight(keep_attrs) - if should_fold: - if ( - not keep_attrs - and self._not_sequential_quantizers() - and self.weight_quantizer.svdquant_lora_a is not None - and self.weight_quantizer.svdquant_lora_b is not None - ): - self.weight.data.copy_( - self.weight - + self.weight_quantizer.svdquant_lora_b @ self.weight_quantizer.svdquant_lora_a - ) - if not keep_attrs: - _attrs = [ - "_svdquant_lora_a", - "_svdquant_lora_b", - ] - for attr in _attrs: - if hasattr(self.weight_quantizer, attr): - delattr(self.weight_quantizer, attr) + if lora_a is not None and lora_b is not None: + self.weight.data.copy_(self.weight + lora_b @ lora_a) class RealQuantLinear(QuantModule): diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py index 8b1d8a4d349..903a44d7b2e 100644 --- a/modelopt/torch/quantization/nn/modules/quant_module.py +++ b/modelopt/torch/quantization/nn/modules/quant_module.py @@ -142,7 +142,7 @@ def iter_weight_quantizers_for_calibration(self): @staticmethod @torch.no_grad() def _fold_weight_quantizer( - quantizer: TensorQuantizer, + quantizer: TensorQuantizer | SequentialQuantizer, weights: Iterable[torch.Tensor], keep_attrs: bool = False, ): @@ -150,23 +150,32 @@ def _fold_weight_quantizer( ``weights`` is an iterable so a single quantizer shared across views (e.g. one per-tensor quantizer over all experts of a fused MoE weight) can be folded view by - view while disabling and dropping its calibration attrs exactly once. + view while disabling and dropping its calibration attrs exactly once. A + :class:`SequentialQuantizer` is folded as its whole chain. """ if not quantizer.fake_quant: return + members = list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer] for weight in weights: - weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype)) + folded = weight + # Each member's output is in the weight dtype, as in the forward pass. + for member in members: + folded = member(folded.float().contiguous()).to(weight.dtype) + weight.data.copy_(folded) quantizer.disable() quantizer.disable_rotate() - if keep_attrs and hasattr(quantizer, "_pre_quant_scale"): - # The scale is already baked into the folded weight. - # Disable pre-quant scaling so it is not applied twice. - quantizer._enable_pre_quant_scale = False - elif not keep_attrs: - for attr_name in ("_pre_quant_scale", "_amax"): - if hasattr(quantizer, attr_name): - delattr(quantizer, attr_name) + for member in members: + if keep_attrs and hasattr(member, "_pre_quant_scale"): + # The scale is already baked into the folded weight. + # Disable pre-quant scaling so it is not applied twice. + member._enable_pre_quant_scale = False + elif not keep_attrs: + # Each of these only re-derives the folded weight (e.g. a pinned GGML payload). + buffers = {name for name, _ in member.named_buffers(recurse=False)} + for attr_name in buffers | {"_pre_quant_scale", "_amax"}: + if hasattr(member, attr_name): + delattr(member, attr_name) def _validate_fold_weight(self): """Validate that weight folding is supported before any module is modified.""" @@ -176,17 +185,17 @@ def fold_weight(self, keep_attrs: bool = False): Every fake-quant weight quantizer is folded, including disabled quantizers whose pre-quant scale or rotation remains active. The folded transform is baked into the stored weight and - then disabled, so subsequent forwards use the stored weight directly. Calibration buffers - (``_pre_quant_scale``, ``_amax``) are dropped unless ``keep_attrs``. A retained pre-quant - scale remains stored but is made inactive because it has already been applied to the folded - weight. + then disabled, so subsequent forwards use the stored weight directly. The quantizer's + buffers (e.g. ``_pre_quant_scale``, ``_amax``) are dropped unless ``keep_attrs``. A + retained pre-quant scale remains stored but is made inactive because it has already been + applied to the folded weight. """ # Handle all attributes that end with _weight_quantizer for name in dir(self): attr = getattr(self, name) if ( name.endswith("weight_quantizer") - and isinstance(attr, TensorQuantizer) + and isinstance(attr, (TensorQuantizer, SequentialQuantizer)) and attr.fake_quant ): # Get the corresponding weight name by removing _weight_quantizer suffix diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index a96e7fd2a33..84fd26cc4a3 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -372,8 +372,12 @@ def _output_layer_untied(config) -> bool: return False -# Statement kinds of the upstream GPTModel.sharded_state_dict body patched below. -_GPT_SSD_STATEMENTS = ["Assign", "Assign", "Assign", "Assert", "Return"] +# Statement kinds of the upstream GPTModel.sharded_state_dict bodies patched below. Newer +# megatron-core also accepts an empty tensor placeholder (one more assignment). +_GPT_SSD_STATEMENTS = ( + ["Assign", "Assign", "Assign", "Assert", "Return"], + ["Assign", "Assign", "Assign", "Assign", "Assert", "Return"], +) def _output_layer_extra_state_has_data(entry: Any) -> bool: @@ -402,7 +406,7 @@ def keep_gpt_output_layer_extra_state() -> bool: body = func.body if isinstance(func, ast.FunctionDef) else [] if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant): body = body[1:] # docstring - if [type(stmt).__name__ for stmt in body] != _GPT_SSD_STATEMENTS: + if [type(stmt).__name__ for stmt in body] not in _GPT_SSD_STATEMENTS: warn_rank_0( "GPTModel.sharded_state_dict is not the version ModelOpt patches; leaving it as is. " "If it does not keep a populated output_layer._extra_state, saving a quantized " diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py index c2ae02a7f93..c7517eb559c 100644 --- a/modelopt/torch/quantization/plugins/transformer_engine.py +++ b/modelopt/torch/quantization/plugins/transformer_engine.py @@ -247,9 +247,12 @@ def iter_weights_for_calibration(self): def fold_weight(self, keep_attrs: bool = False): """Fold each grouped weight with its corresponding fake-quant quantizer.""" - quantizer_weights: dict[TensorQuantizer, list[torch.Tensor]] = {} + quantizer_weights: dict[TensorQuantizer | SequentialQuantizer, list[torch.Tensor]] = {} for weight, quantizer in self.iter_weights_for_calibration(): - if isinstance(quantizer, TensorQuantizer) and quantizer.fake_quant: + if ( + isinstance(quantizer, (TensorQuantizer, SequentialQuantizer)) + and quantizer.fake_quant + ): quantizer_weights.setdefault(quantizer, []).append(weight) for quantizer, weights in quantizer_weights.items(): diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py index b3df875f3ce..e9bc834520f 100644 --- a/tests/examples/megatron_bridge/test_qad.py +++ b/tests/examples/megatron_bridge/test_qad.py @@ -21,6 +21,7 @@ import pytest import torch import torch.distributed.checkpoint as dcp +import yaml from _test_utils.examples.run_command import extend_cmd_parts, run_example_command from _test_utils.torch.export.unified_checkpoint import ( assert_exported_checkpoint_matches, @@ -244,3 +245,31 @@ def test_qad(tmp_path: Path, num_gpus, create_student, lora_rank): if is_moe: # allow_unexpected waives the expert names wholesale; this re-tightens it. assert_per_expert_experts_complete(hf_export_path) + + # Step 4: export for vLLM fake-quant evaluation (test_quantize_export.py covers the rest). + # Only QAD saves a VLM's language model alone; the export must still carry the vision tower. + if is_vlm: + fakequant_path = tmp_path / "qad_vllm_fakequant" + fakequant_cmd = extend_cmd_parts( + [ + "torchrun", + f"--nproc_per_node={num_gpus}", + "export_quantized_megatron_to_hf.py", + "--vllm_fakequant_export", + ], + hf_model_name_or_path=hf_model_path, + megatron_path=distilled_megatron_path, + export_unified_hf_path=fakequant_path, + tp_size=num_gpus, + ) + run_example_command(fakequant_cmd, example_path="megatron_bridge", setup_free_port=True) + assert_exported_checkpoint_matches( + fakequant_path, + hf_model_path, + check_values=False, + bit_exact_prefixes=("model.visual.",), + ) + recipe = yaml.safe_load((fakequant_path / "quant_recipe.yaml").read_text()) + state = torch.load(fakequant_path / "quantizer_state.pth", weights_only=True) + assert any(name.endswith("k_bmm_quantizer") for name in recipe["quantizers"]) + assert {key.rpartition(".")[0] for key in state} <= recipe["quantizers"].keys() diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py index a60fb417acd..d031990e76a 100644 --- a/tests/examples/megatron_bridge/test_quantize_export.py +++ b/tests/examples/megatron_bridge/test_quantize_export.py @@ -17,8 +17,13 @@ from pathlib import Path import pytest +import torch +import yaml from _test_utils.examples.run_command import extend_cmd_parts, run_example_command -from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches +from _test_utils.torch.export.unified_checkpoint import ( + assert_exported_checkpoint_matches, + load_safetensors_dir, +) from _test_utils.torch.megatron.modelopt_state import assert_has_modelopt_state from _test_utils.torch.transformers_models import ( create_tiny_nemotron_h_dir, @@ -41,7 +46,14 @@ @pytest.mark.parametrize( - ("create_model", "model_kwargs", "recipe", "quantize_parallelism", "export_parallelism"), + ( + "create_model", + "model_kwargs", + "recipe", + "quantize_parallelism", + "export_parallelism", + "fakequant_export_parallelism", + ), [ # MoE: routed experts used to be dropped silently from the export. Calibrates with context # parallelism (not on the VLM: M-RoPE CP needs nemo:26.10). @@ -51,17 +63,33 @@ "general/ptq/nvfp4_default-kv_fp8", "cp_size", "pp_size", + # Per-expert HF names, each EP rank owning some of them. + "ep_size", ), # Dense VLM: only the language model is quantized, vision is copied through. Keeps the # FP8 script-to-checkpoint coverage. - (create_tiny_qwen3vl_dir, {}, "general/ptq/fp8_default-kv_fp8", "tp_size", "pp_size"), + ( + create_tiny_qwen3vl_dir, + {}, + "general/ptq/fp8_default-kv_fp8", + "tp_size", + "pp_size", + "tp_size", + ), # Mamba hybrid + MoE with grouped-GEMM experts: exported with its experts sharded # across ranks (EP), resharded from the TP-quantized checkpoint. - (create_tiny_nemotron_h_dir, {}, "general/ptq/nvfp4_default-kv_fp8", "tp_size", "ep_size"), + ( + create_tiny_nemotron_h_dir, + {}, + "general/ptq/nvfp4_default-kv_fp8", + "tp_size", + "ep_size", + "pp_size", + ), ], ids=["qwen3_moe", "qwen3vl", "nemotron_h"], ) -@pytest.mark.timeout(360) # quantize + export in one test; 1-gpu CI exceeds the default 300s +@pytest.mark.timeout(420) # quantize + 2 exports in one test; 1-gpu CI exceeds the default 300s def test_quantize_and_export( tmp_path: Path, num_gpus, @@ -70,6 +98,7 @@ def test_quantize_and_export( recipe, quantize_parallelism, export_parallelism, + fakequant_export_parallelism, fake_hub_checkpoint, ): """Quantize a tiny model via a YAML recipe and export it to a unified HF checkpoint. @@ -134,3 +163,32 @@ def test_quantize_and_export( # ) # outputs = llm.generate(["Hello!"], vllm.SamplingParams(max_tokens=4)) # assert outputs and outputs[0].outputs and outputs[0].outputs[0].text + + # Step 3: export folded weights + quantizer sidecars through Megatron-Bridge for vLLM fake quant + fakequant_path = tmp_path / "vllm_fakequant" + fakequant_cmd = extend_cmd_parts( + [ + "torchrun", + f"--nproc_per_node={num_gpus}", + "export_quantized_megatron_to_hf.py", + "--vllm_fakequant_export", + ], + hf_model_name_or_path=model_arg, + megatron_path=megatron_path, + export_unified_hf_path=fakequant_path, + **{fakequant_export_parallelism: num_gpus}, + ) + run_example_command(fakequant_cmd, example_path="megatron_bridge", setup_free_port=True) + is_vlm = "vision_config" in (hf_model_path / "config.json").read_text() + assert_exported_checkpoint_matches( + fakequant_path, + hf_model_path, + check_values=False, + bit_exact_prefixes=("model.visual.",) if is_vlm else (), + ) + assert not any("quantizer" in name for name in load_safetensors_dir(fakequant_path)) + recipe = yaml.safe_load((fakequant_path / "quant_recipe.yaml").read_text()) + quantizers = recipe["quantizers"] + state = torch.load(fakequant_path / "quantizer_state.pth", weights_only=True) + assert any(name.endswith("k_bmm_quantizer") for name in quantizers) + assert {key.rpartition(".")[0] for key in state} <= quantizers.keys() diff --git a/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py b/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py new file mode 100644 index 00000000000..0eab41a1f63 --- /dev/null +++ b/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. +"""vLLM fake-quant export through Megatron-Bridge matches ModelOpt's Megatron exporter.""" + +import os +from functools import partial +from pathlib import Path + +import pytest +import torch +import yaml +from _test_utils.torch.transformers_models import create_tiny_nemotron_h_dir, create_tiny_qwen3_dir +from megatron.core.utils import unwrap_model +from safetensors.torch import load_file + +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.export import export_mcore_gpt_to_hf_vllm_fq +from modelopt.torch.export.plugins.vllm_fakequant_hf import ( + QUANTIZER_RECIPE_FORMAT, + QUANTIZER_RECIPE_VERSION, +) +from modelopt.torch.export.plugins.vllm_fakequant_megatron import ( + check_mcore_vllm_fakequant_quantizers, +) +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.utils.plugins.mbridge import load_mbridge_model_from_hf +from modelopt.torch.utils.plugins.megatron_generate import megatron_prefill + + +def _load_export(export_dir: Path): + weights = {} + for shard in sorted(export_dir.glob("*.safetensors")): + weights.update(load_file(shard)) + state = torch.load(export_dir / "quantizer_state.pth", weights_only=True) + recipe = yaml.safe_load((export_dir / "quant_recipe.yaml").read_text()) + assert recipe["format"] == QUANTIZER_RECIPE_FORMAT + assert recipe["version"] == QUANTIZER_RECIPE_VERSION + return weights, state, recipe["quantizers"] + + +QUANT_CFGS = { + "nvfp4-kv_fp8": load_recipe("general/ptq/nvfp4_default-kv_fp8").quantize.model_dump(), + # INT4 + FP8 weights: a SequentialQuantizer chain the export folds as a whole. + "w4a8": {**mtq.W4A8_AWQ_BETA_CFG, "algorithm": "max"}, +} + + +def _test_matches_megatron_exporter(source, quant_cfg_name, tmp_path, rank, size): + # Models in a reused worker pick their attention backend through process-wide variables. + for name in ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN"): + os.environ.pop(name, None) + bridge, _, model, unwrapped, _ = load_mbridge_model_from_hf( + hf_model_name_or_path=str(source), provider_overrides={"pipeline_dtype": torch.bfloat16} + ) + torch.manual_seed(0) + input_ids = torch.randint(0, unwrapped.vocab_size, (2, 16), device="cuda") + mtq.quantize(unwrapped, QUANT_CFGS[quant_cfg_name], lambda m: megatron_prefill(m, input_ids)) + reference_dir, export_dir = tmp_path / "megatron_exporter", tmp_path / "bridge" + export_mcore_gpt_to_hf_vllm_fq(unwrapped, str(source), export_dir=str(reference_dir)) + for chunk in model: + check_mcore_vllm_fakequant_quantizers(unwrap_model(chunk)) + mtq.fold_weight(unwrap_model(chunk)) + bridge.save_hf_pretrained(model, str(export_dir)) + if rank != 0: + return + + ref_weights, ref_state, ref_recipe = _load_export(reference_dir) + weights, state, quantizers = _load_export(export_dir) + assert weights.keys() == ref_weights.keys() + for name, weight in weights.items(): + torch.testing.assert_close(weight, ref_weights[name], rtol=0, atol=0, msg=name) + + enabled = {name for name, entry in ref_recipe.items() if not entry["_disabled"]} + assert quantizers.keys() == enabled + assert any(name.endswith("k_bmm_quantizer") for name in quantizers) == quant_cfg_name.endswith( + "kv_fp8" + ) + to_cfg = TensorQuantizer.attribute_config_from_modelopt_state + for name, entry in quantizers.items(): + assert to_cfg(ref_recipe[name]).items() <= to_cfg(entry).items(), name + + assert state.keys() == {key for key in ref_state if key.rpartition(".")[0] in quantizers} + for key, value in state.items(): + torch.testing.assert_close(value, ref_state[key], rtol=0, atol=0, msg=key) + + +@pytest.mark.parametrize( + ("create_source", "quant_cfg_name"), + [(create_tiny_nemotron_h_dir, "nvfp4-kv_fp8"), (create_tiny_qwen3_dir, "w4a8")], +) +def test_matches_megatron_exporter(dist_workers_size_1, tmp_path, create_source, quant_cfg_name): + source = create_source(tmp_path, with_tokenizer=True) + dist_workers_size_1.run( + partial(_test_matches_megatron_exporter, source, quant_cfg_name, tmp_path) + ) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 9e2a7be8a0b..c53a87276e6 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -691,6 +691,7 @@ def _test_fp8_real_quantize_helper(rank, size): assert real_quant_mem < cur_mem +@skip_flaky_on_blackwell def test_fp8_real_quantize(dist_workers): dist_workers.run(_test_fp8_real_quantize_helper) @@ -1725,6 +1726,7 @@ def forward_loop(_model): assert all(m.input_quantizer.amax is not None for m in calibrated) +@skip_flaky_on_blackwell @pytest.mark.parametrize("drive_wrapper", [False, True]) def test_layerwise_calibrate_vlm_text_forward_loop(dist_workers_size_1, drive_wrapper): dist_workers_size_1.run(partial(_test_layerwise_calibrate_vlm_text_forward_loop, drive_wrapper)) diff --git a/tests/unit/torch/quantization/test_ggml_gptq.py b/tests/unit/torch/quantization/test_ggml_gptq.py index 1106d330e3b..fc5a01ca080 100644 --- a/tests/unit/torch/quantization/test_ggml_gptq.py +++ b/tests/unit/torch/quantization/test_ggml_gptq.py @@ -319,6 +319,8 @@ def test_folding_a_bf16_weight_decodes_the_pinned_payload(): assert torch.equal(model.weight_quantizer(model.weight.float()).to(torch.bfloat16), decoded) mtq.fold_weight(model) assert torch.equal(model.weight, decoded) + # The folded weight is the decoding, so the payload is dropped with the other quantizer state. + assert _pin(model) is None def test_pin_survives_save_and_restore(): diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py index 31352e42bb5..81191fd6ae4 100644 --- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py @@ -26,6 +26,7 @@ import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module from modelopt.torch.quantization import QuantModuleRegistry, tensor_quant from modelopt.torch.quantization.config import QuantizerAttributeConfig, RotateConfig +from modelopt.torch.quantization.conversion import create_and_replace_svdquant_linear_on_the_fly from modelopt.torch.quantization.nn import ( SequentialQuantizer, TensorQuantizer, @@ -403,6 +404,47 @@ def test_fold_weight_keep_attrs_keeps_calibration_attrs_inactive(monkeypatch): unregister_quant_backend(backend_name) +def test_fold_weight_folds_svdquant_lora(): + model = torch.nn.Sequential(torch.nn.Linear(4, 3)) + inputs = torch.randn(2, 4) + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + ], + "algorithm": "max", + } + mtq.quantize(model, config, lambda m: m(inputs)) + create_and_replace_svdquant_linear_on_the_fly(model) + model[0].weight_quantizer.svdquant_lora_a = torch.randn(2, 4) + model[0].weight_quantizer.svdquant_lora_b = torch.randn(3, 2) + output_before = model(inputs) + + mtq.fold_weight(model) + + assert not hasattr(model[0].weight_quantizer, "_svdquant_lora_a") + assert torch.allclose(model(inputs), output_before, atol=1e-6) + + +@torch.no_grad() +def test_fold_weight_folds_sequential_weight_quantizer(): + model = SimpleLinear() + inputs = model.get_input() + mtq.quantize(model, {**WINT4INT8_CFG, "algorithm": "max"}, lambda m: m(inputs)) + expected = { + module: module.weight_quantizer(module.weight) + for module in model.modules() + if isinstance(getattr(module, "weight_quantizer", None), SequentialQuantizer) + } + assert expected + + mtq.fold_weight(model) + + for module, weight in expected.items(): + assert torch.equal(module.weight, weight) + assert not any(q.is_enabled or hasattr(q, "_amax") for q in module.weight_quantizer) + + def test_temporarily_fold_weights_restores_after_exception(monkeypatch): calls = [] backend_name = "test_temporary_fold_backend"