From 5fdf518fbba8ce0f684cabc5372c35d425d0b93d Mon Sep 17 00:00:00 2001 From: lvliang-intel Date: Fri, 10 Jul 2026 11:17:25 +0800 Subject: [PATCH 1/4] Support world model Cosmos3 Signed-off-by: lvliang-intel --- auto_round/context/model.py | 12 +- auto_round/special_model_handler.py | 291 ++++++++++++- auto_round/utils/model.py | 14 + test/test_cpu/models/test_cosmos3.py | 610 +++++++++++++++++++++++++++ 4 files changed, 920 insertions(+), 7 deletions(-) create mode 100644 test/test_cpu/models/test_cosmos3.py diff --git a/auto_round/context/model.py b/auto_round/context/model.py index 285f6f5df..ce88c1eca 100644 --- a/auto_round/context/model.py +++ b/auto_round/context/model.py @@ -145,17 +145,17 @@ def device(self, value) -> None: device_manager.device = value def _load_model(self): - if is_mllm_model(self.model, platform=self.platform): + if is_diffusion_model(self.model): + self.is_diffusion = True + self.pipe, self.model = diffusion_load_model( + self.model, platform=self.platform, device="cpu", model_dtype=self.model_dtype + ) + elif is_mllm_model(self.model, platform=self.platform): self.is_mllm = True if isinstance(self.model, str): self.model, self.processor, self.tokenizer, self.image_processor = mllm_load_model( self.model, platform=self.platform, device="cpu", model_dtype=self.model_dtype ) - elif is_diffusion_model(self.model): - self.is_diffusion = True - self.pipe, self.model = diffusion_load_model( - self.model, platform=self.platform, device="cpu", model_dtype=self.model_dtype - ) elif isinstance(self.model, str): config = self.config try: diff --git a/auto_round/special_model_handler.py b/auto_round/special_model_handler.py index 54b7a69a4..ef754b64a 100644 --- a/auto_round/special_model_handler.py +++ b/auto_round/special_model_handler.py @@ -1272,7 +1272,6 @@ def _nextstep_pipeline_fn(pipe, prompts, guidance_scale=7.5, num_inference_steps _PRE_DEFINED_FIXED_ATTR = {"gemma4_unified": {"has_variable_block_shape": True}} - def get_predefined_fixed_attr(model: torch.nn.Module) -> dict | None: """Return fixed compressor attributes for models that need special caching. @@ -1289,3 +1288,293 @@ def get_predefined_fixed_attr(model: torch.nn.Module) -> dict | None: return None attrs = _PRE_DEFINED_FIXED_ATTR.get(config.model_type) return attrs + +# Cosmos3 world-model handler. +def _set_quantizable_bits(module, bits=4): + """Mark all nn.Linear children as quantizable so AutoRound wraps them.""" + import torch.nn as nn + + for _n, _m in module.named_modules(): + if not isinstance(_m, nn.Linear): + continue + _m.bits = getattr(_m, "bits", bits) + _m.group_size = getattr(_m, "group_size", 128) + _m.sym = getattr(_m, "sym", True) + _m.data_type = getattr(_m, "data_type", "int") + _m.scale_dtype = getattr(_m, "scale_dtype", None) + _m.act_bits = getattr(_m, "act_bits", 16) + _m.act_sym = getattr(_m, "act_sym", None) + _m.act_data_type = getattr(_m, "act_data_type", None) + _m.act_dynamic = getattr(_m, "act_dynamic", True) + + +import contextlib +import threading + +_cosmos3_forward_state = threading.local() + +@contextlib.contextmanager +def _cosmos3_forward_mode(): + """Context flag marking a genuine `Cosmos3OmniPipeline.__call__` forward.""" + _cosmos3_forward_state.active = True + try: + yield + finally: + _cosmos3_forward_state.active = False + + +def _wrap_cosmos3_layer_dual_mode(layer): + """Monkey-patch one `Cosmos3VLTextMoTDecoderLayer` in place.""" + if getattr(layer, "_autoround_dual_mode_patched", False): + return + orig_forward = layer.forward + + def _dual_mode_forward(self, und_seq, gen_seq, rotary_emb): + if getattr(_cosmos3_forward_state, "active", False): + return orig_forward(und_seq, gen_seq, rotary_emb) + + und_len = rotary_emb[0].shape[0] + if und_seq.shape[0] != und_len: + combined = und_seq + und_seq = combined[:und_len] + gen_seq = combined[und_len:] + und_out, gen_out = orig_forward(und_seq, gen_seq, rotary_emb) + return torch.cat([und_out, gen_out], dim=0) + + import types + + layer.forward = types.MethodType(_dual_mode_forward, layer) + layer._autoround_dual_mode_patched = True + + +def _get_cosmos3_multimodal_block(model, quant_vision=False): + """Get block names for Cosmos3.""" + if hasattr(model, "layers"): + return [[f"layers.{i}" for i in range(len(model.layers))]] + return [] + + +def _bypass_cosmos3_safety_checker(): + """Patch Cosmos3 safety checker so calibration does not require cosmos_guardrail.""" + try: + import diffusers.pipelines.cosmos.pipeline_cosmos3_omni as pipeline_cosmos3_omni + except ImportError: + return + + safety_checker = getattr(pipeline_cosmos3_omni, "CosmosSafetyChecker", None) + if safety_checker is None or getattr(safety_checker, "_autoround_patched", False): + return + + def _patched_init(self, *args, **kwargs): + torch.nn.Module.__init__(self) + + safety_checker.__init__ = _patched_init + safety_checker.to = lambda self, *args, **kwargs: self + safety_checker.check_text_safety = staticmethod(lambda prompt: True) + safety_checker.check_video_safety = staticmethod(lambda frames: frames) + safety_checker.dtype = property(lambda self: torch.float32) + safety_checker.device = property(lambda self: torch.device("cpu")) + safety_checker._autoround_patched = True + + +def _align_cosmos3_quant_config_to_vllm_omni(transformer_dir): + """Rewrite the fused ``layers`` block name to vllm-omni's split module tree.""" + import json + import os + + def _rewrite_blocks(blocks): + if isinstance(blocks, str): + blocks = [b.strip() for b in blocks.split(",") if b.strip()] + if not isinstance(blocks, list): + return blocks + out = [] + for b in blocks: + if b == "layers": + out.extend(["language_model.layers", "gen_layers"]) + else: + out.append(b) + return out + + config_path = os.path.join(transformer_dir, "config.json") + if not os.path.isfile(config_path): + return + with open(config_path, "r", encoding="utf-8") as f: + config_data = json.load(f) + qcfg = config_data.get("quantization_config") + if not isinstance(qcfg, dict) or "block_name_to_quantize" not in qcfg: + return + rewritten = _rewrite_blocks(qcfg["block_name_to_quantize"]) + qcfg["block_name_to_quantize"] = rewritten + with open(config_path, "w", encoding="utf-8") as f: + json.dump(config_data, f, indent=2, sort_keys=True) + + quant_config_path = os.path.join(transformer_dir, "quantization_config.json") + if os.path.isfile(quant_config_path): + with open(quant_config_path, "r", encoding="utf-8") as f: + quant_config_data = json.load(f) + if "block_name_to_quantize" in quant_config_data: + quant_config_data["block_name_to_quantize"] = rewritten + with open(quant_config_path, "w", encoding="utf-8") as f: + json.dump(quant_config_data, f, indent=2, sort_keys=True) + + +def load_cosmos3_diffusion(pretrained_model_name_or_path, device_str): + """Load Cosmos3 from a diffusers checkpoint for AutoRound quantization. + + Returns a (pipe, model) pair where `model.layers` are the quantizable + blocks. + """ + import json + import os + import shutil + + from diffusers import Cosmos3OmniPipeline + + _bypass_cosmos3_safety_checker() + + diffusers_pipe = Cosmos3OmniPipeline.from_pretrained( + pretrained_model_name_or_path, torch_dtype=torch.bfloat16 + ) + torch.set_grad_enabled(True) + fused_model = diffusers_pipe.transformer + + for layer in fused_model.layers: + _wrap_cosmos3_layer_dual_mode(layer) + _set_quantizable_bits(fused_model.layers) + + # Read transformer config for model dimensions / metadata only. + config_path = os.path.join(pretrained_model_name_or_path, "transformer", "config.json") + with open(config_path, "r") as f: + cfg = json.load(f) + + import torch.nn as nn + from transformers import PretrainedConfig + + class _Cosmos3VllmOmniConfig(PretrainedConfig): + model_type = "cosmos3_vllm_omni" + + def __init__(self, **kwargs): + kwargs.pop("model_type", None) + super().__init__(**kwargs) + + class Cosmos3VllmOmniQuantModel(nn.Module): + def __init__(self): + super().__init__() + self.config = _Cosmos3VllmOmniConfig(**cfg) + self.config._class_name = "Cosmos3VFMTransformer" + self.config.architectures = ["Cosmos3VFMTransformer"] + self.config.name_or_path = pretrained_model_name_or_path + self.layers = fused_model.layers + + @property + def dtype(self): + try: + return next(self.parameters()).dtype + except StopIteration: + return torch.float32 + + @property + def device(self): + try: + return next(self.parameters()).device + except StopIteration: + return torch.device("cpu") + + def save_pretrained(self, save_directory, max_shard_size="5GB", safe_serialization=True, **kwargs): + os.makedirs(save_directory, exist_ok=True) + self.config.save_pretrained(save_directory) + state_dict = {name: tensor.detach().cpu() for name, tensor in self.state_dict().items()} + if safe_serialization: + from safetensors.torch import save_file + + save_file(state_dict, os.path.join(save_directory, "model.safetensors")) + else: + torch.save(state_dict, os.path.join(save_directory, "pytorch_model.bin")) + + model = Cosmos3VllmOmniQuantModel() + logger.info("Cosmos3: quantizing %d decoder-block layers", len(fused_model.layers)) + + class _Cosmos3Pipe: + _autoround_diffusion_pipe = True + + def __init__(self, transformer, diffusers_pipe, pipe_cfg): + self.transformer = transformer + self.model = transformer + self.diffusers_pipe = diffusers_pipe + self.config = pipe_cfg + self.components = {"transformer": transformer} + self.dtype = torch.bfloat16 + + @property + def device(self): + return next(self.transformer.parameters()).device + + def to(self, *args, **kwargs): + self.diffusers_pipe.to(*args, **kwargs) + for component in vars(self.diffusers_pipe).values(): + keep_fp32 = getattr(component, "_keep_in_fp32_modules", None) + if not isinstance(component, torch.nn.Module) or not keep_fp32: + continue + for sub_name, sub_module in component.named_modules(): + if any(sub_name == name or sub_name.startswith(name + ".") for name in keep_fp32): + sub_module.to(torch.float32) + return self + + def save_config(self, save_directory): + os.makedirs(save_directory, exist_ok=True) + model_index_path = os.path.join(pretrained_model_name_or_path, "model_index.json") + with open(model_index_path, "r", encoding="utf-8") as f: + model_index = json.load(f) + + model_index["_class_name"] = "Cosmos3OmniDiffusersPipeline" + model_index["_name_or_path"] = pretrained_model_name_or_path + for component_name, component_spec in model_index.items(): + if component_name.startswith("_") or component_name == "transformer" or not isinstance(component_spec, list): + continue + source = os.path.join(pretrained_model_name_or_path, component_name) + target = os.path.join(save_directory, component_name) + if os.path.isdir(source): + shutil.copytree(source, target, dirs_exist_ok=True) + with open(os.path.join(save_directory, "model_index.json"), "w", encoding="utf-8") as f: + f.write(json.dumps(model_index, indent=2, sort_keys=True) + "\n") + + _SKIP_TOPLEVEL = { + "model_index.json", + "checkpoint.json", + "model.safetensors.index.json", + } + for entry in os.listdir(pretrained_model_name_or_path): + src = os.path.join(pretrained_model_name_or_path, entry) + if not os.path.isfile(src): + continue + if entry in _SKIP_TOPLEVEL: + continue + if entry.endswith((".json", ".txt", ".model")): + shutil.copy2(src, os.path.join(save_directory, entry)) + + # Align the exported quant metadata with vllm-omni's runtime module tree. + _align_cosmos3_quant_config_to_vllm_omni(os.path.join(save_directory, "transformer")) + + def __call__(self, prompts=None, guidance_scale=6.0, num_inference_steps=35, generator=None, **kwargs): + """Calibration forward: drives the actual Cosmos3OmniPipeline.""" + prompt_list = prompts if isinstance(prompts, list) else [prompts] + with _cosmos3_forward_mode(): + for prompt in prompt_list: + self.diffusers_pipe( + prompt=prompt, + num_frames=1, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=generator, + output_type="latent", + enable_safety_check=False, + ) + + pipe = _Cosmos3Pipe(model, diffusers_pipe, dict(diffusers_pipe.config)) + + return pipe, model + + +# Register Cosmos3 block detection. +SPECIAL_MULTIMODAL_BLOCK["cosmos3_vllm_omni"] = _get_cosmos3_multimodal_block +SPECIAL_SHARED_CACHE_KEYS["Cosmos3VllmOmniQuantModel"] = ("rotary_emb",) diff --git a/auto_round/utils/model.py b/auto_round/utils/model.py index cc2c761e1..8334bf770 100644 --- a/auto_round/utils/model.py +++ b/auto_round/utils/model.py @@ -855,6 +855,20 @@ def diffusion_load_model( pipe, model = load_next_step_diffusion(pretrained_model_name_or_path, device_str) return pipe, pipe.model + # A special case for Cosmos3: model_index.json _class_name may not match + # diffusers, and the safety checker requires cosmos_guardrail which may not be installed. + _model_index = None + if isinstance(pretrained_model_name_or_path, str): + _mi_path = os.path.join(pretrained_model_name_or_path, "model_index.json") + if os.path.isfile(_mi_path): + with open(_mi_path, "r", encoding="utf-8") as f: + _model_index = json.load(f) + _pipe_class = (_model_index or {}).get("_class_name", "") + if _pipe_class in ("Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"): + from auto_round.special_model_handler import load_cosmos3_diffusion + + return load_cosmos3_diffusion(pretrained_model_name_or_path, device_str) + pipelines = LazyImport("diffusers.pipelines") if isinstance(pretrained_model_name_or_path, str): model_index = os.path.join(pretrained_model_name_or_path, "model_index.json") diff --git a/test/test_cpu/models/test_cosmos3.py b/test/test_cpu/models/test_cosmos3.py new file mode 100644 index 000000000..066a6d6d2 --- /dev/null +++ b/test/test_cpu/models/test_cosmos3.py @@ -0,0 +1,610 @@ +# Copyright (c) 2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit and integration tests for Cosmos3 diffusion model support.""" + +import json +import os +import shutil +import tempfile +from unittest.mock import MagicMock, patch + +import pytest +import torch + + +class TestCosmos3ForwardMode: + """Test _cosmos3_forward_mode() thread-local context manager.""" + + def test_forward_mode_sets_and_clears_state(self): + from auto_round.special_model_handler import _cosmos3_forward_mode, _cosmos3_forward_state + + assert not getattr(_cosmos3_forward_state, "active", False) + with _cosmos3_forward_mode(): + assert getattr(_cosmos3_forward_state, "active", False) is True + assert not getattr(_cosmos3_forward_state, "active", False) + + def test_forward_mode_cleanup_on_exception(self): + from auto_round.special_model_handler import _cosmos3_forward_mode, _cosmos3_forward_state + + with pytest.raises(RuntimeError): + with _cosmos3_forward_mode(): + raise RuntimeError("simulated error") + assert not getattr(_cosmos3_forward_state, "active", False) + + def test_forward_mode_nested_context(self): + from auto_round.special_model_handler import _cosmos3_forward_mode, _cosmos3_forward_state + + with _cosmos3_forward_mode(): + assert getattr(_cosmos3_forward_state, "active", False) is True + with _cosmos3_forward_mode(): + assert getattr(_cosmos3_forward_state, "active", False) is True + assert not getattr(_cosmos3_forward_state, "active", False) + + +class TestWrapCosmos3LayerDualMode: + """Test _wrap_cosmos3_layer_dual_mode() layer monkey-patching.""" + + def test_patch_idempotent(self): + from auto_round.special_model_handler import _wrap_cosmos3_layer_dual_mode + + mock_layer = MagicMock() + orig_mock_forward = MagicMock(return_value=("und_out", "gen_out")) + mock_layer.forward = orig_mock_forward + mock_layer._autoround_dual_mode_patched = False + + _wrap_cosmos3_layer_dual_mode(mock_layer) + first_forward = mock_layer.forward + _wrap_cosmos3_layer_dual_mode(mock_layer) + assert mock_layer.forward is first_forward + + def test_dual_mode_splits_combined_sequence(self): + from auto_round.special_model_handler import _wrap_cosmos3_layer_dual_mode + + mock_layer = MagicMock() + mock_und_out = torch.randn(2, 4, 8) + mock_gen_out = torch.randn(3, 4, 8) + orig_mock_forward = MagicMock(return_value=(mock_und_out, mock_gen_out)) + mock_layer.forward = orig_mock_forward + mock_layer._autoround_dual_mode_patched = False + + _wrap_cosmos3_layer_dual_mode(mock_layer) + + und_seq = torch.randn(5, 4, 8) + gen_seq = torch.randn(0, 4, 8) + rotary_emb = (torch.randn(2, 4, 8), torch.randn(2, 4, 8)) + + result = mock_layer.forward(und_seq, gen_seq, rotary_emb) + + orig_mock_forward.assert_called_once() + call_args = orig_mock_forward.call_args[0] + und_arg, gen_arg, rot_arg = call_args + + assert und_arg.shape[0] == 2 + assert gen_arg.shape[0] == 3 + assert result.shape[0] == 5 + + def test_active_mode_bypasses_splitting(self): + from auto_round.special_model_handler import _cosmos3_forward_mode, _wrap_cosmos3_layer_dual_mode + + mock_layer = MagicMock() + mock_und_out = torch.randn(2, 4, 8) + mock_gen_out = torch.randn(3, 4, 8) + orig_mock_forward = MagicMock(return_value=(mock_und_out, mock_gen_out)) + mock_layer.forward = orig_mock_forward + mock_layer._autoround_dual_mode_patched = False + + _wrap_cosmos3_layer_dual_mode(mock_layer) + + und_seq = torch.randn(5, 4, 8) + gen_seq = torch.randn(3, 4, 8) + rotary_emb = (torch.randn(2, 4, 8), torch.randn(2, 4, 8)) + + with _cosmos3_forward_mode(): + mock_layer.forward(und_seq, gen_seq, rotary_emb) + + orig_mock_forward.assert_called_once() + call_args = orig_mock_forward.call_args[0] + assert torch.equal(call_args[0], und_seq) + assert torch.equal(call_args[1], gen_seq) + + def test_mismatched_sequence_length_triggers_split(self): + from auto_round.special_model_handler import _wrap_cosmos3_layer_dual_mode + + mock_layer = MagicMock() + mock_layer.forward = MagicMock(return_value=(torch.zeros(1, 4, 8), torch.zeros(1, 4, 8))) + mock_layer._autoround_dual_mode_patched = False + + orig_mock_forward = mock_layer.forward + _wrap_cosmos3_layer_dual_mode(mock_layer) + + und_seq = torch.randn(4, 4, 8) + gen_seq = torch.zeros(0, 4, 8) + rotary_emb = (torch.randn(2, 4, 8), torch.randn(2, 4, 8)) + + mock_layer.forward(und_seq, gen_seq, rotary_emb) + + call_args = orig_mock_forward.call_args[0] + assert call_args[0].shape[0] == 2 + assert call_args[1].shape[0] == 2 + + +class TestGetCosmos3MultimodalBlock: + """Test _get_cosmos3_multimodal_block() block name discovery.""" + + def test_returns_layers_list_with_quant_vision_false(self): + from auto_round.special_model_handler import _get_cosmos3_multimodal_block + + mock_model = MagicMock() + mock_model.layers = [MagicMock() for _ in range(8)] + + result = _get_cosmos3_multimodal_block(mock_model, quant_vision=False) + + assert len(result) == 1 + assert len(result[0]) == 8 + assert result[0][0] == "layers.0" + assert result[0][7] == "layers.7" + + def test_returns_layers_list_with_quant_vision_true(self): + from auto_round.special_model_handler import _get_cosmos3_multimodal_block + + mock_model = MagicMock() + mock_model.layers = [MagicMock() for _ in range(4)] + + result = _get_cosmos3_multimodal_block(mock_model, quant_vision=True) + + # quant_vision is ignored for Cosmos3 — only text layers are quantized + assert len(result) == 1 + assert result[0][0] == "layers.0" + + def test_returns_empty_when_no_layers(self): + from auto_round.special_model_handler import _get_cosmos3_multimodal_block + + mock_model = MagicMock(spec=[]) + del mock_model.layers + + result = _get_cosmos3_multimodal_block(mock_model) + assert result == [] + + +class TestBypassCosmos3SafetyChecker: + """Test _bypass_cosmos3_safety_checker() import safety and patching.""" + + @staticmethod + def _make_real_safety_checker_class(): + return type( + "CosmosSafetyChecker", + (), + { + "_autoround_patched": False, + "__init__": lambda self: None, + "to": lambda self, *a, **kw: self, + }, + ) + + def test_noop_when_diffusers_not_installed(self): + from auto_round.special_model_handler import _bypass_cosmos3_safety_checker + + with patch("builtins.__import__", side_effect=ImportError("No module named 'diffusers'")): + _bypass_cosmos3_safety_checker() + + def test_noop_when_safety_checker_class_not_found(self): + from auto_round.special_model_handler import _bypass_cosmos3_safety_checker + + mock_module = MagicMock() + mock_module.CosmosSafetyChecker = None + with patch.dict("sys.modules", {"diffusers.pipelines.cosmos.pipeline_cosmos3_omni": mock_module}): + _bypass_cosmos3_safety_checker() + + def test_patches_safety_checker_methods(self): + from auto_round.special_model_handler import _bypass_cosmos3_safety_checker + + safety_checker_cls = self._make_real_safety_checker_class() + mock_module = MagicMock() + mock_module.CosmosSafetyChecker = safety_checker_cls + + with patch.dict("sys.modules", {"diffusers.pipelines.cosmos.pipeline_cosmos3_omni": mock_module}): + _bypass_cosmos3_safety_checker() + + assert safety_checker_cls._autoround_patched is True + assert callable(safety_checker_cls.__init__) + assert callable(safety_checker_cls.check_text_safety) + assert callable(safety_checker_cls.check_video_safety) + assert safety_checker_cls.check_text_safety("any prompt") is True + + def test_idempotent_patch(self): + from auto_round.special_model_handler import _bypass_cosmos3_safety_checker + + safety_checker_cls = self._make_real_safety_checker_class() + mock_module = MagicMock() + mock_module.CosmosSafetyChecker = safety_checker_cls + + with patch.dict("sys.modules", {"diffusers.pipelines.cosmos.pipeline_cosmos3_omni": mock_module}): + _bypass_cosmos3_safety_checker() + first_init = safety_checker_cls.__init__ + assert safety_checker_cls._autoround_patched is True + _bypass_cosmos3_safety_checker() + # On second call, the function returns early (idempotent) + assert safety_checker_cls.__init__ is first_init + + +class TestAlignCosmos3QuantConfigToVllmOmni: + """Test _align_cosmos3_quant_config_to_vllm_omni() config rewriting.""" + + def _write_temp_transformer_dir(self, config_data, quant_config_data=None): + tmpdir = tempfile.mkdtemp() + os.makedirs(os.path.join(tmpdir, "transformer")) + + with open(os.path.join(tmpdir, "transformer", "config.json"), "w") as f: + json.dump(config_data, f) + + if quant_config_data is not None: + with open(os.path.join(tmpdir, "transformer", "quantization_config.json"), "w") as f: + json.dump(quant_config_data, f) + + return tmpdir + + def test_rewrites_layers_block_in_config_json(self): + from auto_round.special_model_handler import _align_cosmos3_quant_config_to_vllm_omni + + config = { + "quantization_config": { + "bits": 4, + "block_name_to_quantize": "layers", + } + } + tmpdir = self._write_temp_transformer_dir(config) + + try: + _align_cosmos3_quant_config_to_vllm_omni(os.path.join(tmpdir, "transformer")) + + with open(os.path.join(tmpdir, "transformer", "config.json")) as f: + result = json.load(f) + + assert result["quantization_config"]["block_name_to_quantize"] == [ + "language_model.layers", + "gen_layers", + ] + finally: + shutil.rmtree(tmpdir) + + def test_rewrites_layers_block_in_quantization_config_json(self): + from auto_round.special_model_handler import _align_cosmos3_quant_config_to_vllm_omni + + # Top-level config.json must also have block_name_to_quantize for the + # rewrite cascade to reach the standalone quantization_config.json. + config = { + "quantization_config": { + "bits": 4, + "block_name_to_quantize": "layers", + } + } + quant_config = { + "bits": 4, + "block_name_to_quantize": "layers", + } + tmpdir = self._write_temp_transformer_dir(config, quant_config) + + try: + _align_cosmos3_quant_config_to_vllm_omni(os.path.join(tmpdir, "transformer")) + + with open(os.path.join(tmpdir, "transformer", "quantization_config.json")) as f: + result = json.load(f) + + assert result["block_name_to_quantize"] == ["language_model.layers", "gen_layers"] + finally: + shutil.rmtree(tmpdir) + + def test_leaves_non_layers_blocks_unchanged(self): + from auto_round.special_model_handler import _align_cosmos3_quant_config_to_vllm_omni + + config = { + "quantization_config": { + "bits": 4, + "block_name_to_quantize": ["transformer_blocks", "single_transformer_blocks"], + } + } + tmpdir = self._write_temp_transformer_dir(config) + + try: + _align_cosmos3_quant_config_to_vllm_omni(os.path.join(tmpdir, "transformer")) + + with open(os.path.join(tmpdir, "transformer", "config.json")) as f: + result = json.load(f) + + assert result["quantization_config"]["block_name_to_quantize"] == [ + "transformer_blocks", + "single_transformer_blocks", + ] + finally: + shutil.rmtree(tmpdir) + + def test_handles_string_blocks_format(self): + from auto_round.special_model_handler import _align_cosmos3_quant_config_to_vllm_omni + + config = { + "quantization_config": { + "bits": 4, + "block_name_to_quantize": "layers , single_transformer_blocks", + } + } + tmpdir = self._write_temp_transformer_dir(config) + + try: + _align_cosmos3_quant_config_to_vllm_omni(os.path.join(tmpdir, "transformer")) + + with open(os.path.join(tmpdir, "transformer", "config.json")) as f: + result = json.load(f) + + assert result["quantization_config"]["block_name_to_quantize"] == [ + "language_model.layers", + "gen_layers", + "single_transformer_blocks", + ] + finally: + shutil.rmtree(tmpdir) + + def test_noop_when_config_json_missing(self): + from auto_round.special_model_handler import _align_cosmos3_quant_config_to_vllm_omni + + tmpdir = tempfile.mkdtemp() + try: + _align_cosmos3_quant_config_to_vllm_omni(os.path.join(tmpdir, "transformer")) + finally: + shutil.rmtree(tmpdir) + + def test_noop_when_no_quantization_config(self): + from auto_round.special_model_handler import _align_cosmos3_quant_config_to_vllm_omni + + config = {"model_type": "cosmos3_vllm_omni"} + tmpdir = self._write_temp_transformer_dir(config) + + try: + _align_cosmos3_quant_config_to_vllm_omni(os.path.join(tmpdir, "transformer")) + + with open(os.path.join(tmpdir, "transformer", "config.json")) as f: + result = json.load(f) + + assert "quantization_config" not in result + finally: + shutil.rmtree(tmpdir) + + +class TestCosmos3Registration: + """Test Cosmos3 registration in special model handler registries.""" + + def test_cosmos3_vllm_omni_registered_in_special_multimodal_block(self): + from auto_round.special_model_handler import SPECIAL_MULTIMODAL_BLOCK + + assert "cosmos3_vllm_omni" in SPECIAL_MULTIMODAL_BLOCK + assert callable(SPECIAL_MULTIMODAL_BLOCK["cosmos3_vllm_omni"]) + + def test_cosmos3_quant_model_registered_in_special_shared_cache_keys(self): + from auto_round.special_model_handler import SPECIAL_SHARED_CACHE_KEYS + + assert "Cosmos3VllmOmniQuantModel" in SPECIAL_SHARED_CACHE_KEYS + assert SPECIAL_SHARED_CACHE_KEYS["Cosmos3VllmOmniQuantModel"] == ("rotary_emb",) + + +class TestCosmos3LoadAndQuantize: + """End-to-end tests for ``load_cosmos3_diffusion`` with mocked diffusers pipe. + + The diffusers ``Cosmos3OmniPipeline.from_pretrained`` call is mocked so the + tests exercise the real AutoRound wiring (config construction, layer + patching, save_config flow) without requiring a multi-GB Cosmos3 checkpoint. + """ + + @staticmethod + def _build_fake_diffusers_pipe(n_layers=2): + """Build a minimal stand-in for Cosmos3OmniPipeline with nn.Module layers.""" + import torch.nn as nn + + class _FakeLayer(nn.Module): + def forward(self, *args, **kwargs): + return (args[0] if args else None, args[1] if len(args) > 1 else None) + + class _FakeTransformer(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([_FakeLayer() for _ in range(n_layers)]) + + class _FakePipe: + def __init__(self): + self.transformer = _FakeTransformer() + self.config = {"_class_name": "Cosmos3OmniDiffusersPipeline"} + + def to(self, *args, **kwargs): + return self + + return _FakePipe() + + @pytest.fixture + def tiny_cosmos3_model_path(self): + """Write only the files consumed by ``load_cosmos3_diffusion`` for read-side checks.""" + tmpdir = tempfile.mkdtemp() + + transformer_config = { + "_class_name": "Cosmos3VFMTransformer", + "architectures": ["Cosmos3VFMTransformer"], + "model_type": "cosmos3_vllm_omni", + "num_hidden_layers": 2, + "hidden_size": 64, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "intermediate_size": 128, + "vocab_size": 32000, + "max_position_embeddings": 128, + "rms_norm_eps": 1e-6, + "rope_theta": 10000.0, + "attention_head_dim": 32, + "latent_patch_size": 2, + "max_action_dim": 64, + "num_embodiment_domains": 1, + "guidance_embed": True, + } + + os.makedirs(os.path.join(tmpdir, "transformer")) + with open(os.path.join(tmpdir, "transformer", "config.json"), "w") as f: + json.dump(transformer_config, f) + + model_index = { + "_class_name": "Cosmos3OmniDiffusersPipeline", + "_name_or_path": tmpdir, + "scheduler": ["diffusers", "DDPMScheduler"], + "transformer": ["diffusers", "Cosmos3VFMTransformer"], + } + with open(os.path.join(tmpdir, "model_index.json"), "w") as f: + json.dump(model_index, f) + + yield tmpdir + shutil.rmtree(tmpdir, ignore_errors=True) + + def test_load_cosmos3_diffusion_returns_pipe_and_model(self, tiny_cosmos3_model_path): + from auto_round.special_model_handler import load_cosmos3_diffusion + + fake_pipe = self._build_fake_diffusers_pipe(n_layers=2) + fake_cls = MagicMock() + fake_cls.from_pretrained.return_value = fake_pipe + with patch.dict("sys.modules", {"diffusers": MagicMock(Cosmos3OmniPipeline=fake_cls)}): + pipe, model = load_cosmos3_diffusion(tiny_cosmos3_model_path, "cpu") + + # Pipe shape + assert hasattr(pipe, "transformer") + assert hasattr(pipe, "diffusers_pipe") + assert hasattr(pipe, "save_config") + assert pipe._autoround_diffusion_pipe is True + assert pipe.transformer is model + assert pipe.dtype == torch.bfloat16 + + # Model shape + assert hasattr(model, "config") + assert hasattr(model, "layers") + assert model.config.model_type == "cosmos3_vllm_omni" + assert model.config._class_name == "Cosmos3VFMTransformer" + assert "Cosmos3VFMTransformer" in model.config.architectures + assert len(model.layers) == 2 + + def test_load_cosmos3_patches_each_layer(self, tiny_cosmos3_model_path): + """Each real layer must be wrapped with the dual-mode forward patch.""" + from auto_round.special_model_handler import load_cosmos3_diffusion + + fake_pipe = self._build_fake_diffusers_pipe(n_layers=3) + fake_cls = MagicMock() + fake_cls.from_pretrained.return_value = fake_pipe + with patch.dict("sys.modules", {"diffusers": MagicMock(Cosmos3OmniPipeline=fake_cls)}): + _, model = load_cosmos3_diffusion(tiny_cosmos3_model_path, "cpu") + + for layer in model.layers: + assert getattr(layer, "_autoround_dual_mode_patched", False) is True + + def test_load_cosmos3_pipe_save_config_rewrites_block_names(self, tiny_cosmos3_model_path, tmp_path): + """Calling pipe.save_config() must invoke the vllm-omni config realignment.""" + from auto_round.special_model_handler import load_cosmos3_diffusion + + fake_pipe = self._build_fake_diffusers_pipe(n_layers=2) + fake_cls = MagicMock() + fake_cls.from_pretrained.return_value = fake_pipe + + target_dir = os.path.join(str(tmp_path), "saved") + os.makedirs(os.path.join(target_dir, "transformer"), exist_ok=True) + cfg = { + "_class_name": "Cosmos3VFMTransformer", + "model_type": "cosmos3_vllm_omni", + "quantization_config": {"bits": 4, "block_name_to_quantize": "layers"}, + } + with open(os.path.join(target_dir, "transformer", "config.json"), "w") as f: + json.dump(cfg, f) + with open(os.path.join(target_dir, "transformer", "quantization_config.json"), "w") as f: + json.dump({"bits": 4, "block_name_to_quantize": "layers"}, f) + + with patch.dict("sys.modules", {"diffusers": MagicMock(Cosmos3OmniPipeline=fake_cls)}): + pipe, _ = load_cosmos3_diffusion(tiny_cosmos3_model_path, "cpu") + pipe.save_config(target_dir) + + with open(os.path.join(target_dir, "transformer", "config.json")) as f: + saved_cfg = json.load(f) + assert saved_cfg["quantization_config"]["block_name_to_quantize"] == [ + "language_model.layers", + "gen_layers", + ] + with open(os.path.join(target_dir, "transformer", "quantization_config.json")) as f: + saved_qcfg = json.load(f) + assert saved_qcfg["block_name_to_quantize"] == ["language_model.layers", "gen_layers"] + + +class TestCosmos3DiffusionLoadModel: + """Test the diffusion_load_model dispatch for Cosmos3.""" + + def _write_model_index(self, class_name): + tmpdir = tempfile.mkdtemp() + model_index = {"_class_name": class_name} + with open(os.path.join(tmpdir, "model_index.json"), "w") as f: + json.dump(model_index, f) + return tmpdir + + def test_cosmos3_detected_by_class_name_in_model_index(self): + from auto_round.utils.model import diffusion_load_model + + tmpdir = self._write_model_index("Cosmos3OmniDiffusersPipeline") + try: + with patch( + "auto_round.special_model_handler.load_cosmos3_diffusion" + ) as mock_load: + mock_load.return_value = (MagicMock(), MagicMock()) + result = diffusion_load_model(tmpdir) + + mock_load.assert_called_once() + assert result is mock_load.return_value + finally: + shutil.rmtree(tmpdir) + + def test_cosmos3_detected_by_alternate_class_name(self): + from auto_round.utils.model import diffusion_load_model + + tmpdir = self._write_model_index("Cosmos3OmniPipeline") + try: + with patch( + "auto_round.special_model_handler.load_cosmos3_diffusion" + ) as mock_load: + mock_load.return_value = (MagicMock(), MagicMock()) + result = diffusion_load_model(tmpdir) + + mock_load.assert_called_once() + assert result is mock_load.return_value + finally: + shutil.rmtree(tmpdir) + + def test_non_cosmos3_pipeline_skips_cosmos3_loader(self): + """Routing logic: non-Cosmos class names must not invoke load_cosmos3_diffusion. + + We mock out the actual diffusers pipeline loading to keep this a pure + unit test on the dispatch logic in diffusion_load_model(). + """ + from auto_round.utils.model import diffusion_load_model + + tmpdir = self._write_model_index("DDPMScheduler") + try: + with patch( + "auto_round.special_model_handler.load_cosmos3_diffusion" + ) as mock_cosmos: + with patch( + "auto_round.utils.common.LazyImport", return_value=MagicMock() + ): + try: + diffusion_load_model(tmpdir) + except Exception: + pass + mock_cosmos.assert_not_called() + finally: + shutil.rmtree(tmpdir) From 404b9066a3c124fdace639dc166157eec86a484d Mon Sep 17 00:00:00 2001 From: lvliang-intel Date: Fri, 10 Jul 2026 11:26:50 +0800 Subject: [PATCH 2/4] update readme Signed-off-by: lvliang-intel --- auto_round/compressors/diffusion/README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/auto_round/compressors/diffusion/README.md b/auto_round/compressors/diffusion/README.md index ae78a0567..d44cc467a 100644 --- a/auto_round/compressors/diffusion/README.md +++ b/auto_round/compressors/diffusion/README.md @@ -71,6 +71,8 @@ For diffusion models, currently we validate quantization on the following models | Wan-AI/Wan2.2-I2V-A14B-Diffusers | COCO2014 | - | | Wan-AI/Wan2.2-TI2V-5B-Diffusers | COCO2014 | - | | Wan-AI/Wan2.2-T2V-A14B-Diffusers | COCO2014 | - | +| nvidia/Cosmos3-Nano | COCO2014 | - | +| nvidia/Cosmos3-Super | COCO2014 | - |
Calibration Dataset From a937d09a27789e28faae43bbb4feab699cf8ac85 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 03:31:28 +0000 Subject: [PATCH 3/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- auto_round/special_model_handler.py | 13 +++++++++---- test/test_cpu/models/test_cosmos3.py | 16 ++++------------ 2 files changed, 13 insertions(+), 16 deletions(-) diff --git a/auto_round/special_model_handler.py b/auto_round/special_model_handler.py index ef754b64a..a89a1d71e 100644 --- a/auto_round/special_model_handler.py +++ b/auto_round/special_model_handler.py @@ -1272,6 +1272,7 @@ def _nextstep_pipeline_fn(pipe, prompts, guidance_scale=7.5, num_inference_steps _PRE_DEFINED_FIXED_ATTR = {"gemma4_unified": {"has_variable_block_shape": True}} + def get_predefined_fixed_attr(model: torch.nn.Module) -> dict | None: """Return fixed compressor attributes for models that need special caching. @@ -1289,6 +1290,7 @@ def get_predefined_fixed_attr(model: torch.nn.Module) -> dict | None: attrs = _PRE_DEFINED_FIXED_ATTR.get(config.model_type) return attrs + # Cosmos3 world-model handler. def _set_quantizable_bits(module, bits=4): """Mark all nn.Linear children as quantizable so AutoRound wraps them.""" @@ -1313,6 +1315,7 @@ def _set_quantizable_bits(module, bits=4): _cosmos3_forward_state = threading.local() + @contextlib.contextmanager def _cosmos3_forward_mode(): """Context flag marking a genuine `Cosmos3OmniPipeline.__call__` forward.""" @@ -1432,9 +1435,7 @@ def load_cosmos3_diffusion(pretrained_model_name_or_path, device_str): _bypass_cosmos3_safety_checker() - diffusers_pipe = Cosmos3OmniPipeline.from_pretrained( - pretrained_model_name_or_path, torch_dtype=torch.bfloat16 - ) + diffusers_pipe = Cosmos3OmniPipeline.from_pretrained(pretrained_model_name_or_path, torch_dtype=torch.bfloat16) torch.set_grad_enabled(True) fused_model = diffusers_pipe.transformer @@ -1529,7 +1530,11 @@ def save_config(self, save_directory): model_index["_class_name"] = "Cosmos3OmniDiffusersPipeline" model_index["_name_or_path"] = pretrained_model_name_or_path for component_name, component_spec in model_index.items(): - if component_name.startswith("_") or component_name == "transformer" or not isinstance(component_spec, list): + if ( + component_name.startswith("_") + or component_name == "transformer" + or not isinstance(component_spec, list) + ): continue source = os.path.join(pretrained_model_name_or_path, component_name) target = os.path.join(save_directory, component_name) diff --git a/test/test_cpu/models/test_cosmos3.py b/test/test_cpu/models/test_cosmos3.py index 066a6d6d2..8f0713cbf 100644 --- a/test/test_cpu/models/test_cosmos3.py +++ b/test/test_cpu/models/test_cosmos3.py @@ -558,9 +558,7 @@ def test_cosmos3_detected_by_class_name_in_model_index(self): tmpdir = self._write_model_index("Cosmos3OmniDiffusersPipeline") try: - with patch( - "auto_round.special_model_handler.load_cosmos3_diffusion" - ) as mock_load: + with patch("auto_round.special_model_handler.load_cosmos3_diffusion") as mock_load: mock_load.return_value = (MagicMock(), MagicMock()) result = diffusion_load_model(tmpdir) @@ -574,9 +572,7 @@ def test_cosmos3_detected_by_alternate_class_name(self): tmpdir = self._write_model_index("Cosmos3OmniPipeline") try: - with patch( - "auto_round.special_model_handler.load_cosmos3_diffusion" - ) as mock_load: + with patch("auto_round.special_model_handler.load_cosmos3_diffusion") as mock_load: mock_load.return_value = (MagicMock(), MagicMock()) result = diffusion_load_model(tmpdir) @@ -595,12 +591,8 @@ def test_non_cosmos3_pipeline_skips_cosmos3_loader(self): tmpdir = self._write_model_index("DDPMScheduler") try: - with patch( - "auto_round.special_model_handler.load_cosmos3_diffusion" - ) as mock_cosmos: - with patch( - "auto_round.utils.common.LazyImport", return_value=MagicMock() - ): + with patch("auto_round.special_model_handler.load_cosmos3_diffusion") as mock_cosmos: + with patch("auto_round.utils.common.LazyImport", return_value=MagicMock()): try: diffusion_load_model(tmpdir) except Exception: From d1afcc4defb7bf9e5a98bb2f5c738f62c10c7af4 Mon Sep 17 00:00:00 2001 From: lvliang-intel Date: Fri, 10 Jul 2026 13:56:19 +0800 Subject: [PATCH 4/4] fix pylint Signed-off-by: lvliang-intel --- auto_round/special_model_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/auto_round/special_model_handler.py b/auto_round/special_model_handler.py index ef754b64a..a6fcb6821 100644 --- a/auto_round/special_model_handler.py +++ b/auto_round/special_model_handler.py @@ -1428,7 +1428,7 @@ def load_cosmos3_diffusion(pretrained_model_name_or_path, device_str): import os import shutil - from diffusers import Cosmos3OmniPipeline + from diffusers import Cosmos3OmniPipeline # pylint: disable=import-error,no-name-in-module _bypass_cosmos3_safety_checker()