From 0d5eed58d7f26ea73647a7f17e3ae5ce96513970 Mon Sep 17 00:00:00 2001
From: pengcuo <pzeren@nvidia.com>
Date: Mon, 3 Aug 2026 21:33:34 -0700
Subject: [PATCH 1/5] feat(quantization): add ModelOpt static-FP8 checkpoint
 loading
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Port of the release-synced half of imaginaire4!10639 (i4 main 39845b6510).
`cosmos_framework/utils/generator/quantization.py` is mapped from
`projects/cosmos3/cosmos3/utils/quantization.py` by the cosmos-framework
release pipeline, so this file is byte-identical to what the next scheduled
release will emit. Drop this commit if the release lands first.

Adds:
- `is_modelopt_fp8_checkpoint`: detect `hf_quant_config.json` with
  quant_method=modelopt / quant_algo=FP8, rejecting malformed or
  unsupported quantization configs.
- `apply_modelopt_fp8_checkpoint_inplace`: stream the exported E4M3 weights
  and their static per-tensor weight/input scales out of the safetensors
  shards and install them directly as TorchAO `PrototypeFloat8Tensor`
  weights — no dequantization, calibration, or re-quantization. The whole
  conversion is validated (shapes, duplicate targets, missing scales,
  already-quantized targets) before any module is replaced.
- `_ModelOptFloat8Linear`: works around two TorchAO 0.16 static-FP8 limits —
  zero-row inputs (0 // 0 on the block size) and the rank mismatch between a
  >2D input and its (1, 1) activation scale.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: pengcuo <pzeren@nvidia.com>
---
 .../utils/generator/quantization.py           | 236 ++++++++++++++++++
 1 file changed, 236 insertions(+)

diff --git a/cosmos_framework/utils/generator/quantization.py b/cosmos_framework/utils/generator/quantization.py
index 94b11948..ae215a10 100644
--- a/cosmos_framework/utils/generator/quantization.py
+++ b/cosmos_framework/utils/generator/quantization.py
@@ -10,17 +10,23 @@
 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.nn import functional as F
 
 from cosmos_framework.utils import log
 from cosmos_framework.configs.base.defaults.quantization import QuantizationConfig
@@ -33,6 +39,236 @@
 # quantization is actually requested.
 
 
+class _ModelOptFloat8Linear(nn.Linear):
+    """Work around TorchAO 0.16 static-FP8 limitations for linear inputs."""
+
+    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])
+        flat_outputs = F.linear(flat_inputs, self.weight, self.bias)
+        return flat_outputs.reshape(output_shape)
+
+
+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 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:
+                    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
+                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).to(device=device)
+                weight_scale = exported_scales[source_weight_scale_key].to(device=device).reshape(1, 1)
+                input_scale = exported_scales[source_input_scale_key].to(device=device).reshape(1, 1)
+                quantized_weight = PrototypeFloat8Tensor(
+                    quantized_data,
+                    weight_scale,
+                    act_quant_scale=input_scale,
+                    block_size=list(quantized_data.shape),
+                    mm_config=mm_config,
+                    act_quant_kwargs=activation_quant_kwargs,
+                    dtype=high_precision_dtype,
+                )
+                replacement.weight = nn.Parameter(quantized_weight, requires_grad=False)
+                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()
+    log.info(f"Loaded {len(converted_fqns)} calibrated ModelOpt FP8 weights into TorchAO")
+    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.
 

From 5f6ed58cb2a45aca0f0520492e27e35e41ade570 Mon Sep 17 00:00:00 2001
From: pengcuo <pzeren@nvidia.com>
Date: Mon, 3 Aug 2026 21:33:35 -0700
Subject: [PATCH 2/5] feat(inference): load ModelOpt FP8 checkpoints in
 from_pretrained_dcp
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Port of the CF-owned half of imaginaire4!10639. Unlike quantization.py,
`cosmos_framework/inference/model.py` is not in the release mapping
(see .file_mapping.json), so it has to be maintained here directly.

`_DiffusersLoadPlanner` gains `defer_modelopt_fp8_weights_loading`: when a
ModelOpt FP8 checkpoint is detected, the E4M3 weight tensors are dropped
from the DCP load plan (and excluded from the missing-key check) so no
temporary BF16 copy is ever materialized. `from_pretrained_dcp` then
installs them as TorchAO weights once the rest of the checkpoint is in.

Guards: ModelOpt FP8 checkpoints are rejected for DP-sharded models (the
weights become tensor subclasses, which requires plain-tensor params), when
runtime PTQ is also requested (the checkpoint is already quantized), and
when the checkpoint is not in diffusers layout (the FP8 loader follows the
diffusers weight map).

Note the loader-path difference from i4: this method does not apply runtime
PTQ — that lives in cosmos_framework/utils/generator/model_loader.py — so
ModelOpt FP8 checkpoints are supported on the diffusers path only.

Tests cover the planner skip, the load/install ordering, both guards, and
the quantization-layer conversion (CPU) plus a TorchAO dispatch check on
GPU.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: pengcuo <pzeren@nvidia.com>
---
 cosmos_framework/inference/model.py           |  87 ++++++-
 cosmos_framework/inference/model_test.py      | 100 ++++++++
 .../utils/generator/quantization_test.py      | 228 ++++++++++++++++++
 3 files changed, 408 insertions(+), 7 deletions(-)
 create mode 100644 cosmos_framework/utils/generator/quantization_test.py

diff --git a/cosmos_framework/inference/model.py b/cosmos_framework/inference/model.py
index 5b7c74ce..ae127fd3 100644
--- a/cosmos_framework/inference/model.py
+++ b/cosmos_framework/inference/model.py
@@ -45,6 +45,10 @@
 )
 from cosmos_framework.utils import misc
 from cosmos_framework.utils.flags import SMOKE
+from cosmos_framework.utils.generator.quantization import (
+    apply_modelopt_fp8_checkpoint_inplace,
+    is_modelopt_fp8_checkpoint,
+)
 
 if TYPE_CHECKING:
     from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel
@@ -351,12 +355,25 @@ def read_metadata(self) -> Metadata:
 class _DiffusersLoadPlanner(dcp.DefaultLoadPlanner):
     """Remap diffusers source keys onto the OmniMoTModel.net state dict for DCP load."""
 
