diff --git a/recipes/Boogu/Boogu-Image.md b/recipes/Boogu/Boogu-Image.md index 9a5ce0397..6ff848539 100644 --- a/recipes/Boogu/Boogu-Image.md +++ b/recipes/Boogu/Boogu-Image.md @@ -132,18 +132,19 @@ with `guidance_scale=1.0` remain valid on the same server: every rank evaluates only the positive branch, no negative embeddings are built, and no guidance is applied. -### 1 x RTX 4090 48 GB (FP8) +### 1 x RTX 4090 24 GB / 48 GB (FP8) -Boogu-Image supports two ways to quantize or load the DiT transformer: +Boogu-Image supports the following FP8 loading paths: -| Path | Base model | Edit model | Method | Scheme | -| ------------------ | -------------------------------- | -------------------------------- | ---------------------------- | ----------------------- | -| Serialized TorchAO | `Boogu/Boogu-Image-0.1-Base-fp8` | `Boogu/Boogu-Image-0.1-Edit-fp8` | `torchao_float8_weight_only` | FP8 weight-only (W8A16) | -| Native online FP8 | `Boogu/Boogu-Image-0.1-Base` | `Boogu/Boogu-Image-0.1-Edit` | `fp8` | Dynamic W8A8 | +| Path | Checkpoints | MLLM | DiT | +| --- | --- | --- | --- | +| Pre-quantized FP8 | `Boogu/Boogu-Image-0.1-{Base,Edit}-fp8` | HF checkpoint FP8 configuration | TorchAO FP8 weight-only (W8A16) | +| Online FP8 | `Boogu/Boogu-Image-0.1-{Base,Edit}` | Optional 128x128 block FP8: Triton W8A8 (vLLM 0.28.0); Marlin W8A16 (vLLM 0.29.0) | Native dynamic FP8 (CUTLASS W8A8) | #### Environment -The following configuration was validated: +The following isolated environments were used for Base T2I and Edit I2I +measurements at 512x512: - OS: Linux - Python: 3.12 @@ -152,7 +153,9 @@ The following configuration was validated: - vLLM-Omni version or commit: Use a commit that contains TorchAO FP8 checkpoint loading for diffusion models - TorchAO version: 0.17.0 -- `kernels` / `kernels-data`: 0.16.1 / matching kernels-data +- `kernels` / `kernels-data`: 0.16.1 / 0.16.1 +- Transformers: 5.17.0 +- Tokenizers: 0.23.1 #### Command @@ -167,19 +170,61 @@ vllm serve Boogu/Boogu-Image-0.1-Base-fp8 \ '{"transformer":{"method":"torchao_float8_weight_only"}}' ``` -Native online FP8 (currently applied only to eligible vLLM-quantizable linear -layers in the DiT transformer): +Native online FP8: ```bash +# DiT only vllm serve Boogu/Boogu-Image-0.1-Base \ --omni \ --port 8091 \ --diffusion-quantization-config \ - '{"transformer":{"method":"fp8"}}' + '{"mllm":null,"transformer":"fp8","vae":null}' + +# MLLM + DiT +vllm serve Boogu/Boogu-Image-0.1-Base \ + --omni \ + --port 8091 \ + --quantization fp8 + +# MLLM + DiT (using --diffusion-quantization-config) +vllm serve Boogu/Boogu-Image-0.1-Base \ + --omni \ + --port 8091 \ + --diffusion-quantization-config \ + '{"mllm":"fp8","transformer":"fp8","vae":null}' + ``` -For image editing, use the matching Edit model ID from the table and keep the -corresponding quantization configuration. +For image editing, replace `Base` with `Edit` and keep the corresponding +quantization configuration. + +#### Notes + +- **Memory usage:** + - The serialized TorchAO FP8 path uses approximately 21.1 GiB at 512x512 + and 24.2-24.6 GiB at 1024x1024. + - The DiT-only online FP8 path uses approximately 26.9 GiB at 512x512 + and 29.9-30.9 GiB at 1024x1024. + - The MLLM + DiT online FP8 path uses approximately 20.5-20.6 GiB at + 512x512 and 23.6-24.5 GiB at 1024x1024. + +- **Quantization scope:** For original Base/Edit checkpoints, `--quantization fp8` + quantizes the MLLM's language-model linear layers through per-block FP8 + and the DiT's linear layers through native online FP8. The vision encoder, + embeddings, normalization layers and VAE remain unquantized. + +- **Performance:** In the 48 GB Base/Edit benchmarks at 512x512, both DiT-only + and MLLM + DiT online FP8 reduced E2E latency by over 10 seconds compared with + official FP8 and vLLM-Omni pre-quantized FP8. The main runtime difference lies in the DiT quantization implementation, which switches from TorchAO weight-only FP8 (W8A16) to native CUTLASS + W8A8 kernels. This reduces additional dequantization overhead and enables + FP8 Tensor Core acceleration. + +- **Known Limitation:** + - **Output consistency:** With MLLM + DiT online FP8, same-seed outputs may visibly differ from BF16 for some Base prompts, but these differences do not mean lower image quality. + + - **Import compatibility conflict between TorchAO and Diffusers::** Loading pre-quantized checkpoints requires `torchao>=0.17.0`, which has a compatibility conflict with `diffusers==0.40.0`. This produces a warning when Diffusers is imported, but currently does not affect Boogu-Image weight loading. + + - **Conflict between the default kernels and Transformers versions:** The default Transformers version range (`>=5.10.1,<5.15`) conflicts with `kernels==0.16.1` when loading pre-quantized FP8 MLLM weights through Hugging Face `FP8Linear`. When using this loading path with `kernels==0.16.1`, please upgrade Transformers to `5.17.0`. #### Verification @@ -190,19 +235,6 @@ ID passed to `vllm serve`. For example, use corresponding Edit model. Base text-to-image: -## Fast text-to-image (Boogu-Image-0.1-Turbo) - -Turbo is the distilled text-to-image checkpoint. Its `model_index.json` -declares `BooguImageTurboPipeline`, which runs the four-step DMD path instead -of the scheduler-driven path used by Base. - -### Command - -```bash -vllm serve Boogu/Boogu-Image-0.1-Turbo --omni --port 8091 -``` - -### Verification ```bash curl -X POST http://localhost:8091/v1/images/generations \ @@ -231,22 +263,24 @@ curl -X POST http://localhost:8091/v1/images/edits \ | jq -r '.data[0].b64_json' | base64 -d > edited_fp8.png ``` -#### Notes +## Fast text-to-image (Boogu-Image-0.1-Turbo) -- **Memory usage:** - - The serialized TorchAO FP8 path uses approximately 21.1 GiB at 512x512 - and 24.2-24.6 GiB at 1024x1024. - - The native online FP8 path uses approximately 31.5 GiB for Base and - 30.7 GiB for Edit at 1024x1024. It uses more memory because only eligible - DiT linear layers are quantized online, while the MLLM from the regular - Base/Edit checkpoint remains in BF16. Most of the observed difference - comes from the MLLM precision, although differences in the two DiT FP8 - representations and kernels may also contribute. -- **Quantization scope:** The transformer-scoped configuration applies to the - DiT. With the serialized `-fp8` checkpoints, the MLLM follows its own - checkpoint-declared FP8 configuration. With the regular checkpoints used for - online FP8, the MLLM remains in BF16. The VAE and scheduler are outside this - quantization routing. +Turbo is the distilled text-to-image checkpoint. Its `model_index.json` +declares `BooguImageTurboPipeline`, which runs the four-step DMD path instead +of the scheduler-driven path used by Base. + +### Command + +```bash +vllm serve Boogu/Boogu-Image-0.1-Turbo --omni --port 8091 +``` + +### Verification + +```bash +curl -X POST http://localhost:8091/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ "model": "Boogu/Boogu-Image-0.1-Turbo", "prompt": "A mountain lake at sunset, photorealistic, cinematic lighting", "size": "1024x1024", diff --git a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py index 39c32a4df..4abd7ef9b 100644 --- a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py +++ b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py @@ -13,11 +13,13 @@ Two groups: CFG handling, and reshape logic can be verified numerically on CPU. """ +import os from types import SimpleNamespace import pytest import torch from torch import nn +from transformers import Qwen3VLConfig from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig, TransformerConfig @@ -56,9 +58,17 @@ def mock_dependencies(mocker, monkeypatch): f"{_MODULE}.FlowMatchEulerDiscreteScheduler.from_pretrained", lambda *a, **k: mock_scheduler, ) - monkeypatch.setattr( + mllm_loader = mocker.patch( f"{_MODULE}.Qwen3VLForConditionalGeneration.from_pretrained", - lambda *a, **k: mllm_wrapper, + return_value=mllm_wrapper, + ) + mllm_config_loader = mocker.patch( + f"{_MODULE}.Qwen3VLConfig.from_pretrained", + return_value=Qwen3VLConfig(), + ) + mllm_builder = mocker.patch( + f"{_MODULE}.create_transformers_model_with_vllm_linears", + return_value=nn.Linear(1, 1), ) monkeypatch.setattr( f"{_MODULE}.Qwen3VLProcessor.from_pretrained", @@ -74,12 +84,17 @@ def mock_dependencies(mocker, monkeypatch): mock_transformer_cls.return_value = mock_transformer_instance monkeypatch.setattr(f"{_MODULE}.BooguImageTransformer2DModel", mock_transformer_cls) - # Treat the dummy model id as a local path: skips hub prefetch. - mocker.patch("os.path.exists", return_value=True) + # Treat only the dummy model id as local. Other filesystem checks (for + # example lazy imports in the quantization registry) must remain real. + path_exists = os.path.exists + mocker.patch("os.path.exists", side_effect=lambda path: str(path).startswith("dummy-boogu") or path_exists(path)) return { "inner_encoder": inner_encoder, "mllm_wrapper": mllm_wrapper, + "mllm_loader": mllm_loader, + "mllm_config_loader": mllm_config_loader, + "mllm_builder": mllm_builder, "processor": mock_processor, "vae": mock_vae, "scheduler": mock_scheduler, @@ -128,6 +143,9 @@ def test_constructor_wires_components(boogu_pipeline, mock_dependencies): assert boogu_pipeline.vae_scale_factor == 8 assert boogu_pipeline.default_sample_size == 128 assert hasattr(boogu_pipeline, "load_weights") + assert "quantization_config" not in mock_dependencies["mllm_loader"].call_args.kwargs + mock_dependencies["mllm_config_loader"].assert_not_called() + mock_dependencies["mllm_builder"].assert_not_called() def test_constructor_strips_mllm_lm_head(boogu_pipeline, mock_dependencies): @@ -188,7 +206,7 @@ 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): +def test_constructor_rejects_unsupported_mllm_quantization(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 @@ -204,10 +222,143 @@ def test_constructor_routes_only_transformer_config(mock_dependencies, mocker): {"transformer": transformer_config, "mllm": encoder_config, "vae": None} ), ) + with pytest.raises(ValueError, match="Boogu MLLM only supports FP8 quantization"): + BooguImagePipeline(od_config=od_config) + mock_dependencies["mllm_loader"].assert_not_called() + mock_dependencies["mllm_config_loader"].assert_not_called() + + +@pytest.mark.parametrize( + ("quantization_config", "quantize_mllm"), + [ + pytest.param("fp8", True, id="global-fp8"), + pytest.param({"mllm": {"method": "fp8"}, "transformer": {"method": "fp8"}}, True, id="mllm-and-dit-fp8"), + pytest.param({"mllm": None, "transformer": "fp8"}, False, id="dit-only-fp8"), + ], +) +def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_config, quantize_mllm): + from transformers import AutoModel + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig + from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8Static128BlockSym + + from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline + + od_config = OmniDiffusionConfig( + model="dummy-boogu", + revision="test-revision", + tf_model_config=TransformerConfig(params={}), + dtype=torch.bfloat16, + quantization_config=quantization_config, + ) + pipeline = BooguImagePipeline(od_config=od_config) + + if quantize_mllm: + mock_dependencies["mllm_loader"].assert_not_called() + mock_dependencies["mllm_config_loader"].assert_called_once_with( + "dummy-boogu", subfolder="mllm", local_files_only=True, revision="test-revision" + ) + builder = mock_dependencies["mllm_builder"] + builder.assert_called_once() + auto_cls, hf_config, online_quant_config = builder.call_args.args + assert auto_cls is AutoModel + assert hf_config is mock_dependencies["mllm_config_loader"].return_value + assert isinstance(online_quant_config, OnlineQuantizationConfig) + assert online_quant_config.args.linear.weight == kFp8Static128BlockSym + assert builder.call_args.kwargs == { + "dtype": torch.bfloat16, + "device": pipeline._execution_device, + "prefix": "mllm", + "skip_modules": ("mllm.visual",), + } + assert pipeline.mllm is builder.return_value + assert not pipeline.mllm.training + assert all(not parameter.requires_grad for parameter in pipeline.mllm.parameters()) + transformer_source, mllm_source = pipeline.weights_sources + assert transformer_source.subfolder == "transformer" + assert mllm_source.model_or_path == "dummy-boogu" + assert mllm_source.subfolder == "mllm" + assert mllm_source.prefix == "mllm." + assert mllm_source.revision == "test-revision" + else: + assert "quantization_config" not in mock_dependencies["mllm_loader"].call_args.kwargs + assert pipeline.mllm is mock_dependencies["inner_encoder"] + mock_dependencies["mllm_config_loader"].assert_not_called() + mock_dependencies["mllm_builder"].assert_not_called() + assert [source.subfolder for source in pipeline.weights_sources] == ["transformer"] + assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], Fp8Config) + + +def test_constructor_preserves_mllm_ignored_layers(mock_dependencies): + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + + from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline + + quant_config = Fp8Config( + ignored_layers=["mllm.language_model.layers.0.self_attn.q_proj", "transformer.blocks.0.attn.to_q"] + ) + od_config = OmniDiffusionConfig( + model="dummy-boogu", + tf_model_config=TransformerConfig(params={}), + dtype=torch.bfloat16, + quantization_config=quant_config, + ) BooguImagePipeline(od_config=od_config) - kwargs = mock_dependencies["transformer_cls"].call_args.kwargs - assert kwargs["quant_config"] is transformer_config - assert kwargs["prefix"] == "transformer" + + online_quant_config = mock_dependencies["mllm_builder"].call_args.args[2] + assert online_quant_config.ignored_layers == quant_config.ignored_layers + + +def test_constructor_preserves_serialized_mllm_quantization(mock_dependencies): + from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline + + mock_dependencies["mllm_config_loader"].return_value = Qwen3VLConfig( + quantization_config={"quant_method": "fp8", "modules_to_not_convert": ["model.visual"]} + ) + od_config = OmniDiffusionConfig( + model="dummy-boogu-fp8", + tf_model_config=TransformerConfig(params={}), + dtype=torch.bfloat16, + quantization_config="fp8", + ) + pipeline = BooguImagePipeline(od_config=od_config) + + # No override: HF reads the serialized checkpoint's scales and skip list. + assert "quantization_config" not in mock_dependencies["mllm_loader"].call_args.kwargs + assert pipeline.mllm is mock_dependencies["inner_encoder"] + mock_dependencies["mllm_builder"].assert_not_called() + assert [source.subfolder for source in pipeline.weights_sources] == ["transformer"] + + +def test_load_weights_maps_mllm_and_preserves_transformer(): + from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline + + pipeline = object.__new__(BooguImagePipeline) + nn.Module.__init__(pipeline) + pipeline.mllm = nn.ModuleDict( + { + "language_model": nn.Linear(2, 2, bias=False), + "visual": nn.Linear(2, 2, bias=False), + } + ) + pipeline.transformer = nn.Linear(2, 2, bias=False) + language_weight = torch.full((2, 2), 1.0) + visual_weight = torch.full((2, 2), 2.0) + transformer_weight = torch.full((2, 2), 3.0) + + loaded = pipeline.load_weights( + [ + ("mllm.lm_head.weight", torch.full((2, 2), 9.0)), + ("mllm.model.language_model.weight", language_weight), + ("mllm.model.visual.weight", visual_weight), + ("transformer.weight", transformer_weight), + ] + ) + + assert loaded == {"mllm.language_model.weight", "mllm.visual.weight", "transformer.weight"} + torch.testing.assert_close(pipeline.mllm["language_model"].weight, language_weight) + torch.testing.assert_close(pipeline.mllm["visual"].weight, visual_weight) + torch.testing.assert_close(pipeline.transformer.weight, transformer_weight) @pytest.mark.parametrize( diff --git a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py index c82e6ad7b..b4cd930b2 100644 --- a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py +++ b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py @@ -34,9 +34,13 @@ from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL from diffusers.utils.torch_utils import randn_tensor from torch import nn from torch.nn.attention import SDPBackend, sdpa_kernel -from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor +from transformers import AutoModel, Qwen3VLConfig, Qwen3VLForConditionalGeneration, Qwen3VLModel, Qwen3VLProcessor +from vllm.config.quantization import QuantizationConfigArgs from vllm.logger import init_logger -from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin @@ -54,6 +58,7 @@ from vllm_omni.diffusion.models.boogu_image.scheduling_flow_match_euler_discrete ) from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.models.utils import create_transformers_model_with_vllm_linears from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific @@ -243,6 +248,7 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports self.od_config = od_config self._raise_unsupported_features() transformer_quant_config = resolve_component_quant_config(od_config.quantization_config, "transformer") + mllm_quant_config = resolve_component_quant_config(od_config.quantization_config, "mllm") self.weights_sources = [ DiffusersPipelineLoader.ComponentSource( model_or_path=od_config.model, @@ -273,22 +279,7 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports local_files_only=local_files_only, revision=od_config.revision, ) - - mllm = from_pretrained_with_prefetch( - Qwen3VLForConditionalGeneration.from_pretrained, - model, - subfolder="mllm", - prefetch_list=boogu_subfolders, - local_files_only=local_files_only, - torch_dtype=od_config.dtype, - revision=od_config.revision, - ) - # Upstream reuses the full VLM as an optional instruction rewriter and - # encodes with its inner model (no ``lm_head``); the rewriter is not - # ported, so keep only the inner ``Qwen3VLModel`` as the encoder. - if hasattr(mllm, "lm_head"): - mllm = mllm.model - self.mllm = mllm.to(self._execution_device) + self.mllm = self._load_mllm(model, local_files_only, mllm_quant_config, boogu_subfolders) self.processor = Qwen3VLProcessor.from_pretrained( model, @@ -323,6 +314,68 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports self.SYSTEM_PROMPT_4_TI2I = SYSTEM_PROMPT_4_TI2I_UNIFIED self.SYSTEM_PROMPT_4_I2I = SYSTEM_PROMPT_4_TI2I_UNIFIED + def _load_mllm( + self, + model_path: str, + local_files_only: bool, + quant_config: QuantizationConfig | None, + prefetch_list: list[str], + ) -> Qwen3VLModel: + """Load MLLM through HF, or prepare it for deferred online FP8 loading.""" + config = None + if quant_config is not None: + if not isinstance(quant_config, Fp8Config): + raise ValueError( + "Boogu MLLM only supports FP8 quantization. Set mllm to null to disable online quantization." + ) + config = Qwen3VLConfig.from_pretrained( + model_path, subfolder="mllm", local_files_only=local_files_only, revision=self.od_config.revision + ) + + # With no MLLM quantization override, or with checkpoint quantization + # metadata, preserve HF model initialization and weight loading. + if config is None or getattr(config, "quantization_config", None): + mllm = from_pretrained_with_prefetch( + Qwen3VLForConditionalGeneration.from_pretrained, + model_path, + subfolder="mllm", + prefetch_list=prefetch_list, + local_files_only=local_files_only, + torch_dtype=self.od_config.dtype, + revision=self.od_config.revision, + ) + # Upstream reuses the full VLM as an optional instruction rewriter and + # encodes with its inner model (no ``lm_head``); the rewriter is not + # ported, so keep only the inner ``Qwen3VLModel`` as the encoder. + if hasattr(mllm, "lm_head"): + mllm = mllm.model + return mllm.to(self._execution_device) + + online_quant_config = OnlineQuantizationConfig( + QuantizationConfigArgs(linear="fp8_per_block", ignore=quant_config.ignored_layers) + ) + + mllm = create_transformers_model_with_vllm_linears( + AutoModel, + config, + online_quant_config, + dtype=self.od_config.dtype, + device=self._execution_device, + prefix="mllm", + skip_modules=("mllm.visual",), + ) + + self.weights_sources.append( + DiffusersPipelineLoader.ComponentSource( + model_or_path=model_path, + subfolder="mllm", + revision=self.od_config.revision, + prefix="mllm.", + ) + ) + + return mllm.requires_grad_(False).eval() + def _raise_unsupported_features(self) -> None: """Reject execution modes that do not have Boogu-specific support.""" parallel_config = self.od_config.parallel_config @@ -339,7 +392,8 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + mapper = WeightsMapper(orig_to_new_prefix={"mllm.lm_head.": None, "mllm.model.": "mllm."}) + return loader.load_weights(weights, mapper=mapper) # ------------------------------------------------------------------ # Prompt encoding (upstream ``encode_instruction``, t2i path) diff --git a/vllm_omni/diffusion/models/utils.py b/vllm_omni/diffusion/models/utils.py index 429c34b9d..f55996a76 100644 --- a/vllm_omni/diffusion/models/utils.py +++ b/vllm_omni/diffusion/models/utils.py @@ -5,12 +5,14 @@ from __future__ import annotations import json import os +from collections.abc import Sequence from typing import TYPE_CHECKING, Literal import torch import torch.nn as nn from vllm.model_executor.layers.linear import ( ColumnParallelLinear, + LinearBase, ReplicatedLinear, RowParallelLinear, ) @@ -124,6 +126,28 @@ def recursive_replace_linear(model: nn.Module, od_config: OmniDiffusionConfig): _recursive_replace(model, prefix="") +def recursive_replace_linear_with_quantization_config( + model: nn.Module, + quant_config: QuantizationConfig | None, + prefix: str = "", + skip_modules: Sequence[str] = (), +) -> None: + for child_name, child_module in model.named_children(): + qual_name = maybe_prefix(prefix, child_name) + if any(qual_name == path or qual_name.startswith(path + ".") for path in skip_modules): + continue + if isinstance(child_module, nn.Linear): + setattr( + model, + child_name, + replace_linear_class(child_module, quant_config=quant_config, prefix=qual_name), + ) + else: + recursive_replace_linear_with_quantization_config( + child_module, quant_config, prefix=qual_name, skip_modules=skip_modules + ) + + def init_parameters( module: nn.Module, dtype: torch.dtype | None, @@ -144,6 +168,24 @@ def init_parameters( init_parameters(child, dtype, device) +def init_parameters_preserving_vllm_linear( + module: nn.Module, + dtype: torch.dtype | None, + device: torch.device | None = None, +) -> None: + if isinstance(module, LinearBase): + return + for name, param in module.named_parameters(recurse=False): + if param.is_meta: + setattr( + module, + name, + nn.Parameter(torch.empty_like(param, dtype=dtype, device=device), requires_grad=param.requires_grad), + ) + for child in module.children(): + init_parameters_preserving_vllm_linear(child, dtype, device) + + def create_transformers_model( auto_cls: _BaseAutoModelClass, od_config: OmniDiffusionConfig, @@ -161,6 +203,23 @@ def create_transformers_model( return model +def create_transformers_model_with_vllm_linears( + auto_cls: _BaseAutoModelClass, + hf_config: PretrainedConfig, + quant_config: QuantizationConfig | None, + dtype: torch.dtype, + device: torch.device, + prefix: str = "", + skip_modules: Sequence[str] = (), +) -> PreTrainedModel: + """Build a Transformers model with vLLM linears.""" + with init_on_device_without_buffers("meta"): + model = auto_cls.from_config(hf_config, dtype=dtype) + recursive_replace_linear_with_quantization_config(model, quant_config, prefix=prefix, skip_modules=skip_modules) + init_parameters_preserving_vllm_linear(model, dtype=dtype, device=device) + return model + + def _load_json(model_path: str, filename: str, local_files_only: bool = True) -> dict: """Load a JSON config file from a local path or HuggingFace Hub repo.""" if local_files_only: