[Model] FP8 Support For Boogu-Image (#6925)
Signed-off-by: Dmaner <2663515256@qq.com>
This commit is contained in:
15 files changed
+809
-78
No files matched your search
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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. |
|
||||
|
||||
@@ -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.
|
||||
@@ -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 `<MODEL_ID>` 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": "<MODEL_ID>",
|
||||
"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="<MODEL_ID>" \
|
||||
-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`);
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"),
|
||||
[
|
||||
|
||||
@@ -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
|
||||
@@ -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",
|
||||
|
||||
@@ -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<family>.+)-\d+-of-(?P<count>\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,
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user