From 97a2003db088a81e72080fffd064f9de7b9bd025 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 8 Oct 2026 16:53:42 +0000 Subject: [PATCH 1/6] Add upcast_ggml="bf16" export for evaluating GGML checkpoints in vLLM export_hf_checkpoint(upcast_ggml="bf16") (--upcast_ggml bf16 in hf_ptq.py) writes each GGML-quantized weight as the BF16 values its packed blocks decode to, disables its quantizer, and so leaves the layer out of the quantization config. Runtimes without GGML kernels, such as vLLM, then load the checkpoint as a plain BF16 model. Also drop the payload GPTQ pins to a GGML weight quantizer from the exported state dict: it was written as a stray "*.weight_quantizer._ggml_pinned_" tensor, which loaders reject. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 1 + docs/source/deployment/3_unified_hf.rst | 7 ++ examples/hf_ptq/README.md | 2 + examples/hf_ptq/hf_ptq.py | 21 ++++++ modelopt/torch/export/quant_utils.py | 2 + modelopt/torch/export/unified_export_hf.py | 67 ++++++++++++++++++- tests/unit/torch/export/test_export_weight.py | 48 ++++++++++++- .../torch/export/test_unified_export_hf.py | 53 +++++++++++++++ 8 files changed, 198 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d3cc67d95bf..b3ba815b903 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -14,6 +14,7 @@ Changelog *Quantization* - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. +- Add ``upcast_ggml="bf16"`` to ``export_hf_checkpoint`` (``--upcast_ggml bf16`` in ``examples/hf_ptq/hf_ptq.py``): GGML-quantized weights are exported as the BF16 values their packed blocks decode to, with no quantization config, so runtimes without GGML kernels such as vLLM can evaluate the quantized model. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. - Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body. diff --git a/docs/source/deployment/3_unified_hf.rst b/docs/source/deployment/3_unified_hf.rst index 7791848b0e2..07210aa13d5 100644 --- a/docs/source/deployment/3_unified_hf.rst +++ b/docs/source/deployment/3_unified_hf.rst @@ -89,6 +89,13 @@ The generated configuration records ``quant_method: modelopt``, ``packing: ggml` block size, and the payload byte count. GGML payloads are not represented as compressed-tensors integer ``weights`` groups because all scales and indices are embedded in each packed block. +To evaluate a GGML-quantized model in a runtime without GGML kernels, such as vLLM, pass +``upcast_ggml="bf16"`` to ``export_hf_checkpoint`` (``--upcast_ggml bf16`` in +``examples/hf_ptq/hf_ptq.py``). Each GGML weight is then written as the BF16 values its packed +blocks decode to, and those layers are left out of the quantization configuration, so the +checkpoint loads as an ordinary unquantized one. It is no smaller than the original model: use it +for accuracy evaluation, not deployment. + Each 74-byte IQ2_XS block represents 256 logical weights: * bytes 0--1 are the little-endian FP16 super-block scale ``d``; diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 4be9e35ece3..d610636ad34 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -672,6 +672,8 @@ python hf_ptq.py --pyt_ckpt_path --qformat fp8 --export ``` > *For exporting fake-quantized models for vLLM serving (e.g., for research or kernels not yet supported in real-quant), use the `--vllm_fakequant_export` flag. See [vllm_serve/README.md](../vllm_serve/README.md) for details.* +> +> *To evaluate a GGML-quantized model (IQ formats, Q8_0) in stock vLLM, add `--upcast_ggml bf16`: the GGML weights are exported as the BF16 values their packed blocks decode to, with no quantization config.* ### Hugging Face framework [Script](./scripts/huggingface_example.sh) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index db29d5fbf61..124f763dd42 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -716,6 +716,7 @@ def export_quantized( export_hf_checkpoint( full_model, export_dir=export_path, + upcast_ggml=args.upcast_ggml, ) if args.qformat == "w4a16_nvfp4": @@ -961,6 +962,11 @@ def quantize_main( "layerwise.export_dir is not supported with an AutoQuantize recipe; " "use a PTQ recipe, or drop export_dir and export afterwards." ) + if args.upcast_ggml: + # Fail before calibration rather than after: the layer shards are written packed. + raise NotImplementedError( + "--upcast_ggml is not supported with layerwise.export_dir; drop export_dir." + ) if not args.skip_generate: print("Layerwise export: forcing --skip_generate, the model is left in export form.") args.skip_generate = True @@ -1350,6 +1356,16 @@ def parse_args() -> argparse.Namespace: help="Export as vLLM fake-quant checkpoint (produces vllm_fq_modelopt_state.pth " "for use with vllm_serve_fakequant.py).", ) + parser.add_argument( + "--upcast_ggml", + default=None, + choices=["bf16"], + help=( + "Export GGML-quantized weights (IQ formats, Q8_0) as the values their packed blocks " + "decode to in this dtype, rather than as the blocks, so runtimes without GGML kernels " + "such as vLLM can evaluate the checkpoint." + ), + ) parser.add_argument( "--cast_mxfp4_to_nvfp4", action="store_true", @@ -1432,6 +1448,11 @@ def parse_args() -> argparse.Namespace: parser.error(f"--use_fsdp2 does not support --sparsity_fmt {args.sparsity_fmt}.") if args.use_fsdp2 and args.vllm_fakequant_export: parser.error("--use_fsdp2 does not support --vllm_fakequant_export.") + if args.upcast_ggml and args.vllm_fakequant_export: + parser.error( + "--upcast_ggml applies to the unified HF export; --vllm_fakequant_export already " + "writes fake-quantized weights." + ) if args.use_fsdp2 and args.cast_mxfp4_to_nvfp4: parser.error("--use_fsdp2 does not support --cast_mxfp4_to_nvfp4.") diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 6450c1afc74..eeeccfd43b6 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -1003,6 +1003,8 @@ def from_quantized_weight( "_bias_value", "input_quantizer._pre_quant_scale", "weight_shape", + # The GGML payload GPTQ pins to a weight quantizer; the exported weight already holds it. + "_ggml_pinned_", ) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index f6116f4746a..d7dfee85fc2 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -69,7 +69,7 @@ from modelopt.torch.opt.plugins.huggingface import _MODELOPT_STATE_SAVE_NAME from modelopt.torch.quantization import set_quantizer_by_cfg_context from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY -from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer +from modelopt.torch.quantization.nn import QuantModule, SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper from modelopt.torch.quantization.qtensor.nvfp4_tensor import _cast_per_block_scale_to_fp8 @@ -78,7 +78,10 @@ module_name_maps, quantizer_attr_names, ) -from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload +from modelopt.torch.quantization.utils.core_utils import ( + enable_weight_access_and_writeback, + has_accelerate_offload, +) from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.dataset_utils import _disable_use_cache from modelopt.torch.utils.distributed import is_fsdp2_model @@ -1852,6 +1855,47 @@ def read_unplaced_weights(model: nn.Module, *, keys_only: bool = False) -> dict[ return out +# What export_hf_checkpoint(upcast_ggml=...) can decode GGML weights to. +_GGML_UPCAST_DTYPES: dict[str, torch.dtype] = {"bf16": torch.bfloat16} + + +def _is_ggml_weight_quantizer(quantizer) -> bool: + # getattr: a SequentialQuantizer has no num_bits, and is never GGML. + return getattr(quantizer, "num_bits", None) in GGML_FORMATS and quantizer.is_enabled + + +@torch.no_grad() +def _upcast_ggml_weights(model: nn.Module, dtype: torch.dtype) -> None: + """Decode each GGML-quantized weight to ``dtype`` in place and disable its quantizer. + + The weight becomes what its packed export payload decodes to -- the payload fake quant cached + or GPTQ pinned, when there is one. With its quantizer disabled, the layer then exports as an + unquantized one and stays out of the quantization config, so a runtime without GGML kernels, + such as vLLM, loads the checkpoint as plain weights. + """ + names = module_name_maps(model) + upcast_count = 0 + for module in model.modules(): + # Check the quantizers first, so an offloaded module is materialized only if it has work. + if not isinstance(module, QuantModule) or not any( + map(_is_ggml_weight_quantizer, module.iter_weight_quantizers_for_calibration()) + ): + continue + with enable_weight_access_and_writeback(module, model, names): + for weight, quantizer in module.iter_weights_for_calibration(): + if not _is_ggml_weight_quantizer(quantizer): + continue + ggml_format = GGML_FORMAT_REGISTRY[quantizer.num_bits] + packed_weight = ggml_format.pack(weight, quantizer) + weight.data.copy_( + ggml_format.dequantize(packed_weight, torch.tensor(weight.shape), dtype=dtype) + ) + quantizer.disable() + upcast_count += 1 + if upcast_count == 0: + warnings.warn("upcast_ggml is set, but the model has no GGML-quantized weights.") + + def export_hf_checkpoint( model: Any, dtype: torch.dtype | None = None, @@ -1860,6 +1904,7 @@ def export_hf_checkpoint( components: list[str] | None = None, extra_state_dict: dict[str, torch.Tensor] | None = None, max_shard_size: int | str = "10GB", + upcast_ggml: str | None = None, **kwargs, ): """Export quantized HuggingFace model checkpoint (transformers or diffusers). @@ -1883,6 +1928,12 @@ def export_hf_checkpoint( to export. If None, all quantized components are exported. extra_state_dict: Extra state dictionary to add to the exported model. max_shard_size: Maximum size of each safetensors shard file. Defaults to "10GB". + upcast_ggml: The dtype to upcast GGML-quantized weights to, currently only ``"bf16"``. + Each GGML weight (the IQ formats and Q8_0) is written as the values its packed blocks + decode to in that dtype, rather than as the blocks, and is left out of the + quantization config. Runtimes without GGML kernels, such as vLLM, can then load the + checkpoint to evaluate the quantized model. Not supported for diffusers pipelines or + a ``layerwise.export_dir`` export. **kwargs: Runtime-specific post-processing options forwarded to :func:`_postprocess_safetensors` for diffusion model exports. See its docstring for supported keys. @@ -1916,6 +1967,18 @@ def export_hf_checkpoint( from .layerwise_export import LAYERWISE_EXPORTER_ATTR exporter = getattr(model, LAYERWISE_EXPORTER_ATTR, None) + if upcast_ggml is not None: + if upcast_ggml not in _GGML_UPCAST_DTYPES: + raise ValueError( + f"upcast_ggml must be one of {sorted(_GGML_UPCAST_DTYPES)}, got {upcast_ggml!r}" + ) + if exporter is not None or (HAS_DIFFUSERS and is_diffusers_object(model)): + raise NotImplementedError( + "upcast_ggml supports transformers models only, and not a layerwise.export_dir " + "export, whose layer shards calibration already wrote packed." + ) + _upcast_ggml_weights(model, _GGML_UPCAST_DTYPES[upcast_ggml]) + if exporter is not None: # Per-layer export wrote the shards during calibration; this writes the rest. exporter.finalize(extra_state_dict=extra_state_dict) diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index 8815788e9d6..8b8ad7098ab 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/tests/unit/torch/export/test_export_weight.py @@ -14,6 +14,7 @@ # limitations under the License. +import copy import dataclasses import pytest @@ -22,13 +23,16 @@ from _test_utils.torch.export.utils import ToyModel, partial_fp8_config, partial_w4a8_config import modelopt.torch.quantization as mtq -from modelopt.torch.export.quant_utils import postprocess_state_dict +from modelopt.torch.export.quant_format import QUANTIZATION_NONE +from modelopt.torch.export.quant_utils import get_quantization_format, postprocess_state_dict from modelopt.torch.export.unified_export_hf import ( _export_quantized_weight, _process_quantized_modules, + _upcast_ggml_weights, ) from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY +from modelopt.torch.quantization.ggml.common import pin_packed_weight from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.quantization.utils import quantizer_attr_names @@ -187,6 +191,48 @@ def test_export_repacks_a_weight_changed_since_fake_quant(num_bits): assert torch.equal(linear.weight, expected) +def test_export_drops_the_payload_gptq_pinned(): + """GPTQ pins its payload to the quantizer as a buffer; only the packed weight is exported.""" + linear = _iq_linear("q8_0") + packed, _ = GGML_FORMAT_REGISTRY["q8_0"].quantize(linear.weight.detach()) + pin_packed_weight(linear.weight_quantizer, "q8_0", packed) + + _export_quantized_weight(linear, torch.bfloat16) + state_dict = postprocess_state_dict(linear.state_dict(), maxbound=448, quantization=None) + + assert list(state_dict) == ["weight"] + + +@pytest.mark.parametrize("num_bits", sorted(GGML_FORMAT_REGISTRY)) +def test_upcast_ggml_writes_what_the_packed_export_decodes_to(num_bits): + model = nn.Sequential(nn.Linear(512, 4, bias=False, dtype=torch.bfloat16)) + mtq.quantize( + model, + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": num_bits, "backend": "ggml"}, + "enable": True, + }, + ], + "algorithm": None, + }, + ) + packed = copy.deepcopy(model[0]) + _export_quantized_weight(packed, torch.bfloat16) + + _upcast_ggml_weights(model, torch.bfloat16) + + expected = GGML_FORMAT_REGISTRY[num_bits].dequantize( + packed.weight, torch.tensor((4, 512)), dtype=torch.bfloat16 + ) + assert torch.equal(model[0].weight, expected) + # The layer now exports as an unquantized one, so it stays out of the quantization config. + assert get_quantization_format(model) == QUANTIZATION_NONE + + class QuantMoELinear(nn.Module): def __init__(self): super().__init__() diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 9c9467ba17e..26f98b564bb 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -1075,3 +1075,56 @@ def test_union_degrades_to_the_recorded_set_without_the_loader(tmp_path, monkeyp carried = read_unplaced_weights(model) assert key in carried, "recorded keys must still carry when the structural pass is unavailable" + + +@pytest.mark.parametrize("moe", [False, True]) +def test_upcast_ggml_export_loads_as_a_plain_checkpoint(tmp_path, moe): + """With upcast_ggml="bf16" the checkpoint carries no quantization config, so a runtime without + GGML kernels loads it as an ordinary BF16 model, and every GGML weight holds its decoded payload. + + The MoE model exercises the fused experts, which export splits into per-expert projections. + """ + pytest.importorskip("transformers.core_model_loading") + # Local import: transformers and its test fixtures are optional dependencies. + from _test_utils.torch.transformers_models import get_tiny_qwen3, get_tiny_qwen3_moe + from transformers import AutoModelForCausalLM + + from modelopt.torch.export import export_hf_checkpoint + from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY + + model = (get_tiny_qwen3_moe if moe else get_tiny_qwen3)() + original = {name: param.detach().clone() for name, param in model.named_parameters()} + mtq.quantize( + model, + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*mlp*weight_quantizer", + "cfg": {"num_bits": "q8_0", "backend": "ggml"}, + "enable": True, + }, + ], + "algorithm": None, + }, + ) + + export_hf_checkpoint(model, export_dir=tmp_path, upcast_ggml="bf16") + + assert "quantization_config" not in json.loads((tmp_path / "config.json").read_text()) + assert not (tmp_path / "hf_quant_config.json").exists() + reloaded = AutoModelForCausalLM.from_pretrained(tmp_path, dtype=torch.bfloat16) + q8_0 = GGML_FORMAT_REGISTRY["q8_0"] + for name, param in reloaded.named_parameters(): + expected = original[name] + # The MoE router (mlp.gate) is not a linear, so it is never quantized. + if ".mlp." in name and not name.endswith(".mlp.gate.weight"): + expected = q8_0.dequantize(*q8_0.quantize(expected), dtype=expected.dtype) + assert torch.equal(param, expected), name + + +def test_upcast_ggml_rejects_an_unsupported_dtype(tmp_path): + from modelopt.torch.export import export_hf_checkpoint + + with pytest.raises(ValueError, match="upcast_ggml must be one of"): + export_hf_checkpoint(torch.nn.Linear(2, 2), export_dir=tmp_path, upcast_ggml="fp8") From 1a3840f8b511aaed600a034a64372448a6912ec0 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Thu, 8 Oct 2026 18:18:39 +0000 Subject: [PATCH 2/6] Validate GGML layers before the BF16 upcast and reject ignored routes Address review feedback on upcast_ggml: - Run the packed export's GGML layer checks (backend, enabled input quantizer, AWQ pre_quant_scale) before disabling any weight quantizer, so activation quantization is refused rather than silently dropped. The checks move into validate_ggml_layer, shared with get_quantization_format. - Reject GGML weights not stored in the target dtype: fused-expert slices and the offload/FSDP2 writebacks keep the stored dtype, so FP16/FP32 weights would not export as BF16. Every check runs before any weight changes. - hf_ptq: reject --upcast_ggml with --sparsity_fmt at parse time, and on the TensorRT-LLM and speculative-decoding export routes, which would otherwise ignore it. - Docs: only the GGML layers leave the quantization config; load a fully upcast checkpoint in vLLM without quantization="modelopt". - Tests: refusal cases, offloaded writeback with and without a pinned GPTQ payload, the no-GGML warning and the layerwise rejection. FSDP2 coverage is left as a TODO. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 2 +- docs/source/deployment/3_unified_hf.rst | 9 +- examples/hf_ptq/README.md | 2 +- examples/hf_ptq/hf_ptq.py | 13 +++ modelopt/torch/export/quant_utils.py | 35 ++++--- modelopt/torch/export/unified_export_hf.py | 47 +++++++-- tests/unit/torch/export/test_export_weight.py | 98 ++++++++++++++++--- .../torch/export/test_unified_export_hf.py | 25 +++-- 8 files changed, 179 insertions(+), 52 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index b3ba815b903..7d732089aeb 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -14,7 +14,7 @@ Changelog *Quantization* - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. -- Add ``upcast_ggml="bf16"`` to ``export_hf_checkpoint`` (``--upcast_ggml bf16`` in ``examples/hf_ptq/hf_ptq.py``): GGML-quantized weights are exported as the BF16 values their packed blocks decode to, with no quantization config, so runtimes without GGML kernels such as vLLM can evaluate the quantized model. +- Add ``upcast_ggml="bf16"`` to ``export_hf_checkpoint`` (``--upcast_ggml bf16`` in ``examples/hf_ptq/hf_ptq.py``): GGML-quantized weights are exported as the BF16 values their packed blocks decode to and left out of the quantization config, so runtimes without GGML kernels such as vLLM can evaluate the quantized model. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. - Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body. diff --git a/docs/source/deployment/3_unified_hf.rst b/docs/source/deployment/3_unified_hf.rst index 07210aa13d5..fd247d1822e 100644 --- a/docs/source/deployment/3_unified_hf.rst +++ b/docs/source/deployment/3_unified_hf.rst @@ -91,10 +91,11 @@ integer ``weights`` groups because all scales and indices are embedded in each p To evaluate a GGML-quantized model in a runtime without GGML kernels, such as vLLM, pass ``upcast_ggml="bf16"`` to ``export_hf_checkpoint`` (``--upcast_ggml bf16`` in -``examples/hf_ptq/hf_ptq.py``). Each GGML weight is then written as the BF16 values its packed -blocks decode to, and those layers are left out of the quantization configuration, so the -checkpoint loads as an ordinary unquantized one. It is no smaller than the original model: use it -for accuracy evaluation, not deployment. +``examples/hf_ptq/hf_ptq.py``). Each GGML weight, which must be stored in BF16, is then written +as the BF16 values its packed blocks decode to, and those layers are left out of the quantization +configuration. Layers quantized in other formats, and a quantized KV cache, export as usual; a +model with none loads as an ordinary unquantized checkpoint. The upcast weights are no smaller +than the original ones: use the checkpoint for accuracy evaluation, not deployment. Each 74-byte IQ2_XS block represents 256 logical weights: diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index d610636ad34..da466866d68 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -673,7 +673,7 @@ python hf_ptq.py --pyt_ckpt_path --qformat fp8 --export > *For exporting fake-quantized models for vLLM serving (e.g., for research or kernels not yet supported in real-quant), use the `--vllm_fakequant_export` flag. See [vllm_serve/README.md](../vllm_serve/README.md) for details.* > -> *To evaluate a GGML-quantized model (IQ formats, Q8_0) in stock vLLM, add `--upcast_ggml bf16`: the GGML weights are exported as the BF16 values their packed blocks decode to, with no quantization config.* +> *To evaluate a GGML-quantized model (IQ formats, Q8_0) in stock vLLM, add `--upcast_ggml bf16`: the GGML weights are exported as the BF16 values their packed blocks decode to and left out of the quantization config. If no other layers are quantized, load the checkpoint without `quantization="modelopt"`.* ### Hugging Face framework [Script](./scripts/huggingface_example.sh) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 124f763dd42..2edeb570979 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -648,6 +648,11 @@ def export_quantized( # Early exit for speculative decoding checkpoints # No tokenizer saving needed for spec ckpts if has_spec_opt(full_model): + if args.upcast_ggml: + raise NotImplementedError( + "--upcast_ggml applies to the unified HF export, not a speculative " + "decoding checkpoint." + ) export_speculative_decoding(full_model, export_dir=export_path) args.checkpoint_exported = True print(f"Quantized speculative decoding checkpoint exported to: {export_path}") @@ -666,6 +671,11 @@ def export_quantized( or "int8_smoothquant" in args.qformat ) if is_tensorrt_llm_export: + if args.upcast_ggml: + raise NotImplementedError( + "--upcast_ggml applies to the unified HF export, but this model exports a " + "TensorRT-LLM checkpoint." + ) if ( args.inference_tensor_parallel != 1 or args.inference_pipeline_parallel != 1 ) and args.qformat == "nvfp4_svdquant": @@ -1448,6 +1458,9 @@ def parse_args() -> argparse.Namespace: parser.error(f"--use_fsdp2 does not support --sparsity_fmt {args.sparsity_fmt}.") if args.use_fsdp2 and args.vllm_fakequant_export: parser.error("--use_fsdp2 does not support --vllm_fakequant_export.") + if args.upcast_ggml and args.sparsity_fmt != "dense": + # Sparsity exports a TensorRT-LLM checkpoint, which the upcast does not apply to. + parser.error("--upcast_ggml does not support --sparsity_fmt.") if args.upcast_ggml and args.vllm_fakequant_export: parser.error( "--upcast_ggml applies to the unified HF export; --vllm_fakequant_export already " diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index eeeccfd43b6..0f3756c672b 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -443,6 +443,25 @@ def get_weight_block_size(module: nn.Module, weight_name: str = "weight") -> int return 0 +def validate_ggml_layer(weight_quantizer, input_quantizer) -> None: + """Reject a GGML-quantized weight whose layer export cannot represent. + + Both exporters return before collecting input_scale and before the pre_quant_scale handling, + so an enabled activation quantizer would be dropped without a trace and the checkpoint would + load as weight-only. Refuse instead. + """ + if weight_quantizer.backend != "ggml": + raise ValueError("GGML formats require the built-in 'ggml' quantization backend") + if input_quantizer is not None and input_quantizer.is_enabled: + raise NotImplementedError( + "GGML export is weight-only, but this layer has an enabled input " + "quantizer. The GGML block payload carries no activation scale, so the " + "activation quantization would be silently lost." + ) + if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"): + raise NotImplementedError("GGML export does not support an AWQ-style pre_quant_scale.") + + def uses_iq_quantization(module) -> bool: """Whether any weight quantizer in ``module`` or its children targets a GGML format. @@ -501,21 +520,7 @@ def _get_quantization_from_layer(layer, quantizer_attr_names: QuantizerAttrNames # Handle individual num_bits cases if weight_quantizer.num_bits in GGML_FORMATS: - if weight_quantizer.backend != "ggml": - raise ValueError("GGML formats require the built-in 'ggml' quantization backend") - # Both exporters return before collecting input_scale and before the pre_quant_scale - # handling below, so an enabled activation quantizer would be dropped without a trace - # and the checkpoint would load as weight-only. Refuse instead. - if input_quantizer is not None and input_quantizer.is_enabled: - raise NotImplementedError( - "GGML export is weight-only, but this layer has an enabled input " - "quantizer. The GGML block payload carries no activation scale, so the " - "activation quantization would be silently lost." - ) - if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"): - raise NotImplementedError( - "GGML export does not support an AWQ-style pre_quant_scale." - ) + validate_ggml_layer(weight_quantizer, input_quantizer) return weight_quantizer.num_bits if weight_quantizer.num_bits == 4: diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index d7dfee85fc2..a67d980fed3 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -77,6 +77,8 @@ fsdp2_aware_weight_update, module_name_maps, quantizer_attr_names, + representative_weight_quantizer, + weight_attr_names, ) from modelopt.torch.quantization.utils.core_utils import ( enable_weight_access_and_writeback, @@ -134,6 +136,7 @@ preprocess_linear_fusion, sync_tied_input_amax, to_quantized_weight, + validate_ggml_layer, ) from .registry import ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry @@ -1860,6 +1863,7 @@ def read_unplaced_weights(model: nn.Module, *, keys_only: bool = False) -> dict[ def _is_ggml_weight_quantizer(quantizer) -> bool: + """Whether ``quantizer`` is an enabled GGML weight quantizer.""" # getattr: a SequentialQuantizer has no num_bits, and is never GGML. return getattr(quantizer, "num_bits", None) in GGML_FORMATS and quantizer.is_enabled @@ -1874,13 +1878,39 @@ def _upcast_ggml_weights(model: nn.Module, dtype: torch.dtype) -> None: such as vLLM, loads the checkpoint as plain weights. """ names = module_name_maps(model) - upcast_count = 0 - for module in model.modules(): - # Check the quantizers first, so an offloaded module is materialized only if it has work. - if not isinstance(module, QuantModule) or not any( - map(_is_ggml_weight_quantizer, module.iter_weight_quantizers_for_calibration()) - ): - continue + # Read only the quantizers here, so an offloaded module is materialized only if it has work. + ggml_modules = [ + module + for module in model.modules() + if isinstance(module, QuantModule) + and any(map(_is_ggml_weight_quantizer, module.iter_weight_quantizers_for_calibration())) + ] + if not ggml_modules: + warnings.warn("upcast_ggml is set, but the model has no GGML-quantized weights.") + return + + # Check every layer before changing any, so a refusal leaves the model as it was. + for module in ggml_modules: + for weight_name in weight_attr_names(module): + weight_quantizer = representative_weight_quantizer(module, weight_name) + if not _is_ggml_weight_quantizer(weight_quantizer): + continue + # Disabling the quantizer would skip the checks the packed export runs on the layer. + input_quantizer = getattr( + module, quantizer_attr_names(weight_name).input_quantizer, None + ) + validate_ggml_layer(weight_quantizer, input_quantizer) + # Fused-expert slices and the offload and FSDP2 writebacks keep the stored dtype, so + # the decoded values can only be written to a weight already stored in ``dtype``. + weight_dtype = getattr(module, weight_name).dtype + if weight_dtype != dtype: + raise ValueError( + f"upcast_ggml decodes to {dtype}, but " + f"{names.module_to_name.get(id(module), '')}.{weight_name} is stored as " + f"{weight_dtype}. Load the model in {dtype} before quantizing it." + ) + + for module in ggml_modules: with enable_weight_access_and_writeback(module, model, names): for weight, quantizer in module.iter_weights_for_calibration(): if not _is_ggml_weight_quantizer(quantizer): @@ -1891,9 +1921,6 @@ def _upcast_ggml_weights(model: nn.Module, dtype: torch.dtype) -> None: ggml_format.dequantize(packed_weight, torch.tensor(weight.shape), dtype=dtype) ) quantizer.disable() - upcast_count += 1 - if upcast_count == 0: - warnings.warn("upcast_ggml is set, but the model has no GGML-quantized weights.") def export_hf_checkpoint( diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index 8b8ad7098ab..f35527d095a 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/tests/unit/torch/export/test_export_weight.py @@ -203,23 +203,26 @@ def test_export_drops_the_payload_gptq_pinned(): assert list(state_dict) == ["weight"] -@pytest.mark.parametrize("num_bits", sorted(GGML_FORMAT_REGISTRY)) -def test_upcast_ggml_writes_what_the_packed_export_decodes_to(num_bits): - model = nn.Sequential(nn.Linear(512, 4, bias=False, dtype=torch.bfloat16)) - mtq.quantize( - model, +def _ggml_model(num_bits, dtype=torch.bfloat16, input_cfg=None): + """A one-linear model whose weight quantizer is ``num_bits``; ``input_cfg`` enables its input.""" + model = nn.Sequential(nn.Linear(512, 4, bias=False, dtype=dtype)) + quant_cfg = [ + {"quantizer_name": "*", "enable": False}, { - "quant_cfg": [ - {"quantizer_name": "*", "enable": False}, - { - "quantizer_name": "*weight_quantizer", - "cfg": {"num_bits": num_bits, "backend": "ggml"}, - "enable": True, - }, - ], - "algorithm": None, + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": num_bits, "backend": "ggml"}, + "enable": True, }, - ) + ] + if input_cfg is not None: + quant_cfg.append({"quantizer_name": "*input_quantizer", "cfg": input_cfg, "enable": True}) + return mtq.quantize(model, {"quant_cfg": quant_cfg, "algorithm": None}) + + +@pytest.mark.parametrize("num_bits", sorted(GGML_FORMAT_REGISTRY)) +def test_upcast_ggml_writes_what_the_packed_export_decodes_to(num_bits): + """The upcast weight is exactly the decoding of the payload the packed export writes.""" + model = _ggml_model(num_bits) packed = copy.deepcopy(model[0]) _export_quantized_weight(packed, torch.bfloat16) @@ -233,6 +236,71 @@ def test_upcast_ggml_writes_what_the_packed_export_decodes_to(num_bits): assert get_quantization_format(model) == QUANTIZATION_NONE +@pytest.mark.parametrize( + ("dtype", "input_cfg", "error"), + [ + (torch.bfloat16, {"num_bits": (4, 3)}, "enabled input quantizer"), + (torch.bfloat16, "pre_quant_scale", "pre_quant_scale"), + (torch.float16, None, "stored as torch.float16"), + (torch.float32, None, "stored as torch.float32"), + ], + ids=["input_quantizer", "pre_quant_scale", "fp16_weight", "fp32_weight"], +) +def test_upcast_ggml_refuses_a_layer_before_changing_any(dtype, input_cfg, error): + """Activation quantization the upcast would drop, or a weight not stored in BF16, is refused. + + The packed export refuses the first two too; disabling the quantizer must not skip that. The + model is left as it was, so a refusal never leaves it half upcast. + """ + model = _ggml_model("q8_0", dtype, None if input_cfg == "pre_quant_scale" else input_cfg) + if input_cfg == "pre_quant_scale": + model[0].input_quantizer.pre_quant_scale = torch.ones(512) + weight = model[0].weight.detach().clone() + + with pytest.raises((NotImplementedError, ValueError), match=error): + _upcast_ggml_weights(model, torch.bfloat16) + + assert torch.equal(model[0].weight, weight) + assert model[0].weight_quantizer.is_enabled + + +@pytest.mark.parametrize("pinned", [False, True]) +def test_upcast_ggml_writes_back_an_offloaded_weight(monkeypatch, pinned): + """The decoded values reach the offloaded copy; a payload GPTQ pinned is decoded, not redone.""" + pytest.importorskip("accelerate") + from accelerate.hooks import AlignDevicesHook, add_hook_to_module + from accelerate.utils import set_module_tensor_to_device + + model = _ggml_model("q8_0") + linear = model[0] + q8_0 = GGML_FORMAT_REGISTRY["q8_0"] + payload, shape = q8_0.quantize(linear.weight.detach()) + expected = q8_0.dequantize(payload, shape, dtype=torch.bfloat16) + assert not torch.equal(linear.weight, expected) + if pinned: + # How GPTQ leaves a layer: its payload pinned, and the weight equal to the decoding. + with torch.no_grad(): + linear.weight.copy_(expected) + pin_packed_weight(linear.weight_quantizer, "q8_0", payload) + _without_search(monkeypatch, "q8_0") + weights_map = {"weight": linear.weight.detach().clone()} + add_hook_to_module( + linear, AlignDevicesHook(execution_device="cpu", offload=True, weights_map=weights_map) + ) + set_module_tensor_to_device(linear, "weight", "meta") + + _upcast_ggml_weights(model, torch.bfloat16) + + assert linear.weight.is_meta + assert torch.equal(weights_map["weight"], expected) + + +def test_upcast_ggml_warns_when_there_is_nothing_to_upcast(): + """Asking for an upcast of a model with no GGML weights is most likely a mistake.""" + with pytest.warns(UserWarning, match="no GGML-quantized weights"): + _upcast_ggml_weights(nn.Sequential(nn.Linear(4, 4)), torch.bfloat16) + + class QuantMoELinear(nn.Module): def __init__(self): super().__init__() diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 26f98b564bb..45738445015 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -28,6 +28,7 @@ from safetensors.torch import save_file import modelopt.torch.quantization as mtq +from modelopt.torch.export.layerwise_export import LAYERWISE_EXPORTER_ATTR from modelopt.torch.export.model_utils import ( TiedWeightMap, get_language_model_from_vl, @@ -38,7 +39,12 @@ postprocess_state_dict, sync_tied_input_amax, ) -from modelopt.torch.export.unified_export_hf import _resolve_export_dtype, read_unplaced_weights +from modelopt.torch.export.unified_export_hf import ( + _resolve_export_dtype, + export_hf_checkpoint, + read_unplaced_weights, +) +from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.nn import TensorQuantizer @@ -1077,6 +1083,8 @@ def test_union_degrades_to_the_recorded_set_without_the_loader(tmp_path, monkeyp assert key in carried, "recorded keys must still carry when the structural pass is unavailable" +# TODO: cover the FSDP2 export path too, with a multi-rank upcast test in +# tests/gpu/torch/export/test_fsdp2_export.py. @pytest.mark.parametrize("moe", [False, True]) def test_upcast_ggml_export_loads_as_a_plain_checkpoint(tmp_path, moe): """With upcast_ggml="bf16" the checkpoint carries no quantization config, so a runtime without @@ -1089,9 +1097,6 @@ def test_upcast_ggml_export_loads_as_a_plain_checkpoint(tmp_path, moe): from _test_utils.torch.transformers_models import get_tiny_qwen3, get_tiny_qwen3_moe from transformers import AutoModelForCausalLM - from modelopt.torch.export import export_hf_checkpoint - from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY - model = (get_tiny_qwen3_moe if moe else get_tiny_qwen3)() original = {name: param.detach().clone() for name, param in model.named_parameters()} mtq.quantize( @@ -1124,7 +1129,15 @@ def test_upcast_ggml_export_loads_as_a_plain_checkpoint(tmp_path, moe): def test_upcast_ggml_rejects_an_unsupported_dtype(tmp_path): - from modelopt.torch.export import export_hf_checkpoint - + """Only the dtypes the upcast knows are accepted.""" with pytest.raises(ValueError, match="upcast_ggml must be one of"): export_hf_checkpoint(torch.nn.Linear(2, 2), export_dir=tmp_path, upcast_ggml="fp8") + + +def test_upcast_ggml_rejects_a_layerwise_export(tmp_path): + """A layerwise export wrote its layer shards packed during calibration, too early to upcast.""" + model = torch.nn.Linear(2, 2) + setattr(model, LAYERWISE_EXPORTER_ATTR, object()) + + with pytest.raises(NotImplementedError, match=r"layerwise\.export_dir"): + export_hf_checkpoint(model, export_dir=tmp_path, upcast_ggml="bf16") From 48276120a629ea384cae52ebe1e5f392938f4d25 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 9 Oct 2026 18:52:22 +0000 Subject: [PATCH 3/6] Refuse grouped GGML layouts in the upcast and test it under FSDP2 Address review feedback on upcast_ggml: - Refuse a GGML weight not stored under the name weight_attr_names reports (TEGroupedLinear keeps weight0..N) with NotImplementedError before changing anything, instead of an AttributeError. - Add a two-rank FSDP2 test: the sharded model's upcast, gathered again after the reshard, equals what a single-process packed export decodes to, with and without GPTQ-pinned payloads, and the GGML layers leave the quantization config. It calls the upcast directly because the FSDP2 streaming writer deadlocks at world > 1 on main (the existing test_fsdp2_streaming_export_matches_reference times out too). - Note why the accelerate imports stay local in the offload test. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/export/unified_export_hf.py | 17 +++- tests/gpu/torch/export/test_fsdp2_export.py | 91 +++++++++++++++++++ tests/unit/torch/export/test_export_weight.py | 25 ++++- 3 files changed, 127 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index a67d980fed3..e82dbe88a69 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1900,14 +1900,21 @@ def _upcast_ggml_weights(model: nn.Module, dtype: torch.dtype) -> None: module, quantizer_attr_names(weight_name).input_quantizer, None ) validate_ggml_layer(weight_quantizer, input_quantizer) + name = f"{names.module_to_name.get(id(module), '')}.{weight_name}" + weight = getattr(module, weight_name, None) + if not isinstance(weight, torch.Tensor): + # A grouped layout, such as TEGroupedLinear's weight0..N; the packed export + # supports only a standard weight attribute too. + raise NotImplementedError( + f"upcast_ggml needs each GGML weight stored under its own name, but {name} " + f"of {type(module).__name__} is not a tensor." + ) # Fused-expert slices and the offload and FSDP2 writebacks keep the stored dtype, so # the decoded values can only be written to a weight already stored in ``dtype``. - weight_dtype = getattr(module, weight_name).dtype - if weight_dtype != dtype: + if weight.dtype != dtype: raise ValueError( - f"upcast_ggml decodes to {dtype}, but " - f"{names.module_to_name.get(id(module), '')}.{weight_name} is stored as " - f"{weight_dtype}. Load the model in {dtype} before quantizing it." + f"upcast_ggml decodes to {dtype}, but {name} is stored as {weight.dtype}. " + f"Load the model in {dtype} before quantizing it." ) for module in ggml_modules: diff --git a/tests/gpu/torch/export/test_fsdp2_export.py b/tests/gpu/torch/export/test_fsdp2_export.py index 939fe3e581f..84870a630b4 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/tests/gpu/torch/export/test_fsdp2_export.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import copy +import dataclasses import json from functools import partial from pathlib import Path @@ -33,12 +34,16 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.layer_utils import is_quantlinear from modelopt.torch.export.model_utils import TiedWeightMap +from modelopt.torch.export.quant_utils import get_quant_config from modelopt.torch.export.unified_export_hf import ( _export_quantized_weight, _export_transformers_checkpoint, + _upcast_ggml_weights, requantize_resmooth_fused_llm_layers, ) from modelopt.torch.export.unified_export_hf_streaming import _export_fsdp2_checkpoint_streaming +from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY +from modelopt.torch.quantization.ggml.common import pin_packed_weight from modelopt.torch.quantization.utils import ( enable_weight_access_and_writeback, fsdp2_aware_weight_update, @@ -463,3 +468,89 @@ def test_fsdp2_streaming_export_matches_reference(dist_workers, tmp_path, quant_ quant_config=quant_config, ) ) + + +_Q8_0_MLP_CFG = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*mlp*weight_quantizer", + "cfg": {"num_bits": "q8_0", "block_sizes": {-1: 32}, "backend": "ggml"}, + "enable": True, + }, + ], + "algorithm": None, +} + + +def _pin_ggml_payloads(model): + """Leave each GGML layer as GPTQ does: a payload pinned, and the weight equal to its decoding.""" + q8_0 = GGML_FORMAT_REGISTRY["q8_0"] + names = module_name_maps(model) + for module in list(model.modules()): + quantizer = getattr(module, "weight_quantizer", None) + if not (is_quantlinear(module) and quantizer is not None and quantizer.is_enabled): + continue + with enable_weight_access_and_writeback(module, model, names): + payload, shape = q8_0.quantize(module.weight.flip(-1)) + module.weight.copy_(q8_0.dequantize(payload, shape, dtype=module.weight.dtype)) + pin_packed_weight(quantizer, "q8_0", payload) + + +def _no_search(*args, **kwargs): + raise AssertionError("the upcast re-ran the q8_0 search instead of decoding the pinned payload") + + +def _upcast_ggml_matches_reference_test(rank, size, pinned): + """A sharded model's upcast holds what a single-process packed export decodes to. + + The upcast decodes each GGML weight inside a gather window and writes it back to the shards, + so after a reshard every rank must gather the decoded payload, and the GGML layers must have + dropped out of the quantization config. With ``pinned``, every GGML layer carries a payload + pinned as GPTQ pins one, and the encoder is disabled to prove the pinned codes are decoded. + + This calls the upcast directly rather than ``export_hf_checkpoint``, whose FSDP2 streaming + writer currently deadlocks at world > 1 (see test_fsdp2_streaming_export_matches_reference). + """ + q8_0 = GGML_FORMAT_REGISTRY["q8_0"] + with patch_fsdp_mp_dtypes(), torch.no_grad(): + model = get_tiny_llama(num_hidden_layers=2).to("cuda").eval() + reference = get_tiny_llama(num_hidden_layers=2).to("cuda").eval() + for layer in model.model.layers: + fully_shard(layer) + fully_shard(model) + mtq.quantize(model, _Q8_0_MLP_CFG) + mtq.quantize(reference, _Q8_0_MLP_CFG) + if pinned: + _pin_ggml_payloads(model) + _pin_ggml_payloads(reference) + # The pool reuses its workers across tests, so the patch is undone below. + GGML_FORMAT_REGISTRY["q8_0"] = dataclasses.replace(q8_0, quantize=_no_search) + try: + packed, _ = _export_transformers_checkpoint(reference, torch.bfloat16) + _upcast_ggml_weights(model, torch.bfloat16) + finally: + GGML_FORMAT_REGISTRY["q8_0"] = q8_0 + + assert get_quant_config(model)["quantization"]["quant_algo"] is None + ggml_keys = [key for key, tensor in packed.items() if tensor.dtype == torch.uint8] + assert ggml_keys, "the reference export packed no GGML weight" + names = module_name_maps(model) + for key in ggml_keys: + payload = packed[key] + expected = q8_0.dequantize( + payload, + torch.tensor((*payload.shape[:-2], payload.shape[-2] * q8_0.block_size)), + dtype=torch.bfloat16, + ) + module = names.name_to_module[key.removesuffix(".weight")] + # Gather again after the upcast's window closed, so this reads what reached the shards. + with enable_weight_access_and_writeback(module, model, names, writeback=False): + assert torch.equal(module.weight, expected.to(module.weight.device)), key + + +@pytest.mark.parametrize("pinned", [False, True], ids=["encoded", "gptq_pinned"]) +def test_fsdp2_upcast_ggml_matches_reference(dist_workers, pinned): + if torch.cuda.device_count() < 2: + pytest.skip("needs >=2 GPUs for the weights to actually be sharded") + dist_workers.run(partial(_upcast_ggml_matches_reference_test, pinned=pinned)) diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index f35527d095a..1917e2897aa 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/tests/unit/torch/export/test_export_weight.py @@ -33,7 +33,7 @@ from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.ggml.common import pin_packed_weight -from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.nn import GroupedQuantizer, TensorQuantizer from modelopt.torch.quantization.utils import quantizer_attr_names @@ -268,6 +268,7 @@ def test_upcast_ggml_refuses_a_layer_before_changing_any(dtype, input_cfg, error def test_upcast_ggml_writes_back_an_offloaded_weight(monkeypatch, pinned): """The decoded values reach the offloaded copy; a payload GPTQ pinned is decoded, not redone.""" pytest.importorskip("accelerate") + # Local, so importorskip can skip this test when the optional accelerate is not installed. from accelerate.hooks import AlignDevicesHook, add_hook_to_module from accelerate.utils import set_module_tensor_to_device @@ -295,6 +296,28 @@ def test_upcast_ggml_writes_back_an_offloaded_weight(monkeypatch, pinned): assert torch.equal(weights_map["weight"], expected) +def test_upcast_ggml_refuses_a_grouped_layout(): + """A grouped layer keeps its GGML weights as weight0..N, not under the name it reports. + + Stands in for TEGroupedLinear: one GroupedQuantizer, which weight_attr_names reports as + "weight", over expert weights that only its own iter_weights_for_calibration yields. + """ + linear = _ggml_model("q8_0")[0] + linear.weight0 = linear.weight + del linear.weight + linear.weight_quantizer = GroupedQuantizer(linear.weight_quantizer) + linear.iter_weights_for_calibration = lambda: iter( + [(linear.weight0, linear.weight_quantizer[0])] + ) + weight = linear.weight0.detach().clone() + + with pytest.raises(NotImplementedError, match="stored under its own name"): + _upcast_ggml_weights(nn.Sequential(linear), torch.bfloat16) + + assert torch.equal(linear.weight0, weight) + assert linear.weight_quantizer[0].is_enabled + + def test_upcast_ggml_warns_when_there_is_nothing_to_upcast(): """Asking for an upcast of a model with no GGML weights is most likely a mistake.""" with pytest.warns(UserWarning, match="no GGML-quantized weights"): From 61e93eb4de39acc66b22a9b32ab87562d2aa0bd2 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Sun, 11 Oct 2026 16:39:09 +0000 Subject: [PATCH 4/6] Replace upcast_ggml with exact folding of GPTQ'd GGML weights Review asked why GGML needed its own upcast rather than weight folding. Folding is the general mechanism, and --vllm_fakequant_export already writes a plain BF16 Hugging Face checkpoint (no quantization config) that a stock vLLM image serves. It was only wrong for GGML weights that GPTQ calibrated: both fold paths quantize a float32 copy of the weight, which never matched the payload GPTQ pinned against the BF16 weight, so they re-encoded it and folded in different codes. On Qwen3-0.6B with the iq2_xxs recipe this raised vLLM perplexity from 42.2 to 72.0. - The pinned-payload lookup also matches a float32 weight against the decoding rounded through BF16 or FP16, so folding a GPTQ'd weight decodes GPTQ's codes instead of re-encoding. - Remove upcast_ggml and --upcast_ggml with their docs and tests, and document --vllm_fakequant_export for evaluating GGML checkpoints in vLLM. - Keep dropping the pinned payload from the packed export's state dict. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 1 - docs/source/deployment/3_unified_hf.rst | 13 +- examples/hf_ptq/README.md | 4 +- examples/hf_ptq/hf_ptq.py | 34 ----- modelopt/torch/export/quant_utils.py | 35 ++--- modelopt/torch/export/unified_export_hf.py | 101 +------------- modelopt/torch/quantization/ggml/common.py | 29 ++-- tests/gpu/torch/export/test_fsdp2_export.py | 91 ------------- tests/unit/torch/export/test_export_weight.py | 128 +----------------- .../torch/export/test_unified_export_hf.py | 68 +--------- .../unit/torch/quantization/test_ggml_gptq.py | 20 +++ 11 files changed, 65 insertions(+), 459 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7d732089aeb..d3cc67d95bf 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -14,7 +14,6 @@ Changelog *Quantization* - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. -- Add ``upcast_ggml="bf16"`` to ``export_hf_checkpoint`` (``--upcast_ggml bf16`` in ``examples/hf_ptq/hf_ptq.py``): GGML-quantized weights are exported as the BF16 values their packed blocks decode to and left out of the quantization config, so runtimes without GGML kernels such as vLLM can evaluate the quantized model. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. - Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body. diff --git a/docs/source/deployment/3_unified_hf.rst b/docs/source/deployment/3_unified_hf.rst index fd247d1822e..d3dc840cc62 100644 --- a/docs/source/deployment/3_unified_hf.rst +++ b/docs/source/deployment/3_unified_hf.rst @@ -89,13 +89,12 @@ The generated configuration records ``quant_method: modelopt``, ``packing: ggml` block size, and the payload byte count. GGML payloads are not represented as compressed-tensors integer ``weights`` groups because all scales and indices are embedded in each packed block. -To evaluate a GGML-quantized model in a runtime without GGML kernels, such as vLLM, pass -``upcast_ggml="bf16"`` to ``export_hf_checkpoint`` (``--upcast_ggml bf16`` in -``examples/hf_ptq/hf_ptq.py``). Each GGML weight, which must be stored in BF16, is then written -as the BF16 values its packed blocks decode to, and those layers are left out of the quantization -configuration. Layers quantized in other formats, and a quantized KV cache, export as usual; a -model with none loads as an ordinary unquantized checkpoint. The upcast weights are no smaller -than the original ones: use the checkpoint for accuracy evaluation, not deployment. +To evaluate a GGML-quantized model in a runtime without GGML kernels, such as vLLM, export it with +``--vllm_fakequant_export`` in ``examples/hf_ptq/hf_ptq.py`` instead. Each GGML weight is folded into +the values its blocks decode to, including the codes GPTQ chose, and the checkpoint carries no +quantization configuration, so a stock vLLM image serves it as an ordinary model. The extra +``vllm_fq_modelopt_state.pth`` is read only by the fake-quant launcher in ``examples/vllm_serve``, so +activation or KV-cache quantization elsewhere in the model applies only when served through it. Each 74-byte IQ2_XS block represents 256 logical weights: diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index da466866d68..2252f5643e7 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -671,9 +671,7 @@ with torch.inference_mode(): python hf_ptq.py --pyt_ckpt_path --qformat fp8 --export_path --trust_remote_code ``` -> *For exporting fake-quantized models for vLLM serving (e.g., for research or kernels not yet supported in real-quant), use the `--vllm_fakequant_export` flag. See [vllm_serve/README.md](../vllm_serve/README.md) for details.* -> -> *To evaluate a GGML-quantized model (IQ formats, Q8_0) in stock vLLM, add `--upcast_ggml bf16`: the GGML weights are exported as the BF16 values their packed blocks decode to and left out of the quantization config. If no other layers are quantized, load the checkpoint without `quantization="modelopt"`.* +> *For exporting fake-quantized models for vLLM serving (e.g., for research or kernels not yet supported in real-quant), use the `--vllm_fakequant_export` flag. See [vllm_serve/README.md](../vllm_serve/README.md) for details. For weight-only formats such as the GGML formats (IQ, Q8_0), the export is a plain BF16 checkpoint that a stock vLLM image serves directly.* ### Hugging Face framework [Script](./scripts/huggingface_example.sh) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 2edeb570979..db29d5fbf61 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -648,11 +648,6 @@ def export_quantized( # Early exit for speculative decoding checkpoints # No tokenizer saving needed for spec ckpts if has_spec_opt(full_model): - if args.upcast_ggml: - raise NotImplementedError( - "--upcast_ggml applies to the unified HF export, not a speculative " - "decoding checkpoint." - ) export_speculative_decoding(full_model, export_dir=export_path) args.checkpoint_exported = True print(f"Quantized speculative decoding checkpoint exported to: {export_path}") @@ -671,11 +666,6 @@ def export_quantized( or "int8_smoothquant" in args.qformat ) if is_tensorrt_llm_export: - if args.upcast_ggml: - raise NotImplementedError( - "--upcast_ggml applies to the unified HF export, but this model exports a " - "TensorRT-LLM checkpoint." - ) if ( args.inference_tensor_parallel != 1 or args.inference_pipeline_parallel != 1 ) and args.qformat == "nvfp4_svdquant": @@ -726,7 +716,6 @@ def export_quantized( export_hf_checkpoint( full_model, export_dir=export_path, - upcast_ggml=args.upcast_ggml, ) if args.qformat == "w4a16_nvfp4": @@ -972,11 +961,6 @@ def quantize_main( "layerwise.export_dir is not supported with an AutoQuantize recipe; " "use a PTQ recipe, or drop export_dir and export afterwards." ) - if args.upcast_ggml: - # Fail before calibration rather than after: the layer shards are written packed. - raise NotImplementedError( - "--upcast_ggml is not supported with layerwise.export_dir; drop export_dir." - ) if not args.skip_generate: print("Layerwise export: forcing --skip_generate, the model is left in export form.") args.skip_generate = True @@ -1366,16 +1350,6 @@ def parse_args() -> argparse.Namespace: help="Export as vLLM fake-quant checkpoint (produces vllm_fq_modelopt_state.pth " "for use with vllm_serve_fakequant.py).", ) - parser.add_argument( - "--upcast_ggml", - default=None, - choices=["bf16"], - help=( - "Export GGML-quantized weights (IQ formats, Q8_0) as the values their packed blocks " - "decode to in this dtype, rather than as the blocks, so runtimes without GGML kernels " - "such as vLLM can evaluate the checkpoint." - ), - ) parser.add_argument( "--cast_mxfp4_to_nvfp4", action="store_true", @@ -1458,14 +1432,6 @@ def parse_args() -> argparse.Namespace: parser.error(f"--use_fsdp2 does not support --sparsity_fmt {args.sparsity_fmt}.") if args.use_fsdp2 and args.vllm_fakequant_export: parser.error("--use_fsdp2 does not support --vllm_fakequant_export.") - if args.upcast_ggml and args.sparsity_fmt != "dense": - # Sparsity exports a TensorRT-LLM checkpoint, which the upcast does not apply to. - parser.error("--upcast_ggml does not support --sparsity_fmt.") - if args.upcast_ggml and args.vllm_fakequant_export: - parser.error( - "--upcast_ggml applies to the unified HF export; --vllm_fakequant_export already " - "writes fake-quantized weights." - ) if args.use_fsdp2 and args.cast_mxfp4_to_nvfp4: parser.error("--use_fsdp2 does not support --cast_mxfp4_to_nvfp4.") diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 0f3756c672b..eeeccfd43b6 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -443,25 +443,6 @@ def get_weight_block_size(module: nn.Module, weight_name: str = "weight") -> int return 0 -def validate_ggml_layer(weight_quantizer, input_quantizer) -> None: - """Reject a GGML-quantized weight whose layer export cannot represent. - - Both exporters return before collecting input_scale and before the pre_quant_scale handling, - so an enabled activation quantizer would be dropped without a trace and the checkpoint would - load as weight-only. Refuse instead. - """ - if weight_quantizer.backend != "ggml": - raise ValueError("GGML formats require the built-in 'ggml' quantization backend") - if input_quantizer is not None and input_quantizer.is_enabled: - raise NotImplementedError( - "GGML export is weight-only, but this layer has an enabled input " - "quantizer. The GGML block payload carries no activation scale, so the " - "activation quantization would be silently lost." - ) - if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"): - raise NotImplementedError("GGML export does not support an AWQ-style pre_quant_scale.") - - def uses_iq_quantization(module) -> bool: """Whether any weight quantizer in ``module`` or its children targets a GGML format. @@ -520,7 +501,21 @@ def _get_quantization_from_layer(layer, quantizer_attr_names: QuantizerAttrNames # Handle individual num_bits cases if weight_quantizer.num_bits in GGML_FORMATS: - validate_ggml_layer(weight_quantizer, input_quantizer) + if weight_quantizer.backend != "ggml": + raise ValueError("GGML formats require the built-in 'ggml' quantization backend") + # Both exporters return before collecting input_scale and before the pre_quant_scale + # handling below, so an enabled activation quantizer would be dropped without a trace + # and the checkpoint would load as weight-only. Refuse instead. + if input_quantizer is not None and input_quantizer.is_enabled: + raise NotImplementedError( + "GGML export is weight-only, but this layer has an enabled input " + "quantizer. The GGML block payload carries no activation scale, so the " + "activation quantization would be silently lost." + ) + if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"): + raise NotImplementedError( + "GGML export does not support an AWQ-style pre_quant_scale." + ) return weight_quantizer.num_bits if weight_quantizer.num_bits == 4: diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index e82dbe88a69..f6116f4746a 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -69,7 +69,7 @@ from modelopt.torch.opt.plugins.huggingface import _MODELOPT_STATE_SAVE_NAME from modelopt.torch.quantization import set_quantizer_by_cfg_context from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY -from modelopt.torch.quantization.nn import QuantModule, SequentialQuantizer, TensorQuantizer +from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper from modelopt.torch.quantization.qtensor.nvfp4_tensor import _cast_per_block_scale_to_fp8 @@ -77,13 +77,8 @@ fsdp2_aware_weight_update, module_name_maps, quantizer_attr_names, - representative_weight_quantizer, - weight_attr_names, -) -from modelopt.torch.quantization.utils.core_utils import ( - enable_weight_access_and_writeback, - has_accelerate_offload, ) +from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.dataset_utils import _disable_use_cache from modelopt.torch.utils.distributed import is_fsdp2_model @@ -136,7 +131,6 @@ preprocess_linear_fusion, sync_tied_input_amax, to_quantized_weight, - validate_ggml_layer, ) from .registry import ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry @@ -1858,78 +1852,6 @@ def read_unplaced_weights(model: nn.Module, *, keys_only: bool = False) -> dict[ return out -# What export_hf_checkpoint(upcast_ggml=...) can decode GGML weights to. -_GGML_UPCAST_DTYPES: dict[str, torch.dtype] = {"bf16": torch.bfloat16} - - -def _is_ggml_weight_quantizer(quantizer) -> bool: - """Whether ``quantizer`` is an enabled GGML weight quantizer.""" - # getattr: a SequentialQuantizer has no num_bits, and is never GGML. - return getattr(quantizer, "num_bits", None) in GGML_FORMATS and quantizer.is_enabled - - -@torch.no_grad() -def _upcast_ggml_weights(model: nn.Module, dtype: torch.dtype) -> None: - """Decode each GGML-quantized weight to ``dtype`` in place and disable its quantizer. - - The weight becomes what its packed export payload decodes to -- the payload fake quant cached - or GPTQ pinned, when there is one. With its quantizer disabled, the layer then exports as an - unquantized one and stays out of the quantization config, so a runtime without GGML kernels, - such as vLLM, loads the checkpoint as plain weights. - """ - names = module_name_maps(model) - # Read only the quantizers here, so an offloaded module is materialized only if it has work. - ggml_modules = [ - module - for module in model.modules() - if isinstance(module, QuantModule) - and any(map(_is_ggml_weight_quantizer, module.iter_weight_quantizers_for_calibration())) - ] - if not ggml_modules: - warnings.warn("upcast_ggml is set, but the model has no GGML-quantized weights.") - return - - # Check every layer before changing any, so a refusal leaves the model as it was. - for module in ggml_modules: - for weight_name in weight_attr_names(module): - weight_quantizer = representative_weight_quantizer(module, weight_name) - if not _is_ggml_weight_quantizer(weight_quantizer): - continue - # Disabling the quantizer would skip the checks the packed export runs on the layer. - input_quantizer = getattr( - module, quantizer_attr_names(weight_name).input_quantizer, None - ) - validate_ggml_layer(weight_quantizer, input_quantizer) - name = f"{names.module_to_name.get(id(module), '')}.{weight_name}" - weight = getattr(module, weight_name, None) - if not isinstance(weight, torch.Tensor): - # A grouped layout, such as TEGroupedLinear's weight0..N; the packed export - # supports only a standard weight attribute too. - raise NotImplementedError( - f"upcast_ggml needs each GGML weight stored under its own name, but {name} " - f"of {type(module).__name__} is not a tensor." - ) - # Fused-expert slices and the offload and FSDP2 writebacks keep the stored dtype, so - # the decoded values can only be written to a weight already stored in ``dtype``. - if weight.dtype != dtype: - raise ValueError( - f"upcast_ggml decodes to {dtype}, but {name} is stored as {weight.dtype}. " - f"Load the model in {dtype} before quantizing it." - ) - - for module in ggml_modules: - with enable_weight_access_and_writeback(module, model, names): - for weight, quantizer in module.iter_weights_for_calibration(): - if not _is_ggml_weight_quantizer(quantizer): - continue - ggml_format = GGML_FORMAT_REGISTRY[quantizer.num_bits] - packed_weight = ggml_format.pack(weight, quantizer) - weight.data.copy_( - ggml_format.dequantize(packed_weight, torch.tensor(weight.shape), dtype=dtype) - ) - quantizer.disable() - - def export_hf_checkpoint( model: Any, dtype: torch.dtype | None = None, @@ -1938,7 +1860,6 @@ def export_hf_checkpoint( components: list[str] | None = None, extra_state_dict: dict[str, torch.Tensor] | None = None, max_shard_size: int | str = "10GB", - upcast_ggml: str | None = None, **kwargs, ): """Export quantized HuggingFace model checkpoint (transformers or diffusers). @@ -1962,12 +1883,6 @@ def export_hf_checkpoint( to export. If None, all quantized components are exported. extra_state_dict: Extra state dictionary to add to the exported model. max_shard_size: Maximum size of each safetensors shard file. Defaults to "10GB". - upcast_ggml: The dtype to upcast GGML-quantized weights to, currently only ``"bf16"``. - Each GGML weight (the IQ formats and Q8_0) is written as the values its packed blocks - decode to in that dtype, rather than as the blocks, and is left out of the - quantization config. Runtimes without GGML kernels, such as vLLM, can then load the - checkpoint to evaluate the quantized model. Not supported for diffusers pipelines or - a ``layerwise.export_dir`` export. **kwargs: Runtime-specific post-processing options forwarded to :func:`_postprocess_safetensors` for diffusion model exports. See its docstring for supported keys. @@ -2001,18 +1916,6 @@ def export_hf_checkpoint( from .layerwise_export import LAYERWISE_EXPORTER_ATTR exporter = getattr(model, LAYERWISE_EXPORTER_ATTR, None) - if upcast_ggml is not None: - if upcast_ggml not in _GGML_UPCAST_DTYPES: - raise ValueError( - f"upcast_ggml must be one of {sorted(_GGML_UPCAST_DTYPES)}, got {upcast_ggml!r}" - ) - if exporter is not None or (HAS_DIFFUSERS and is_diffusers_object(model)): - raise NotImplementedError( - "upcast_ggml supports transformers models only, and not a layerwise.export_dir " - "export, whose layer shards calibration already wrote packed." - ) - _upcast_ggml_weights(model, _GGML_UPCAST_DTYPES[upcast_ggml]) - if exporter is not None: # Per-layer export wrote the shards during calibration; this writes the rest. exporter.finalize(extra_state_dict=extra_state_dict) diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index 8f2aa927ea5..28cabbe52d2 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -305,8 +305,10 @@ def _pinned_rows(self, quantizer, weight: torch.Tensor) -> torch.Tensor | None: """The pinned payload of ``weight``, a view of the pinned weight or rows taken from it. Fake quant hands over the pinned weight itself in 256-value rows; export may hand over row - slices of it, such as the gate and up halves of a fused projection. The comparison syncs, - but it only runs for weights the identity cache has not seen. + slices of it, such as the gate and up halves of a fused projection. Folding hands over a + float32 copy of a narrower weight, which holds the decoding rounded to that narrower dtype, + so a float32 weight also matches the decoding rounded through BF16 or FP16. The comparison + syncs, but it only runs for weights the identity cache has not seen. """ pinned = pinned_packed_weight(quantizer, self.name) width = weight.shape[-1] // self.block_size @@ -317,17 +319,22 @@ def _pinned_rows(self, quantizer, weight: torch.Tensor) -> torch.Tensor | None: ): return None pinned = pinned.to(weight.device).reshape(-1, width, self.block_bytes) - table = self.dequantize( + decoded = self.dequantize( pinned, torch.tensor((pinned.shape[0], weight.shape[-1])), dtype=torch.float32 - ).to(weight.dtype) + ) rows = weight.reshape(-1, weight.shape[-1]) - if rows.shape == table.shape and torch.equal(rows, table): - index = slice(None) - else: - index = _matching_rows(table, rows) - if index is None: - return None - return pinned[index].reshape(*weight.shape[:-1], width, self.block_bytes) + storage_dtypes = [weight.dtype] + if weight.dtype == torch.float32: + storage_dtypes += [torch.bfloat16, torch.float16] + for storage_dtype in storage_dtypes: + table = decoded.to(storage_dtype).to(weight.dtype) + if rows.shape == table.shape and torch.equal(rows, table): + index = slice(None) + else: + index = _matching_rows(table, rows) + if index is not None: + return pinned[index].reshape(*weight.shape[:-1], width, self.block_bytes) + return None def narrow_to_float32(blocks: torch.Tensor) -> torch.Tensor: diff --git a/tests/gpu/torch/export/test_fsdp2_export.py b/tests/gpu/torch/export/test_fsdp2_export.py index 84870a630b4..939fe3e581f 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/tests/gpu/torch/export/test_fsdp2_export.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import copy -import dataclasses import json from functools import partial from pathlib import Path @@ -34,16 +33,12 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.layer_utils import is_quantlinear from modelopt.torch.export.model_utils import TiedWeightMap -from modelopt.torch.export.quant_utils import get_quant_config from modelopt.torch.export.unified_export_hf import ( _export_quantized_weight, _export_transformers_checkpoint, - _upcast_ggml_weights, requantize_resmooth_fused_llm_layers, ) from modelopt.torch.export.unified_export_hf_streaming import _export_fsdp2_checkpoint_streaming -from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY -from modelopt.torch.quantization.ggml.common import pin_packed_weight from modelopt.torch.quantization.utils import ( enable_weight_access_and_writeback, fsdp2_aware_weight_update, @@ -468,89 +463,3 @@ def test_fsdp2_streaming_export_matches_reference(dist_workers, tmp_path, quant_ quant_config=quant_config, ) ) - - -_Q8_0_MLP_CFG = { - "quant_cfg": [ - {"quantizer_name": "*", "enable": False}, - { - "quantizer_name": "*mlp*weight_quantizer", - "cfg": {"num_bits": "q8_0", "block_sizes": {-1: 32}, "backend": "ggml"}, - "enable": True, - }, - ], - "algorithm": None, -} - - -def _pin_ggml_payloads(model): - """Leave each GGML layer as GPTQ does: a payload pinned, and the weight equal to its decoding.""" - q8_0 = GGML_FORMAT_REGISTRY["q8_0"] - names = module_name_maps(model) - for module in list(model.modules()): - quantizer = getattr(module, "weight_quantizer", None) - if not (is_quantlinear(module) and quantizer is not None and quantizer.is_enabled): - continue - with enable_weight_access_and_writeback(module, model, names): - payload, shape = q8_0.quantize(module.weight.flip(-1)) - module.weight.copy_(q8_0.dequantize(payload, shape, dtype=module.weight.dtype)) - pin_packed_weight(quantizer, "q8_0", payload) - - -def _no_search(*args, **kwargs): - raise AssertionError("the upcast re-ran the q8_0 search instead of decoding the pinned payload") - - -def _upcast_ggml_matches_reference_test(rank, size, pinned): - """A sharded model's upcast holds what a single-process packed export decodes to. - - The upcast decodes each GGML weight inside a gather window and writes it back to the shards, - so after a reshard every rank must gather the decoded payload, and the GGML layers must have - dropped out of the quantization config. With ``pinned``, every GGML layer carries a payload - pinned as GPTQ pins one, and the encoder is disabled to prove the pinned codes are decoded. - - This calls the upcast directly rather than ``export_hf_checkpoint``, whose FSDP2 streaming - writer currently deadlocks at world > 1 (see test_fsdp2_streaming_export_matches_reference). - """ - q8_0 = GGML_FORMAT_REGISTRY["q8_0"] - with patch_fsdp_mp_dtypes(), torch.no_grad(): - model = get_tiny_llama(num_hidden_layers=2).to("cuda").eval() - reference = get_tiny_llama(num_hidden_layers=2).to("cuda").eval() - for layer in model.model.layers: - fully_shard(layer) - fully_shard(model) - mtq.quantize(model, _Q8_0_MLP_CFG) - mtq.quantize(reference, _Q8_0_MLP_CFG) - if pinned: - _pin_ggml_payloads(model) - _pin_ggml_payloads(reference) - # The pool reuses its workers across tests, so the patch is undone below. - GGML_FORMAT_REGISTRY["q8_0"] = dataclasses.replace(q8_0, quantize=_no_search) - try: - packed, _ = _export_transformers_checkpoint(reference, torch.bfloat16) - _upcast_ggml_weights(model, torch.bfloat16) - finally: - GGML_FORMAT_REGISTRY["q8_0"] = q8_0 - - assert get_quant_config(model)["quantization"]["quant_algo"] is None - ggml_keys = [key for key, tensor in packed.items() if tensor.dtype == torch.uint8] - assert ggml_keys, "the reference export packed no GGML weight" - names = module_name_maps(model) - for key in ggml_keys: - payload = packed[key] - expected = q8_0.dequantize( - payload, - torch.tensor((*payload.shape[:-2], payload.shape[-2] * q8_0.block_size)), - dtype=torch.bfloat16, - ) - module = names.name_to_module[key.removesuffix(".weight")] - # Gather again after the upcast's window closed, so this reads what reached the shards. - with enable_weight_access_and_writeback(module, model, names, writeback=False): - assert torch.equal(module.weight, expected.to(module.weight.device)), key - - -@pytest.mark.parametrize("pinned", [False, True], ids=["encoded", "gptq_pinned"]) -def test_fsdp2_upcast_ggml_matches_reference(dist_workers, pinned): - if torch.cuda.device_count() < 2: - pytest.skip("needs >=2 GPUs for the weights to actually be sharded") - dist_workers.run(partial(_upcast_ggml_matches_reference_test, pinned=pinned)) diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index 1917e2897aa..acbe2f758b8 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/tests/unit/torch/export/test_export_weight.py @@ -14,7 +14,6 @@ # limitations under the License. -import copy import dataclasses import pytest @@ -23,17 +22,15 @@ from _test_utils.torch.export.utils import ToyModel, partial_fp8_config, partial_w4a8_config import modelopt.torch.quantization as mtq -from modelopt.torch.export.quant_format import QUANTIZATION_NONE -from modelopt.torch.export.quant_utils import get_quantization_format, postprocess_state_dict +from modelopt.torch.export.quant_utils import postprocess_state_dict from modelopt.torch.export.unified_export_hf import ( _export_quantized_weight, _process_quantized_modules, - _upcast_ggml_weights, ) from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.ggml.common import pin_packed_weight -from modelopt.torch.quantization.nn import GroupedQuantizer, TensorQuantizer +from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.quantization.utils import quantizer_attr_names @@ -203,127 +200,6 @@ def test_export_drops_the_payload_gptq_pinned(): assert list(state_dict) == ["weight"] -def _ggml_model(num_bits, dtype=torch.bfloat16, input_cfg=None): - """A one-linear model whose weight quantizer is ``num_bits``; ``input_cfg`` enables its input.""" - model = nn.Sequential(nn.Linear(512, 4, bias=False, dtype=dtype)) - quant_cfg = [ - {"quantizer_name": "*", "enable": False}, - { - "quantizer_name": "*weight_quantizer", - "cfg": {"num_bits": num_bits, "backend": "ggml"}, - "enable": True, - }, - ] - if input_cfg is not None: - quant_cfg.append({"quantizer_name": "*input_quantizer", "cfg": input_cfg, "enable": True}) - return mtq.quantize(model, {"quant_cfg": quant_cfg, "algorithm": None}) - - -@pytest.mark.parametrize("num_bits", sorted(GGML_FORMAT_REGISTRY)) -def test_upcast_ggml_writes_what_the_packed_export_decodes_to(num_bits): - """The upcast weight is exactly the decoding of the payload the packed export writes.""" - model = _ggml_model(num_bits) - packed = copy.deepcopy(model[0]) - _export_quantized_weight(packed, torch.bfloat16) - - _upcast_ggml_weights(model, torch.bfloat16) - - expected = GGML_FORMAT_REGISTRY[num_bits].dequantize( - packed.weight, torch.tensor((4, 512)), dtype=torch.bfloat16 - ) - assert torch.equal(model[0].weight, expected) - # The layer now exports as an unquantized one, so it stays out of the quantization config. - assert get_quantization_format(model) == QUANTIZATION_NONE - - -@pytest.mark.parametrize( - ("dtype", "input_cfg", "error"), - [ - (torch.bfloat16, {"num_bits": (4, 3)}, "enabled input quantizer"), - (torch.bfloat16, "pre_quant_scale", "pre_quant_scale"), - (torch.float16, None, "stored as torch.float16"), - (torch.float32, None, "stored as torch.float32"), - ], - ids=["input_quantizer", "pre_quant_scale", "fp16_weight", "fp32_weight"], -) -def test_upcast_ggml_refuses_a_layer_before_changing_any(dtype, input_cfg, error): - """Activation quantization the upcast would drop, or a weight not stored in BF16, is refused. - - The packed export refuses the first two too; disabling the quantizer must not skip that. The - model is left as it was, so a refusal never leaves it half upcast. - """ - model = _ggml_model("q8_0", dtype, None if input_cfg == "pre_quant_scale" else input_cfg) - if input_cfg == "pre_quant_scale": - model[0].input_quantizer.pre_quant_scale = torch.ones(512) - weight = model[0].weight.detach().clone() - - with pytest.raises((NotImplementedError, ValueError), match=error): - _upcast_ggml_weights(model, torch.bfloat16) - - assert torch.equal(model[0].weight, weight) - assert model[0].weight_quantizer.is_enabled - - -@pytest.mark.parametrize("pinned", [False, True]) -def test_upcast_ggml_writes_back_an_offloaded_weight(monkeypatch, pinned): - """The decoded values reach the offloaded copy; a payload GPTQ pinned is decoded, not redone.""" - pytest.importorskip("accelerate") - # Local, so importorskip can skip this test when the optional accelerate is not installed. - from accelerate.hooks import AlignDevicesHook, add_hook_to_module - from accelerate.utils import set_module_tensor_to_device - - model = _ggml_model("q8_0") - linear = model[0] - q8_0 = GGML_FORMAT_REGISTRY["q8_0"] - payload, shape = q8_0.quantize(linear.weight.detach()) - expected = q8_0.dequantize(payload, shape, dtype=torch.bfloat16) - assert not torch.equal(linear.weight, expected) - if pinned: - # How GPTQ leaves a layer: its payload pinned, and the weight equal to the decoding. - with torch.no_grad(): - linear.weight.copy_(expected) - pin_packed_weight(linear.weight_quantizer, "q8_0", payload) - _without_search(monkeypatch, "q8_0") - weights_map = {"weight": linear.weight.detach().clone()} - add_hook_to_module( - linear, AlignDevicesHook(execution_device="cpu", offload=True, weights_map=weights_map) - ) - set_module_tensor_to_device(linear, "weight", "meta") - - _upcast_ggml_weights(model, torch.bfloat16) - - assert linear.weight.is_meta - assert torch.equal(weights_map["weight"], expected) - - -def test_upcast_ggml_refuses_a_grouped_layout(): - """A grouped layer keeps its GGML weights as weight0..N, not under the name it reports. - - Stands in for TEGroupedLinear: one GroupedQuantizer, which weight_attr_names reports as - "weight", over expert weights that only its own iter_weights_for_calibration yields. - """ - linear = _ggml_model("q8_0")[0] - linear.weight0 = linear.weight - del linear.weight - linear.weight_quantizer = GroupedQuantizer(linear.weight_quantizer) - linear.iter_weights_for_calibration = lambda: iter( - [(linear.weight0, linear.weight_quantizer[0])] - ) - weight = linear.weight0.detach().clone() - - with pytest.raises(NotImplementedError, match="stored under its own name"): - _upcast_ggml_weights(nn.Sequential(linear), torch.bfloat16) - - assert torch.equal(linear.weight0, weight) - assert linear.weight_quantizer[0].is_enabled - - -def test_upcast_ggml_warns_when_there_is_nothing_to_upcast(): - """Asking for an upcast of a model with no GGML weights is most likely a mistake.""" - with pytest.warns(UserWarning, match="no GGML-quantized weights"): - _upcast_ggml_weights(nn.Sequential(nn.Linear(4, 4)), torch.bfloat16) - - class QuantMoELinear(nn.Module): def __init__(self): super().__init__() diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 45738445015..9c9467ba17e 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -28,7 +28,6 @@ from safetensors.torch import save_file import modelopt.torch.quantization as mtq -from modelopt.torch.export.layerwise_export import LAYERWISE_EXPORTER_ATTR from modelopt.torch.export.model_utils import ( TiedWeightMap, get_language_model_from_vl, @@ -39,12 +38,7 @@ postprocess_state_dict, sync_tied_input_amax, ) -from modelopt.torch.export.unified_export_hf import ( - _resolve_export_dtype, - export_hf_checkpoint, - read_unplaced_weights, -) -from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY +from modelopt.torch.export.unified_export_hf import _resolve_export_dtype, read_unplaced_weights from modelopt.torch.quantization.nn import TensorQuantizer @@ -1081,63 +1075,3 @@ def test_union_degrades_to_the_recorded_set_without_the_loader(tmp_path, monkeyp carried = read_unplaced_weights(model) assert key in carried, "recorded keys must still carry when the structural pass is unavailable" - - -# TODO: cover the FSDP2 export path too, with a multi-rank upcast test in -# tests/gpu/torch/export/test_fsdp2_export.py. -@pytest.mark.parametrize("moe", [False, True]) -def test_upcast_ggml_export_loads_as_a_plain_checkpoint(tmp_path, moe): - """With upcast_ggml="bf16" the checkpoint carries no quantization config, so a runtime without - GGML kernels loads it as an ordinary BF16 model, and every GGML weight holds its decoded payload. - - The MoE model exercises the fused experts, which export splits into per-expert projections. - """ - pytest.importorskip("transformers.core_model_loading") - # Local import: transformers and its test fixtures are optional dependencies. - from _test_utils.torch.transformers_models import get_tiny_qwen3, get_tiny_qwen3_moe - from transformers import AutoModelForCausalLM - - model = (get_tiny_qwen3_moe if moe else get_tiny_qwen3)() - original = {name: param.detach().clone() for name, param in model.named_parameters()} - mtq.quantize( - model, - { - "quant_cfg": [ - {"quantizer_name": "*", "enable": False}, - { - "quantizer_name": "*mlp*weight_quantizer", - "cfg": {"num_bits": "q8_0", "backend": "ggml"}, - "enable": True, - }, - ], - "algorithm": None, - }, - ) - - export_hf_checkpoint(model, export_dir=tmp_path, upcast_ggml="bf16") - - assert "quantization_config" not in json.loads((tmp_path / "config.json").read_text()) - assert not (tmp_path / "hf_quant_config.json").exists() - reloaded = AutoModelForCausalLM.from_pretrained(tmp_path, dtype=torch.bfloat16) - q8_0 = GGML_FORMAT_REGISTRY["q8_0"] - for name, param in reloaded.named_parameters(): - expected = original[name] - # The MoE router (mlp.gate) is not a linear, so it is never quantized. - if ".mlp." in name and not name.endswith(".mlp.gate.weight"): - expected = q8_0.dequantize(*q8_0.quantize(expected), dtype=expected.dtype) - assert torch.equal(param, expected), name - - -def test_upcast_ggml_rejects_an_unsupported_dtype(tmp_path): - """Only the dtypes the upcast knows are accepted.""" - with pytest.raises(ValueError, match="upcast_ggml must be one of"): - export_hf_checkpoint(torch.nn.Linear(2, 2), export_dir=tmp_path, upcast_ggml="fp8") - - -def test_upcast_ggml_rejects_a_layerwise_export(tmp_path): - """A layerwise export wrote its layer shards packed during calibration, too early to upcast.""" - model = torch.nn.Linear(2, 2) - setattr(model, LAYERWISE_EXPORTER_ATTR, object()) - - with pytest.raises(NotImplementedError, match=r"layerwise\.export_dir"): - export_hf_checkpoint(model, export_dir=tmp_path, upcast_ggml="bf16") diff --git a/tests/unit/torch/quantization/test_ggml_gptq.py b/tests/unit/torch/quantization/test_ggml_gptq.py index d32c6ae89cb..bc14dd51fdf 100644 --- a/tests/unit/torch/quantization/test_ggml_gptq.py +++ b/tests/unit/torch/quantization/test_ggml_gptq.py @@ -215,6 +215,26 @@ def test_pin_outlives_the_weight_tensor(): assert not torch.equal(IQ1_S.quantize(model.weight.detach())[0], _pin(model)) +def test_folding_a_bf16_weight_decodes_the_pinned_payload(): + """Folding fake-quantizes a float32 copy of the weight; it must still decode GPTQ's codes. + + A BF16 weight holds the pinned payload's decoding rounded to BF16, so the float32 copy that + ``mtq.fold_weight`` and the vLLM fake-quant export hand to the quantizer is not that decoding + in float32. Re-encoding it instead would fold in different codes than GPTQ chose. + """ + torch.manual_seed(0) + model = torch.nn.Linear(512, 8, bias=False, dtype=torch.bfloat16) + inputs = _inputs().to(torch.bfloat16) + mtq.quantize(model, _iq_config(GPTQ), forward_loop=lambda m: m(inputs)) + decoded = _decoded(_pin(model), model.weight).to(torch.bfloat16) + assert torch.equal(model.weight, decoded) + + # The vLLM fake-quant export folds this way. + assert torch.equal(model.weight_quantizer(model.weight.float()).to(torch.bfloat16), decoded) + mtq.fold_weight(model) + assert torch.equal(model.weight, decoded) + + def test_pin_survives_save_and_restore(): model, inputs = _gptq_model() buffer = io.BytesIO() From 7331bdf706942e72806906aabaed258070b2dec0 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Sun, 11 Oct 2026 17:08:59 +0000 Subject: [PATCH 5/6] Bound the pinned-payload lookup's memory and test every GGML format - Decode the pinned payload straight into each candidate storage dtype, holding one table at a time instead of a float32 decode beside each rounding, and run every whole-weight comparison before falling back to row matching, whose int64 keys are several times the weight's size and which only row slices need. For a 17408x5120 weight the lookup's peak extra GPU memory drops from 596 to 256 MiB for a BF16 weight and from 1700 to 510 MiB for the float32 copy folding passes (510 and 1700 MiB before this PR). - Check that the packed export drops the pinned payload for every GGML format, not only Q8_0. - Scope the README's stock-vLLM note to models whose quantizers are all weight quantizers. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- examples/hf_ptq/README.md | 2 +- modelopt/torch/quantization/ggml/common.py | 22 ++++++++++++------- tests/unit/torch/export/test_export_weight.py | 9 ++++---- 3 files changed, 20 insertions(+), 13 deletions(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 2252f5643e7..beb4f0b0872 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -671,7 +671,7 @@ with torch.inference_mode(): python hf_ptq.py --pyt_ckpt_path --qformat fp8 --export_path --trust_remote_code ``` -> *For exporting fake-quantized models for vLLM serving (e.g., for research or kernels not yet supported in real-quant), use the `--vllm_fakequant_export` flag. See [vllm_serve/README.md](../vllm_serve/README.md) for details. For weight-only formats such as the GGML formats (IQ, Q8_0), the export is a plain BF16 checkpoint that a stock vLLM image serves directly.* +> *For exporting fake-quantized models for vLLM serving (e.g., for research or kernels not yet supported in real-quant), use the `--vllm_fakequant_export` flag. See [vllm_serve/README.md](../vllm_serve/README.md) for details. When every quantizer is a weight quantizer, as in the GGML recipes (IQ, Q8_0), the export is a plain BF16 checkpoint that a stock vLLM image serves directly.* ### Hugging Face framework [Script](./scripts/huggingface_example.sh) diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index 28cabbe52d2..1e21abf59df 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -319,19 +319,25 @@ def _pinned_rows(self, quantizer, weight: torch.Tensor) -> torch.Tensor | None: ): return None pinned = pinned.to(weight.device).reshape(-1, width, self.block_bytes) - decoded = self.dequantize( - pinned, torch.tensor((pinned.shape[0], weight.shape[-1])), dtype=torch.float32 - ) + shape = torch.tensor((pinned.shape[0], weight.shape[-1])) rows = weight.reshape(-1, weight.shape[-1]) storage_dtypes = [weight.dtype] if weight.dtype == torch.float32: storage_dtypes += [torch.bfloat16, torch.float16] + + def decoding(storage_dtype): + # Decoding straight into the storage dtype rounds once, as GPTQ did writing the weight, + # and holds one table at a time rather than a float32 copy beside each rounding. + return self.dequantize(pinned, shape, dtype=storage_dtype).to(weight.dtype) + + # Every whole-weight comparison first: matching rows builds int64 keys several times the + # weight's size, and only row slices need it. + if rows.shape[0] == pinned.shape[0]: + for storage_dtype in storage_dtypes: + if torch.equal(rows, decoding(storage_dtype)): + return pinned.reshape(*weight.shape[:-1], width, self.block_bytes) for storage_dtype in storage_dtypes: - table = decoded.to(storage_dtype).to(weight.dtype) - if rows.shape == table.shape and torch.equal(rows, table): - index = slice(None) - else: - index = _matching_rows(table, rows) + index = _matching_rows(decoding(storage_dtype), rows) if index is not None: return pinned[index].reshape(*weight.shape[:-1], width, self.block_bytes) return None diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index acbe2f758b8..c29f25379c2 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/tests/unit/torch/export/test_export_weight.py @@ -188,11 +188,12 @@ def test_export_repacks_a_weight_changed_since_fake_quant(num_bits): assert torch.equal(linear.weight, expected) -def test_export_drops_the_payload_gptq_pinned(): +@pytest.mark.parametrize("num_bits", sorted(GGML_FORMAT_REGISTRY)) +def test_export_drops_the_payload_gptq_pinned(num_bits): """GPTQ pins its payload to the quantizer as a buffer; only the packed weight is exported.""" - linear = _iq_linear("q8_0") - packed, _ = GGML_FORMAT_REGISTRY["q8_0"].quantize(linear.weight.detach()) - pin_packed_weight(linear.weight_quantizer, "q8_0", packed) + linear = _iq_linear(num_bits) + packed, _ = GGML_FORMAT_REGISTRY[num_bits].quantize(linear.weight.detach()) + pin_packed_weight(linear.weight_quantizer, num_bits, packed) _export_quantized_weight(linear, torch.bfloat16) state_dict = postprocess_state_dict(linear.state_dict(), maxbound=448, quantization=None) From 72303f4383df501715d817d08cf2d1b94c2ac78c Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Mon, 12 Oct 2026 16:53:49 +0000 Subject: [PATCH 6/6] Hash pinned-payload rows a column chunk at a time Matching row slices of a pinned weight, as Megatron's gate/up split does, built int64 copies of the whole decoded table and of its products at once: about 9x the weight. Sum the keys 256 columns at a time instead. Integer sums wrap and are associative, so the keys are bit-identical. For the gate half of a fused 34816x5120 weight, peak extra GPU memory drops from 3060 to 597 MiB (BF16) and from 3400 to 1107 MiB (float32 export dtype), at the same speed. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/ggml/common.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index 1e21abf59df..5e83b9c79c8 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -26,6 +26,8 @@ GGML_BLOCK_SIZE = 256 _SIGNED_INT_OF_SIZE = {2: torch.int16, 4: torch.int32, 8: torch.int64} +# Columns hashed at a time, bounding _row_keys' int64 temporaries to this slice of the weight. +_ROW_KEY_CHUNK = 256 class _CacheKey(NamedTuple): @@ -168,13 +170,18 @@ def _row_keys(x: torch.Tensor) -> torch.Tensor: """A 64-bit hash of each row's bits; equal rows always get equal keys. Integer sums wrap and are associative, so unlike a floating-point projection the key does not - depend on how the reduction is ordered, which differs between differently shaped products. + depend on how the reduction is ordered, which differs between differently shaped products. That + is also what lets the columns be summed a chunk at a time. """ multipliers = torch.randint( -(2**62), 2**62, (x.shape[-1],), generator=torch.Generator().manual_seed(0) - ) - bits = x.contiguous().view(_SIGNED_INT_OF_SIZE[x.element_size()]).to(torch.int64) - return (bits * multipliers.to(x.device)).sum(-1) + ).to(x.device) + bits = x.contiguous().view(_SIGNED_INT_OF_SIZE[x.element_size()]) + keys = torch.zeros(x.shape[:-1], dtype=torch.int64, device=x.device) + for start in range(0, x.shape[-1], _ROW_KEY_CHUNK): + columns = slice(start, start + _ROW_KEY_CHUNK) + keys += (bits[..., columns].to(torch.int64) * multipliers[columns]).sum(-1) + return keys def fake_quantize_with_cache(