cosmos_framework/utils/generator/quantization.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

"""Low-precision quantization helpers for the Cosmos3 VFM MoT network.

Quantization is applied via torchao's :func:`apply_quantization_inplace`, which
uses the ``quantize_`` path to replace each selected weight with a quantized
tensor subclass in place. Because the live parameter becomes a tensor subclass,
this only works on unsharded (plain-tensor) params and is therefore restricted
to replicated inference (``data_parallel_shard_degree == 1``); it cannot be
applied to an FSDP-sharded model whose params are ``DTensor`` shards.

ModelOpt FP8 checkpoint loading installs the exported E4M3 weights and static
scales into TorchAO tensor subclasses without calibration or re-quantization.

This is an inference-only path: the ``quantize_`` PTQ configs have no backward
support. Module selection is delegated to the filter built by
:func:`_get_filter_fn`.
"""

import gc
import json
import re
from collections.abc import Callable
from pathlib import Path

import torch
from torch import nn
from torch.distributed.tensor import DTensor, distribute_tensor
from torch.nn import functional as F

from cosmos_framework.utils import log
from cosmos_framework.configs.base.defaults.quantization import QuantizationConfig

# NOTE: ``torchao`` is imported lazily inside the functions below rather than at
# module top level. These two helpers are the only torchao consumers, but this
# module is imported transitively by the model package (e.g. during tests that
# never quantize). Keeping the imports lazy means importing this module does not
# require torchao to be installed; the imports only run — and only fail — when
# quantization is actually requested.


class _ModelOptFloat8Linear(nn.Linear):
    """Work around TorchAO 0.16 static-FP8 limitations for linear inputs."""

    # Compute dtype the FP8 weight dequantizes to. Recorded when the module is
    # swapped in on meta, where the plain E4M3 placeholder no longer carries it.
    _modelopt_high_precision_dtype: torch.dtype | None = None
    # True once real checkpoint data has been installed.
    _modelopt_weight_loaded: bool = False
    # Mixed-precision diffusion steps (see utils/generator/mixed_precision.py).
    # None until install_mixed_precision_runtime tags this module.
    _mixed_precision_runtime = None
    _mixed_precision_path: str = "generation"
    # A (N, K) dense view staged by the block provider for the current decoder
    # layer, or None. Read here, written only by the provider hooks.
    _mixed_precision_staged_weight: torch.Tensor | None = None

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        output_shape = (*inputs.shape[:-1], self.out_features)
        # PrototypeFloat8Tensor divides each input dimension by its block size, so
        # a zero-row input fails with 0 // 0 instead of producing an empty output.
        if inputs.numel() == 0:
            return inputs.new_empty(output_shape)
        # PrototypeFloat8Tensor requires the input and its static (1, 1) activation
        # scale to have equal rank, so flatten token dimensions before dispatch.
        flat_inputs = inputs.reshape(-1, inputs.shape[-1])
        runtime = self._mixed_precision_runtime
        if runtime is not None and runtime.use_high_precision(self._mixed_precision_path):
            # W8A16: dense GEMM against the dequantized FP8 weight; the
            # activation stays in the compute dtype. The resolved weight is
            # in the runtime's configured activation_dtype, which normally
            # matches the live activations -- but defensively cast both
            # operands to the weight's dtype and cast the result back to the
            # input dtype when they differ (e.g. a fp16-activation model run
            # against a bf16-configured runtime), instead of crashing.
            weight = runtime.resolve_w8a16_weight(self)
            if weight.dtype != flat_inputs.dtype:
                output = F.linear(flat_inputs.to(weight.dtype), weight, self.bias).to(inputs.dtype)
            else:
                output = F.linear(flat_inputs, weight, self.bias)
            return output.reshape(output_shape)
        flat_outputs = F.linear(flat_inputs, self.weight, self.bias)
        return flat_outputs.reshape(output_shape)


_torchao_fsdp_support_installed = False


def _rewrap_float8(source, qdata: torch.Tensor):
    """Rebuild a static-FP8 tensor around new ``qdata``, carrying the scales over.

    Valid only because ModelOpt exports PerTensor granularity: one scalar scale
    covers the whole weight, so reshaping or slicing the data never redistributes
    it. ``block_size`` is therefore always "one block spanning everything".
    """
    return source.__class__(
        qdata,
        source.scale,
        act_quant_scale=source.act_quant_scale,
        output_act_quant_scale=source.output_act_quant_scale,
        block_size=list(qdata.shape),
        mm_config=source.mm_config,
        act_quant_kwargs=source.act_quant_kwargs,
        kernel_preference=source.kernel_preference,
        dtype=source.dtype,
        output_act_quant_kwargs=source.output_act_quant_kwargs,
    )


