[Distributed] [model_free_ptq] Eliminate reindexing step via fine-grained parallelized partial reads (#2498)

## Purpose
Eliminates the `reindex_fused_weights` preprocessing step for microscale
schemes (NVFP4, MXFP4) by enabling each shard to be processed
independently
with full parallelism, even when fused weight sets (q/k/v, gate/up) span
multiple shards.

## Approach
Instead of grouping shards together (which reduces parallelism), each
shard
process fetches only the specific fused partner tensors it needs from
other
shards via targeted partial safetensors reads, computes the fused global
scale locally, and writes only its own output shard. No cross-process
coordination or file locking required.

## Changes

### `helpers.py`
Added `build_tensor_file_index()` — reads `index.json` once at startup
and
builds a flat mapping of `tensor_name → resolved_file_path`. This gives
each
worker process an O(1) lookup to find which file contains any fused
partner
tensor, without re-scanning headers at runtime.

### `process.py`
Updated `process_file_microscale_scheme()` with an optional
`tensor_file_index` parameter. When provided:
- `_fetch_fused_partners()` is called to identify any fused set members
missing from the current shard, then fetches only those specific tensors
via partial safetensors reads (headers + target tensors only, not full
files)
- Fused global scale is computed locally using all members of the fused
set
- `_belongs_to_shard()` ensures only native tensors are written to the
output
shard — fetched partner tensors are used for scale computation only and
  never written to the wrong shard

### `__init__.py`
Simplified back to one job per shard — full parallelism restored. For
microscale schemes, builds the `tensor_file_index` once from
`index.json`
and passes it to each job. No union-find, no grouping logic needed.

### `validate.py`
Removed `NotImplementedError` for cross-shard fused weights — the case
is
now handled natively. Replaced with `logger.debug` noting that partner
tensors will be resolved via partial reads.

## Latest Updates: Eliminate reindexing step via inverse_weights_map
with unified job signatures

## Approach
Each shard job receives a precomputed `inverse_weights_map` specifying
exactly
which tensors to load from which files. For cross-shard fused weights,
only the
shard owning the **primary** tensor (q_proj, gate_proj) fetches its
partners —
preventing double reads. All jobs share a unified signature for both
standard
and microscale schemes.

## Changes

### `microscale.py`
- Refactor `DEFAULT_FUSED_MAPPINGS` from a list of lists to
`{primary_pattern: [partner_templates]}` — only the primary-owning shard
fetches its partners, preventing double reads for cross-shard fused
weights
- Move `build_inverse_weights_map()` here — uses regex match on primary
  patterns to construct partner names and locate them in other shards

### `process.py`
- **Unified signature** for `validate_file`, `process_file`, and
  `process_file_microscale_scheme`:
  `(inverse_weights_map, save_path, scheme, ignore, device, converter)`
- All functions use `safe_open` + `f.get_tensor()` for true partial
reads
- Partner tensors re-saved into requesting shard's output; caller
updates
  safetensors index to reflect new locations

### `__init__.py`
- Single `_get_weights_map()` helper handles both single-file and
multi-file
models (reads `safetensors.index.json` or scans file headers via
`safe_open`)
- Single `_build_quantization_jobs()` replaces separate
standard/microscale
  builders — one job per shard with identical tuple structure for both
- Validate jobs use `*job[1:]` for full future-proofing

### `helpers.py`
- Removed `build_weights_map` and `build_inverse_weights_map` (moved to
  `microscale.py`)

### `validate.py`
- Removed `NotImplementedError` for cross-shard fused weights — handled
natively
- Updated to reflect `inverse_weights_map`-based approach

## Testing
- `pytest tests/llmcompressor/entrypoints/model_free/` — all passing
locally
- `make style && make quality` — all checks pass

Signed-off-by: David Zheng <dqzheng1996@gmail.com>

Closes #2497
Related to #2448

Signed-off-by: David Zheng <dqzheng1996@gmail.com>

Author: dzhengAP

src/llmcompressor/entrypoints/model_free/__init__.py

@@ -1,3 +1,4 @@
+import json
 import os
 import shutil
 from pathlib import Path
@@ -12,12 +13,15 @@
 from compressed_tensors.utils.safetensors_load import (
     get_checkpoint_files,
     is_weights_file,
-    update_safetensors_index,
 )
 from loguru import logger
 
-from llmcompressor.entrypoints.model_free.helpers import gpu_if_available
+from llmcompressor.entrypoints.model_free.helpers import (
+    find_safetensors_index_file,
+    gpu_if_available,
+)
 from llmcompressor.entrypoints.model_free.microscale import (
+    build_inverse_weights_map,
     is_microscale_scheme,
 )
 from llmcompressor.entrypoints.model_free.process import (
@@ -46,52 +50,55 @@ def model_free_ptq(
     converter: Converter | None = None,
 ):
     """
-    Quantize a model without the need for a model definition. This function operates on
-    a model stub or folder containing weights saved in safetensors files
+    Quantize a model without the need for a model definition. This function
+    operates on a model stub or folder containing weights saved in safetensors
+    files.
+
+    For microscale schemes (NVFP4, MXFP4), fused weight sets (q/k/v, gate/up)
+    are handled correctly even when split across shards. Each shard job receives
+    a precomputed inverse_weights_map specifying exactly which tensors to load
+    from which files — enabling true partial reads with no runtime discovery
+    and no redundant tensor reads.
 
     :param model_stub: huggingface model hub or path to local weights files
+    :param save_directory: directory to save quantized weights to
     :param scheme: weight quantization scheme or preset scheme name
-    :param ignore: modules to ignore. Modules ending with "norm" are automatically
-        ignored
+    :param ignore: modules to ignore. Modules ending with "norm" are
+        automatically ignored
     :param max_workers: number of worker threads to process files with
     :param device: gpu device to accelerate quantization with
-    :param converter: optional converter to apply to the checkpoint to convert it to
-        compressed-tensors format before running model-free PTQ
-        e.g. conversion of some layers from modelopt format to compressed-tensors
-        See compressed-tensors convert_checkpoint entrypoint for more information
+    :param converter: optional converter to apply to the checkpoint to convert
+        it to compressed-tensors format before running model-free PTQ
     """
     # validate arguments
     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 = []
-    job_fn = (
-        process_file
-        if not is_microscale_scheme(scheme)
-        else process_file_microscale_scheme
-    )
+    # copy non-safetensors files (configs, tokenizers, etc.)
     for file_path, resolved_path in model_files.items():
-        save_path = Path(save_directory) / file_path
-
-        if file_path.endswith("safetensors"):
-            jobs.append(
-                (job_fn, resolved_path, save_path, scheme, ignore, device, converter)
-            )
-
-        else:
+        if not file_path.endswith("safetensors"):
+            save_path = Path(save_directory) / file_path
             if is_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}")
+            logger.info(f"Copying {file_path} -> {save_path}")
             shutil.copyfile(resolved_path, save_path)
 
-    # 1. validate quantizable tensors fail fast before long-running quantization
-    exec_jobs(
-        [(validate_file, *job[1:]) for job in jobs], max_workers, desc="Validating"
-    )
+    # build quantization jobs
+    if is_microscale_scheme(scheme):
+        jobs = _build_microscale_jobs(
+            model_files, save_directory, scheme, ignore, device, converter
+        )
+    else:
+        jobs = _build_standard_jobs(
+            model_files, save_directory, scheme, ignore, device, converter
+        )
+
+    # 1. validate quantizable tensors — fail fast before long-running quantization
+    validate_jobs = _build_validate_jobs(jobs)
+    exec_jobs(validate_jobs, max_workers, desc="Validating")
 
     # 2-5. quantize and compress weights
     total_size = 0
@@ -101,6 +108,175 @@ def model_free_ptq(
         total_size += _total_size
         weight_map.update(_weight_map)
 
-    # 5. update config and safetensors index
+    # 6. update config and safetensors index
+    # weight_map may contain tensors re-located to new shards (partner tensors
+    # re-saved alongside the shard that needed them for fused scale computation)
     update_config(save_directory, scheme_name, scheme, ignore, converter)
-    update_safetensors_index(save_directory, total_size, weight_map)
+
+
+def _build_standard_jobs(
+    model_files: dict[str, str],
+    save_directory: str | os.PathLike,
+    scheme: QuantizationScheme,
+    ignore: Iterable[str],
+    device: torch.device,
+    converter: Converter | None,
+    job_fn=None,
+) -> list[tuple]:
+    """Build one job per safetensors file using the given processing function."""
+    if job_fn is None:
+        job_fn = process_file
+    jobs = []
+    for file_path, resolved_path in model_files.items():
+        if file_path.endswith("safetensors"):
+            save_path = Path(save_directory) / file_path
+            jobs.append(
+                (job_fn, resolved_path, save_path, scheme, ignore, device, converter)
+            )
+    return jobs
+
+
+def _build_microscale_jobs(
+    model_files: dict[str, str],
+    save_directory: str | os.PathLike,
+    scheme: QuantizationScheme,
+    ignore: Iterable[str],
+    device: torch.device,
+    converter: Converter | None,
+) -> list[tuple]:
+    """
+    Build microscale jobs with precomputed inverse_weights_map per shard.
+
+    For each output shard, build_inverse_weights_map() determines exactly which
+    tensors to load from which source files — including any fused partner tensors
+    from other shards. This avoids runtime fused-partner discovery inside the
+    process function and eliminates redundant tensor reads.
+
+    Job tuple format:
+        (process_file_microscale_scheme, inverse_weights_map, save_path,
+         scheme, ignore, device, converter)
+    """
+    index_file = find_safetensors_index_file(model_files)
+
+    if index_file is None:
+        # Single-file model — no cross-shard fused weights possible,
+        # Create inverse_weights_map dict format for process_file_microscale_scheme
+        jobs = []
+        for file_path, resolved_path in model_files.items():
+            if file_path.endswith("safetensors"):
+                save_path = Path(save_directory) / file_path
+                # Wrap as inverse_weights_map: {source_file: None}
+                # means load all tensors
+                inverse_weights_map = {resolved_path: []}
+                jobs.append(
+                    (
+                        process_file_microscale_scheme,
+                        inverse_weights_map,
+                        save_path,
+                        scheme,
+                        ignore,
+                        device,
+                        converter,
+                    )
+                )
+        return jobs
+
+    # Read weight map from safetensors.index.json
+    with open(index_file, "r") as f:
+        weight_map: dict[str, str] = json.load(f)["weight_map"]
+
+    jobs = []
+    for shard_name, resolved_path in model_files.items():
+        if not shard_name.endswith("safetensors"):
+            continue
+
+        save_path = Path(save_directory) / shard_name
+
+        # Precompute exactly which tensors to load from which files for this shard,
+        # including fused partner tensors that live in other shards
+        inverse_weights_map = build_inverse_weights_map(
+            shard_name=shard_name,
+            weight_map=weight_map,
+            model_files=model_files,
+        )
+
+        if len(inverse_weights_map) > 1:
+            partner_shards = [s for s in inverse_weights_map if s != resolved_path]
+            logger.info(
+                f"{shard_name}: will fetch fused partners from "
+                f"{[os.path.basename(s) for s in partner_shards]}"
+            )
+
+        jobs.append(
+            (
+                process_file_microscale_scheme,
+                inverse_weights_map,
+                save_path,
+                scheme,
+                ignore,
+                device,
+                converter,
+            )
+        )
+
+    return jobs
+
+
+def _build_validate_jobs(jobs: list[tuple]) -> list[tuple]:
+    """
+    Build validation jobs from processing jobs.
+
+    Handles both job formats:
+    - Standard/fallback: (proc_fn, file_path_str, save_path, scheme, ignore, device, \
+        converter)
+    - Microscale with index: (proc_fn, inverse_weights_map_dict, save_path, scheme, \
+        ignore, device, converter)
+    """
+    validate_jobs = []
+    for job in jobs:
+        # job[0] is the processing function
+        # Check if second element is a dict (microscale with index)
+        # or string (standard/fallback)
+        second_arg = job[1]
+
+        if isinstance(second_arg, dict):
+            # Microscale job with inverse_weights_map dict
+            _, inverse_weights_map, save_path, scheme, ignore, device, converter = job
+            # Use first source file path from inverse_weights_map for validation
+            source_file = next(iter(inverse_weights_map.keys()))
+            validate_jobs.append(
+                (
+                    validate_file,
+                    source_file,
+                    save_path,
+                    scheme,
+                    ignore,
+                    device,
+                    converter,
+                    inverse_weights_map,
+                )
+            )
+        else:
+            # Standard job or microscale fallback: second_arg is file_path string
+            _, file_path, save_path, scheme, ignore, device, converter = job
+            validate_jobs.append(
+                (
+                    validate_file,
+                    file_path,
+                    save_path,
+                    scheme,
+                    ignore,
+                    device,
+                    converter,
+                    None,
+                )
+            )
+    return validate_jobs
+
+
+def _get_all_tensor_names(file_path: str) -> list[str]:
+    """Get all tensor names from a safetensors file without loading tensors."""
+    from safetensors import safe_open
+
+    with safe_open(file_path, framework="pt", device="cpu") as f:
+        return list(f.keys())

src/llmcompressor/entrypoints/model_free/helpers.py

@@ -96,3 +96,23 @@ def invert_mapping(
         inverse[value].append(key)
 
     return inverse
+
+
+def build_weights_map(
+    weight_map: dict[str, str],
+    model_files: dict[str, str],
+) -> dict[str, str]:
+    """
+    Build a mapping of tensor name -> resolved file path from the model's
+    weight_map (index.json). This allows any process to locate fused partner
+    tensors from other shards without loading entire files.
+
+    :param weight_map: mapping of tensor name -> shard filename (from index.json)
+    :param model_files: mapping of shard filename -> resolved absolute path
+    :return: mapping of tensor name -> resolved absolute path
+    """
+    return {
+        tensor_name: model_files[shard_name]
+        for tensor_name, shard_name in weight_map.items()
+        if shard_name in model_files
+    }