[Model] FP8 Support For Boogu-Image (#6925)

Signed-off-by: Dmaner <2663515256@qq.com>

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -17,6 +17,7 @@ from huggingface_hub import snapshot_download
 from safetensors.torch import save_file
 from vllm.config.load import LoadConfig
 
+import vllm_omni.diffusion.model_loader.diffusers_loader as loader_module
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.lora.manager import LoRABackend
@@ -30,6 +31,7 @@ from vllm_omni.diffusion.model_loader.host_weights import source_identity as sou
 from vllm_omni.diffusion.models.helios import HeliosPipeline
 from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
 from vllm_omni.diffusion.registry import initialize_model
+from vllm_omni.quantization.component_config import ComponentQuantizationConfig
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
@@ -165,6 +167,43 @@ def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoade
     return loader
 
 
+def test_serialized_torchao_component_uses_pytorch_iterator(mocker):
+    torchao_config = SimpleNamespace(
+        get_name=lambda: "torchao",
+        is_checkpoint_torchao_serialized=True,
+    )
+    loader = _make_loader_with_weights([])
+    loader.quant_config = ComponentQuantizationConfig({"transformer": torchao_config})
+    loader.load_config.pt_load_map_location = "cpu"
+    files = ["/weights/model-10.bin", "/weights/model-2.bin"]
+    mocker.patch.object(loader, "_prepare_weights", return_value=("/weights", files, False))
+    pt_iterator = mocker.patch.object(
+        loader_module,
+        "pt_weights_iterator",
+        return_value=iter([("block.weight", torch.ones(1))]),
+    )
+    safetensors_iterator = mocker.patch.object(loader_module, "safetensors_weights_iterator")
+    multithread_iterator = mocker.patch.object(loader_module, "multi_thread_safetensors_weights_iterator")
+    source = DiffusersPipelineLoader.ComponentSource(
+        model_or_path="/weights",
+        subfolder=None,
+        revision=None,
+        prefix="transformer.",
+    )
+
+    weights = list(loader._get_weights_iterator(source))
+
+    pt_iterator.assert_called_once_with(
+        ["/weights/model-2.bin", "/weights/model-10.bin"],
+        loader.load_config.use_tqdm_on_load,
+        "cpu",
+    )
+    safetensors_iterator.assert_not_called()
+    multithread_iterator.assert_not_called()
+    assert weights[0][0] == "transformer.block.weight"
+    assert torch.equal(weights[0][1], torch.ones(1))
+
+
 def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
     tmp_path: Path,
     monkeypatch,
@@ -434,7 +473,7 @@ def test_prepare_weights_honors_component_index_and_explicit_override(tmp_path,
     assert folder == str(transformer)
     assert [str(transformer / filename) for filename in indexed_files] == files
     assert use_safetensors
-    assert hub_api.hf_hub_download.call_count == 2
+    assert hub_api.hf_hub_download.call_count == len(loader_mod.SAFETENSORS_INDEX_FILES)
     hub_api.hf_hub_download.assert_any_call(
         repo_id="org/model",
         filename="transformer/diffusion_pytorch_model.safetensors.index.json",
@@ -493,6 +532,45 @@ def test_prepare_local_weights_honors_component_index(tmp_path):
     assert use_safetensors
 
 
+def test_prepare_local_bin_index_is_authoritative_and_safetensors_remains_preferred(tmp_path):
+    transformer = tmp_path / "transformer"
+    transformer.mkdir()
+    indexed_files = [f"diffusion_pytorch_model-{index:05d}-of-00002.bin" for index in (1, 2)]
+    stale_file = "diffusion_pytorch_model-00001-of-00008.bin"
+    (transformer / "diffusion_pytorch_model.bin.index.json").write_text(
+        json.dumps({"weight_map": {f"weight.{index}": filename for index, filename in enumerate(indexed_files)}})
+    )
+    for filename in (*indexed_files, stale_file):
+        (transformer / filename).touch()
+
+    _, files, use_safetensors = _make_loader_with_weights([])._prepare_weights(
+        tmp_path,
+        subfolder="transformer",
+        revision=None,
+        fall_back_to_pt=True,
+        allow_patterns_overrides=None,
+    )
+
+    assert files == [str(transformer / filename) for filename in indexed_files]
+    assert not use_safetensors
+
+    safetensors_file = "diffusion_pytorch_model.safetensors"
+    (transformer / "diffusion_pytorch_model.safetensors.index.json").write_text(
+        json.dumps({"weight_map": {"weight": safetensors_file}})
+    )
+    (transformer / safetensors_file).touch()
+    _, files, use_safetensors = _make_loader_with_weights([])._prepare_weights(
+        tmp_path,
+        subfolder="transformer",
+        revision=None,
+        fall_back_to_pt=True,
+        allow_patterns_overrides=None,
+    )
+
+    assert files == [str(transformer / safetensors_file)]
+    assert use_safetensors
+
+
 def test_prepare_weights_rejects_polluted_offline_cache_without_index(tmp_path, mocker):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
 
@@ -519,7 +597,9 @@ def test_prepare_weights_rejects_polluted_offline_cache_without_index(tmp_path,
             allow_patterns_overrides=None,
         )
 
-    assert hub_api.hf_hub_download.call_count == 2
+    assert hub_api.hf_hub_download.call_count == len(loader_mod.SAFETENSORS_INDEX_FILES) + len(
+        loader_mod.PT_INDEX_FILES
+    )
     assert all(call.kwargs["local_files_only"] for call in hub_api.hf_hub_download.call_args_list)
 
 

tests/diffusion/models/boogu_image/test_boogu_image_transformer.py

@@ -330,6 +330,74 @@ def test_transformer_instantiates():
     assert model.x_embedder.out_features == HIDDEN_SIZE
 
 
+@pytest.fixture
+def recording_quant_config():
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+    class RecordingQuantConfig(QuantizationConfig):
+        def __init__(self):
+            super().__init__()
+            self.prefixes: list[str] = []
+
+        @classmethod
+        def get_name(cls):
+            return "recording"
+
+        @classmethod
+        def get_supported_act_dtypes(cls):
+            return [torch.float32]
+
+        @classmethod
+        def get_min_capability(cls):
+            return 0
+
+        @classmethod
+        def get_config_filenames(cls):
+            return []
+
+        @classmethod
+        def from_config(cls, config):
+            return cls()
+
+        def get_quant_method(self, layer, prefix):
+            self.prefixes.append(prefix)
+            return UnquantizedLinearMethod()
+
+    return RecordingQuantConfig()
+
+
+def test_transformer_propagates_quant_config_and_prefix(recording_quant_config):
+    from vllm.model_executor.layers.linear import LinearBase
+
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageTransformer2DModel,
+    )
+
+    model = BooguImageTransformer2DModel(
+        od_config=_tiny_od_config(),
+        quant_config=recording_quant_config,
+        prefix="transformer",
+    )
+
+    configured_prefixes = set()
+    for name, module in model.named_modules():
+        if isinstance(module, LinearBase):
+            qualified_name = f"transformer.{name}"
+            assert module.prefix == qualified_name
+            assert module.quant_config is recording_quant_config
+            configured_prefixes.add(qualified_name)
+
+    assert set(recording_quant_config.prefixes) == configured_prefixes
+    assert {name for name, module in model.named_modules() if isinstance(module, torch.nn.Linear)} == {
+        "x_embedder",
+        "ref_image_patch_embedder",
+        "time_caption_embed.timestep_embedder.linear_1",
+        "time_caption_embed.timestep_embedder.linear_2",
+        "time_caption_embed.caption_embedder.1",
+    }
+
+
 def test_transformer_preprocesses_multiple_instruction_feature_layers():
     from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
         BooguImageTransformer2DModel,

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -83,6 +83,7 @@ def mock_dependencies(mocker, monkeypatch):
         "processor": mock_processor,
         "vae": mock_vae,
         "scheduler": mock_scheduler,
+        "transformer_cls": mock_transformer_cls,
         "transformer": mock_transformer_instance,
     }
 
@@ -187,6 +188,28 @@ def test_constructor_weights_sources(boogu_pipeline):
     assert source.fall_back_to_pt is True
 
 
+def test_constructor_routes_only_transformer_config(mock_dependencies, mocker):
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+    from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+
+    transformer_config = mocker.MagicMock(spec=QuantizationConfig)
+    encoder_config = mocker.MagicMock(spec=QuantizationConfig)
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.bfloat16,
+        quantization_config=ComponentQuantizationConfig(
+            {"transformer": transformer_config, "mllm": encoder_config, "vae": None}
+        ),
+    )
+    BooguImagePipeline(od_config=od_config)
+    kwargs = mock_dependencies["transformer_cls"].call_args.kwargs
+    assert kwargs["quant_config"] is transformer_config
+    assert kwargs["prefix"] == "transformer"
+
+
 @pytest.mark.parametrize(
     ("parallel_config", "cache_backend", "message"),
     [

tests/diffusion/quantization/test_torchao_config.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.quantization import build_quant_config
+from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize(
+    "transformer_spec",
+    [
+        {
+            "method": "torchao",
+            "quant_type": {
+                "default": {
+                    "_type": "Float8WeightOnlyConfig",
+                    "_version": 2,
+                    "_data": {
+                        "weight_dtype": {
+                            "_type": "torch.dtype",
+                            "_data": "float8_e4m3fn",
+                        },
+                        "set_inductor_config": False,
+                    },
+                }
+            },
+        },
+        {"method": "torchao_float8_weight_only"},
+    ],
+    ids=["serialized-json", "checkpoint-shorthand"],
+)
+def test_build_quant_config_torchao_checkpoint(transformer_spec):
+    quantization = pytest.importorskip("torchao.quantization")
+    result = build_quant_config({"transformer": transformer_spec})
+
+    assert isinstance(result, ComponentQuantizationConfig)
+    transformer = result.resolve("transformer")
+    assert transformer.get_name() == "torchao"
+    assert transformer.is_checkpoint_torchao_serialized is True
+    assert isinstance(transformer.torchao_config, quantization.Float8WeightOnlyConfig)
+    assert transformer.torchao_config.version == 2
+    assert transformer.torchao_config.weight_dtype is torch.float8_e4m3fn
+    assert transformer.torchao_config.set_inductor_config is False
+
+
+def test_build_quant_config_torchao_runtime():
+    from vllm.model_executor.layers.quantization.torchao import TorchAOConfig
+
+    torchao_config = object()
+    result = build_quant_config("torchao", torchao_config=torchao_config)
+
+    assert isinstance(result, TorchAOConfig)
+    assert result.torchao_config is torchao_config
+    assert result.is_checkpoint_torchao_serialized is False