[model_free_ptq] NVFP4A16 (#1988)

## Purpose ##
* Support NVFP4A16 for `model_free_ptq`

```bash
llmcompressor.reindex_fused_weights \
    unsloth/Kimi-K2-Thinking-BF16 \
    Kimi-K2-Thinking-BF16-reindexed \
    --num_workers=10
```

```python
model_free_ptq(
    model_stub="Kimi-K2-Thinking-BF16-reindexed",
    save_directory="Kimi-K2-Thinking-BF16-NVFP4A16",
    scheme="NVFP4A16",
    ignore=[
        "re:.*gate$",
        "lm_head",
        "re:.*kv_a_proj_with_mqa$",
        "re:.*q_a_proj$",
        "model.embed_tokens",
    ],
    max_workers=15,
    device="cuda:0",
)
```

## Changes ##
* Restructure files
  * Move `validate_scheme` to `validate.py`
* Move `find_safetensors_index_path`, `find_config_path`,
`find_safetensors_index_file` to `helpers.py`
  * Move `process_file` to `process.py`
  * Move `validate_scheme` to `validate.py`
* Break `calibrate_weights` into `calibrate_global_scale` and
`calibrate_scale_zp`
* Add extra utility functions
  * `match_names_set_eager`
  * `invert_mapping`
* Add microscale/fused module utility functions
  * `is_microscale_scheme`
  * `get_fused_names`
* Add `process_file_microscale_scheme` to separate the fp4 lifecycle
from the regular lifecycle (this script should be very trustworthy. By
separating the functions, an FP8 user does not have to trust anything
about FP4)
* Add `llm.compressor.reindex_fused_weights` script which reindexes a
model's weights so that fused modules are in the same files.
* Fix
[bug](https://github.com/vllm-project/llm-compressor/pull/1988/files#diff-8d11f284a49f6c4e559617aaf7750f3437a074cd526ee94dbefe86866f250a42R80-R82)
where safetensors index metadata was not being saved correctly

## Testing ##
* Add NVFP4A16 to `test_model_free_ptq_matches_oneshot`
* Regression tested large mistral model e2e with FP8_BLOCK
* Tested large mistral model e2e with NVFP4A16

## Mistral 3 ##
This branch was used to quantize Mistral 3

1. Quantize to W4A16
```python3
from llmcompressor import model_free_ptq

model_free_ptq(
    "mistralai/Mistral-Large-3-675B-Instruct-2512",
    "Mistral-Large-3-675B-Instruct-2512-FP8_BLOCK",
    scheme="NVFP4_A16",
    ignore=[
        "tok_embeddings",  # embeddings
        "re:patch_merger.*",  # patch merger
        "re:vision_encoder.*",  # vision tower
        "re:vision_language_adapter.*",  # vision adapter
        "re:.*attention$",  # sensitive to quantization
        "re:.*gate$",  # sensitive to quantization
        "output",  # lm head
    ],
    max_workers=10,  # 10 = 52Gb
    device="cuda:0",
)
```

2. Update ignore list to use vLLM checkpoint format
```
[
  "model.embed_tokens",
  "re:patch_merger.*",
  "re:vision_encoder.*",
  "re:vision_language_adapter.*",
  "lm_head",
  "re:.*self_attn.*",
  "re:.*gate$"
]
```

3. Add observers to vLLM model definition and run for 100 samples from
ultrachat
4. Save model checkpoint, making sure to reduce values from shards
For more information on how observers were added to vLLM, please reach
out to @kylesayrs

---------

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Co-authored-by: Dipika Sikka <dipikasikka1@gmail.com>

Author: kylesayrs

examples/model_free_ptq/README.md

@@ -13,3 +13,53 @@
 In `kimi_k2_thinking_fp8_block.py`, we call `model_free_ptq` by providing a `scheme` and `ignore` list, similar to how we provide reicpes to `oneshot` calls. In the case of Kimi-K2 Thinking, we apply the `FP8_BLOCK` scheme and ignore layers that are incompatible with a block_size of 128 (specifically, `kv_a_proj_with_mqa` and `q_a_proj`).
 
 In contrast to `oneshot`, we expect the model stub or pathway string to be directly passed in, as opposed to first being loaded through transformers. Once complete, the model is compressed using compressed-tensors and saved to `SAVE_DIR`.
+
+To get started, simply call `model_free_ptq` with your desired model stub and save directory
+```python
+model_free_ptq(
+    model_stub="unsloth/Kimi-K2-Thinking-BF16",
+    save_directory="Kimi-K2-Thinking-FP8-BLOCK",
+    scheme="FP8_BLOCK",
+    ignore=[
+        "re:.*gate$",
+        "lm_head",
+        "re:.*kv_a_proj_with_mqa$",
+        "re:.*q_a_proj$",
+        "model.embed_tokens",
+    ],
+    max_workers=15,
+    device="cuda:0",
+)
+
+```
+
+
+# Quantizing models to NVFP4A16/ MXFP4A16
+
+Using model_free_ptq to quantize models with microscale schemes (NVFP4/MXFP4) is the same as quantizing models with non-microscale schemes, except for one additional step. That extra step is that the safetensors in the model files must be reindexed to ensure that fused modules (qkv, gate_up) end up in the same safetensors files, which allows model_free_ptq to fuse global scales.
+
+First, apply `llmcompressor.reindex_fused_weights` from the command line entrypoint
+```bash
+llmcompressor.reindex_fused_weights \
+    unsloth/Kimi-K2-Thinking-BF16 \
+    Kimi-K2-Thinking-BF16-reindexed \
+    --num_workers=10
+```
+
+Then, call `model_free_ptq` on the reindex files
+```python
+model_free_ptq(
+    model_stub="Kimi-K2-Thinking-BF16-reindexed",
+    save_directory="Kimi-K2-Thinking-BF16-NVFP4A16",
+    scheme="NVFP4A16",
+    ignore=[
+        "re:.*gate$",
+        "lm_head",
+        "re:.*kv_a_proj_with_mqa$",
+        "re:.*q_a_proj$",
+        "model.embed_tokens",
+    ],
+    max_workers=15,
+    device="cuda:0",
+)
+```

examples/model_free_ptq/kimi_k2_thinking_fp8_block.py

@@ -1,7 +1,7 @@
 from llmcompressor import model_free_ptq
 
 MODEL_ID = "unsloth/Kimi-K2-Thinking-BF16"
-SAVE_DIR = "Kimi-K2-Thinking-FP8-Block"
+SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-FP8-BLOCK"
 
 # Apply FP8-Block to the model
 # Once quantized, the model is saved

examples/model_free_ptq/kimi_k2_thinking_nvfp4a16.py

@@ -0,0 +1,36 @@
+"""
+NOTE: Please run the following script before using `model_free_ptq`
+
+This script is used to reindex the safetensors files of a model such that all fused
+modules (gate_up, qkv) are in the same safetensors file. This is required by
+model_free_ptq for microscale schemes (NVFP4A16, MXFP4A16)
+
+llmcompressor.reindex_fused_weights \
+    unsloth/Kimi-K2-Thinking-BF16 \
+    Kimi-K2-Thinking-BF16-reindexed \
+    --num_workers=10
+"""
+
+from llmcompressor import model_free_ptq
+
+MODEL_ID = "unsloth/Kimi-K2-Thinking-BF16"
+REINDEX_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-reindexed"
+SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4A16"
+
+# See above notice pertaining to safetensors reindexing
+# After running `llmcompressor.reindex_fused_weights`,
+# use `model_free_ptq` to apply NVFP4A16 quantization
+model_free_ptq(
+    model_stub=REINDEX_DIR,
+    save_directory=SAVE_DIR,
+    scheme="NVFP4A16",
+    ignore=[
+        "re:.*gate$",
+        "lm_head",
+        "re:.*kv_a_proj_with_mqa$",
+        "re:.*q_a_proj$",
+        "model.embed_tokens",
+    ],
+    max_workers=15,
+    device="cuda:0",
+)

setup.py

@@ -188,6 +188,7 @@ def localversion_func(version: ScmVersion) -> str:
     entry_points={
         "console_scripts": [
             "llmcompressor.trace=llmcompressor.transformers.tracing.debug:main",
+            "llmcompressor.reindex_fused_weights=llmcompressor.entrypoints.model_free.reindex_fused_weights:main",
         ]
     },
     python_requires=">=3.10",

src/llmcompressor/entrypoints/model_free/__init__.py

@@ -7,27 +7,28 @@
 import torch
 import tqdm
 from compressed_tensors.quantization import QuantizationScheme
-from compressed_tensors.utils.match import _match_name
 from loguru import logger
-from safetensors.torch import load_file, save_file
 
-from llmcompressor.entrypoints.model_free.helpers import (
-    gpu_if_available,
-    validate_scheme,
-)
-from llmcompressor.entrypoints.model_free.lifecycle import (
-    calibrate_weights,
-    compress_module,
-    initialize_quantized_linear,
+from llmcompressor.entrypoints.model_free.helpers import gpu_if_available
+from llmcompressor.entrypoints.model_free.microscale import (
+    is_microscale_scheme,
 )
 from llmcompressor.entrypoints.model_free.model_utils import (
     get_checkpoint_files,
     is_weights_file,
 )
+from llmcompressor.entrypoints.model_free.process import (
+    process_file,
+    process_file_microscale_scheme,
+)
 from llmcompressor.entrypoints.model_free.save_utils import (
     update_config,
     update_safetensors_index,
 )
+from llmcompressor.entrypoints.model_free.validate import (
+    validate_safetensors_index,
+    validate_scheme,
+)
 
 __all__ = ["model_free_ptq"]
 
@@ -55,20 +56,24 @@ def model_free_ptq(
     model_files = get_checkpoint_files(model_stub)
     scheme_name, scheme = validate_scheme(scheme)
     device = gpu_if_available(device)
+    validate_safetensors_index(model_files, scheme)
 
     # 0. collect safetensors files, copy files
     jobs = []
-    for file_path, resolved_path in model_files:
+    job_fn = (
+        process_file
+        if not is_microscale_scheme(scheme)
+        else process_file_microscale_scheme
+    )
+    for file_path, resolved_path in model_files.items():
         save_path = Path(save_directory) / file_path
 
         if file_path.endswith("safetensors"):
-            jobs.append(
-                (_process_file, resolved_path, save_path, scheme, ignore, device)
-            )
+            jobs.append((job_fn, resolved_path, save_path, scheme, ignore, device))
 
         else:
             if is_weights_file(file_path):
-                logger.warning(f"Skipping weights file {file_path}")
+                logger.warning(f"Skip processing for weights file {file_path}")
             save_path.parent.mkdir(parents=True, exist_ok=True)
             logger.info(f"Copying {file_path} {save_path}")
             shutil.copyfile(resolved_path, save_path)
@@ -89,50 +94,3 @@ def model_free_ptq(
     # 5. update config and safetensors index
     update_config(save_directory, scheme_name, scheme, ignore)
     update_safetensors_index(save_directory, total_size, weight_map)
-
-
-def _process_file(
-    file_path: str | os.PathLike,
-    save_path: str | os.PathLike,
-    scheme: QuantizationScheme,
-    ignore: str | list[str],
-    device: str | torch.device,
-) -> tuple[int, dict[str, str]]:
-    """
-    Quantize and compress tensors in a given safetensors file
-
-    :param file_path: safetensors file to process
-    :param save_path: save path of file with quantized weights
-    :param scheme: quantization scheme to apply to tensors
-    :param ignore: modules to ignore. Modules ending with "norm" are automatically
-        ignored
-    :param device: device used to quantize and compress weights
-    """
-    tensors = load_file(file_path)
-
-    for name in list(tensors.keys()):
-        module_name, param_name = name.rsplit(".", 1)
-        is_linear_weight = param_name == "weight" and not module_name.endswith("norm")
-        is_ignored = any(_match_name(module_name, ign) for ign in ignore)
-        if not is_linear_weight or is_ignored:
-            continue
-
-        # 1. initialize module with qparams (on device)
-        module = initialize_quantized_linear(tensors[name], scheme, device)
-
-        # 2. calibrate weight qparams
-        calibrate_weights(module)
-
-        # 3. compress module using qparams
-        compress_module(module)
-
-        # 4. save compressed data (on cpu)
-        del tensors[name]
-        prefix = module_name + "."
-        for key, value in module.state_dict(prefix=prefix).items():
-            tensors[key] = value.to("cpu")
-
-    save_file(tensors, save_path)
-    total_size = sum(tensor.nbytes for tensor in tensors.values())
-    weight_map = {key: os.path.basename(save_path) for key in tensors.keys()}
-    return total_size, weight_map

src/llmcompressor/entrypoints/model_free/helpers.py

@@ -1,51 +1,25 @@
-from typing import Optional
+import os
+from collections import defaultdict
+from typing import Mapping, TypeVar
 
 import torch
-from compressed_tensors.quantization import QuantizationScheme, preset_name_to_scheme
-from compressed_tensors.utils import getattr_chain
 from compressed_tensors.utils.match import _match_name
 from loguru import logger
+from transformers.file_utils import CONFIG_NAME
 
-__all__ = ["validate_scheme", "gpu_if_available", "is_match_name"]
+__all__ = [
+    "gpu_if_available",
+    "find_safetensors_index_path",
+    "find_config_path",
+    "find_safetensors_index_file",
+    "match_names_set_eager",
+    "MatchedNamesSet",
+    "invert_mapping",
+]
 
-
-def validate_scheme(scheme: QuantizationScheme) -> tuple[str, QuantizationScheme]:
-    # treat strings as preset schemes
-    if isinstance(scheme, str):
-        scheme_name, scheme = scheme, preset_name_to_scheme(scheme, [])
-    else:
-        scheme_name = "config_group_0"
-
-    # weight quantization must be provided
-    if scheme.weights is None:
-        raise ValueError(
-            "Must provide a weights quanitization scheme to perform weights-only PTQ"
-        )
-
-    # activation quantization must be dynamic
-    input_dynamic = getattr_chain(scheme, "input_activations.dynamic", True)
-    output_dynamic = getattr_chain(scheme, "output_activations.dynamic", True)
-    if input_dynamic is not True or output_dynamic is not True:
-        raise ValueError(
-            "Model Free PTQ cannot calibrate activations. "
-            "Please use `oneshot` instead."
-        )
-
-    # override with static observers
-    # Remove after https://github.com/vllm-project/compressed-tensors/pull/489
-    if scheme.weights.observer in ("minmax", "mse"):
-        new_observer = f"static_{scheme.weights.observer}"
-        logger.warning(
-            f"Scheme uses {scheme.weights.observer} weight observer. "
-            f"Using {new_observer} instead"
-        )
-        scheme.weights.observer = new_observer
-
-    # target all modules; filter by ignore list
-    # technically this should be "re:.*", but vllm's
-    # ct moe layer has a hard coded check for "Linear"
-    scheme.targets = ["Linear"]
-    return scheme_name, scheme
+KeyType = TypeVar("K")
+ValueType = TypeVar("V")
+MatchedNamesSet = dict[str, str | None]
 
 
 def gpu_if_available(device: torch.device | str | None) -> torch.device:
@@ -63,13 +37,70 @@ def gpu_if_available(device: torch.device | str | None) -> torch.device:
         return torch.device("cpu")
 
 
-def is_match_name(
-    name: str, targets: list[str], ignore: Optional[str | list[str]] = None
-) -> bool:
-    targets = targets if isinstance(targets, list) else [targets]
-    ignore = ignore if isinstance(ignore, list) else [ignore]
+def find_safetensors_index_path(save_directory: str | os.PathLike) -> str | None:
+    for file_name in os.listdir(save_directory):
+        if file_name.endswith("safetensors.index.json"):
+            return os.path.join(save_directory, file_name)
+
+    return None
+
+
+def find_config_path(save_directory: str | os.PathLike) -> str | None:
+    for file_name in os.listdir(save_directory):
+        if file_name in (CONFIG_NAME, "params.json"):
+            return os.path.join(save_directory, file_name)
+
+    return None
+
+
+def find_safetensors_index_file(model_files: dict[str, str]) -> str | None:
+    for file_path, resolved_path in model_files.items():
+        if file_path.endswith("safetensors.index.json"):
+            return resolved_path
+
+    return None
+
+
+def match_names_set_eager(
+    names: set[str] | list[str],
+    targets: set[str] | list[str],
+    return_unmatched: bool = True,
+) -> list[MatchedNamesSet] | tuple[list[MatchedNamesSet], MatchedNamesSet]:
+    matched_sets = []
+    matches = dict.fromkeys(targets, None)
+
+    for name in names:
+        # match until we get a full set
+        for target in targets:
+            if _match_name(name, target):
+                if matches[target] is None:
+                    matches[target] = name
+                else:
+                    # matched target twice without completing a set
+                    raise ValueError(
+                        f"Matched a {target} twice before "
+                        f"completing set ({matches[target]}, {name})"
+                    )
+
+        # once we have a full set, yield and reset
+        if all((matches[target] is not None for target in targets)):
+            matched_sets.append(matches)
+            matches = dict.fromkeys(targets, None)
+
+    unmatched_set = matches if any((v is not None for v in matches.values())) else None
+
+    if return_unmatched:
+        return matched_sets, unmatched_set
+    else:
+        return matched_sets
+
+
+def invert_mapping(
+    mapping: Mapping[KeyType, ValueType],
+) -> dict[ValueType, list[KeyType]]:
+    inverse = defaultdict(list)
 
-    matches_target = any(_match_name(name, target) for target in targets)
-    matches_ignore = any(_match_name(name, ign) for ign in ignore)
+    for key, value in mapping.items():
+        inverse[value].append(key)
 
-    return matches_target and not matches_ignore
+    return inverse