feat: load ModelOpt static-FP8 checkpoints via TorchAO

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()
 

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)

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)
 

cosmos_framework/inference/model.py

@@ -45,6 +45,11 @@
 )
 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,
+    plan_modelopt_fp8_targets,
+)
 
 if TYPE_CHECKING:
     from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel
@@ -351,12 +356,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 +383,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,16 +431,36 @@ 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
         # them rather than erroring on unmapped projector/visual keys.
         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 +474,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 +578,30 @@ 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. 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 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
         # lazy ``_ensure_vision_tower``): checkpoints that bundle a
@@ -555,9 +627,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)

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,90 @@ 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.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(
+            "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_diffusers_weight_map_registered_checkpoint():
     checkpoint_path = Path(_CHECKPOINTS["Cosmos3-Nano"].hf.download())
 

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(

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(

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.