From 06541d1f523dc25f13c4258732dcd0ed4cb248ff Mon Sep 17 00:00:00 2001 From: DMAN <32004742+Dmaner@users.noreply.github.com> Date: Sun, 6 Sep 2026 12:48:38 +0800 Subject: [PATCH] [Model] FP8 Support For Boogu-Image (#6925) Signed-off-by: Dmaner <2663515256@qq.com> --- docs/.nav.yml | 1 + docs/design/feature/quantization.md | 4 +- docs/user_guide/quantization/online.md | 16 +- docs/user_guide/quantization/overview.md | 41 +-- docs/user_guide/quantization/torchao.md | 94 ++++++ recipes/Boogu/Boogu-Image.md | 103 ++++++ .../model_loader/test_diffusers_loader.py | 84 ++++- .../test_boogu_image_transformer.py | 68 ++++ .../boogu_image/test_pipeline_boogu_image.py | 23 ++ .../quantization/test_torchao_config.py | 58 ++++ tests/entrypoints/test_serve.py | 20 ++ .../model_loader/diffusers_loader.py | 40 ++- .../boogu_image/boogu_image_transformer.py | 304 +++++++++++++++--- .../boogu_image/pipeline_boogu_image.py | 8 +- vllm_omni/quantization/factory.py | 23 ++ 15 files changed, 809 insertions(+), 78 deletions(-) create mode 100644 docs/user_guide/quantization/torchao.md create mode 100644 tests/diffusion/quantization/test_torchao_config.py diff --git a/docs/.nav.yml b/docs/.nav.yml index 6477323c4..8a03c7a4f 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -66,6 +66,7 @@ nav: - ModelOpt: user_guide/quantization/modelopt.md - AutoRound: user_guide/quantization/autoround.md - msModelSlim: user_guide/quantization/msmodelslim.md + - TorchAO: user_guide/quantization/torchao.md - Diffusion Acceleration: - Overview and Compatibility: - Overview: user_guide/diffusion_features.md diff --git a/docs/design/feature/quantization.md b/docs/design/feature/quantization.md index 84c9641b3..f99263221 100644 --- a/docs/design/feature/quantization.md +++ b/docs/design/feature/quantization.md @@ -84,8 +84,8 @@ dictionary narrows this global scope through runtime layer-prefix routing. The factory uses Omni-specific overrides for methods whose diffusion behavior or platform dispatch is not provided by the vLLM registry. The current override -set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/ -AutoRound. Other methods, such as FP8, GGUF, and ModelOpt, are resolved through +set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, INC/AutoRound, +and TorchAO. Other methods, such as FP8, GGUF, and ModelOpt, are resolved through vLLM's registry when their configuration is compatible. ## Per-component routing diff --git a/docs/user_guide/quantization/online.md b/docs/user_guide/quantization/online.md index 6103be9a0..c3329bf4c 100644 --- a/docs/user_guide/quantization/online.md +++ b/docs/user_guide/quantization/online.md @@ -7,16 +7,16 @@ loading the model. Use it when you want memory savings without preparing a separate quantized checkpoint. This mode is different from pre-quantized checkpoint formats such as GGUF, -AutoRound, msModelSlim, or serialized Int8 checkpoints. Those formats are -prepared before serving and are documented in their method-specific guides. -For MXFP8 and MXFP4, use this page for load-time quantization from BF16 -checkpoints, and use the method-specific pages for offline checkpoints produced -by msModelSlim and the merge tools. +AutoRound, msModelSlim, serialized TorchAO checkpoints, or serialized Int8 +checkpoints. Those formats are prepared before serving and are documented in +their method-specific guides. For MXFP8 and MXFP4, use this page for load-time +quantization from BF16 checkpoints, and use the method-specific pages for +offline checkpoints produced by msModelSlim and the merge tools. ## Hardware Support | Device | FP8 W8A8 | Int8 W8A8 | MXFP8 W8A8 | MXFP4 W4A4 | -|--------|----------|-----------|------------|------------| +| -------- | ---------- | ----------- | ------------ | ------------ | | NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ⭕ | ⭕ | | NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ | | NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | @@ -33,7 +33,7 @@ MXFP4 are documented for the Ascend NPU path. ### Diffusion Model (Qwen-Image, Wan2.2) | Method | Guide | Example models | Status | -|--------|-------|----------------|--------| +| -------- | ------- | ---------------- | -------- | | FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models | | Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image | | MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Validated on Ascend NPU and Intel XPU | @@ -88,7 +88,7 @@ config = build_quant_config({ ## Parameters | Parameter | Methods | Description | -|-----------|---------|-------------| +| ----------- | --------- | ------------- | | `method` | FP8, Int8, MXFP8, MXFP4 | Quantization method: `"fp8"`, `"int8"`, `"mxfp8"`, `"mxfp4"`, or `"mxfp4_dualscale"` | | `ignored_layers` | FP8, Int8, MXFP8, MXFP4 | Layer name patterns to keep in BF16/FP16 | | `activation_scheme` | FP8, Int8 | The runtime value `"dynamic"` selects online activation scaling | diff --git a/docs/user_guide/quantization/overview.md b/docs/user_guide/quantization/overview.md index 9ceadda5c..8319077a2 100644 --- a/docs/user_guide/quantization/overview.md +++ b/docs/user_guide/quantization/overview.md @@ -11,21 +11,21 @@ For the internal architecture and backend extension points, see the ## Quantization Modes | Mode | Guide | Description | Methods | -|------|-------|-------------|---------| +| ------ | ------- | ------------- | --------- | | Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, BitsAndBytes W4, MXFP8 W8A8, MXFP4 W4A4 | | Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA | -| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale | +| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, AutoRound, TorchAO, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale | ## Hardware Support -| Device | FP8 W8A8 | Int8 W8A8 | BitsAndBytes W4 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | AutoRound | msModelSlim | -|--------|----------|-----------|------------------|----------|------------|------------|-----------|-------------| -| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ | -| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ | -| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ❌ | -| AMD ROCm | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ✅ | ⭕ | ❌ | -| Intel XPU | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ⭕ | ✅ | ❌ | -| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | +| Device | FP8 W8A8 | Int8 W8A8 | BitsAndBytes W4 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | AutoRound | msModelSlim | TorchAO FP8 weight-only | +| -------- | ---------- | ----------- | ------------------ | ---------- | ------------ | ------------ | ----------- | ------------- | --------- | +| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ | ⭕ | +| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ | ✅ | +| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ❌ | ⭕ | +| AMD ROCm | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ✅ | ⭕ | ❌ | ⭕ | +| Intel XPU | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ⭕ | ✅ | ❌ | ⭕ | +| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ⭕ | Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. FP8 on Ampere may use a weight-only path where available. @@ -40,7 +40,7 @@ encoder, and VAE stay on the base checkpoint unless a method guide says otherwise. | Method | Guide | Mode | Example models | Status | -|--------|-------|------|----------------|--------| +| -------- | ------- | ------ | ---------------- | -------- | | FP8 W8A8 | [FP8](fp8.md) | Online W8A8 or checkpoint FP8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models | | Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image | | BitsAndBytes W4 | [BitsAndBytes](bitsandbytes.md) | Online W4 weight-only | Z-Image-Turbo; Qwen-Image/Wan2.2 not validated | Validated for Z-Image-Turbo | @@ -49,6 +49,7 @@ otherwise. | MXFP4 W4A4 | [MXFP4](mxfp4.md) | `mxfp4`: online single-scale only; `mxfp4_dualscale`: online or offline dual-scale (offline recommended) | Wan2.2-T2V-A14B, I2V-A14B | Ascend NPU only; validated for Wan2.2 A14B cascade models; TI2V-5B not supported; offline `mxfp4_dualscale` uses calibrated `mul_scale` for best accuracy | | AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven | | msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path | +| TorchAO | [TorchAO](torchao.md) | Pre-quantized FP8 weight-only (W8A16) | Boogu-Image Base/Edit `-fp8` | Validated for the Boogu-Image diffusion transformer | ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) @@ -58,13 +59,14 @@ checkpoint contains a supported `quantization_config`; the non-AR stages stay in BF16 unless the model guide explicitly adds support. | Method | Guide | Scope | Example models | Status | -|--------|-------|-------|----------------|--------| +| -------- | ------- | ------- | ---------------- | -------- | | ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path | | Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints | +| TorchAO | [TorchAO](torchao.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | ### Multi-Stage Diffusion Model (BAGEL, GLM-Image) @@ -73,7 +75,7 @@ These models split generation across multiple stages. Quantization must be attached to the intended stage rather than applied globally. | Method | Guide | Scope | Example models | Status | -|--------|-------|-------|----------------|--------| +| -------- | ------- | ------- | ---------------- | -------- | | FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation | | Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation | | BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Stage-specific transformer or DiT module | BAGEL, GLM-Image | Not validated | @@ -81,6 +83,7 @@ attached to the intended stage rather than applied globally. | MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated | | MXFP4 | [MXFP4](mxfp4.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated | | AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed | +| TorchAO | [TorchAO](torchao.md) | Stage-specific serialized checkpoint | BAGEL, GLM-Image | Not validated | | msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation | !!! note @@ -105,8 +108,8 @@ config = build_quant_config({ ``` | Component | Default quantized? | Notes | -|-----------|--------------------|-------| -| Diffusion transformer | Yes | Primary target for FP8, Int8, BitsAndBytes, ModelOpt, MXFP8, MXFP4, AutoRound, and msModelSlim | +| ----------- | -------------------- | ------- | +| Diffusion transformer | Yes | Primary target for FP8, Int8, BitsAndBytes, ModelOpt, MXFP8, MXFP4, AutoRound, TorchAO, and msModelSlim | | Text encoder | No | Keep BF16 unless a method-specific guide documents support | | VAE | No | Keep BF16; storage-only paths are method-specific | | Scheduler/tokenizer | No | Loaded from the base model repository | @@ -114,7 +117,7 @@ config = build_quant_config({ ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) | Component | Default quantized? | Notes | -|-----------|--------------------|-------| +| ----------- | -------------------- | ------- | | Thinker or AR language model | Yes, when checkpoint config is supported | ModelOpt FP8/NVFP4 or AutoRound checkpoint config | | Audio encoder | No | BF16 | | Vision encoder | No | BF16 | @@ -124,7 +127,7 @@ config = build_quant_config({ ### Multi-Stage Diffusion Model (BAGEL, GLM-Image) | Component | Default quantized? | Notes | -|-----------|--------------------|-------| +| ----------- | -------------------- | ------- | | Selected diffusion or transformer stage | Method-specific | Must be routed to the intended stage | | Other generation stages | No | Keep BF16 unless separately validated | | VAE, tokenizer, scheduler | No | Loaded from the base checkpoint | @@ -204,7 +207,7 @@ The output JSON includes `output_metrics`, `reference_generation`, and `candidate_generation`. | Metric | Direction | Meaning | -|--------|-----------|---------| +| -------- | ----------- | --------- | | `cosine_similarity` | Higher is better | Vector direction similarity between output pixels or frames. Useful as a broad sanity check. | | `mae` | Lower is better | Mean absolute pixel or frame error. For decoded outputs, values are in uint8 pixel units. | | `mse` / `rmse` | Lower is better | Squared error and its square root. These penalize localized large differences more than `mae`. | @@ -217,7 +220,7 @@ The output JSON includes `output_metrics`, `reference_generation`, and Recommended starting thresholds for same-seed diffusion comparisons: | Metric | Smoke threshold | Stricter target | Notes | -|--------|-----------------|-----------------|-------| +| -------- | ----------------- | ----------------- | ------- | | `psnr_db` | `>= 20.0` | `>= 25.0` | Good for quick image or frame regression checks. | | `mae` | `<= 12.0` | `<= 6.0` | Interpreted in decoded uint8 pixel units. | | `cosine_similarity` | `>= 0.98` | `>= 0.995` | Less sensitive to global scale than L2-style metrics. | diff --git a/docs/user_guide/quantization/torchao.md b/docs/user_guide/quantization/torchao.md new file mode 100644 index 000000000..133fd1c23 --- /dev/null +++ b/docs/user_guide/quantization/torchao.md @@ -0,0 +1,94 @@ +# TorchAO Quantization + +## Overview + +[TorchAO](https://github.com/pytorch/ao) provides quantization tools for +PyTorch and supports both pre-quantized checkpoint loading and runtime +quantization. + +## Hardware Support + +| Device | Support | +| -------- | --------- | +| NVIDIA Blackwell GPU (SM 100+) | ⭕ | +| NVIDIA Ada GPU (SM 89) | ✅ | +| NVIDIA Hopper GPU (SM 90) | ⭕ | +| NVIDIA Ampere GPU (SM 80+) | ⭕ | +| AMD ROCm | ⭕ | +| Intel XPU | ⭕ | +| Ascend NPU | ⭕ | + +Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this +guide. + +## Model Type Support + +### Diffusion Model (Boogu-Image) + +| Model | Checkpoint | Scope | Scheme | +| --------------------- | ------------------------------------------------------------------ | -------------------------- | ----------------------- | +| Boogu-Image Base/Edit | `Boogu/Boogu-Image-0.1-Base-fp8`, `Boogu/Boogu-Image-0.1-Edit-fp8` | Diffusion transformer only | FP8 weight-only (W8A16) | + +## Configuration + +Using Boogu-Image as an example: + +```bash +vllm serve Boogu/Boogu-Image-0.1-Base-fp8 \ + --omni \ + --port 8091 \ + --diffusion-quantization-config \ + '{"transformer":{"method":"torchao_float8_weight_only"}}' +``` + +The equivalent command using the complete serialized TorchAO configuration is: + +```bash +vllm serve Boogu/Boogu-Image-0.1-Base-fp8 \ + --omni \ + --port 8091 \ + --diffusion-quantization-config \ + '{ + "transformer": { + "method": "torchao", + "quant_type": { + "default": { + "_type": "Float8WeightOnlyConfig", + "_version": 2, + "_data": { + "weight_dtype": { + "_type": "torch.dtype", + "_data": "float8_e4m3fn" + }, + "set_inductor_config": false + } + } + } + } + }' +``` + +## Parameters + +| Parameter | Type | Default | Description | +| ------------ | ---- | ------- | ----------------------------------------------------------------------------------------------------------------------------------- | +| `method` | str | - | Use `torchao_float8_weight_only` for the serialized FP8 weight-only shorthand, or `torchao` with a full `quant_type` configuration. | +| `quant_type` | dict | - | Complete serialized TorchAO configuration used with `method: "torchao"`. | + +See TorchAO's [vLLM integration guide](https://docs.pytorch.org/ao/0.17/eager_tutorials/torchao_vllm_integration.html#configuration-system) +for the serialized form and [workflow configs](https://docs.pytorch.org/ao/0.17/api_reference/api_ref_quantization.html#workflow-configs) +for available configuration classes and parameters. + +## Validation and Notes + +At load time, vLLM-Omni builds a `TorchAOConfig` from +`--diffusion-quantization-config`. For the Boogu-Image checkpoints documented +here, indexed PyTorch `.bin` shards are loaded through `pt_weights_iterator`. + +When a component map is used, only the components included in the +configuration use TorchAO. Other components keep their own checkpoint and +runtime settings. + +Use either the `torchao_float8_weight_only` shorthand or the equivalent full +`quant_type` configuration shown above. The checkpoint must already contain +weights quantized with TorchAO. diff --git a/recipes/Boogu/Boogu-Image.md b/recipes/Boogu/Boogu-Image.md index 741ea606b..ecc2c1d0f 100644 --- a/recipes/Boogu/Boogu-Image.md +++ b/recipes/Boogu/Boogu-Image.md @@ -128,6 +128,109 @@ 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) + +Boogu-Image supports two ways to quantize or load the DiT transformer: + +| 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 | + +#### Environment + +The following configuration was validated: + +- OS: Linux +- Python: 3.12 +- Driver / runtime: NVIDIA CUDA environment with one RTX 4090 48 GB GPU +- vLLM version: Match the vLLM-Omni checkout used for deployment +- 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.15.2 / 0.16.1 + +#### Command + +Serve the official Base FP8 checkpoint using its pre-quantized TorchAO FP8 +weight-only transformer weights: + +```bash +vllm serve Boogu/Boogu-Image-0.1-Base-fp8 \ + --omni \ + --port 8091 \ + --diffusion-quantization-config \ + '{"transformer":{"method":"torchao_float8_weight_only"}}' +``` + +Native online FP8 (currently applied only to eligible vLLM-quantizable linear +layers in the DiT transformer): + +```bash +vllm serve Boogu/Boogu-Image-0.1-Base \ + --omni \ + --port 8091 \ + --diffusion-quantization-config \ + '{"transformer":{"method":"fp8"}}' +``` + +For image editing, use the matching Edit model ID from the table and keep the +corresponding quantization configuration. + +#### Verification + +The request format is shared by both FP8 paths. Set `` to the model +ID passed to `vllm serve`. For example, use +`Boogu/Boogu-Image-0.1-Base-fp8` for serialized TorchAO FP8 or +`Boogu/Boogu-Image-0.1-Base` for native online FP8. For image editing, use the +corresponding Edit model. + +Base text-to-image: + +```bash +curl -X POST http://localhost:8091/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "model": "", + "prompt": "A mountain lake at sunset, photorealistic, cinematic lighting", + "size": "1024x1024", + "num_inference_steps": 28, + "guidance_scale": 4.0, + "seed": 42 + }' | jq -r '.data[0].b64_json' | base64 -d > output_fp8.png +``` + +Image editing: + +```bash +curl -X POST http://localhost:8091/v1/images/edits \ + -F model="" \ + -F image="@input.png" \ + -F prompt="Change the style to a colored pencil drawing." \ + -F num_inference_steps=28 \ + -F guidance_scale=4.0 \ + -F guidance_scale_2=1.0 \ + -F seed=42 \ + | jq -r '.data[0].b64_json' | base64 -d > edited_fp8.png +``` + +#### Notes + +- **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. + ## Image editing (Boogu-Image-0.1-Edit) The Edit checkpoint is served by the same native pipeline (`BooguImagePipeline`); diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 4f60e379b..555cb19bb 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/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) diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py index 9daff16fe..4296c1cde 100644 --- a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py +++ b/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, 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 d59544c3b..39c32a4df 100644 --- a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py +++ b/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"), [ diff --git a/tests/diffusion/quantization/test_torchao_config.py b/tests/diffusion/quantization/test_torchao_config.py new file mode 100644 index 000000000..10b2c98aa --- /dev/null +++ b/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 diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 01d3e7ddb..29090c425 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -90,6 +90,26 @@ def test_serve_parser_accepts_ulysses_a2a_permute() -> None: assert args.get_explicit_kwargs_dict()["ulysses_a2a_permute"] is True +def test_serve_parser_accepts_diffusion_quantization_config() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + expected = {"transformer": {"method": "torchao_float8_weight_only"}} + + args = parser.parse_args( + [ + "serve", + "Boogu/Boogu-Image-0.1-Base-fp8", + "--omni", + "--diffusion-quantization-config", + '{"transformer":{"method":"torchao_float8_weight_only"}}', + ] + ) + + assert args.diffusion_quantization_config == expected + assert args.get_explicit_kwargs_dict()["diffusion_quantization_config"] == expected + + def _make_headless_args(**kwargs) -> TrackingNamespace: defaults = { "model": "fake-model", diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index ec0c1b14e..27774aa37 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -22,6 +22,7 @@ from vllm.model_executor.model_loader.weight_utils import ( filter_files_not_needed_for_inference, maybe_download_from_modelscope, multi_thread_safetensors_weights_iterator, + pt_weights_iterator, safetensors_weights_iterator, ) from vllm.transformers_utils.repo_utils import hf_api @@ -98,7 +99,9 @@ def _natural_sort_key(filepath: str) -> list: DIFFUSION_MODEL_WEIGHTS_INDEX = "diffusion_pytorch_model.safetensors.index.json" TRANSFORMER_WEIGHTS_INDEX = "model.safetensors.index.json" -INDEX_FILES = [DIFFUSION_MODEL_WEIGHTS_INDEX, TRANSFORMER_WEIGHTS_INDEX] +DIFFUSION_MODEL_BIN_WEIGHTS_INDEX = "diffusion_pytorch_model.bin.index.json" +SAFETENSORS_INDEX_FILES = [DIFFUSION_MODEL_WEIGHTS_INDEX, TRANSFORMER_WEIGHTS_INDEX] +PT_INDEX_FILES = [DIFFUSION_MODEL_BIN_WEIGHTS_INDEX] SHARDED_SAFETENSORS_PATTERN = re.compile(r"^(?P.+)-\d+-of-(?P\d+)\.safetensors$") @@ -191,11 +194,12 @@ class DiffusersPipelineLoader(HWRLoaderMixin): model_name_or_path: Path | str, subfolder: str | None, revision: str | None, + index_files: Sequence[str], ) -> list[str] | None: """Resolve an index and return its authoritative shard manifest.""" is_local = os.path.isdir(model_name_or_path) index_paths: list[tuple[str, Path]] = [] - for index_file in INDEX_FILES: + for index_file in index_files: repo_index_path = self._repo_relative_path(subfolder, index_file) if is_local: index_path = Path(model_name_or_path) / repo_index_path @@ -249,11 +253,17 @@ class DiffusersPipelineLoader(HWRLoaderMixin): is_local = os.path.isdir(model_name_or_path) load_format = self.load_config.load_format use_safetensors = False - indexed_weight_files = ( - self._resolve_weight_index(model_name_or_path, subfolder, revision) - if allow_patterns_overrides is None - else None - ) + indexed_weight_files = None + if allow_patterns_overrides is None: + for index_files in (SAFETENSORS_INDEX_FILES, PT_INDEX_FILES): + indexed_weight_files = self._resolve_weight_index( + model_name_or_path, + subfolder, + revision, + index_files, + ) + if indexed_weight_files is not None: + break # only hf is supported currently if load_format == "auto": @@ -332,12 +342,19 @@ class DiffusersPipelineLoader(HWRLoaderMixin): source.fall_back_to_pt, source.allow_patterns_overrides, ) - + quant_config = self._get_source_quant_config(source) use_multithread = ( use_safetensors and getattr(self.od_config, "enable_multithread_weight_load", False) and self.load_config.safetensors_load_strategy != "torchao" ) + use_torchao = ( + not use_safetensors + and quant_config is not None + and hasattr(quant_config, "get_name") + and quant_config.get_name() == "torchao" + and getattr(quant_config, "is_checkpoint_torchao_serialized", False) + ) if use_multithread: num_threads = getattr(self.od_config, "num_weight_load_threads", 4) # Keep deterministic shard order before passing to vLLM helper. @@ -347,6 +364,13 @@ class DiffusersPipelineLoader(HWRLoaderMixin): self.load_config.use_tqdm_on_load, max_workers=num_threads, ) + elif use_torchao: + sorted_hf_weights_files = sorted(hf_weights_files, key=_natural_sort_key) + weights_iterator = pt_weights_iterator( + sorted_hf_weights_files, + self.load_config.use_tqdm_on_load, + self.load_config.pt_load_map_location, + ) else: weights_iterator = safetensors_weights_iterator( hf_weights_files, diff --git a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py index 7016b4f69..2025859ed 100644 --- a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py +++ b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # # Native vLLM-Omni port of the Boogu-Image transformer. # @@ -16,6 +16,7 @@ import itertools from collections.abc import Iterable +from typing import TYPE_CHECKING import torch import torch.nn as nn @@ -36,12 +37,19 @@ from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.platforms import current_omni_platform +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + logger = init_logger(__name__) RotaryEmbedding = tuple[torch.Tensor, torch.Tensor] RotaryFrequencyTables = list[RotaryEmbedding] +def _join_prefix(prefix: str, name: str) -> str: + return f"{prefix}.{name}" if prefix else name + + def apply_rotary_emb(x: torch.Tensor, rotary_emb: RotaryEmbedding) -> torch.Tensor: """Apply rotary embeddings with real-valued cosine and sine tensors. @@ -82,10 +90,23 @@ class TimestepEmbedding(nn.Module): class LuminaRMSNormZero(nn.Module): """AdaRMS modulation: projects `temb` into scale/gate terms.""" - def __init__(self, embedding_dim: int, norm_eps: float) -> None: + def __init__( + self, + embedding_dim: int, + norm_eps: float, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", + ) -> None: super().__init__() self.silu = nn.SiLU() - self.linear = nn.Linear(min(embedding_dim, 1024), 4 * embedding_dim, bias=True) + self.linear = ReplicatedLinear( + min(embedding_dim, 1024), + 4 * embedding_dim, + bias=True, + return_bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "linear"), + ) self.norm = RMSNorm(embedding_dim, eps=norm_eps) def forward( @@ -108,12 +129,32 @@ class LuminaLayerNormContinuous(nn.Module): eps: float = 1e-5, bias: bool = True, out_dim: int | None = None, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ) -> None: super().__init__() self.silu = nn.SiLU() - self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias) + self.linear_1 = ReplicatedLinear( + conditioning_embedding_dim, + embedding_dim, + bias=bias, + return_bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "linear_1"), + ) self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias) - self.linear_2 = nn.Linear(embedding_dim, out_dim, bias=bias) if out_dim is not None else None + self.linear_2 = ( + ReplicatedLinear( + embedding_dim, + out_dim, + bias=bias, + return_bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "linear_2"), + ) + if out_dim is not None + else None + ) def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor: scale = self.linear_1(self.silu(conditioning_embedding).to(x.dtype)) @@ -132,15 +173,35 @@ class LuminaFeedForward(nn.Module): inner_dim: int, multiple_of: int = 256, ffn_dim_multiplier: float | None = None, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ) -> None: super().__init__() if ffn_dim_multiplier is not None: inner_dim = int(ffn_dim_multiplier * inner_dim) inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of) - self.linear_1 = ColumnParallelLinear(dim, inner_dim, bias=False) # gate - self.linear_3 = ColumnParallelLinear(dim, inner_dim, bias=False) # input - self.linear_2 = RowParallelLinear(inner_dim, dim, bias=False) + self.linear_1 = ColumnParallelLinear( + dim, + inner_dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "linear_1"), + ) # gate + self.linear_3 = ColumnParallelLinear( + dim, + inner_dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "linear_3"), + ) # input + self.linear_2 = RowParallelLinear( + inner_dim, + dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "linear_2"), + ) def forward(self, x: torch.Tensor) -> torch.Tensor: h1, _ = self.linear_1(x) @@ -427,17 +488,48 @@ class BooguImageSelfAttention(nn.Module): SP/KV-cache concerns. """ - def __init__(self, dim: int, num_attention_heads: int, num_kv_heads: int) -> None: + def __init__( + self, + dim: int, + num_attention_heads: int, + num_kv_heads: int, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", + ) -> None: super().__init__() self.head_dim = dim // num_attention_heads kv_dim = self.head_dim * num_kv_heads - self.to_q = ColumnParallelLinear(dim, dim, bias=False) - self.to_k = ColumnParallelLinear(dim, kv_dim, bias=False) - self.to_v = ColumnParallelLinear(dim, kv_dim, bias=False) + self.to_q = ColumnParallelLinear( + dim, + dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "to_q"), + ) + self.to_k = ColumnParallelLinear( + dim, + kv_dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "to_k"), + ) + self.to_v = ColumnParallelLinear( + dim, + kv_dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "to_v"), + ) self.norm_q = RMSNorm(self.head_dim, eps=1e-5) self.norm_k = RMSNorm(self.head_dim, eps=1e-5) - self.to_out = RowParallelLinear(dim, dim, bias=False) + self.to_out = RowParallelLinear( + dim, + dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "to_out"), + ) self.num_local_heads = self.to_q.output_size_per_partition // self.head_dim self.num_local_kv_heads = self.to_k.output_size_per_partition // self.head_dim @@ -492,28 +584,89 @@ class BooguImageJointAttention(nn.Module): `img_instruct_attn.to_out[0]`. """ - def __init__(self, dim: int, num_attention_heads: int, num_kv_heads: int) -> None: + def __init__( + self, + dim: int, + num_attention_heads: int, + num_kv_heads: int, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", + ) -> None: super().__init__() self.head_dim = dim // num_attention_heads kv_dim = self.head_dim * num_kv_heads - self.img_to_q = ColumnParallelLinear(dim, dim, bias=False) - self.img_to_k = ColumnParallelLinear(dim, kv_dim, bias=False) - self.img_to_v = ColumnParallelLinear(dim, kv_dim, bias=False) + self.img_to_q = ColumnParallelLinear( + dim, + dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "img_to_q"), + ) + self.img_to_k = ColumnParallelLinear( + dim, + kv_dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "img_to_k"), + ) + self.img_to_v = ColumnParallelLinear( + dim, + kv_dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "img_to_v"), + ) - self.instruct_to_q = ColumnParallelLinear(dim, dim, bias=False) - self.instruct_to_k = ColumnParallelLinear(dim, kv_dim, bias=False) - self.instruct_to_v = ColumnParallelLinear(dim, kv_dim, bias=False) + self.instruct_to_q = ColumnParallelLinear( + dim, + dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "instruct_to_q"), + ) + self.instruct_to_k = ColumnParallelLinear( + dim, + kv_dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "instruct_to_k"), + ) + self.instruct_to_v = ColumnParallelLinear( + dim, + kv_dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "instruct_to_v"), + ) self.norm_q = RMSNorm(self.head_dim, eps=1e-5) self.norm_k = RMSNorm(self.head_dim, eps=1e-5) # Per-stream output projections (attention output is head-sharded # under TP, hence row-parallel). - self.instruct_out = RowParallelLinear(dim, dim, bias=False) - self.img_out = RowParallelLinear(dim, dim, bias=False) + self.instruct_out = RowParallelLinear( + dim, + dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "instruct_out"), + ) + self.img_out = RowParallelLinear( + dim, + dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "img_out"), + ) # Final joint projection applied to the merged full-dim sequence. - self.to_out = ReplicatedLinear(dim, dim, bias=False) + self.to_out = ReplicatedLinear( + dim, + dim, + bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "to_out"), + ) self.num_local_heads = self.img_to_q.output_size_per_partition // self.head_dim self.num_local_kv_heads = self.img_to_k.output_size_per_partition // self.head_dim @@ -594,21 +747,36 @@ class BooguImageTransformerBlock(nn.Module): ffn_dim_multiplier: float | None, norm_eps: float, modulation: bool = True, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ) -> None: super().__init__() self.head_dim = dim // num_attention_heads self.modulation = modulation - self.attn = BooguImageSelfAttention(dim, num_attention_heads, num_kv_heads) + self.attn = BooguImageSelfAttention( + dim, + num_attention_heads, + num_kv_heads, + quant_config=quant_config, + prefix=_join_prefix(prefix, "attn"), + ) self.feed_forward = LuminaFeedForward( dim=dim, inner_dim=4 * dim, multiple_of=multiple_of, ffn_dim_multiplier=ffn_dim_multiplier, + quant_config=quant_config, + prefix=_join_prefix(prefix, "feed_forward"), ) if modulation: - self.norm1 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps) + self.norm1 = LuminaRMSNormZero( + embedding_dim=dim, + norm_eps=norm_eps, + quant_config=quant_config, + prefix=_join_prefix(prefix, "norm1"), + ) else: self.norm1 = RMSNorm(dim, eps=norm_eps) @@ -671,6 +839,8 @@ class BooguImageDoubleStreamTransformerBlock(nn.Module): ffn_dim_multiplier: float | None, norm_eps: float, modulation: bool = True, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ) -> None: super().__init__() self.head_dim = dim // num_attention_heads @@ -678,21 +848,50 @@ class BooguImageDoubleStreamTransformerBlock(nn.Module): self.modulation = modulation self.hidden_size = dim - self.img_instruct_attn = BooguImageJointAttention(dim, num_attention_heads, num_kv_heads) - self.img_self_attn = BooguImageSelfAttention(dim, num_attention_heads, num_kv_heads) + self.img_instruct_attn = BooguImageJointAttention( + dim, + num_attention_heads, + num_kv_heads, + quant_config=quant_config, + prefix=_join_prefix(prefix, "img_instruct_attn"), + ) + self.img_self_attn = BooguImageSelfAttention( + dim, + num_attention_heads, + num_kv_heads, + quant_config=quant_config, + prefix=_join_prefix(prefix, "img_self_attn"), + ) self.img_feed_forward = LuminaFeedForward( dim=dim, inner_dim=4 * dim, multiple_of=multiple_of, ffn_dim_multiplier=ffn_dim_multiplier, + quant_config=quant_config, + prefix=_join_prefix(prefix, "img_feed_forward"), ) if modulation: # Image modulation terms: cross-attn, MLP, self-attn. - self.img_norm1 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps) - self.img_norm2 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps) - self.img_norm3 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps) + self.img_norm1 = LuminaRMSNormZero( + embedding_dim=dim, + norm_eps=norm_eps, + quant_config=quant_config, + prefix=_join_prefix(prefix, "img_norm1"), + ) + self.img_norm2 = LuminaRMSNormZero( + embedding_dim=dim, + norm_eps=norm_eps, + quant_config=quant_config, + prefix=_join_prefix(prefix, "img_norm2"), + ) + self.img_norm3 = LuminaRMSNormZero( + embedding_dim=dim, + norm_eps=norm_eps, + quant_config=quant_config, + prefix=_join_prefix(prefix, "img_norm3"), + ) else: self.img_norm1 = RMSNorm(dim, eps=norm_eps) self.img_norm2 = RMSNorm(dim, eps=norm_eps) @@ -708,12 +907,24 @@ class BooguImageDoubleStreamTransformerBlock(nn.Module): inner_dim=4 * dim, multiple_of=multiple_of, ffn_dim_multiplier=ffn_dim_multiplier, + quant_config=quant_config, + prefix=_join_prefix(prefix, "instruct_feed_forward"), ) if modulation: # Instruction modulation terms: cross-attn, MLP. - self.instruct_norm1 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps) - self.instruct_norm2 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps) + self.instruct_norm1 = LuminaRMSNormZero( + embedding_dim=dim, + norm_eps=norm_eps, + quant_config=quant_config, + prefix=_join_prefix(prefix, "instruct_norm1"), + ) + self.instruct_norm2 = LuminaRMSNormZero( + embedding_dim=dim, + norm_eps=norm_eps, + quant_config=quant_config, + prefix=_join_prefix(prefix, "instruct_norm2"), + ) else: self.instruct_norm1 = RMSNorm(dim, eps=norm_eps) self.instruct_norm2 = RMSNorm(dim, eps=norm_eps) @@ -859,7 +1070,12 @@ class BooguImageTransformer2DModel(nn.Module): ] _layerwise_offload_blocks_attrs = ["single_stream_layers", "double_stream_layers"] - def __init__(self, od_config: OmniDiffusionConfig) -> None: + def __init__( + self, + od_config: OmniDiffusionConfig, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", + ) -> None: super().__init__() self.od_config = od_config cfg = od_config.tf_model_config @@ -946,8 +1162,10 @@ class BooguImageTransformer2DModel(nn.Module): ffn_dim_multiplier, norm_eps, modulation=True, + quant_config=quant_config, + prefix=_join_prefix(prefix, f"noise_refiner.{i}"), ) - for _ in range(num_refiner_layers) + for i in range(num_refiner_layers) ] ) @@ -961,8 +1179,10 @@ class BooguImageTransformer2DModel(nn.Module): ffn_dim_multiplier, norm_eps, modulation=True, + quant_config=quant_config, + prefix=_join_prefix(prefix, f"ref_image_refiner.{i}"), ) - for _ in range(num_refiner_layers) + for i in range(num_refiner_layers) ] ) @@ -976,8 +1196,10 @@ class BooguImageTransformer2DModel(nn.Module): ffn_dim_multiplier, norm_eps, modulation=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, f"context_refiner.{i}"), ) - for _ in range(num_refiner_layers) + for i in range(num_refiner_layers) ] ) @@ -992,8 +1214,10 @@ class BooguImageTransformer2DModel(nn.Module): ffn_dim_multiplier, norm_eps, modulation=True, + quant_config=quant_config, + prefix=_join_prefix(prefix, f"double_stream_layers.{i}"), ) - for _ in range(num_double_stream_layers) + for i in range(num_double_stream_layers) ] ) @@ -1007,8 +1231,10 @@ class BooguImageTransformer2DModel(nn.Module): ffn_dim_multiplier, norm_eps, modulation=True, + quant_config=quant_config, + prefix=_join_prefix(prefix, f"single_stream_layers.{i}"), ) - for _ in range(self.num_single_stream_layers) + for i in range(self.num_single_stream_layers) ] ) @@ -1019,6 +1245,8 @@ class BooguImageTransformer2DModel(nn.Module): eps=1e-6, bias=True, out_dim=patch_size * patch_size * self.out_channels, + quant_config=quant_config, + prefix=_join_prefix(prefix, "norm_out"), ) # Distinguish multiple reference images (max 5). 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 257e4e21a..c82e6ad7b 100644 --- a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py +++ b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py @@ -57,6 +57,7 @@ from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin 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 +from vllm_omni.quantization.component_config import resolve_component_quant_config logger = init_logger(__name__) @@ -241,6 +242,7 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports super().__init__() self.od_config = od_config self._raise_unsupported_features() + transformer_quant_config = resolve_component_quant_config(od_config.quantization_config, "transformer") self.weights_sources = [ DiffusersPipelineLoader.ComponentSource( model_or_path=od_config.model, @@ -304,7 +306,11 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports revision=od_config.revision, ).to(self._execution_device) - self.transformer = BooguImageTransformer2DModel(od_config=od_config) + self.transformer = BooguImageTransformer2DModel( + od_config=od_config, + quant_config=transformer_quant_config, + prefix="transformer", + ) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 self.default_sample_size = 128 diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index f3b6e8cd2..8f15aaf50 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -146,6 +146,27 @@ def _build_inc(**kw: Any) -> QuantizationConfig: return OmniINCConfig(**filtered) +def _build_torchao(**kw: Any) -> QuantizationConfig: + """Build a TorchAO runtime or serialized-checkpoint config.""" + from vllm.model_executor.layers.quantization.torchao import TorchAOConfig + + if "quant_type" in kw: + return TorchAOConfig.from_config({**kw, "quant_method": "torchao"}) + return TorchAOConfig(**kw) + + +def _build_torchao_float8_weight_only(**kw: Any) -> QuantizationConfig: + """Build the serialized TorchAO FP8 weight-only checkpoint config.""" + from torchao.quantization import Float8WeightOnlyConfig + + return _build_torchao( + torchao_config=Float8WeightOnlyConfig( + set_inductor_config=False, + ), + is_checkpoint_torchao_serialized=True, + ) + + _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = { "int8": _build_int8, "bitsandbytes": _build_bitsandbytes, @@ -156,6 +177,8 @@ _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = { "inc": _build_inc, "auto-round": _build_inc, "auto_round": _build_inc, + "torchao": _build_torchao, + "torchao_float8_weight_only": _build_torchao_float8_weight_only, }