def install_torchao_float8_fsdp_support() -> None:
    """Teach TorchAO's static-FP8 tensor the ops and hooks FSDP2 needs.

    TorchAO 0.16 ships this tensor for single-device inference: it implements no
    ``fsdp_pre_all_gather`` / ``fsdp_post_all_gather``, and its shape-changing ops
    derive a scale layout from ``block_size`` assuming a fixed rank, so FSDP2's
    reshape/split/allocate path fails on a 2D linear weight. Sharding it is
    nonetheless well defined for a PerTensor scale — the scalar is identical on
    every shard, so only the FP8 bytes ever need to move.

    Registered once, globally, since the dispatch table is keyed by class.
    """
    global _torchao_fsdp_support_installed
    if _torchao_fsdp_support_installed:
        return
    from torchao.prototype.quantization.float8_static_quant.prototype_float8_tensor import (
        PrototypeFloat8Tensor,
    )
    from torchao.utils import return_and_correct_aliasing

    aten = torch.ops.aten
    implements = PrototypeFloat8Tensor.implements

    @implements([aten.view.default, aten._unsafe_view.default, aten.reshape.default])
    def _(func, types, args, kwargs):
        self, size = args[0], args[1]
        return return_and_correct_aliasing(func, args, kwargs, _rewrap_float8(self, self.qdata.reshape(*size)))

    @implements(aten.split.Tensor)
    def _(func, types, args, kwargs):
        self, size = args[0], args[1]
        dim = args[2] if len(args) > 2 else 0
        return [_rewrap_float8(self, part) for part in self.qdata.split(size, dim)]

    @implements(aten.slice.Tensor)
    def _(func, types, args, kwargs):
        self = args[0]
        return return_and_correct_aliasing(
            func, args, kwargs, _rewrap_float8(self, aten.slice.Tensor(self.qdata, *args[1:]))
        )

    @implements(aten.as_strided.default)
    def _(func, types, args, kwargs):
        self, size, stride = args[0], args[1], args[2]
        offset = args[3] if len(args) > 3 else 0
        return return_and_correct_aliasing(
            func, args, kwargs, _rewrap_float8(self, self.qdata.as_strided(size, stride, offset))
        )

    @implements(aten.new_zeros.default)
    def _(func, types, args, kwargs):
        self, size = args[0], args[1]
        return _rewrap_float8(self, self.qdata.new_zeros(size))

    @implements([aten.empty_like.default, aten.zeros_like.default])
    def _(func, types, args, kwargs):
        self = args[0]
        device = kwargs.get("device")
        qdata = func(self.qdata, **{key: value for key, value in kwargs.items() if key != "dtype"})
        # `to_empty` moves storage without data, so the scales are allocated empty
        # on the target device as well: `.to(device)` would try to copy out of a
        # meta tensor. Real scale values arrive with the checkpoint weights.
        scale = torch.empty_like(self.scale, device=device) if device is not None else self.scale
        act_quant_scale = (
            torch.empty_like(self.act_quant_scale, device=device) if device is not None else self.act_quant_scale
        )
        return self.__class__(
            qdata,
            scale,
            act_quant_scale=act_quant_scale,
            output_act_quant_scale=self.output_act_quant_scale,
            block_size=list(qdata.shape),
            mm_config=self.mm_config,
            act_quant_kwargs=self.act_quant_kwargs,
            kernel_preference=self.kernel_preference,
            dtype=self.dtype,
            output_act_quant_kwargs=self.output_act_quant_kwargs,
        )

    @implements([aten.detach.default, aten.clone.default])
    def _(func, types, args, kwargs):
        self = args[0]
        return return_and_correct_aliasing(func, args, kwargs, _rewrap_float8(self, func(self.qdata)))

    @implements(aten._to_copy.default)
    def _(func, types, args, kwargs):
        # Device moves must go through, but a dtype request must not: casting the
        # E4M3 payload to the compute dtype would silently undo the quantization.
        self = args[0]
        forwarded = {key: value for key, value in kwargs.items() if key != "dtype"}
        return return_and_correct_aliasing(func, args, kwargs, _rewrap_float8(self, func(self.qdata, **forwarded)))

    @implements(aten.copy_.default)
    def _(func, types, args, kwargs):
        destination, source = args[0], args[1]
        if isinstance(source, PrototypeFloat8Tensor):
            destination.qdata.copy_(source.qdata)
            destination.scale.copy_(source.scale)
            destination.act_quant_scale.copy_(source.act_quant_scale)
        else:
            destination.qdata.copy_(source)
        return destination

    def fsdp_pre_all_gather(self, mesh):
        # Only the FP8 bytes travel. The static scales are already identical on
        # every rank, so keeping them out of the collective saves the traffic.
        return (self.qdata,), (self.scale, self.act_quant_scale, self.dtype)

    def fsdp_post_all_gather(self, all_gather_outputs, metadata, param_dtype, *, out=None):
        (qdata,) = all_gather_outputs
        scale, act_quant_scale, dtype = metadata
        if out is not None:
            return None
        gathered = self.__class__(
            qdata,
            scale,
            act_quant_scale=act_quant_scale,
            output_act_quant_scale=self.output_act_quant_scale,
            block_size=list(qdata.shape),
            mm_config=self.mm_config,
            act_quant_kwargs=self.act_quant_kwargs,
            kernel_preference=self.kernel_preference,
            dtype=param_dtype or dtype,
            output_act_quant_kwargs=self.output_act_quant_kwargs,
        )
        return gathered, (qdata,)

    PrototypeFloat8Tensor.fsdp_pre_all_gather = fsdp_pre_all_gather
    PrototypeFloat8Tensor.fsdp_post_all_gather = fsdp_post_all_gather
    _torchao_fsdp_support_installed = True
    log.info("Installed TorchAO static-FP8 FSDP support (8 aten ops + all-gather hooks)")