-    def __init__(self, checkpoint_path: Path) -> None:
+    def __init__(
+        self,
+        checkpoint_path: Path,
+        *,
+        defer_modelopt_fp8_weights_loading: bool = False,
+    ) -> None:
+        """Initialize the planner.
+
+        When ``defer_modelopt_fp8_weights_loading`` is enabled, ModelOpt FP8
+        weight tensors are omitted from the DCP load so they can be installed
+        directly as TorchAO weights after the remaining checkpoint is loaded.
+        """
         super().__init__()
         self.checkpoint_path = checkpoint_path
         self.weight_map = _diffusers_weight_map(checkpoint_path)
         self.files_to_keys = _diffusers_files_to_keys(self.weight_map)
         self.has_vision_weights = any(rel_path.startswith("vision_encoder/") for rel_path in self.files_to_keys)
+        self.defer_modelopt_fp8_weights_loading = defer_modelopt_fp8_weights_loading
+        self.skipped_source_keys: set[str] = set()
 
     def set_up_planner(
         self,
@@ -365,9 +382,19 @@ def set_up_planner(
         is_coordinator: bool = False,
     ) -> None:
         target_state_dict = self._normalize_target_state_dict(state_dict)
-        remapped_state_dict, loaded_keys = self._build_remapped_state_dict(target_state_dict)
-
-        missing_keys = set(target_state_dict) - loaded_keys
+        if self.defer_modelopt_fp8_weights_loading:
+            if metadata is None:
+                raise ValueError("Checkpoint metadata is required to identify ModelOpt FP8 weights.")
+            self.skipped_source_keys = {
+                key
+                for key, tensor_metadata in metadata.state_dict_metadata.items()
+                if isinstance(tensor_metadata, TensorStorageMetadata)
+                and tensor_metadata.properties.dtype == torch.float8_e4m3fn
+                and key.endswith(".weight")
+            }
+        remapped_state_dict, loaded_keys, skipped_target_keys = self._build_remapped_state_dict(target_state_dict)
+
+        missing_keys = set(target_state_dict) - loaded_keys - skipped_target_keys
         if not self.has_vision_weights:
             missing_keys = {key for key in missing_keys if not key.startswith("language_model.visual.")}
         # Task-specialized checkpoints (e.g. Text2Image, Image2Video) omit the
@@ -403,9 +430,25 @@ def _normalize_target_state_dict(state_dict: dict[str, Any]) -> dict[str, Any]:
             target_state_dict[net_key] = tensor
         return target_state_dict
 
-    def _build_remapped_state_dict(self, target_state_dict: dict[str, Any]) -> tuple[dict[str, Any], set[str]]:
+    def _build_remapped_state_dict(
+        self, target_state_dict: dict[str, Any]
+    ) -> tuple[dict[str, Any], set[str], set[str]]:
+        """Build the state dict used to load diffusers weights into the model.
+
+        Args:
+            target_state_dict: Model state dict keyed by normalized Cosmos3 target names.
+
+        Returns:
+            A tuple containing:
+            - A state dict keyed by diffusers checkpoint names whose values reference
+              the corresponding target model tensors.
+            - Target keys that will be populated by the DCP load.
+            - Target keys whose ModelOpt FP8 weights were deferred for direct TorchAO
+              installation.
+        """
         remapped_state_dict: dict[str, Any] = {}
         loaded_keys: set[str] = set()
+        skipped_target_keys: set[str] = set()
         # When the model is built without a visual tower (e.g. Cosmos3-Edge t2i with
         # include_visual disabled), its state dict has no `language_model.visual.*`
         # targets, so the checkpoint's vision_encoder weights have nowhere to go — skip
@@ -413,6 +456,10 @@ def _build_remapped_state_dict(self, target_state_dict: dict[str, Any]) -> tuple
         has_visual_target = any(key.startswith("language_model.visual.") for key in target_state_dict)
         for diff_key, rel_path in sorted(self.weight_map.items()):
             net_key = _diffusers_to_net_key(diff_key, rel_path)
+            if diff_key in self.skipped_source_keys:
+                if net_key in target_state_dict:
+                    skipped_target_keys.add(net_key)
+                continue
             if net_key is None:
                 if _is_diffusers_model_weight_path(rel_path):
                     if rel_path.startswith("vision_encoder/") and not has_visual_target:
@@ -426,7 +473,7 @@ def _build_remapped_state_dict(self, target_state_dict: dict[str, Any]) -> tuple
                 raise KeyError(f"Multiple diffusers keys map to target model key {net_key!r}.")
             remapped_state_dict[diff_key] = target_tensor
             loaded_keys.add(net_key)
-        return remapped_state_dict, loaded_keys
+        return remapped_state_dict, loaded_keys, skipped_target_keys
 
 
 class Cosmos3OmniConfig(transformers.PretrainedConfig):
@@ -530,6 +577,20 @@ def from_pretrained_dcp(
         config.parallelism = attrs.asdict(parallelism_config)
         config.compile = attrs.asdict(compile_config)
         config.quantization = attrs.asdict(quantization_config)
+
+        # ModelOpt FP8 checkpoints ship already-quantized E4M3 weights plus static
+        # scales. They are installed as TorchAO tensor subclasses after the rest of
+        # the checkpoint loads, which only works on unsharded (plain-tensor) params.
+        modelopt_checkpoint = is_modelopt_fp8_checkpoint(checkpoint_path)
+        if modelopt_checkpoint and parallelism_config.data_parallel_shard_degree > 1:
+            raise ValueError("ModelOpt FP8 checkpoints are not supported for DP sharded models.")
+        if modelopt_checkpoint and quantization_config.method is not None:
+            raise ValueError(
+                "A ModelOpt FP8 checkpoint is already quantized; do not also request runtime quantization."
+            )
+        if modelopt_checkpoint and not _is_diffusers_checkpoint(checkpoint_path):
+            raise ValueError(f"ModelOpt FP8 loading requires a diffusers-format checkpoint layout: {checkpoint_path}")
+
         model = cls(config)
         # Thread the local checkpoint dir to the reasoner LM (consumed by Edge's
         # lazy ``_ensure_vision_tower``): checkpoints that bundle a
@@ -555,9 +616,21 @@ def from_pretrained_dcp(
                     dcp.load(
                         state_dict=state_dict,
                         storage_reader=_DiffusersHuggingFaceStorageReader(checkpoint_path),
-                        planner=_DiffusersLoadPlanner(checkpoint_path),
+                        planner=_DiffusersLoadPlanner(
+                            checkpoint_path,
+                            defer_modelopt_fp8_weights_loading=modelopt_checkpoint,
+                        ),
                         no_dist=no_dist,
                     )
+                    if modelopt_checkpoint:
+                        # The FP8 weights were skipped by the planner above; install
+                        # them straight from the checkpoint as TorchAO weights.
+                        apply_modelopt_fp8_checkpoint_inplace(
+                            model.model.net,
+                            checkpoint_path,
+                            key_mapper=_diffusers_to_net_key,
+                            weight_map=_diffusers_weight_map(checkpoint_path),
+                        )
                     return model
                 state_dict = get_model_state_dict(model)
                 _raise_on_missing_vision_keys(checkpoint_path, state_dict)
diff --git a/cosmos_framework/inference/model_test.py b/cosmos_framework/inference/model_test.py
index ca2c8c9c..9a3ac889 100644
--- a/cosmos_framework/inference/model_test.py
+++ b/cosmos_framework/inference/model_test.py
@@ -3,20 +3,27 @@
 
 import json
 from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import patch
 
 import attrs
 import hydra
+import pytest
 import safetensors.torch
 import torch
 import torch.distributed.checkpoint as dcp
+from torch.distributed.checkpoint.metadata import Metadata, TensorProperties, TensorStorageMetadata
 
 from cosmos_framework.configs.base.defaults.compile import CompileConfig
 from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig
+from cosmos_framework.configs.base.defaults.quantization import QuantizationConfig
 from cosmos_framework.inference.args import _CHECKPOINTS, DEFAULT_CHECKPOINT
 from cosmos_framework.inference.common.args import CheckpointType
 from cosmos_framework.inference.common.config import structure_config
 from cosmos_framework.inference.model import (
     Cosmos3OmniConfig,
+    Cosmos3OmniModel,
     _diffusers_to_net_key,
     _diffusers_weight_map,
     _DiffusersHuggingFaceStorageReader,
@@ -146,6 +153,99 @@ def test_diffusers_dcp_load_remaps_nested_safetensors(tmp_path: Path):
     torch.testing.assert_close(target["model.net._orig_mod.vae2llm.weight"], source)
 
 
+class LightweightDcpModel(Cosmos3OmniModel):
+    """Skip the hydra model build so the checkpoint-loading plumbing can be tested on CPU."""
+
+    model: Any
+
+    def __init__(self, config: Cosmos3OmniConfig, *args: object, **kwargs: object) -> None:
+        object.__setattr__(self, "model", SimpleNamespace(net=object()))
+
+
+def test_diffusers_load_planner_skips_modelopt_fp8_weights(tmp_path: Path) -> None:
+    checkpoint_path = tmp_path / "checkpoint"
+    checkpoint_path.mkdir()
+    source_key = "transformer.time_embedder.linear_1.weight"
+    target_key = "time_embedder.mlp.0.weight"
+    (checkpoint_path / "model.safetensors.index.json").write_text(
+        json.dumps({"weight_map": {source_key: "transformer/model.safetensors"}}),
+        encoding="utf-8",
+    )
+    metadata = Metadata(
+        state_dict_metadata={
+            source_key: TensorStorageMetadata(
+                properties=TensorProperties(dtype=torch.float8_e4m3fn),
+                size=torch.Size([1]),
+                chunks=[],
+            )
+        }
+    )
+    planner = _DiffusersLoadPlanner(checkpoint_path, defer_modelopt_fp8_weights_loading=True)
+
+    planner.set_up_planner({target_key: torch.empty(1)}, metadata)
+    plan = planner.create_local_plan()
+
+    assert planner.state_dict == {}
+    assert planner.skipped_source_keys == {source_key}
+    assert plan.items == []
+
+
+def test_from_pretrained_dcp_installs_modelopt_fp8_after_load(tmp_path: Path) -> None:
+    events: list[str] = []
+    converted_targets: list[object] = []
+    merged_weight_map = {"selected.weight": "transformer/model.safetensors"}
+
+    def record_load(**kwargs: Any) -> None:
+        assert kwargs["planner"].defer_modelopt_fp8_weights_loading
+        events.append("load")
+
+    def record_modelopt_conversion(
+        model: object, checkpoint_path: Path, *, key_mapper: object, weight_map: dict[str, str]
+    ) -> list[str]:
+        del checkpoint_path, key_mapper
+        events.append("modelopt")
+        converted_targets.append(model)
+        assert weight_map is merged_weight_map
+        return ["selected"]
+
+    with (
+        patch.object(CheckpointType, "from_path", return_value=CheckpointType.HF),
+        patch("cosmos_framework.inference.model._is_diffusers_checkpoint", return_value=True),
+        patch("cosmos_framework.inference.model.is_modelopt_fp8_checkpoint", return_value=True),
+        patch("cosmos_framework.inference.model._diffusers_weight_map", return_value=merged_weight_map),
+        patch("cosmos_framework.inference.model.get_model_state_dict", return_value={}),
+        patch("cosmos_framework.inference.model.dcp.load", side_effect=record_load),
+        patch(
+            "cosmos_framework.inference.model.apply_modelopt_fp8_checkpoint_inplace",
+            side_effect=record_modelopt_conversion,
+        ),
+    ):
+        model = LightweightDcpModel.from_pretrained_dcp(checkpoint_path=tmp_path, config=Cosmos3OmniConfig())
+
+    assert events == ["load", "modelopt"]
+    assert converted_targets == [model.model.net]
+
+
+def test_from_pretrained_dcp_rejects_modelopt_fp8_with_runtime_quantization(tmp_path: Path) -> None:
+    with patch("cosmos_framework.inference.model.is_modelopt_fp8_checkpoint", return_value=True):
+        with pytest.raises(ValueError, match="already quantized"):
+            LightweightDcpModel.from_pretrained_dcp(
+                checkpoint_path=tmp_path,
+                config=Cosmos3OmniConfig(),
+                quantization_config=QuantizationConfig(method="fp8"),
+            )
+
+
+def test_from_pretrained_dcp_rejects_modelopt_fp8_with_dp_sharding(tmp_path: Path) -> None:
+    with patch("cosmos_framework.inference.model.is_modelopt_fp8_checkpoint", return_value=True):
+        with pytest.raises(ValueError, match="not supported for DP sharded models"):
+            LightweightDcpModel.from_pretrained_dcp(
+                checkpoint_path=tmp_path,
+                config=Cosmos3OmniConfig(),
+                parallelism_config=ParallelismConfig(data_parallel_shard_degree=2),
+            )
+
+
 def test_diffusers_weight_map_registered_checkpoint():
     checkpoint_path = Path(_CHECKPOINTS["Cosmos3-Nano"].hf.download())
 
diff --git a/cosmos_framework/utils/generator/quantization_test.py b/cosmos_framework/utils/generator/quantization_test.py
new file mode 100644
index 00000000..3fb9f720
--- /dev/null
+++ b/cosmos_framework/utils/generator/quantization_test.py
@@ -0,0 +1,228 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+import json
+from pathlib import Path
+
+import pytest
+import safetensors.torch
+import torch
+from torch import nn
+
+from cosmos_framework.utils.generator.quantization import (
+    apply_modelopt_fp8_checkpoint_inplace,
+    is_modelopt_fp8_checkpoint,
+)
+
+
+class TinyLinearModel(nn.Module):
+    selected: nn.Linear
+    unselected: nn.Linear
+
+    def __init__(self, device: torch.device | str = "cuda") -> None:
+        super().__init__()
+        self.selected = nn.Linear(16, 16, bias=False, device=device, dtype=torch.bfloat16)
+        self.unselected = nn.Linear(16, 16, bias=False, device=device, dtype=torch.bfloat16)
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:  # inputs: [B,D], returns: [B,D]
+        return self.selected(inputs) + self.unselected(inputs)  # [B,D]
+
+
+def _write_modelopt_fp8_checkpoint(
+    checkpoint_path: Path,
+    quantized_weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_scale: torch.Tensor,
+    additional_module: str | None = None,
+    additional_module_has_scales: bool = True,
+) -> dict[str, str]:
+    """Write a minimal ModelOpt-style FP8 export with weights and scales in separate shards."""
+    checkpoint_path.mkdir()
+    (checkpoint_path / "hf_quant_config.json").write_text(
+        json.dumps({"quant_method": "modelopt", "quant_algo": "FP8"}),
+        encoding="utf-8",
+    )
+    weight_shard_name = "model-00001-of-00002.safetensors"
+    scale_shard_name = "model-00002-of-00002.safetensors"
+    module_names = ["selected"]
+    if additional_module is not None:
+        module_names.append(additional_module)
+    scale_module_names = module_names if additional_module_has_scales else ["selected"]
+    weight_tensors = {f"{module_name}.weight": quantized_weight.clone() for module_name in module_names}
+    scale_tensors = {
+        key: value
+        for module_name in scale_module_names
+        for key, value in (
+            (f"{module_name}.input_scale", input_scale.clone()),
+            (f"{module_name}.weight_scale", weight_scale.clone()),
+        )
+    }
+    safetensors.torch.save_file(weight_tensors, checkpoint_path / weight_shard_name)
+    safetensors.torch.save_file(scale_tensors, checkpoint_path / scale_shard_name)
+    weight_map = {
+        **{f"{module_name}.weight": weight_shard_name for module_name in module_names},
+        **{
+            f"{module_name}.{scale_name}": scale_shard_name
+            for module_name in scale_module_names
+            for scale_name in ("input_scale", "weight_scale")
+        },
+    }
+    (checkpoint_path / "model.safetensors.index.json").write_text(
+        json.dumps({"weight_map": weight_map}),
+        encoding="utf-8",
+    )
+    return weight_map
+
+
+def _identity_key_mapper(source_key: str, shard_path: str) -> str:
+    del shard_path
+    return source_key
+
+
+def _ignore_extra_key_mapper(source_key: str, shard_path: str) -> str | None:
+    del shard_path
+    return None if source_key.startswith("ignored.") else source_key
+
+
+def test_is_modelopt_fp8_checkpoint_without_config(tmp_path: Path) -> None:
+    assert not is_modelopt_fp8_checkpoint(tmp_path)
+
+
+def test_is_modelopt_fp8_checkpoint_rejects_malformed_config(tmp_path: Path) -> None:
+    (tmp_path / "hf_quant_config.json").write_text("{", encoding="utf-8")
+
+    with pytest.raises(ValueError, match="Unable to read a valid checkpoint quantization config"):
+        is_modelopt_fp8_checkpoint(tmp_path)
+
+
+@pytest.mark.parametrize(
+    "quant_config",
+    [
+        pytest.param([], id="non-object"),
+        pytest.param({"quant_method": "other", "quant_algo": "FP8"}, id="unsupported-method"),
+        pytest.param({"quant_method": "modelopt", "quant_algo": "OTHER"}, id="unsupported-algorithm"),
+    ],
+)
+def test_is_modelopt_fp8_checkpoint_rejects_unsupported_config(tmp_path: Path, quant_config: object) -> None:
+    (tmp_path / "hf_quant_config.json").write_text(json.dumps(quant_config), encoding="utf-8")
+
+    with pytest.raises(ValueError, match="Unsupported checkpoint quantization configuration"):
+        is_modelopt_fp8_checkpoint(tmp_path)
+
+
+def test_apply_modelopt_fp8_checkpoint_preserves_exported_tensors(tmp_path: Path) -> None:
+    model = TinyLinearModel(device="cpu").eval()
+    original_selected = model.selected
+    quantized_weight = torch.arange(256, dtype=torch.uint8).view(torch.float8_e4m3fn).reshape(16, 16)
+    weight_scale = torch.tensor(0.125, dtype=torch.float32)
+    input_scale = torch.tensor(0.25, dtype=torch.float32)
+    weight_map = _write_modelopt_fp8_checkpoint(tmp_path / "checkpoint", quantized_weight, weight_scale, input_scale)
+
+    converted = apply_modelopt_fp8_checkpoint_inplace(
+        model,
+        tmp_path / "checkpoint",
+        key_mapper=_identity_key_mapper,
+        weight_map=weight_map,
+    )
+
+    assert is_modelopt_fp8_checkpoint(tmp_path / "checkpoint")
+    assert converted == ["selected"]
+    assert isinstance(model.selected, nn.Linear)
+    assert model.selected is not original_selected
+    assert type(original_selected) is nn.Linear
+    assert type(model.selected.weight).__name__ == "PrototypeFloat8Tensor"
+    assert torch.equal(model.selected.weight.qdata.view(torch.uint8), quantized_weight.view(torch.uint8))
+    assert torch.equal(model.selected.weight.scale, weight_scale.reshape(1, 1))
+    assert torch.equal(model.selected.weight.act_quant_scale, input_scale.reshape(1, 1))
+    assert type(model.unselected.weight) is nn.Parameter
+
+
+def test_apply_modelopt_fp8_checkpoint_skips_ignored_keys(tmp_path: Path) -> None:
+    model = TinyLinearModel(device="cpu").eval()
+    quantized_weight = torch.arange(256, dtype=torch.uint8).view(torch.float8_e4m3fn).reshape(16, 16)
+    weight_scale = torch.tensor(0.125, dtype=torch.float32)
+    input_scale = torch.tensor(0.25, dtype=torch.float32)
+    weight_map = _write_modelopt_fp8_checkpoint(
+        tmp_path / "checkpoint",
+        quantized_weight,
+        weight_scale,
+        input_scale,
+        additional_module="ignored",
+        additional_module_has_scales=False,
+    )
+
+    converted = apply_modelopt_fp8_checkpoint_inplace(
+        model,
+        tmp_path / "checkpoint",
+        key_mapper=_ignore_extra_key_mapper,
+        weight_map=weight_map,
+    )
+
+    assert converted == ["selected"]
+
+
+def test_apply_modelopt_fp8_checkpoint_skips_absent_target(tmp_path: Path) -> None:
+    model = TinyLinearModel(device="cpu").eval()
+    quantized_weight = torch.arange(256, dtype=torch.uint8).view(torch.float8_e4m3fn).reshape(16, 16)
+    weight_scale = torch.tensor(0.125, dtype=torch.float32)
+    input_scale = torch.tensor(0.25, dtype=torch.float32)
+    weight_map = _write_modelopt_fp8_checkpoint(
+        tmp_path / "checkpoint",
+        quantized_weight,
+        weight_scale,
+        input_scale,
+        additional_module="missing",
+        additional_module_has_scales=False,
+    )
+
+    converted = apply_modelopt_fp8_checkpoint_inplace(
+        model,
+        tmp_path / "checkpoint",
+        key_mapper=_identity_key_mapper,
+        weight_map=weight_map,
+    )
+
+    assert converted == ["selected"]
+
+
+@pytest.mark.gpus(1)
+def test_apply_modelopt_fp8_checkpoint_uses_torchao_linear_dispatch(tmp_path: Path) -> None:
+    if torch.cuda.get_device_capability() < (8, 9):
+        pytest.skip("requires an Ada or newer GPU")
+    pytest.importorskip("torchao")
+
+    model = TinyLinearModel().eval()
+    original_weight = model.selected.weight.detach().cpu()
+    weight_scale = original_weight.abs().amax().float() / torch.finfo(torch.float8_e4m3fn).max
+    quantized_weight = (original_weight / weight_scale).to(torch.float8_e4m3fn)
+    input_scale = torch.tensor(0.025, dtype=torch.float32)
+    weight_map = _write_modelopt_fp8_checkpoint(tmp_path / "checkpoint", quantized_weight, weight_scale, input_scale)
+    apply_modelopt_fp8_checkpoint_inplace(
+        model,
+        tmp_path / "checkpoint",
+        key_mapper=_identity_key_mapper,
+        weight_map=weight_map,
+    )
+
+    inputs = torch.randn((6, 16), device="cuda", dtype=torch.bfloat16)
+    output = model.selected(inputs)
+    reasoning_inputs = torch.randn((2, 3, 16), device="cuda", dtype=torch.bfloat16)
+    reasoning_output = model.selected(reasoning_inputs)
+    output_after_reasoning = model.selected(inputs)
+    compiled_model = torch.compile(model.selected, dynamic=True)
+    compiled_output = compiled_model(inputs)
+    compiled_reasoning_output = compiled_model(reasoning_inputs)
+    empty_output = model.selected(inputs[:0])
+
+    assert output.shape == (6, 16)
+    assert torch.isfinite(output).all()
+    assert reasoning_output.shape == (2, 3, 16)
+    assert torch.isfinite(reasoning_output).all()
+    assert output_after_reasoning.shape == (6, 16)
+    assert torch.isfinite(output_after_reasoning).all()
+    assert compiled_output.shape == (6, 16)
+    assert torch.isfinite(compiled_output).all()
+    assert compiled_reasoning_output.shape == (2, 3, 16)
+    assert torch.isfinite(compiled_reasoning_output).all()
+    assert empty_output.shape == (0, 16)
+    assert model.selected.weight.act_quant_scale.shape == (1, 1)

From 1bc2690b846548cf5be3fcf3c430ec41f3464b34 Mon Sep 17 00:00:00 2001
From: pengcuo <pzeren@nvidia.com>
Date: Thu, 6 Aug 2026 23:27:30 -0700
Subject: [PATCH 3/5] feat(quantization): support FSDP-sharded ModelOpt FP8
 checkpoint loading

Lifts the data_parallel_shard_degree > 1 guard added in #157 by teaching
TorchAO's PrototypeFloat8Tensor the ops and hooks FSDP2 needs, and by
swapping ModelOpt FP8 linears before the network is parallelized so
peak memory follows FP8 shapes rather than bf16.

- utils/generator/quantization.py:
  * install_torchao_float8_fsdp_support() registers 8 aten ops
    (view/reshape/split/slice/as_strided/new_zeros/empty_like/
    detach/clone/_to_copy/copy_) and fsdp_pre_all_gather /
    fsdp_post_all_gather on PrototypeFloat8Tensor. Only the E4M3
    qdata travels; the PerTensor scale is identical on every rank.
  * plan_modelopt_fp8_targets() derives the target FQN list from the
    checkpoint index alone, so build_net can consume it without
    reaching back into the loader.
  * swap_modelopt_fp8_linears_on_meta() replaces target linears with
    _ModelOptFloat8Linear + PrototypeFloat8Tensor placeholders on
    meta before FSDP wrap -- FSDP registers the pre-all-gather
    extension at wrap time, so a post-hoc swap wouldn't be picked up.
  * apply_modelopt_fp8_checkpoint_inplace() now handles the DTensor-
    sharded (per-rank slice via distribute_tensor) and pre-swapped
    meta paths in addition to the legacy replicated path.
- configs/base/defaults/quantization.py: adds modelopt_fp8_checkpoint_path
  and modelopt_fp8_target_fqns so the meta swap is driven from config.
- inference/common/config.py: undo_config_replacements normalizes
  checkpoints that ship un-rewritten internal-tree paths (e.g.
  cosmos3-super-i2v-fp8-14072026) by running the forward table first.
- inference/model.py: drops the DP-shard reject; pre-computes the
  target FQN list and threads it into config.quantization.
- model/generator/omni_mot_model.py: swaps ModelOpt FP8 linears in
  build_net before parallelize_vfm_network, gated on
  config.quantization.modelopt_fp8_checkpoint_path.
- model/generator/reasoner/qwen3_vl/qwen3_vl.py: _init_weights early-
  returns on _ModelOptFloat8Linear (E4M3 has no normal_ kernel; random
  init is pointless for a checkpoint-filled weight).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 .../configs/base/defaults/quantization.py     |  14 +
 cosmos_framework/inference/common/config.py   |   9 +
 cosmos_framework/inference/model.py           |  19 +-
 .../model/generator/omni_mot_model.py         |   9 +
 .../generator/reasoner/qwen3_vl/qwen3_vl.py   |   7 +
 .../utils/generator/quantization.py           | 434 ++++++++++++++++--
 6 files changed, 456 insertions(+), 36 deletions(-)

diff --git a/cosmos_framework/configs/base/defaults/quantization.py b/cosmos_framework/configs/base/defaults/quantization.py
index f07d32f7..64cf04ef 100644
--- a/cosmos_framework/configs/base/defaults/quantization.py
+++ b/cosmos_framework/configs/base/defaults/quantization.py
@@ -44,3 +44,17 @@ class QuantizationConfig:
     # considered as excluded.
     include_regex: list[str] = attrs.field(factory=list)
     exclude_regex: list[str] = attrs.field(factory=list)
+
+    # Local root of a ModelOpt static-FP8 diffusers checkpoint. When set, the
+    # linears named by ``modelopt_fp8_target_fqns`` are swapped to FP8 modules on
+    # the meta device *before* the network is parallelized and materialized, so
+    # peak memory follows the FP8 weights rather than their bf16 shapes. This is
+    # independent of ``method``, which selects runtime (post-training)
+    # quantization; a ModelOpt checkpoint arrives already quantized.
+    modelopt_fp8_checkpoint_path: str | None = attrs.field(default=None)
+
+    # Target module FQNs (relative to the VFM network) that the ModelOpt FP8
+    # checkpoint carries quantized weights for. Computed from the checkpoint
+    # index by the loader, which knows the diffusers key mapping; passed through
+    # the config so the meta-device swap in ``build_net`` needs no mapper.
+    modelopt_fp8_target_fqns: list[str] = attrs.field(factory=list)
diff --git a/cosmos_framework/inference/common/config.py b/cosmos_framework/inference/common/config.py
index 80c3eec9..ff73a51e 100644
--- a/cosmos_framework/inference/common/config.py
+++ b/cosmos_framework/inference/common/config.py
@@ -569,6 +569,15 @@ def apply_config_replacements(config_str: str) -> str:
 
 def undo_config_replacements(config_str: str) -> str:
     """Undo config replacements to a config string."""
+    # Some published checkpoints (e.g. cosmos3-super-i2v-fp8-14072026) carry a
+    # config.json dumped straight from the internal tree, so it still names
+    # ``projects.cosmos3.vfm.*`` / ``imaginaire.*`` instead of the ``cosmos3._src.*``
+    # form the release rewrite produces. Normalize to that form first so the
+    # inverse table below sees what it expects; already-rewritten configs are
+    # unaffected (the patterns cannot match ``cosmos3._src.*``).
+    for pattern, repl in CONFIG_REPLACEMENTS:
+        config_str = re.sub(pattern, repl, config_str)
+
     for pattern, repl in CONFIG_REPLACEMENTS_INVERSE:
         config_str = re.sub(pattern, repl, config_str)
 
diff --git a/cosmos_framework/inference/model.py b/cosmos_framework/inference/model.py
index ae127fd3..7fd5eee8 100644
--- a/cosmos_framework/inference/model.py
+++ b/cosmos_framework/inference/model.py
@@ -48,6 +48,7 @@
 from cosmos_framework.utils.generator.quantization import (
     apply_modelopt_fp8_checkpoint_inplace,
     is_modelopt_fp8_checkpoint,
+    plan_modelopt_fp8_targets,
 )
 
 if TYPE_CHECKING:
@@ -579,17 +580,27 @@ def from_pretrained_dcp(
         config.quantization = attrs.asdict(quantization_config)
 
         # ModelOpt FP8 checkpoints ship already-quantized E4M3 weights plus static
-        # scales. They are installed as TorchAO tensor subclasses after the rest of
-        # the checkpoint loads, which only works on unsharded (plain-tensor) params.
+        # scales. The target linears are swapped to FP8 modules on the meta device
+        # during `build_net` — before FSDP wrap and materialization — so both the
+        # sharded and replicated paths work and peak memory follows the FP8 weights.
         modelopt_checkpoint = is_modelopt_fp8_checkpoint(checkpoint_path)
-        if modelopt_checkpoint and parallelism_config.data_parallel_shard_degree > 1:
-            raise ValueError("ModelOpt FP8 checkpoints are not supported for DP sharded models.")
         if modelopt_checkpoint and quantization_config.method is not None:
             raise ValueError(
                 "A ModelOpt FP8 checkpoint is already quantized; do not also request runtime quantization."
             )
         if modelopt_checkpoint and not _is_diffusers_checkpoint(checkpoint_path):
             raise ValueError(f"ModelOpt FP8 loading requires a diffusers-format checkpoint layout: {checkpoint_path}")
+        if modelopt_checkpoint:
+            # Resolve which linears the checkpoint quantizes here, where the
+            # diffusers key mapping lives, and hand the plain FQN list to the
+            # model config so `build_net` can do the meta-device swap without
+            # reaching back into the loader.
+            config.quantization["modelopt_fp8_checkpoint_path"] = str(checkpoint_path)
+            config.quantization["modelopt_fp8_target_fqns"] = plan_modelopt_fp8_targets(
+                checkpoint_path,
+                _diffusers_to_net_key,
+                weight_map=_diffusers_weight_map(checkpoint_path),
+            )
 
         model = cls(config)
         # Thread the local checkpoint dir to the reasoner LM (consumed by Edge's
diff --git a/cosmos_framework/model/generator/omni_mot_model.py b/cosmos_framework/model/generator/omni_mot_model.py
index 2b568b0e..4ff66bb5 100644
--- a/cosmos_framework/model/generator/omni_mot_model.py
+++ b/cosmos_framework/model/generator/omni_mot_model.py
@@ -81,6 +81,7 @@
 from cosmos_framework.utils.generator.dtensor_helper import DTensorFastEmaModelUpdater
 from cosmos_framework.utils.generator.model_weights_stats import WeightTrainingStat
 from cosmos_framework.utils.generator.parallelism import ParallelDims
+from cosmos_framework.utils.generator.quantization import swap_modelopt_fp8_linears_on_meta
 
 
 class OmniMoTModel(ImaginaireModel):
@@ -266,6 +267,14 @@ def build_net(self, dtype: torch.dtype, *, lora_enabled: bool | None = None) ->
                     lora_target_modules=self.config.lora_target_modules,
                 )
 
+        # Swap ModelOpt FP8 linears BEFORE FSDP wrap and materialization, for the
+        # same reason LoRA is injected early. `fully_shard` wraps the whole network
+        # into one parameter group, so a linear replaced afterwards leaves the group
+        # holding a stale parameter; and `to_empty` below is what sets peak memory,
+        # which at bf16 shapes exceeds a single 80 GB device for a Super-class model.
+        if self.config.quantization.modelopt_fp8_checkpoint_path:
+            swap_modelopt_fp8_linears_on_meta(net, self.config.quantization.modelopt_fp8_target_fqns)
+
         self.install_attention_dispatch(net)
 
         net = parallelize_vfm_network(
diff --git a/cosmos_framework/model/generator/reasoner/qwen3_vl/qwen3_vl.py b/cosmos_framework/model/generator/reasoner/qwen3_vl/qwen3_vl.py
index d442752a..0eb93a30 100644
--- a/cosmos_framework/model/generator/reasoner/qwen3_vl/qwen3_vl.py
+++ b/cosmos_framework/model/generator/reasoner/qwen3_vl/qwen3_vl.py
@@ -53,6 +53,8 @@ def _default_rope_init(config, device=None, **kwargs):
     get_rope_index as _get_rope_index,
 )
 
+from cosmos_framework.utils.generator.quantization import _ModelOptFloat8Linear
+
 from .configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig
 
 TransformersKwargs = Any
@@ -608,6 +610,11 @@ class Qwen3VLPreTrainedModel(PreTrainedModel):
 
     def _init_weights(self, module: nn.Module, buffer_device: torch.device | None) -> None:
         """Initialize the weights."""
+        # A ModelOpt FP8 linear carries an E4M3 weight that the checkpoint fills
+        # wholesale, and `normal_` has no float8 kernel. Random init would be both
+        # unimplemented and pointless here.
+        if isinstance(module, _ModelOptFloat8Linear):
+            return
         super()._init_weights(module)
 
         if isinstance(
diff --git a/cosmos_framework/utils/generator/quantization.py b/cosmos_framework/utils/generator/quantization.py
index ae215a10..56adbbed 100644
--- a/cosmos_framework/utils/generator/quantization.py
+++ b/cosmos_framework/utils/generator/quantization.py
@@ -26,6 +26,7 @@
 
 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
@@ -42,6 +43,12 @@
 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
+
     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
@@ -55,6 +62,169 @@ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         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"
@@ -73,6 +243,145 @@ def is_modelopt_fp8_checkpoint(checkpoint_path: str | Path) -> bool:
     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,
@@ -167,7 +476,18 @@ def apply_modelopt_fp8_checkpoint_inplace(
                         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:
+                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):
@@ -227,36 +547,80 @@ def apply_modelopt_fp8_checkpoint_inplace(
                 source_input_scale_key = f"{source_module_key}.input_scale"
                 module = model.get_submodule(target_module_fqn)
                 device = module.weight.device
-                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).to(device=device)
-                weight_scale = exported_scales[source_weight_scale_key].to(device=device).reshape(1, 1)
-                input_scale = exported_scales[source_input_scale_key].to(device=device).reshape(1, 1)
-                quantized_weight = PrototypeFloat8Tensor(
-                    quantized_data,
-                    weight_scale,
-                    act_quant_scale=input_scale,
-                    block_size=list(quantized_data.shape),
-                    mm_config=mm_config,
-                    act_quant_kwargs=activation_quant_kwargs,
-                    dtype=high_precision_dtype,
-                )
-                replacement.weight = nn.Parameter(quantized_weight, requires_grad=False)
+
+                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)
 
@@ -264,7 +628,13 @@ def apply_modelopt_fp8_checkpoint_inplace(
     if missing_conversions:
         raise ValueError(f"ModelOpt FP8 weights were not converted: {sorted(missing_conversions)}")
     converted_fqns.sort()
-    log.info(f"Loaded {len(converted_fqns)} calibrated ModelOpt FP8 weights into TorchAO")
+    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
 

From 4877f276e33a6c2ab11c889e26cefb16392d2ae8 Mon Sep 17 00:00:00 2001
From: pengcuo <pzeren@nvidia.com>
Date: Thu, 6 Aug 2026 23:51:51 -0700
Subject: [PATCH 4/5] fix(quantization): mirror quantization onto
 OmniMoTModelConfig for build_net

The FSDP FP8 commit reads self.config.quantization from OmniMoTModel.build_net,
but OmniMoTModelConfig had no quantization field, so every path that went
through build_net -- including plain bf16 training and convert_model_to_dcp --
crashed with "Key 'quantization' not in 'OmniMoTModelConfig'".

Add quantization: QuantizationConfig on OmniMoTModelConfig (default is the
disabled-quantization instance, so bf16 flows through unchanged). This
mirrors how parallelism and compile are already threaded from the outer
Cosmos3OmniConfig property setter to the inner model schema.

Also update the tests:
- test_from_pretrained_dcp_installs_modelopt_fp8_after_load: mock the new
  plan_modelopt_fp8_targets call so the fake checkpoint fixture isn't asked
  to open real safetensors shards.
- test_from_pretrained_dcp_rejects_modelopt_fp8_with_dp_sharding: removed,
  DP sharding is now supported by the FSDP FP8 commit.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 .../configs/base/defaults/model_config.py             |  7 +++++++
 cosmos_framework/inference/model_test.py              | 11 +----------
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/cosmos_framework/configs/base/defaults/model_config.py b/cosmos_framework/configs/base/defaults/model_config.py
index be411313..a46e45fe 100644
--- a/cosmos_framework/configs/base/defaults/model_config.py
+++ b/cosmos_framework/configs/base/defaults/model_config.py
@@ -10,6 +10,7 @@
 from cosmos_framework.configs.base.defaults.compile import CompileConfig
 from cosmos_framework.configs.base.defaults.ema import EMAConfig
 from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig
+from cosmos_framework.configs.base.defaults.quantization import QuantizationConfig
 from cosmos_framework.configs.base.defaults.reasoner import VLMConfig
 
 
@@ -161,6 +162,12 @@ class OmniMoTModelConfig:
     # torch.compile knobs (enabled, compiled_region, dynamic, ...).
     compile: CompileConfig = CompileConfig()
 
+    # Post-training quantization + ModelOpt FP8 checkpoint metadata. Mirrored
+    # from Cosmos3OmniConfig.quantization (see ``inference/model.py``) so
+    # ``build_net`` can read modelopt_fp8_checkpoint_path / target_fqns without
+    # reaching outside the model config schema.
+    quantization: QuantizationConfig = QuantizationConfig()
+
     # Activation-checkpointing policy (trade-off between memory and speed).
     activation_checkpointing: ActivationCheckpointingConfig = ActivationCheckpointingConfig()
 
diff --git a/cosmos_framework/inference/model_test.py b/cosmos_framework/inference/model_test.py
index 9a3ac889..02a7a897 100644
--- a/cosmos_framework/inference/model_test.py
+++ b/cosmos_framework/inference/model_test.py
@@ -213,6 +213,7 @@ def record_modelopt_conversion(
         patch("cosmos_framework.inference.model._is_diffusers_checkpoint", return_value=True),
         patch("cosmos_framework.inference.model.is_modelopt_fp8_checkpoint", return_value=True),
         patch("cosmos_framework.inference.model._diffusers_weight_map", return_value=merged_weight_map),
+        patch("cosmos_framework.inference.model.plan_modelopt_fp8_targets", return_value=["selected"]),
         patch("cosmos_framework.inference.model.get_model_state_dict", return_value={}),
         patch("cosmos_framework.inference.model.dcp.load", side_effect=record_load),
         patch(
@@ -236,16 +237,6 @@ def test_from_pretrained_dcp_rejects_modelopt_fp8_with_runtime_quantization(tmp_
             )
 
 
-def test_from_pretrained_dcp_rejects_modelopt_fp8_with_dp_sharding(tmp_path: Path) -> None:
-    with patch("cosmos_framework.inference.model.is_modelopt_fp8_checkpoint", return_value=True):
-        with pytest.raises(ValueError, match="not supported for DP sharded models"):
-            LightweightDcpModel.from_pretrained_dcp(
-                checkpoint_path=tmp_path,
-                config=Cosmos3OmniConfig(),
-                parallelism_config=ParallelismConfig(data_parallel_shard_degree=2),
-            )
-
-
 def test_diffusers_weight_map_registered_checkpoint():
     checkpoint_path = Path(_CHECKPOINTS["Cosmos3-Nano"].hf.download())
 

From 144553e95c66c6776656b14f5473dfd919efd503 Mon Sep 17 00:00:00 2001
From: pengcuo <pzeren@nvidia.com>
Date: Fri, 7 Aug 2026 00:07:37 -0700
Subject: [PATCH 5/5] fix(export): teach diffusers-export validator about
 ModelOpt FP8 defaults

The FSDP FP8 commit added modelopt_fp8_checkpoint_path and
modelopt_fp8_target_fqns to QuantizationConfig, but the
convert_model_to_diffusers validator compared the incoming
quantization dict against a hardcoded {exclude_regex, include_regex,
method}. After the previous commit mirrored quantization onto
OmniMoTModelConfig, the exported config now carries the full
quantization sub-config, so even a fully-default QuantizationConfig()
tripped the "non-default internal quantization" ValueError and broke
convert_model_to_diffusers for every trained bf16 model.

Add both fields (None / [] defaults) to disabled_quantization so a
disabled config still matches. Non-default values (an actual
modelopt_fp8_checkpoint_path or a populated target_fqns list) still
raise, which is what we want: exported diffusers checkpoints should
not carry a rank-local ModelOpt path.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 cosmos_framework/scripts/convert_model_to_diffusers.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/cosmos_framework/scripts/convert_model_to_diffusers.py b/cosmos_framework/scripts/convert_model_to_diffusers.py
index f18edc95..188bad67 100644
--- a/cosmos_framework/scripts/convert_model_to_diffusers.py
+++ b/cosmos_framework/scripts/convert_model_to_diffusers.py
@@ -134,6 +134,8 @@ def _build_public_export_model_config(model_dict: dict[str, Any]) -> dict[str, A
             "exclude_regex": [],
             "include_regex": [],
             "method": None,
+            "modelopt_fp8_checkpoint_path": None,
+            "modelopt_fp8_target_fqns": [],
         }
         # ``fp8_granularity`` is inert when quantization is disabled (method=None);
         # ignore it so a disabled config still matches the expected default.