diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index dd01a287..976fcd19 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -9,7 +9,7 @@ # The six GPU jobs then run (one at a time on the single runner): # * training-smoke — Nano SFT pipeline (convert -> train 5 -> export -> t2i) # * generator-training-regression — vision_sft_nano loss vs goldens (4-GPU subset) -# * generator-inference-smoke — Nano + Edge multi-modality inference (t2vs/t2v + policy + forward_dynamics) +# * generator-inference-smoke — Nano + Edge multi-modality inference (t2vs/t2v + policy + forward_dynamics) + Nano ModelOpt FP8 # * distilled-inference-smoke — Super 4-step T2I + I2V inference (4-GPU subset) # * reasoner-inference-smoke — Nano reasoner golden + Edge reasoner smoke + Qwen conversion coverage (4-GPU) # * reasoner-training-regression — llava_ov loss vs goldens (4-GPU subset) @@ -18,7 +18,9 @@ # * a self-hosted runner labelled [self-hosted, gpu, h200] with 8 GPUs, # NVIDIA drivers, and `uv` on PATH; # * an `HF_TOKEN` repository secret (gated dataset/model downloads, incl. the -# streamed LLaVA-OneVision-Data dataset). +# streamed LLaVA-OneVision-Data dataset). The Nano ModelOpt FP8 cases also +# need it to read nvidia/Cosmos3-Experimental, which is not public; they skip +# (they do not fail) when it cannot. # # Inputs/checkpoints download to examples/ + the HF cache and are reused across # runs (the h100 goldens are reused on H200 — see _detect_arch). @@ -111,7 +113,9 @@ jobs: generator-inference-smoke: needs: pre-commit runs-on: [self-hosted, gpu, h200] - timeout-minutes: 60 + # 90 (not 60) since the Nano step gained the two ModelOpt FP8 cases: a 20 GB + # checkpoint on a cold cache plus two more 8-GPU runs. + timeout-minutes: 90 env: HF_TOKEN: ${{ secrets.HF_TOKEN }} HF_HUB_DISABLE_XET: "1" @@ -124,9 +128,14 @@ jobs: run: uv sync --all-extras --group=cu128-train # One inference call over t2vs (+sound), action policy, and forward_dynamics; checks each output. + # Plus the ModelOpt static-FP8 Nano checkpoint in both parallelism layouts + # (FSDP-sharded and replicated). The FP8 checkpoint lives in the + # access-controlled nvidia/Cosmos3-Experimental repo, so those two cases SKIP + # unless HF_TOKEN can read it — check the log for "no access to + # nvidia/Cosmos3-Experimental" before trusting a green run to have covered FP8. # MAX_GPUS defaults to 8. -s streams the live process log. # Reuse the same input-asset cache dir as the unittest job. - - name: Nano inference smoke (t2vs + action policy + forward_dynamics, 8 GPU) + - name: Nano inference smoke (t2vs + action policy + forward_dynamics + FP8, 8 GPU) run: | export LD_LIBRARY_PATH= export COSMOS_DOWNLOAD_CACHE_DIR="$RUNNER_WORKSPACE/cosmos_input_cache" 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/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 5b7c74ce..7fd5eee8 100644 --- a/cosmos_framework/inference/model.py +++ b/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,9 +431,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 +457,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 +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) diff --git a/cosmos_framework/inference/model_test.py b/cosmos_framework/inference/model_test.py index ca2c8c9c..02a7a897 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,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()) 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/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. diff --git a/cosmos_framework/utils/generator/quantization.py b/cosmos_framework/utils/generator/quantization.py index 94b11948..56adbbed 100644 --- a/cosmos_framework/utils/generator/quantization.py +++ b/cosmos_framework/utils/generator/quantization.py @@ -10,17 +10,24 @@ to replicated inference (``data_parallel_shard_degree == 1``); it cannot be applied to an FSDP-sharded model whose params are ``DTensor`` shards. +ModelOpt FP8 checkpoint loading installs the exported E4M3 weights and static +scales into TorchAO tensor subclasses without calibration or re-quantization. + This is an inference-only path: the ``quantize_`` PTQ configs have no backward support. Module selection is delegated to the filter built by :func:`_get_filter_fn`. """ import gc +import json import re from collections.abc import Callable +from pathlib import Path import torch from torch import nn +from torch.distributed.tensor import DTensor, distribute_tensor +from torch.nn import functional as F from cosmos_framework.utils import log from cosmos_framework.configs.base.defaults.quantization import QuantizationConfig @@ -33,6 +40,605 @@ # quantization is actually requested. +class _ModelOptFloat8Linear(nn.Linear): + """Work around TorchAO 0.16 static-FP8 limitations for linear inputs.""" + + # Compute dtype the FP8 weight dequantizes to. Recorded when the module is + # swapped in on meta, where the plain E4M3 placeholder no longer carries it. + _modelopt_high_precision_dtype: torch.dtype | None = None + # True once real checkpoint data has been installed. + _modelopt_weight_loaded: bool = False + + 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) + + +_torchao_fsdp_support_installed = False + + +def _rewrap_float8(source, qdata: torch.Tensor): + """Rebuild a static-FP8 tensor around new ``qdata``, carrying the scales over. + + Valid only because ModelOpt exports PerTensor granularity: one scalar scale + covers the whole weight, so reshaping or slicing the data never redistributes + it. ``block_size`` is therefore always "one block spanning everything". + """ + return source.__class__( + qdata, + source.scale, + act_quant_scale=source.act_quant_scale, + output_act_quant_scale=source.output_act_quant_scale, + block_size=list(qdata.shape), + mm_config=source.mm_config, + act_quant_kwargs=source.act_quant_kwargs, + kernel_preference=source.kernel_preference, + dtype=source.dtype, + output_act_quant_kwargs=source.output_act_quant_kwargs, + ) + + +def install_torchao_float8_fsdp_support() -> None: + """Teach TorchAO's static-FP8 tensor the ops and hooks FSDP2 needs. + + TorchAO 0.16 ships this tensor for single-device inference: it implements no + ``fsdp_pre_all_gather`` / ``fsdp_post_all_gather``, and its shape-changing ops + derive a scale layout from ``block_size`` assuming a fixed rank, so FSDP2's + reshape/split/allocate path fails on a 2D linear weight. Sharding it is + nonetheless well defined for a PerTensor scale — the scalar is identical on + every shard, so only the FP8 bytes ever need to move. + + Registered once, globally, since the dispatch table is keyed by class. + """ + global _torchao_fsdp_support_installed + if _torchao_fsdp_support_installed: + return + from torchao.prototype.quantization.float8_static_quant.prototype_float8_tensor import ( + PrototypeFloat8Tensor, + ) + from torchao.utils import return_and_correct_aliasing + + aten = torch.ops.aten + implements = PrototypeFloat8Tensor.implements + + @implements([aten.view.default, aten._unsafe_view.default, aten.reshape.default]) + def _(func, types, args, kwargs): + self, size = args[0], args[1] + return return_and_correct_aliasing( + func, args, kwargs, _rewrap_float8(self, self.qdata.reshape(*size)) + ) + + @implements(aten.split.Tensor) + def _(func, types, args, kwargs): + self, size = args[0], args[1] + dim = args[2] if len(args) > 2 else 0 + return [_rewrap_float8(self, part) for part in self.qdata.split(size, dim)] + + @implements(aten.slice.Tensor) + def _(func, types, args, kwargs): + self = args[0] + return return_and_correct_aliasing( + func, args, kwargs, _rewrap_float8(self, aten.slice.Tensor(self.qdata, *args[1:])) + ) + + @implements(aten.as_strided.default) + def _(func, types, args, kwargs): + self, size, stride = args[0], args[1], args[2] + offset = args[3] if len(args) > 3 else 0 + return return_and_correct_aliasing( + func, args, kwargs, _rewrap_float8(self, self.qdata.as_strided(size, stride, offset)) + ) + + @implements(aten.new_zeros.default) + def _(func, types, args, kwargs): + self, size = args[0], args[1] + return _rewrap_float8(self, self.qdata.new_zeros(size)) + + @implements([aten.empty_like.default, aten.zeros_like.default]) + def _(func, types, args, kwargs): + self = args[0] + device = kwargs.get("device") + qdata = func(self.qdata, **{key: value for key, value in kwargs.items() if key != "dtype"}) + # `to_empty` moves storage without data, so the scales are allocated empty + # on the target device as well: `.to(device)` would try to copy out of a + # meta tensor. Real scale values arrive with the checkpoint weights. + scale = torch.empty_like(self.scale, device=device) if device is not None else self.scale + act_quant_scale = ( + torch.empty_like(self.act_quant_scale, device=device) + if device is not None + else self.act_quant_scale + ) + return self.__class__( + qdata, + scale, + act_quant_scale=act_quant_scale, + output_act_quant_scale=self.output_act_quant_scale, + block_size=list(qdata.shape), + mm_config=self.mm_config, + act_quant_kwargs=self.act_quant_kwargs, + kernel_preference=self.kernel_preference, + dtype=self.dtype, + output_act_quant_kwargs=self.output_act_quant_kwargs, + ) + + @implements([aten.detach.default, aten.clone.default]) + def _(func, types, args, kwargs): + self = args[0] + return return_and_correct_aliasing(func, args, kwargs, _rewrap_float8(self, func(self.qdata))) + + @implements(aten._to_copy.default) + def _(func, types, args, kwargs): + # Device moves must go through, but a dtype request must not: casting the + # E4M3 payload to the compute dtype would silently undo the quantization. + self = args[0] + forwarded = {key: value for key, value in kwargs.items() if key != "dtype"} + return return_and_correct_aliasing( + func, args, kwargs, _rewrap_float8(self, func(self.qdata, **forwarded)) + ) + + @implements(aten.copy_.default) + def _(func, types, args, kwargs): + destination, source = args[0], args[1] + if isinstance(source, PrototypeFloat8Tensor): + destination.qdata.copy_(source.qdata) + destination.scale.copy_(source.scale) + destination.act_quant_scale.copy_(source.act_quant_scale) + else: + destination.qdata.copy_(source) + return destination + + def fsdp_pre_all_gather(self, mesh): + # Only the FP8 bytes travel. The static scales are already identical on + # every rank, so keeping them out of the collective saves the traffic. + return (self.qdata,), (self.scale, self.act_quant_scale, self.dtype) + + def fsdp_post_all_gather(self, all_gather_outputs, metadata, param_dtype, *, out=None): + (qdata,) = all_gather_outputs + scale, act_quant_scale, dtype = metadata + if out is not None: + return None + gathered = self.__class__( + qdata, + scale, + act_quant_scale=act_quant_scale, + output_act_quant_scale=self.output_act_quant_scale, + block_size=list(qdata.shape), + mm_config=self.mm_config, + act_quant_kwargs=self.act_quant_kwargs, + kernel_preference=self.kernel_preference, + dtype=param_dtype or dtype, + output_act_quant_kwargs=self.output_act_quant_kwargs, + ) + return gathered, (qdata,) + + PrototypeFloat8Tensor.fsdp_pre_all_gather = fsdp_pre_all_gather + PrototypeFloat8Tensor.fsdp_post_all_gather = fsdp_post_all_gather + _torchao_fsdp_support_installed = True + log.info("Installed TorchAO static-FP8 FSDP support (8 aten ops + all-gather hooks)") + + +def is_modelopt_fp8_checkpoint(checkpoint_path: str | Path) -> bool: + """Return whether a local checkpoint declares ModelOpt FP8 quantization.""" + quant_config_path = Path(checkpoint_path) / "hf_quant_config.json" + if not quant_config_path.is_file(): + return False + try: + quant_config = json.loads(quant_config_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"Unable to read a valid checkpoint quantization config: {quant_config_path}") from error + if ( + not isinstance(quant_config, dict) + or quant_config.get("quant_method") != "modelopt" + or quant_config.get("quant_algo") != "FP8" + ): + raise ValueError(f"Unsupported checkpoint quantization configuration: {quant_config_path}") + return True + + +def collect_modelopt_fp8_params(model: nn.Module) -> set[nn.Parameter]: + """Return the FP8 weights of every ModelOpt linear currently in ``model``. + + Collected by module type immediately before use rather than captured earlier, + because parallelization rebuilds parts of the network (notably the + ``language_model`` subtree) and hands back fresh parameter objects, so any + set gathered ahead of that no longer matches by identity. + """ + return {module.weight for module in model.modules() if isinstance(module, _ModelOptFloat8Linear)} + + +def plan_modelopt_fp8_targets( + checkpoint_path: str | Path, + key_mapper: Callable[[str, str], str | None], + *, + weight_map: dict[str, str], +) -> list[str]: + """Return the network-relative FQNs a ModelOpt FP8 checkpoint quantizes. + + Derived from the checkpoint index alone — no model instance is required — so + the result can be threaded through the model config and consumed by + ``build_net`` before the network exists. Only the E4M3 tensors are reported; + weights ModelOpt intentionally left in high precision (``proj_in``, + ``lm_head``, the vision tower, ...) are skipped. + + Args: + checkpoint_path: Local root of the ModelOpt diffusers checkpoint. + key_mapper: Maps a diffusers source key and shard path to a model state key. + weight_map: Merged checkpoint weight map. + + Returns: + Sorted target module FQNs, relative to the VFM network. + """ + from safetensors import safe_open + + checkpoint_path = Path(checkpoint_path) + weights_by_shard: dict[str, list[str]] = {} + for source_key, shard_path in weight_map.items(): + if source_key.endswith(".weight"): + weights_by_shard.setdefault(shard_path, []).append(source_key) + + target_fqns: set[str] = set() + for shard_path, source_weight_keys in sorted(weights_by_shard.items()): + with safe_open(checkpoint_path / shard_path, framework="pt", device="cpu") as shard: + shard_keys = set(shard.keys()) + for source_weight_key in sorted(source_weight_keys): + if source_weight_key not in shard_keys: + continue + if shard.get_slice(source_weight_key).get_dtype() != "F8_E4M3": + continue + target_weight_key = key_mapper(source_weight_key, shard_path) + if target_weight_key is None or not target_weight_key.endswith(".weight"): + continue + target_fqns.add(target_weight_key.removesuffix(".weight")) + return sorted(target_fqns) + + +def swap_modelopt_fp8_linears_on_meta(model: nn.Module, target_fqns: list[str]) -> list[str]: + """Swap the given linears to FP8 modules while the network is still on meta. + + This must run *before* the network is parallelized and materialized. FSDP2 + wraps the whole network into a single parameter group, so a linear replaced + after ``fully_shard`` would leave the group holding a stale parameter; and + ``to_empty`` on the original bf16 shapes is what sets peak memory, which for + a Super-class model exceeds a single 80 GB device. Swapping here means FSDP + shards — and ``to_empty`` materializes — one byte per element instead of two. + + The weight is created as a real (meta) TorchAO static-FP8 tensor rather than a + plain E4M3 placeholder, because FSDP decides at wrap time whether a parameter + carries the ``fsdp_pre_all_gather`` extension. A plain tensor swapped to a + subclass afterwards is already registered as an ordinary parameter and fails + on the first all-gather. Actual weight bytes and scales arrive later, in + :func:`apply_modelopt_fp8_checkpoint_inplace`. + + Args: + model: Meta-device VFM network. + target_fqns: Module FQNs to convert, as planned from the checkpoint. + + Returns: + Sorted FQNs actually swapped; targets absent from this model variant are + skipped, matching the loader's tolerance for task-specialized exports. + """ + install_torchao_float8_fsdp_support() + from torchao.float8.inference import Float8MMConfig + from torchao.prototype.quantization.float8_static_quant.prototype_float8_tensor import ( + PrototypeFloat8Tensor, + ) + from torchao.quantization import PerTensor + from torchao.quantization.quantize_.workflows import QuantizeTensorToFloat8Kwargs + + mm_config = Float8MMConfig(use_fast_accum=True) + activation_quant_kwargs = QuantizeTensorToFloat8Kwargs( + float8_dtype=torch.float8_e4m3fn, + granularity=PerTensor(), + mm_config=mm_config, + ) + swapped_fqns: list[str] = [] + for target_module_fqn in sorted(target_fqns): + try: + module = model.get_submodule(target_module_fqn) + except AttributeError: + continue + if not isinstance(module, nn.Linear): + raise KeyError(f"ModelOpt FP8 target {target_module_fqn!r} is not an nn.Linear module") + + replacement = _ModelOptFloat8Linear( + module.in_features, + module.out_features, + bias=module.bias is not None, + device="meta", + dtype=module.weight.dtype, + ) + replacement._modelopt_high_precision_dtype = module.weight.dtype + quantized_shape = tuple(module.weight.shape) + replacement.weight = nn.Parameter( + PrototypeFloat8Tensor( + torch.empty(quantized_shape, dtype=torch.float8_e4m3fn, device="meta"), + torch.empty(1, 1, dtype=torch.float32, device="meta"), + act_quant_scale=torch.empty(1, 1, dtype=torch.float32, device="meta"), + block_size=list(quantized_shape), + mm_config=mm_config, + act_quant_kwargs=activation_quant_kwargs, + dtype=module.weight.dtype, + ), + requires_grad=False, + ) + if module.bias is not None: + replacement.bias = module.bias + replacement.train(module.training) + + parent_fqn, _, child_name = target_module_fqn.rpartition(".") + parent = model.get_submodule(parent_fqn) if parent_fqn else model + setattr(parent, child_name, replacement) + swapped_fqns.append(target_module_fqn) + + log.info(f"Swapped {len(swapped_fqns)} linears to meta-device ModelOpt FP8 modules") + return swapped_fqns + + +def apply_modelopt_fp8_checkpoint_inplace( + model: nn.Module, + checkpoint_path: str | Path, + key_mapper: Callable[[str, str], str | None], + *, + weight_map: dict[str, str], +) -> list[str]: + """Install ModelOpt static-FP8 checkpoint tensors as TorchAO weights. + + ModelOpt's HF export stores the already-quantized E4M3 bytes under ``weight`` + and its per-tensor dequantization scales under ``weight_scale`` and + ``input_scale``. This adapter streams those tensors from the checkpoint and + constructs TorchAO ``PrototypeFloat8Tensor`` weights directly, avoiding any + dequantization, recalibration, or re-quantization. + + Args: + model: Unsharded Cosmos3 network whose linears will be converted in place. + checkpoint_path: Local root of the ModelOpt diffusers checkpoint. + key_mapper: Maps a diffusers source key and shard path to a model state key. + weight_map: Merged checkpoint weight map. + + Returns: + Sorted fully qualified names of converted linear modules. + """ + checkpoint_path = Path(checkpoint_path) + if not is_modelopt_fp8_checkpoint(checkpoint_path): + raise ValueError(f"Not a ModelOpt FP8/FP8 checkpoint: {checkpoint_path}") + + from safetensors import safe_open + from torchao.float8.inference import Float8MMConfig + from torchao.prototype.quantization.float8_static_quant.prototype_float8_tensor import ( + PrototypeFloat8Tensor, + ) + from torchao.quantization import PerTensor + from torchao.quantization.quantize_.workflows import QuantizeTensorToFloat8Kwargs + + source_weights_by_shard: dict[str, list[str]] = {} + for source_key, shard_path in weight_map.items(): + if source_key.endswith(".weight"): + source_weights_by_shard.setdefault(shard_path, []).append(source_key) + + mm_config = Float8MMConfig(use_fast_accum=True) + activation_quant_kwargs = QuantizeTensorToFloat8Kwargs( + float8_dtype=torch.float8_e4m3fn, + granularity=PerTensor(), + mm_config=mm_config, + ) + # Target module FQNs successfully replaced with static-FP8 linear modules. + converted_fqns: list[str] = [] + # Source FP8 weights mapped to linear modules present in this model variant. + applicable_source_weight_keys: set[str] = set() + # Source FP8 weights actually installed; compared with the applicable set + # after conversion to detect incomplete plans. + converted_source_weight_keys: set[str] = set() + # Weight shard -> validated (source weight key, target module FQN) conversions. + conversion_plan_by_shard: dict[str, list[tuple[str, str]]] = {} + # Scale shard -> scale keys to preload; scales may be separate from their weights. + source_scales_by_shard: dict[str, list[str]] = {} + # Target module FQNs reserved by the plan, used to reject duplicate mappings. + planned_target_fqns: set[str] = set() + + # Identify and validate every applicable FP8 conversion before replacing any + # modules. The checkpoint may contain intentionally ignored weights or + # optional heads absent from this model variant; those must be filtered + # before requiring their scales. + for shard_path, source_weight_keys in sorted(source_weights_by_shard.items()): + full_shard_path = checkpoint_path / shard_path + with safe_open(full_shard_path, framework="pt", device="cpu") as shard: + shard_keys = set(shard.keys()) + missing_weight_keys = set(source_weight_keys) - shard_keys + if missing_weight_keys: + raise KeyError( + f"ModelOpt weight tensor(s) indexed in {full_shard_path} are missing: {sorted(missing_weight_keys)}" + ) + for source_weight_key in sorted(source_weight_keys): + tensor_slice = shard.get_slice(source_weight_key) + if tensor_slice.get_dtype() != "F8_E4M3": + continue + + target_weight_key = key_mapper(source_weight_key, shard_path) + if target_weight_key is None: + continue + if not target_weight_key.endswith(".weight"): + raise KeyError(f"ModelOpt FP8 key {source_weight_key!r} has no Cosmos3 weight mapping") + target_module_fqn = target_weight_key.removesuffix(".weight") + try: + module = model.get_submodule(target_module_fqn) + except AttributeError: + continue + if not isinstance(module, nn.Linear): + raise KeyError( + f"ModelOpt FP8 key {source_weight_key!r} mapped to {target_module_fqn!r}, " + "which is not an nn.Linear module" + ) + if type(module.weight) is not nn.Parameter and not isinstance(module, _ModelOptFloat8Linear): + # Distinguish the two ways a weight stops being a plain + # parameter: FSDP replaced it with a shard, or something + # already quantized it. Reporting the former as the latter + # sends debugging in entirely the wrong direction. A module + # already swapped by `swap_modelopt_fp8_linears_on_meta` is + # neither — its weight may legitimately be a DTensor shard. + if type(module.weight).__name__ == "DTensor": + raise ValueError( + f"ModelOpt FP8 target is an FSDP shard, not a plain parameter: {target_module_fqn}. " + "The FP8 swap must run before the network is parallelized." + ) + raise ValueError(f"ModelOpt FP8 target is already quantized: {target_module_fqn}") + checkpoint_shape = tuple(tensor_slice.get_shape()) + if checkpoint_shape != tuple(module.weight.shape): + raise ValueError( + f"Shape mismatch for {target_module_fqn}: checkpoint has {checkpoint_shape}, " + f"model expects {tuple(module.weight.shape)}" + ) + if target_module_fqn in planned_target_fqns: + raise ValueError(f"Multiple ModelOpt FP8 weights map to target: {target_module_fqn}") + + source_module_key = source_weight_key.removesuffix(".weight") + source_scale_keys = ( + f"{source_module_key}.weight_scale", + f"{source_module_key}.input_scale", + ) + missing_scale_keys = set(source_scale_keys) - weight_map.keys() + if missing_scale_keys: + raise KeyError( + f"ModelOpt FP8 tensor {source_weight_key!r} is missing scale tensor(s) " + f"from the checkpoint index: {sorted(missing_scale_keys)}" + ) + for source_scale_key in source_scale_keys: + source_scales_by_shard.setdefault(weight_map[source_scale_key], []).append(source_scale_key) + + applicable_source_weight_keys.add(source_weight_key) + planned_target_fqns.add(target_module_fqn) + conversion_plan_by_shard.setdefault(shard_path, []).append((source_weight_key, target_module_fqn)) + + if not applicable_source_weight_keys: + raise ValueError(f"No ModelOpt FP8 linear weights were found in {checkpoint_path}") + + # ModelOpt's consolidated HF export can place the scalar scales in a + # different shard from the corresponding FP8 weight. Read only those tiny + # tensors up front, following the root index rather than assuming locality. + exported_scales: dict[str, torch.Tensor] = {} + for shard_path, source_scale_keys in sorted(source_scales_by_shard.items()): + full_shard_path = checkpoint_path / shard_path + with safe_open(full_shard_path, framework="pt", device="cpu") as shard: + shard_keys = set(shard.keys()) + missing_scale_keys = set(source_scale_keys) - shard_keys + if missing_scale_keys: + raise KeyError( + f"ModelOpt scale tensor(s) indexed in {full_shard_path} are missing: {sorted(missing_scale_keys)}" + ) + for source_scale_key in source_scale_keys: + scale = shard.get_tensor(source_scale_key) + if scale.numel() != 1: + raise ValueError(f"ModelOpt scale tensor {source_scale_key!r} must contain exactly one value") + exported_scales[source_scale_key] = scale + + for shard_path, conversions in sorted(conversion_plan_by_shard.items()): + full_shard_path = checkpoint_path / shard_path + with safe_open(full_shard_path, framework="pt", device="cpu") as shard: + for source_weight_key, target_module_fqn in conversions: + source_module_key = source_weight_key.removesuffix(".weight") + source_weight_scale_key = f"{source_module_key}.weight_scale" + source_input_scale_key = f"{source_module_key}.input_scale" + module = model.get_submodule(target_module_fqn) + device = module.weight.device + + if isinstance(module, _ModelOptFloat8Linear): + # Already swapped on meta before parallelization; the weight + # is an E4M3 placeholder, so the compute dtype comes from the + # dtype recorded at swap time rather than from the parameter. + replacement = module + high_precision_dtype = module._modelopt_high_precision_dtype + if high_precision_dtype is None: + raise ValueError(f"ModelOpt FP8 module has no recorded compute dtype: {target_module_fqn}") + else: + high_precision_dtype = module.weight.dtype + replacement = _ModelOptFloat8Linear( + module.in_features, + module.out_features, + bias=False, + device="meta", + dtype=high_precision_dtype, + ) + replacement.bias = module.bias + replacement.train(module.training) + + parent_fqn, _, child_name = target_module_fqn.rpartition(".") + parent = model.get_submodule(parent_fqn) if parent_fqn else model + setattr(parent, child_name, replacement) + del module + + quantized_data = shard.get_tensor(source_weight_key) + weight_scale = exported_scales[source_weight_scale_key].reshape(1, 1) + input_scale = exported_scales[source_input_scale_key].reshape(1, 1) + + existing_param = replacement.weight + if isinstance(existing_param, DTensor): + # FSDP already sharded the placeholder, so only this rank's + # slice of the checkpoint tensor is written. The scales are + # per-tensor and therefore identical on every rank. + local_weight = existing_param.to_local() + quantized_data = distribute_tensor( + quantized_data.to(device=device), + existing_param.device_mesh, + existing_param.placements, + ).to_local() + elif isinstance(existing_param.data, PrototypeFloat8Tensor): + local_weight = existing_param.data + else: + local_weight = None + + if local_weight is not None: + # Fill the parameter that FSDP is already tracking; rebinding + # `replacement.weight` here would detach it from the parameter + # group and strand the all-gather on a stale tensor. + if tuple(local_weight.qdata.shape) != tuple(quantized_data.shape): + raise ValueError( + f"Shard mismatch for {target_module_fqn}: model shard " + f"{tuple(local_weight.qdata.shape)}, checkpoint slice {tuple(quantized_data.shape)}" + ) + local_weight.qdata.copy_(quantized_data) + local_weight.scale.copy_(weight_scale) + local_weight.act_quant_scale.copy_(input_scale) + else: + # Unswapped legacy path: the module still holds a plain bf16 + # weight, so build the quantized tensor and bind it directly. + replacement.weight = nn.Parameter( + PrototypeFloat8Tensor( + quantized_data.to(device=device), + weight_scale.to(device=device), + act_quant_scale=input_scale.to(device=device), + block_size=list(quantized_data.shape), + mm_config=mm_config, + act_quant_kwargs=activation_quant_kwargs, + dtype=high_precision_dtype, + ), + requires_grad=False, + ) + replacement._modelopt_weight_loaded = True + converted_fqns.append(target_module_fqn) + converted_source_weight_keys.add(source_weight_key) + + missing_conversions = applicable_source_weight_keys - converted_source_weight_keys + if missing_conversions: + raise ValueError(f"ModelOpt FP8 weights were not converted: {sorted(missing_conversions)}") + converted_fqns.sort() + sharded_count = sum( + 1 for fqn in converted_fqns if isinstance(model.get_submodule(fqn).weight.data, DTensor) + ) + log.info( + f"Loaded {len(converted_fqns)} calibrated ModelOpt FP8 weights into TorchAO " + f"({sharded_count} sharded / {len(converted_fqns) - sharded_count} replicated)" + ) + log.debug(f"ModelOpt FP8 matched_fqns={converted_fqns}") + return converted_fqns + + def _get_filter_fn(quantization_config: QuantizationConfig) -> Callable[[nn.Module, str], bool]: """Build a module-selection predicate from the quantization config. 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) diff --git a/tests/nano_inference_smoke_test.py b/tests/nano_inference_smoke_test.py index 098d8510..c63a4db1 100644 --- a/tests/nano_inference_smoke_test.py +++ b/tests/nano_inference_smoke_test.py @@ -30,6 +30,13 @@ action, the t2vs sample an audio track, and the transfer sample exercises the control-guidance branch. +3. Two ``text2video`` calls against the ModelOpt static-FP8 Nano checkpoint, one + per parallelism layout (FSDP-sharded and replicated) -> a non-degenerate + ``vision.mp4`` each. Covers the FP8 checkpoint path end to end: detection, + the meta-device linear swap, the TorchAO weight install, and — in the sharded + layout — the FP8 all-gather. Skipped when the checkpoint is not reachable + (see ``_download_fp8_checkpoint``). + Smoke-level only (output validity, not numeric goldens). The checkpoint + its tokenizers download from the HF Hub on first run and are reused afterward. @@ -44,6 +51,7 @@ import json import os +import re import shutil import socket import subprocess @@ -149,6 +157,60 @@ "emphasize_control_in_prompt": False, } +# ModelOpt static-FP8 Cosmos3-Nano checkpoint. It is not published under its own +# repository yet, so it cannot be a ``--checkpoint-path`` registry name (see +# ``_CHECKPOINTS`` in ``cosmos_framework/inference/args.py``): it lives in a +# subdirectory of the access-controlled nvidia/Cosmos3-Experimental repo, pinned to +# the revision the FP8 loader was validated against. The test downloads that one +# subdirectory and passes the resulting local path to the CLI. Once the checkpoint +# is released under its own name, register it and drop ``_download_fp8_checkpoint``. +_FP8_REPOSITORY = "nvidia/Cosmos3-Experimental" +_FP8_REVISION = "f0cdb8ea37360e8510e2c0caf84c0f9f3e8751c8" +_FP8_SUBDIRECTORY = "cosmos3-nano-fp8-14072026" + +# Emitted by ``swap_modelopt_fp8_linears_on_meta`` / ``install_torchao_float8_fsdp_support`` +# (``cosmos_framework/utils/generator/quantization.py``). The swap count is parsed rather +# than string-matched: a checkpoint whose FP8 targets failed to resolve would swap zero +# linears, run the whole model in bf16, and otherwise produce a perfectly valid video. +_FP8_SWAP_LOG = re.compile(r"Swapped (\d+) linears to meta-device ModelOpt FP8 modules") +_FP8_FSDP_LOG = "Installed TorchAO static-FP8 FSDP support" + +# Downscaled generation overrides shared by both FP8 layouts (480p / 29 frames / +# 10 steps, one chunk). The FP8 path under test is per-tensor weight quantization, +# which is independent of resolution and step count, so a short clip exercises it +# just as well as the checkpoint's 720p/189-frame defaults at a fraction of the +# runtime — the same trade-off ``_TRANSFER_SPEC`` above makes. +_FP8_GENERATION_ARGS = ( + "--resolution=480", + "--aspect-ratio=16,9", + "--fps=30", + "--num-frames=29", + "--num-steps=10", + "--seed=0", +) + +# One entry per parallelism layout. ``sharded`` is the layout that matters most +# here: FSDP2 all-gathers the FP8 weights through the TorchAO tensor-subclass +# hooks, which is the path that does not exist upstream, and it is the only +# layout a Super-class FP8 model fits in. ``replicated`` guards the +# single-device-weights path that worked before those hooks were added. +_FP8_LAYOUTS = { + "sharded": ( + "--parallelism-preset=throughput", + "--dp-shard-size=8", + "--dp-replicate-size=1", + "--cp-size=1", + "--cfgp-size=1", + ), + "replicated": ( + "--parallelism-preset=latency", + "--dp-shard-size=1", + "--dp-replicate-size=1", + "--cp-size=8", + "--cfgp-size=1", + ), +} + # Audio sanity thresholds for the muxed sound track. _RMS_SILENCE_FLOOR = 1e-4 # below this the track is effectively silence _PEAK_SANITY_CEIL = 1.5 # decoded float audio should sit within ~[-1, 1] @@ -279,6 +341,40 @@ def _assert_valid_action(content: dict, where: str) -> None: assert np.all(np.isfinite(arr)), f"action output has NaN/Inf ({where})" +def _download_fp8_checkpoint() -> Path: + """Download the pinned ModelOpt FP8 Nano checkpoint and return its local root. + + Skips the test — rather than failing it — when the repository is unreachable + for a credentials reason (no ``HF_TOKEN``, or a token without access to + nvidia/Cosmos3-Experimental), so a fork PR without the runner secret does not + go red. Any other failure (a deleted revision, a broken download) still fails + loudly: a silently-skipping FP8 job would look green while testing nothing. + """ + from huggingface_hub import snapshot_download + from huggingface_hub.errors import GatedRepoError, HfHubHTTPError, RepositoryNotFoundError + + try: + repo_root = snapshot_download( + repo_id=_FP8_REPOSITORY, + revision=_FP8_REVISION, + allow_patterns=[f"{_FP8_SUBDIRECTORY}/*"], + ) + except (GatedRepoError, RepositoryNotFoundError) as error: + pytest.skip(f"no access to {_FP8_REPOSITORY} (needs an HF_TOKEN with read access): {error!r}") + except HfHubHTTPError as error: + status_code = getattr(error.response, "status_code", None) + if status_code in (401, 403): + pytest.skip(f"no access to {_FP8_REPOSITORY} (HTTP {status_code}): {error!r}") + raise + + checkpoint_path = Path(repo_root) / _FP8_SUBDIRECTORY + assert (checkpoint_path / "hf_quant_config.json").is_file(), ( + f"{checkpoint_path} is not a ModelOpt FP8 checkpoint (no hf_quant_config.json); " + f"revision {_FP8_REVISION} may have changed" + ) + return checkpoint_path + + @pytest.fixture(scope="module", autouse=True) def _require_8_gpus() -> None: """Skip the module unless we can launch an 8-GPU run here.""" @@ -470,3 +566,58 @@ def test_nano_inference_multi_control_transfer(tmp_path: Path) -> None: video = so.parent / "vision.mp4" assert video.is_file(), f"multi-control run produced no vision.mp4 ({so})" _assert_video_has_content(video) + + @pytest.mark.level(2) + @pytest.mark.gpus(8) + @pytest.mark.parametrize("layout", sorted(_FP8_LAYOUTS)) + def test_nano_fp8_inference(tmp_path: Path, layout: str) -> None: + """text2video from the ModelOpt static-FP8 Nano checkpoint, once per layout. + + The FP8 checkpoint ships already-quantized E4M3 weights plus static + per-tensor scales, so this run covers a path the bf16 cases above never + touch: ``is_modelopt_fp8_checkpoint`` detection, the meta-device swap of the + target linears to TorchAO FP8 modules (before FSDP wrap, so peak memory + follows the FP8 shapes), the deferred weight install, and the FP8 forward. + + ``sharded`` additionally covers the FSDP2 path — the TorchAO static-FP8 + tensor upstream implements neither the all-gather hooks nor the shape ops + FSDP2 needs, so without the local support shim the run dies on the first + all-gather rather than producing a degraded video. Completing the run *is* + the assertion there; ``_assert_video_has_content`` then catches the + numerically-broken-but-still-running case (wrong scales -> collapsed clip). + """ + checkpoint_path = _download_fp8_checkpoint() + out_dir = tmp_path / f"out_fp8_{layout}" + cmd = [ + "torchrun", + "--nproc_per_node=8", + f"--master_port={_free_port()}", + "-m", + "cosmos_framework.scripts.inference", + *_FP8_LAYOUTS[layout], + "-i", + "inputs/omni/t2v.json", + "-o", + str(out_dir), + "--checkpoint-path", + str(checkpoint_path), + *_FP8_GENERATION_ARGS, + ] + log = _run(cmd, tmp_path / f"inference_fp8_{layout}.log") + + # The checkpoint was recognized as ModelOpt FP8 and its linears really were + # swapped. Without the count check a checkpoint whose targets failed to + # resolve would run entirely in bf16 and still pass every output assertion. + swap_match = _FP8_SWAP_LOG.search(log) + assert swap_match is not None, f"no ModelOpt FP8 linear swap in the {layout} run; FP8 path never engaged" + assert int(swap_match.group(1)) > 0, f"ModelOpt FP8 swap matched 0 linears in the {layout} run" + assert _FP8_FSDP_LOG in log, f"TorchAO static-FP8 FSDP support was not installed in the {layout} run" + + results = sorted(out_dir.rglob("sample_outputs.json")) + assert len(results) == 1, f"expected 1 FP8 sample_outputs.json, found {[str(p) for p in results]}" + so = results[0] + args = json.loads(so.read_text()).get("args", {}) + assert args.get("model_mode") == "text2video", f"expected a text2video sample, got {args.get('model_mode')}" + video = so.parent / "vision.mp4" + assert video.is_file(), f"FP8 {layout} run produced no vision.mp4 ({so})" + _assert_video_has_content(video)