def is_modelopt_fp8_checkpoint(checkpoint_path: str | Path) -> bool:
    """Return whether a local checkpoint declares ModelOpt FP8 quantization."""
    quant_config_path = Path(checkpoint_path) / "hf_quant_config.json"
    if not quant_config_path.is_file():
        return False
    try:
        quant_config = json.loads(quant_config_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as error:
        raise ValueError(f"Unable to read a valid checkpoint quantization config: {quant_config_path}") from error
    if (
        not isinstance(quant_config, dict)
        or quant_config.get("quant_method") != "modelopt"
        or quant_config.get("quant_algo") != "FP8"
    ):
        raise ValueError(f"Unsupported checkpoint quantization configuration: {quant_config_path}")
    return True


def collect_modelopt_fp8_params(model: nn.Module) -> set[nn.Parameter]:
    """Return the FP8 weights of every ModelOpt linear currently in ``model``.

    Collected by module type immediately before use rather than captured earlier,
    because parallelization rebuilds parts of the network (notably the
    ``language_model`` subtree) and hands back fresh parameter objects, so any
    set gathered ahead of that no longer matches by identity.
    """
    return {module.weight for module in model.modules() if isinstance(module, _ModelOptFloat8Linear)}


def plan_modelopt_fp8_targets(
    checkpoint_path: str | Path,
    key_mapper: Callable[[str, str], str | None],
    *,
    weight_map: dict[str, str],
) -> list[str]:
    """Return the network-relative FQNs a ModelOpt FP8 checkpoint quantizes.

    Derived from the checkpoint index alone — no model instance is required — so
    the result can be threaded through the model config and consumed by
    ``build_net`` before the network exists. Only the E4M3 tensors are reported;
    weights ModelOpt intentionally left in high precision (``proj_in``,
    ``lm_head``, the vision tower, ...) are skipped.

    Args:
        checkpoint_path: Local root of the ModelOpt diffusers checkpoint.
        key_mapper: Maps a diffusers source key and shard path to a model state key.
        weight_map: Merged checkpoint weight map.

    Returns:
        Sorted target module FQNs, relative to the VFM network.
    """
    from safetensors import safe_open

    checkpoint_path = Path(checkpoint_path)
    weights_by_shard: dict[str, list[str]] = {}
    for source_key, shard_path in weight_map.items():
        if source_key.endswith(".weight"):
            weights_by_shard.setdefault(shard_path, []).append(source_key)

    target_fqns: set[str] = set()
    for shard_path, source_weight_keys in sorted(weights_by_shard.items()):
        with safe_open(checkpoint_path / shard_path, framework="pt", device="cpu") as shard:
            shard_keys = set(shard.keys())
            for source_weight_key in sorted(source_weight_keys):
                if source_weight_key not in shard_keys:
                    continue
                if shard.get_slice(source_weight_key).get_dtype() != "F8_E4M3":
                    continue
                target_weight_key = key_mapper(source_weight_key, shard_path)
                if target_weight_key is None or not target_weight_key.endswith(".weight"):
                    continue
                target_fqns.add(target_weight_key.removesuffix(".weight"))
    return sorted(target_fqns)


def swap_modelopt_fp8_linears_on_meta(model: nn.Module, target_fqns: list[str]) -> list[str]:
    """Swap the given linears to FP8 modules while the network is still on meta.

    This must run *before* the network is parallelized and materialized. FSDP2
    wraps the whole network into a single parameter group, so a linear replaced
    after ``fully_shard`` would leave the group holding a stale parameter; and
    ``to_empty`` on the original bf16 shapes is what sets peak memory, which for
    a Super-class model exceeds a single 80 GB device. Swapping here means FSDP
    shards — and ``to_empty`` materializes — one byte per element instead of two.

    The weight is created as a real (meta) TorchAO static-FP8 tensor rather than a
    plain E4M3 placeholder, because FSDP decides at wrap time whether a parameter
    carries the ``fsdp_pre_all_gather`` extension. A plain tensor swapped to a
    subclass afterwards is already registered as an ordinary parameter and fails
    on the first all-gather. Actual weight bytes and scales arrive later, in
    :func:`apply_modelopt_fp8_checkpoint_inplace`.

    Args:
        model: Meta-device VFM network.
        target_fqns: Module FQNs to convert, as planned from the checkpoint.

    Returns:
        Sorted FQNs actually swapped; targets absent from this model variant are
        skipped, matching the loader's tolerance for task-specialized exports.
    """
    install_torchao_float8_fsdp_support()
    from torchao.float8.inference import Float8MMConfig
    from torchao.prototype.quantization.float8_static_quant.prototype_float8_tensor import (
        PrototypeFloat8Tensor,
    )
    from torchao.quantization import PerTensor
    from torchao.quantization.quantize_.workflows import QuantizeTensorToFloat8Kwargs

    mm_config = Float8MMConfig(use_fast_accum=True)
    activation_quant_kwargs = QuantizeTensorToFloat8Kwargs(
        float8_dtype=torch.float8_e4m3fn,
        granularity=PerTensor(),
        mm_config=mm_config,
    )
    swapped_fqns: list[str] = []
    for target_module_fqn in sorted(target_fqns):
        try:
            module = model.get_submodule(target_module_fqn)
        except AttributeError:
            continue
        if not isinstance(module, nn.Linear):
            raise KeyError(f"ModelOpt FP8 target {target_module_fqn!r} is not an nn.Linear module")

        replacement = _ModelOptFloat8Linear(
            module.in_features,
            module.out_features,
            bias=module.bias is not None,
            device="meta",
            dtype=module.weight.dtype,
        )
        replacement._modelopt_high_precision_dtype = module.weight.dtype
        quantized_shape = tuple(module.weight.shape)
        replacement.weight = nn.Parameter(
            PrototypeFloat8Tensor(
                torch.empty(quantized_shape, dtype=torch.float8_e4m3fn, device="meta"),
                torch.empty(1, 1, dtype=torch.float32, device="meta"),
                act_quant_scale=torch.empty(1, 1, dtype=torch.float32, device="meta"),
                block_size=list(quantized_shape),
                mm_config=mm_config,
                act_quant_kwargs=activation_quant_kwargs,
                dtype=module.weight.dtype,
            ),
            requires_grad=False,
        )
        if module.bias is not None:
            replacement.bias = module.bias
        replacement.train(module.training)

        parent_fqn, _, child_name = target_module_fqn.rpartition(".")
        parent = model.get_submodule(parent_fqn) if parent_fqn else model
        setattr(parent, child_name, replacement)
        swapped_fqns.append(target_module_fqn)

    log.info(f"Swapped {len(swapped_fqns)} linears to meta-device ModelOpt FP8 modules")
    return swapped_fqns


def apply_modelopt_fp8_checkpoint_inplace(
    model: nn.Module,
    checkpoint_path: str | Path,
    key_mapper: Callable[[str, str], str | None],
    *,
    weight_map: dict[str, str],
) -> list[str]:
    """Install ModelOpt static-FP8 checkpoint tensors as TorchAO weights.

    ModelOpt's HF export stores the already-quantized E4M3 bytes under ``weight``
    and its per-tensor dequantization scales under ``weight_scale`` and
    ``input_scale``. This adapter streams those tensors from the checkpoint and
    constructs TorchAO ``PrototypeFloat8Tensor`` weights directly, avoiding any
    dequantization, recalibration, or re-quantization.

    Args:
        model: Unsharded Cosmos3 network whose linears will be converted in place.
        checkpoint_path: Local root of the ModelOpt diffusers checkpoint.
        key_mapper: Maps a diffusers source key and shard path to a model state key.
        weight_map: Merged checkpoint weight map.

    Returns:
        Sorted fully qualified names of converted linear modules.
    """
    checkpoint_path = Path(checkpoint_path)
    if not is_modelopt_fp8_checkpoint(checkpoint_path):
        raise ValueError(f"Not a ModelOpt FP8/FP8 checkpoint: {checkpoint_path}")

    from safetensors import safe_open
    from torchao.float8.inference import Float8MMConfig
    from torchao.prototype.quantization.float8_static_quant.prototype_float8_tensor import (
        PrototypeFloat8Tensor,
    )
    from torchao.quantization import PerTensor
    from torchao.quantization.quantize_.workflows import QuantizeTensorToFloat8Kwargs

    source_weights_by_shard: dict[str, list[str]] = {}
    for source_key, shard_path in weight_map.items():
        if source_key.endswith(".weight"):
            source_weights_by_shard.setdefault(shard_path, []).append(source_key)

    mm_config = Float8MMConfig(use_fast_accum=True)
    activation_quant_kwargs = QuantizeTensorToFloat8Kwargs(
        float8_dtype=torch.float8_e4m3fn,
        granularity=PerTensor(),
        mm_config=mm_config,
    )
    # Target module FQNs successfully replaced with static-FP8 linear modules.
    converted_fqns: list[str] = []
    # Source FP8 weights mapped to linear modules present in this model variant.
    applicable_source_weight_keys: set[str] = set()
    # Source FP8 weights actually installed; compared with the applicable set
    # after conversion to detect incomplete plans.
    converted_source_weight_keys: set[str] = set()
    # Weight shard -> validated (source weight key, target module FQN) conversions.
    conversion_plan_by_shard: dict[str, list[tuple[str, str]]] = {}
    # Scale shard -> scale keys to preload; scales may be separate from their weights.
    source_scales_by_shard: dict[str, list[str]] = {}
    # Target module FQNs reserved by the plan, used to reject duplicate mappings.
    planned_target_fqns: set[str] = set()

    # Identify and validate every applicable FP8 conversion before replacing any
    # modules. The checkpoint may contain intentionally ignored weights or
    # optional heads absent from this model variant; those must be filtered
    # before requiring their scales.
    for shard_path, source_weight_keys in sorted(source_weights_by_shard.items()):
        full_shard_path = checkpoint_path / shard_path
        with safe_open(full_shard_path, framework="pt", device="cpu") as shard:
            shard_keys = set(shard.keys())
            missing_weight_keys = set(source_weight_keys) - shard_keys
            if missing_weight_keys:
                raise KeyError(
                    f"ModelOpt weight tensor(s) indexed in {full_shard_path} are missing: {sorted(missing_weight_keys)}"
                )
            for source_weight_key in sorted(source_weight_keys):
                tensor_slice = shard.get_slice(source_weight_key)
                if tensor_slice.get_dtype() != "F8_E4M3":
                    continue

                target_weight_key = key_mapper(source_weight_key, shard_path)
                if target_weight_key is None:
                    continue
                if not target_weight_key.endswith(".weight"):
                    raise KeyError(f"ModelOpt FP8 key {source_weight_key!r} has no Cosmos3 weight mapping")
                target_module_fqn = target_weight_key.removesuffix(".weight")
                try:
                    module = model.get_submodule(target_module_fqn)
                except AttributeError:
                    continue
                if not isinstance(module, nn.Linear):
                    raise KeyError(
                        f"ModelOpt FP8 key {source_weight_key!r} mapped to {target_module_fqn!r}, "
                        "which is not an nn.Linear module"
                    )
                if type(module.weight) is not nn.Parameter and not isinstance(module, _ModelOptFloat8Linear):
                    # Distinguish the two ways a weight stops being a plain
                    # parameter: FSDP replaced it with a shard, or something
                    # already quantized it. Reporting the former as the latter
                    # sends debugging in entirely the wrong direction. A module
                    # already swapped by `swap_modelopt_fp8_linears_on_meta` is
                    # neither — its weight may legitimately be a DTensor shard.
                    if type(module.weight).__name__ == "DTensor":
                        raise ValueError(
                            f"ModelOpt FP8 target is an FSDP shard, not a plain parameter: {target_module_fqn}. "
                            "The FP8 swap must run before the network is parallelized."
                        )
                    raise ValueError(f"ModelOpt FP8 target is already quantized: {target_module_fqn}")
                checkpoint_shape = tuple(tensor_slice.get_shape())
                if checkpoint_shape != tuple(module.weight.shape):
                    raise ValueError(
                        f"Shape mismatch for {target_module_fqn}: checkpoint has {checkpoint_shape}, "
                        f"model expects {tuple(module.weight.shape)}"
                    )
                if target_module_fqn in planned_target_fqns:
                    raise ValueError(f"Multiple ModelOpt FP8 weights map to target: {target_module_fqn}")

                source_module_key = source_weight_key.removesuffix(".weight")
                source_scale_keys = (
                    f"{source_module_key}.weight_scale",
                    f"{source_module_key}.input_scale",
                )
                missing_scale_keys = set(source_scale_keys) - weight_map.keys()
                if missing_scale_keys:
                    raise KeyError(
                        f"ModelOpt FP8 tensor {source_weight_key!r} is missing scale tensor(s) "
                        f"from the checkpoint index: {sorted(missing_scale_keys)}"
                    )
                for source_scale_key in source_scale_keys:
                    source_scales_by_shard.setdefault(weight_map[source_scale_key], []).append(source_scale_key)

                applicable_source_weight_keys.add(source_weight_key)
                planned_target_fqns.add(target_module_fqn)
                conversion_plan_by_shard.setdefault(shard_path, []).append((source_weight_key, target_module_fqn))

    if not applicable_source_weight_keys:
        raise ValueError(f"No ModelOpt FP8 linear weights were found in {checkpoint_path}")

    # ModelOpt's consolidated HF export can place the scalar scales in a
    # different shard from the corresponding FP8 weight. Read only those tiny
    # tensors up front, following the root index rather than assuming locality.
    exported_scales: dict[str, torch.Tensor] = {}
    for shard_path, source_scale_keys in sorted(source_scales_by_shard.items()):
        full_shard_path = checkpoint_path / shard_path
        with safe_open(full_shard_path, framework="pt", device="cpu") as shard:
            shard_keys = set(shard.keys())
            missing_scale_keys = set(source_scale_keys) - shard_keys
            if missing_scale_keys:
                raise KeyError(
                    f"ModelOpt scale tensor(s) indexed in {full_shard_path} are missing: {sorted(missing_scale_keys)}"
                )
            for source_scale_key in source_scale_keys:
                scale = shard.get_tensor(source_scale_key)
                if scale.numel() != 1:
                    raise ValueError(f"ModelOpt scale tensor {source_scale_key!r} must contain exactly one value")
                exported_scales[source_scale_key] = scale

    for shard_path, conversions in sorted(conversion_plan_by_shard.items()):
        full_shard_path = checkpoint_path / shard_path
        with safe_open(full_shard_path, framework="pt", device="cpu") as shard:
            for source_weight_key, target_module_fqn in conversions:
                source_module_key = source_weight_key.removesuffix(".weight")
                source_weight_scale_key = f"{source_module_key}.weight_scale"
                source_input_scale_key = f"{source_module_key}.input_scale"
                module = model.get_submodule(target_module_fqn)
                device = module.weight.device

                if isinstance(module, _ModelOptFloat8Linear):
                    # Already swapped on meta before parallelization; the weight
                    # is an E4M3 placeholder, so the compute dtype comes from the
                    # dtype recorded at swap time rather than from the parameter.
                    replacement = module
                    high_precision_dtype = module._modelopt_high_precision_dtype
                    if high_precision_dtype is None:
                        raise ValueError(f"ModelOpt FP8 module has no recorded compute dtype: {target_module_fqn}")
                else:
                    high_precision_dtype = module.weight.dtype
                    replacement = _ModelOptFloat8Linear(
                        module.in_features,
                        module.out_features,
                        bias=False,
                        device="meta",
                        dtype=high_precision_dtype,
                    )
                    replacement.bias = module.bias
                    replacement.train(module.training)

                    parent_fqn, _, child_name = target_module_fqn.rpartition(".")
                    parent = model.get_submodule(parent_fqn) if parent_fqn else model
                    setattr(parent, child_name, replacement)
                    del module

                quantized_data = shard.get_tensor(source_weight_key)
                weight_scale = exported_scales[source_weight_scale_key].reshape(1, 1)
                input_scale = exported_scales[source_input_scale_key].reshape(1, 1)

                existing_param = replacement.weight
                if isinstance(existing_param, DTensor):
                    # FSDP already sharded the placeholder, so only this rank's
                    # slice of the checkpoint tensor is written. The scales are
                    # per-tensor and therefore identical on every rank.
                    local_weight = existing_param.to_local()
                    quantized_data = distribute_tensor(
                        quantized_data.to(device=device),
                        existing_param.device_mesh,
                        existing_param.placements,
                    ).to_local()
                elif isinstance(existing_param.data, PrototypeFloat8Tensor):
                    local_weight = existing_param.data
                else:
                    local_weight = None

                if local_weight is not None:
                    # Fill the parameter that FSDP is already tracking; rebinding
                    # `replacement.weight` here would detach it from the parameter
                    # group and strand the all-gather on a stale tensor.
                    if tuple(local_weight.qdata.shape) != tuple(quantized_data.shape):
                        raise ValueError(
                            f"Shard mismatch for {target_module_fqn}: model shard "
                            f"{tuple(local_weight.qdata.shape)}, checkpoint slice {tuple(quantized_data.shape)}"
                        )
                    local_weight.qdata.copy_(quantized_data)
                    local_weight.scale.copy_(weight_scale)
                    local_weight.act_quant_scale.copy_(input_scale)
                else:
                    # Unswapped legacy path: the module still holds a plain bf16
                    # weight, so build the quantized tensor and bind it directly.
                    replacement.weight = nn.Parameter(
                        PrototypeFloat8Tensor(
                            quantized_data.to(device=device),
                            weight_scale.to(device=device),
                            act_quant_scale=input_scale.to(device=device),
                            block_size=list(quantized_data.shape),
                            mm_config=mm_config,
                            act_quant_kwargs=activation_quant_kwargs,
                            dtype=high_precision_dtype,
                        ),
                        requires_grad=False,
                    )
                replacement._modelopt_weight_loaded = True
                converted_fqns.append(target_module_fqn)
                converted_source_weight_keys.add(source_weight_key)

    missing_conversions = applicable_source_weight_keys - converted_source_weight_keys
    if missing_conversions:
        raise ValueError(f"ModelOpt FP8 weights were not converted: {sorted(missing_conversions)}")
    converted_fqns.sort()
    sharded_count = sum(1 for fqn in converted_fqns if isinstance(model.get_submodule(fqn).weight.data, DTensor))
    log.info(
        f"Loaded {len(converted_fqns)} calibrated ModelOpt FP8 weights into TorchAO "
        f"({sharded_count} sharded / {len(converted_fqns) - sharded_count} replicated)"
    )
    log.debug(f"ModelOpt FP8 matched_fqns={converted_fqns}")
    return converted_fqns


def _get_filter_fn(quantization_config: QuantizationConfig) -> Callable[[nn.Module, str], bool]:
    """Build a module-selection predicate from the quantization config.

    The returned closure captures ``include_regex`` / ``exclude_regex`` and
    implements the selection policy documented on :class:`QuantizationConfig`.
    Each key is treated as a regular expression and matched against a module's
    FQN with :func:`re.search` (a plain substring remains a valid pattern, so
    existing substring-style keys keep working). It is passed to torchao as
    ``filter_fn`` (for ``quantize_``), which expects a
    ``(module, fqn) -> bool`` signature.

    Args:
        quantization_config: Config carrying the include/exclude key lists.

    Returns:
        A predicate suitable for torchao's ``filter_fn`` / ``module_filter_fn``.
    """

    include_patterns: list[re.Pattern[str]] = []
    for pattern in quantization_config.include_regex:
        try:
            include_patterns.append(re.compile(pattern))
        except re.error as error:
            raise ValueError(f"Invalid include_regex pattern {pattern!r}: {error}") from error

    exclude_patterns: list[re.Pattern[str]] = []
    for pattern in quantization_config.exclude_regex:
        try:
            exclude_patterns.append(re.compile(pattern))
        except re.error as error:
            raise ValueError(f"Invalid exclude_regex pattern {pattern!r}: {error}") from error

    def _filter_fn(mod: nn.Module, name: str) -> bool:
        """Decide whether a single module should be quantized.

        Used by preflight and torchao as each walks the model recursively. A
        module is selected only when ALL of the following hold:

        1. It is an ``nn.Linear`` (the only layer type these recipes support).
        2. ``include_regex`` is empty (include everything) OR the module's FQN
           matches at least one include pattern.
        3. The module's FQN matches none of the ``exclude_regex`` patterns.

        Each include/exclude key is treated as a regular expression and matched
        against the FQN with :func:`re.search`, so the pattern can match anywhere
        in the name (a plain substring is still a valid regex, preserving the
        previous substring-match behavior, while enabling anchors like ``^``/``$``,
        alternation, character classes, etc.).

        Note the parenthesization around the include check: ``and`` binds tighter
        than ``or`` in Python, so without it the ``nn.Linear`` and exclude
        checks would not apply across both include branches.

        Args:
            mod (torch.nn.Module): The module that is being processed.
            name (str): A fully qualified name of the module that is being processed.

        Return:
            True if the module should be quantized, False otherwise.
        """
        # torch.compile inserts `_orig_mod` into FQNs; hide it from user-facing regex matching.
        canonical_name = ".".join(part for part in name.split(".") if part != "_orig_mod")
        return (
            isinstance(mod, nn.Linear)
            and (not include_patterns or any(pattern.search(canonical_name) for pattern in include_patterns))
            and not any(pattern.search(canonical_name) for pattern in exclude_patterns)
        )

    return _filter_fn


def _get_validated_quantization_fqns(model: nn.Module, filter_fn: Callable[[nn.Module, str], bool]) -> list[str]:
    """Validate the selected modules and return their sorted FQNs."""
    matched_modules = sorted(
        ((name, module) for name, module in model.named_modules() if filter_fn(module, name)),
        key=lambda item: item[0],
    )
    if not matched_modules:
        raise ValueError("No nn.Linear modules matched the quantization selection")
    matched_fqns = [name for name, _ in matched_modules]
    already_quantized_fqns = [name for name, module in matched_modules if type(module.weight) is not nn.Parameter]
    if already_quantized_fqns:
        raise ValueError(f"Quantization targets are already quantized: {', '.join(already_quantized_fqns)}")
    return matched_fqns


def apply_quantization_inplace(model: nn.Module, quantization_config: QuantizationConfig) -> list[str]:
    """Apply quantization in place via ``quantize_`` (replaces weights with quantized tensors).

    This is the replication path. ``quantize_`` replaces each weight with a
    quantized tensor subclass as the live parameter, which only works when the
    parameters are plain tensors. It therefore cannot be applied to an already
    FSDP-sharded model (the params are ``DTensor`` shards), so it is restricted
    to replicated inference (``data_parallel_shard_degree == 1``).

    These configs (``MXDynamicActivationMXWeightConfig`` /
    ``NVFP4DynamicActivationNVFP4WeightConfig`` /
    ``Float8DynamicActivationFloat8WeightConfig``) are inference-only (PTQ) and
    have no backward support. For the sharded case use ``apply_quantization``
    (the module-swap path) instead; both functions are currently inference
    paths, selected by whether FSDP is sharding the model.

    Returns:
        Sorted fully qualified names of the matched modules.
    """
    # No-op when quantization is disabled.
    if quantization_config.method is None:
        return []

    filter_fn = _get_filter_fn(quantization_config)
    matched_fqns = _get_validated_quantization_fqns(model, filter_fn)

    from torchao.prototype.mx_formats import (
        MXDynamicActivationMXWeightConfig,
        NVFP4DynamicActivationNVFP4WeightConfig,
    )
    from torchao.quantization import (
        Float8DynamicActivationFloat8WeightConfig,
        PerRow,
        PerTensor,
        quantize_,
    )

    def _reclaim_and_log_gpu_memory(tag: str) -> None:
        # synchronize() + gc.collect() + empty_cache() are load-bearing, not passive
        # observability: torch.mem_get_info reports memory free at the CUDA-driver level.
        # synchronize() first drains any in-flight device work (e.g. async copies from the
        # checkpoint load path) so their allocations retire before we measure; otherwise the
        # baseline could over-report free memory. PyTorch's caching allocator then holds
        # freed blocks in its own pool without returning them to the driver, so empty_cache()
        # flushes that pool back to the driver.
        torch.cuda.synchronize()
        gc.collect()
        torch.cuda.empty_cache()
        free, total = torch.cuda.mem_get_info(torch.device("cuda", torch.cuda.current_device()))
        log.info(f"GPU memory ({tag}): {free / 2**30:.2f} GiB free / {total / 2**30:.2f} GiB total")

    _reclaim_and_log_gpu_memory("before quantization")

    if quantization_config.method == "mxfp8":
        quantize_(
            model,
            config=MXDynamicActivationMXWeightConfig(),
            filter_fn=filter_fn,
        )

    elif quantization_config.method == "nvfp4":
        # use_triton_kernel=False avoids torchao's fused NVFP4 Triton kernel, which
        # requires the external `mslk` package. Prebuilt mslk wheels are linked
        # against upstream torch and fail to load against NVIDIA's NGC custom torch
        # builds (ABI mismatch on `torch::Library::_def`), so we use torchao's
        # built-in NVFP4 path instead.
        quantize_(
            model,
            config=NVFP4DynamicActivationNVFP4WeightConfig(use_triton_kernel=False),
            filter_fn=filter_fn,
        )

    elif quantization_config.method == "fp8":
        # Hopper-compatible FP8. Unlike mxfp8 / nvfp4 (block-scaled MX / NVFP4
        # formats whose accelerated kernels require Blackwell sm_100 tensor
        # cores), this is plain e4m3 dynamic-activation + fp8-weight quantization
        # executed via torch._scaled_mm, which is supported on Hopper (sm_90) and
        # Ada (sm_89). Scaling granularity is user-selectable: PerRow (rowwise,
        # better accuracy) or PerTensor (single scale, slightly faster); both are
        # supported on Hopper.
        granularity = PerRow() if quantization_config.fp8_granularity == "per_row" else PerTensor()
        quantize_(
            model,
            config=Float8DynamicActivationFloat8WeightConfig(granularity=granularity),
            filter_fn=filter_fn,
        )

    else:
        raise ValueError(f"Unsupported quantization method: {quantization_config.method}")

    _reclaim_and_log_gpu_memory("after quantization")
    log.info(f"Applied runtime PTQ method={quantization_config.method}, matched_count={len(matched_fqns)}")
    log.debug(f"Runtime PTQ matched_fqns={matched_fqns}")
    return matched_fqns