fix(quant): export modelopt_state.pth after quantization

cosmos_framework/inference/modelopt_diffusers.py

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+"""Compatibility for loading serialized ModelOpt graphs with native Diffusers.
+
+Diffusers 0.39's ModelOpt quantizer assigns every checkpoint tensor to
+``module._parameters``. That leaves quantizer buffers on meta and discards
+QTensorWrapper metadata. Enable this alongside ModelOpt HF checkpointing before
+loading a pre-quantized Diffusers model.
+"""
+
+
+def enable_modelopt_diffusers_checkpointing() -> None:
+    """Enable ModelOpt restoration, preserving quantizer buffers and FP8 wrappers."""
+    import modelopt.torch.opt as mto
+    from diffusers.quantizers.modelopt.modelopt_quantizer import NVIDIAModelOptQuantizer
+    from diffusers.utils import get_module_from_name
+    from modelopt.torch.quantization.qtensor import QTensorWrapper
+
+    mto.enable_huggingface_checkpointing()
+    original = NVIDIAModelOptQuantizer.create_quantized_param
+    if getattr(original, "_cosmos_preserves_modelopt_state", False):
+        return
+    original_preprocess = NVIDIAModelOptQuantizer._process_model_before_weight_loading
+
+    def preprocess(self, model, device_map, keep_in_fp32_modules=None, **kwargs):
+        if keep_in_fp32_modules is None:
+            keep_in_fp32_modules = []
+        if self.pre_quantized:
+            # Diffusers casts tensors before create_quantized_param is called.
+            # Extend the shared loader list so calibrated scales retain FP32.
+            for role in ("input_quantizer", "weight_quantizer"):
+                if role not in keep_in_fp32_modules:
+                    keep_in_fp32_modules.append(role)
+        return original_preprocess(self, model, device_map, keep_in_fp32_modules, **kwargs)
+
+    def create_quantized_param(self, model, param_value, param_name, target_device, *args, **kwargs):
+        if self.pre_quantized:
+            module, tensor_name = get_module_from_name(model, param_name)
+            if tensor_name in module._buffers:
+                module._buffers[tensor_name] = param_value.to(device=target_device)
+                return
+            parameter = module._parameters.get(tensor_name)
+            if isinstance(parameter, QTensorWrapper):
+                module._parameters[tensor_name] = QTensorWrapper(
+                    qtensor=param_value.to(device=target_device),
+                    metadata=parameter.get_state()["metadata"],
+                )
+                return
+        return original(self, model, param_value, param_name, target_device, *args, **kwargs)
+
+    create_quantized_param._cosmos_preserves_modelopt_state = True
+    NVIDIAModelOptQuantizer._process_model_before_weight_loading = preprocess
+    NVIDIAModelOptQuantizer.create_quantized_param = create_quantized_param

cosmos_framework/quantization/__init__.py

@@ -10,7 +10,7 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Literal
 
@@ -26,6 +26,7 @@
     resolve_checkpoint_path,
     save_sharded_safetensors,
 )
+from .metadata import collect_quantization_metadata, write_quantization_metadata
 
 FP8_E4M3_MAX = 448.0
 
@@ -126,6 +127,7 @@ def quantize_fp8_checkpoint(
     i2v_cond_dataset: str | None = None,
     keep_model: bool = False,
     calibration_behavior: Literal["framework", "legacy"] = "framework",
+    mixed_precision_steps: int = 3,
 ):
     """Quantize a diffusers-layout Cosmos3 checkpoint to FP8 and write a drop-in dir.
 
@@ -140,6 +142,11 @@ def quantize_fp8_checkpoint(
     recipe that restores the historical scheduler, numeric bridges, and attention
     behavior to pursue legacy-equivalent FP8 calibration. Returns the assembled
     ``output_dir`` (and, if ``keep_model``, the calibrated model for inspection).
+
+    ``mixed_precision_steps`` controls the symmetric first/last denoising-step
+    A16 policy written to ``transformer/config.json`` for vLLM-Omni.
+    ``quantization_metadata.json`` records the producing environment, source
+    identity, and calibration recipe independently of inherited source configs.
     """
     input_dir, output_dir = resolve_checkpoint_path(model_name_or_path), Path(output_dir)
     if profile not in ("t2v", "t2i", "i2v"):
@@ -180,6 +187,32 @@ def quantize_fp8_checkpoint(
             shape.num_frames,
         )
 
+    explicit_sigmas = sampler.explicit_sigmas or _calib.resolve_distilled_sigmas(input_dir)
+    metadata = collect_quantization_metadata(
+        source=model_name_or_path,
+        input_dir=input_dir,
+        recipe={
+            "format": "fp8",
+            "profile": profile,
+            "shape": asdict(shape),
+            "sampler": asdict(sampler),
+            "scheduler_class": f"{type(scheduler).__module__}.{type(scheduler).__qualname__}",
+            "resolved_explicit_sigmas": explicit_sigmas,
+            "num_samples": num_samples,
+            "quant_algo": quant_algo,
+            "seed": seed,
+            "negative_prompt": negative_prompt,
+            "fps": fps,
+            "quantize_language_model": quantize_language_model,
+            "calibration_behavior": calibration_behavior,
+            "mixed_precision_steps": mixed_precision_steps,
+            "i2v_cond_dir": str(i2v_cond_dir) if i2v_cond_dir is not None else None,
+            "i2v_cond_dataset": i2v_cond_dataset,
+        },
+        prompts=prompts,
+    )
+    metadata["environment"]["torch_cuda"] = torch.version.cuda
+
     forward_loop = _calib.make_forward_loop(
         vae=vae,
         scheduler=scheduler,
@@ -198,7 +231,7 @@ def quantize_fp8_checkpoint(
         use_flow_sigmas=sampler.use_flow_sigmas,
         i2v=i2v,
         cond_latents=cond_latents,
-        explicit_sigmas=sampler.explicit_sigmas or _calib.resolve_distilled_sigmas(input_dir),
+        explicit_sigmas=explicit_sigmas,
         use_legacy_scheduler=legacy_behavior,
     )
 
@@ -214,8 +247,14 @@ def quantize_fp8_checkpoint(
     mtq.print_quant_summary(model)
 
     staging_dir = output_dir.parent / f".quantized_transformer_{profile}.tmp"
-    _export.export_quantized_transformer(model, staging_dir, transformer_dir)
+    _export.export_quantized_transformer(
+        model,
+        staging_dir,
+        transformer_dir,
+        mixed_precision_steps=mixed_precision_steps,
+    )
     _export.assemble_output_dir(input_dir, output_dir, staging_dir)
+    write_quantization_metadata(output_dir, metadata)
 
     print(f"[done] {output_dir}")
     return (output_dir, model) if keep_model else output_dir

cosmos_framework/quantization/calibration.py

@@ -47,7 +47,7 @@
     "norm_moe_gen|rotary_emb|action2llm|llm2action|sound2llm|llm2sound"
 )
 _BASE_SKIP_TOKENS_NO_LM = (
-    "time_embedder|vae2llm|llm2vae|embed_tokens|"
+    "visual|lm_head|time_embedder|vae2llm|llm2vae|embed_tokens|"
     "norm_q|norm_k|input_layernorm|post_attention_layernorm|"
     "norm_moe_gen|rotary_emb|action2llm|llm2action|sound2llm|llm2sound"
 )
@@ -137,9 +137,10 @@ def apply_legacy_attention(model, *, text_width: int = 512) -> None:
 
     from cosmos_framework.data.generator.sequence_packing.runtime import (
         from_mode_splits,
-        get_causal_seq_padded,
-        get_full_only_seq_padded,
+        get_causal_seq,
+        get_full_only_seq,
         get_gen_seq,
+        get_num_real_samples,
         get_num_real_tokens,
         get_und_seq,
     )
@@ -150,7 +151,7 @@ def legacy_dispatch(query_pack, key_pack, value_pack, attention_mask, **kwargs):
             or attention_mask.control_stream_token_ranges is not None
             or attention_mask.flex_block_mask is not None
             or kwargs.get("memory_value") is not None
-            or query_pack["sample_offsets"].shape[0] != 2
+            or get_num_real_samples(query_pack) != 1
         ):
             raise ValueError(
                 "legacy calibration behavior supports only single-sample, "
@@ -163,9 +164,9 @@ def heads_first(tensor: torch.Tensor) -> torch.Tensor:
         def flatten_heads(tensor: torch.Tensor) -> torch.Tensor:
             return tensor.squeeze(0).transpose(0, 1).flatten(-2, -1)
 
-        causal_q, _, _ = get_causal_seq_padded(query_pack)
-        causal_k, _, _ = get_causal_seq_padded(key_pack)
-        causal_v, _, _ = get_causal_seq_padded(value_pack)
+        causal_q, _ = get_causal_seq(query_pack)
+        causal_k, _ = get_causal_seq(key_pack)
+        causal_v, _ = get_causal_seq(value_pack)
         causal_output = F.scaled_dot_product_attention(
             heads_first(causal_q),
             heads_first(causal_k),
@@ -184,7 +185,7 @@ def flatten_heads(tensor: torch.Tensor) -> torch.Tensor:
         key_gen = get_gen_seq(generation_key_pack)[:num_gen]
         value_und = get_und_seq(value_pack)[:num_und]
         value_gen = get_gen_seq(value_pack)[:num_gen]
-        full_q, _, _ = get_full_only_seq_padded(query_pack)
+        full_q, _ = get_full_only_seq(query_pack)
         zero_k = key_und.new_zeros((padding_rows, *key_und.shape[1:]))
         zero_v = value_und.new_zeros((padding_rows, *value_und.shape[1:]))
         generation_output = F.scaled_dot_product_attention(