[diffusion] feat: admit compatible quantized native encoders (#35962)
This commit is contained in:
@@ -139,7 +139,7 @@ or a native pipeline's registered module name.
|
|||||||
For a native text encoder:
|
For a native text encoder:
|
||||||
|
|
||||||
- `--component-paths.text_encoder {MODEL}` replaces the text-encoder checkpoint; `--text-encoder-path {MODEL}` is its shorter alias
|
- `--component-paths.text_encoder {MODEL}` replaces the text-encoder checkpoint; `--text-encoder-path {MODEL}` is its shorter alias
|
||||||
- Quantization metadata is auto-detected from that checkpoint. Each native encoder must explicitly support the serialized format; this is not blanket quantization support for every component, and unsupported combinations fail before weight loading.
|
- Quantization metadata is auto-detected from that checkpoint. The native loader accepts compatible serialized formats without a model-name allowlist and rejects implementations that do not construct the required quantized layers.
|
||||||
|
|
||||||
The same contract applies to every weighted component: path routing is generic,
|
The same contract applies to every weighted component: path routing is generic,
|
||||||
while quantized materialization is capability-based. Native auxiliary loaders
|
while quantized materialization is capability-based. Native auxiliary loaders
|
||||||
|
|||||||
@@ -70,7 +70,7 @@ paths:
|
|||||||
| Component path | Quantized checkpoint behavior |
|
| Component path | Quantized checkpoint behavior |
|
||||||
| --- | --- |
|
| --- | --- |
|
||||||
| `transformer`, `transformer_2`, `unconditional_transformer`, `audio_dit`, `video_dit` | Uses the SGLang transformer quantization adapters documented below. |
|
| `transformer`, `transformer_2`, `unconditional_transformer`, `audio_dit`, `video_dit` | Uses the SGLang transformer quantization adapters documented below. |
|
||||||
| `text_encoder*`, `image_encoder*` | Requires the native encoder class to declare support for the detected format. MiniMax-H3 FP8 and model-managed integrations such as Ideogram are supported; unknown combinations fail closed. |
|
| `text_encoder*`, `image_encoder*` | Passes detected metadata to the registered native encoder. Loading continues only when that implementation constructs compatible quantized layers; unknown or incompatible combinations fail closed. |
|
||||||
| `vae`, `video_vae`, `audio_vae` | A standard top-level Diffusers `quantization_config` is delegated to `AutoModel.from_pretrained`. Native-only VAEs and nested/compression metadata fail closed. |
|
| `vae`, `video_vae`, `audio_vae` | A standard top-level Diffusers `quantization_config` is delegated to `AutoModel.from_pretrained`. Native-only VAEs and nested/compression metadata fail closed. |
|
||||||
| Library-managed Transformers or Diffusers components | Delegates to the upstream `from_pretrained` path and inherits its format support and validation behavior. The local PE model uses this path; compatible formats remain model-specific. |
|
| Library-managed Transformers or Diffusers components | Delegates to the upstream `from_pretrained` path and inherits its format support and validation behavior. The local PE model uses this path; compatible formats remain model-specific. |
|
||||||
| Native auxiliary components that load raw state dicts | Quantized checkpoints are rejected before model construction until that component has a quantized materialization implementation. This includes connectors, duration heads, bridges, diffusion decoders, sound tokenizers, spatial upsamplers, and vocoders. |
|
| Native auxiliary components that load raw state dicts | Quantized checkpoints are rejected before model construction until that component has a quantized materialization implementation. This includes connectors, duration heads, bridges, diffusion decoders, sound tokenizers, spatial upsamplers, and vocoders. |
|
||||||
@@ -402,27 +402,6 @@ sglang generate \
|
|||||||
|
|
||||||
Each pattern is matched against the full layer prefix (e.g. `layers.0.attention.to_q`). A layer is skipped and left unquantized if its prefix contains any of the given patterns.
|
Each pattern is matched against the full layer prefix (e.g. `layers.0.attention.to_q`). A layer is skipped and left unquantized if its prefix contains any of the given patterns.
|
||||||
|
|
||||||
## MiniMax-H3 Text Encoder FP8
|
|
||||||
|
|
||||||
MiniMax-H3 can load a serialized FP8 checkpoint for the language linear layers
|
|
||||||
in its native Qwen3-VL text encoder independently of the DiT. Embeddings,
|
|
||||||
normalization layers, and the Qwen vision tower remain in BF16.
|
|
||||||
|
|
||||||
```bash Command
|
|
||||||
sglang serve \
|
|
||||||
--model-path MiniMaxAI/MiniMax-H3 \
|
|
||||||
--model-variant fl2va \
|
|
||||||
--component-paths.text_encoder Qwen/Qwen3-VL-32B-Instruct-FP8 \
|
|
||||||
--num-gpus 4 \
|
|
||||||
--port 30010
|
|
||||||
```
|
|
||||||
|
|
||||||
`--text-encoder-path` is accepted as a shorter alias. No quantization flag is
|
|
||||||
needed: SGLang detects the checkpoint metadata and only enables formats that
|
|
||||||
the native encoder explicitly supports. Text-encoder FP8 is approximate, is
|
|
||||||
not enabled by default, and is rejected by MiniMax-H3's strict
|
|
||||||
`quality="high"` deployment contract.
|
|
||||||
|
|
||||||
## Transformers Component BnB4
|
## Transformers Component BnB4
|
||||||
|
|
||||||
Model components that already have a native Transformers loading path can load
|
Model components that already have a native Transformers loading path can load
|
||||||
|
|||||||
+20
-21
@@ -129,27 +129,6 @@ def _configure_encoder_quantization(
|
|||||||
f"got {model_cls.__name__}"
|
f"got {model_cls.__name__}"
|
||||||
)
|
)
|
||||||
|
|
||||||
capability = model_cls.checkpoint_quantization_capability
|
|
||||||
if capability is None:
|
|
||||||
raise ComponentCheckpointUnsupportedError(
|
|
||||||
f"{model_cls.__name__} does not support quantized checkpoints for "
|
|
||||||
f"{component_name!r}: no checkpoint quantization capability is declared"
|
|
||||||
)
|
|
||||||
if capability.backend != "diffusion":
|
|
||||||
raise ComponentCheckpointUnsupportedError(
|
|
||||||
f"{model_cls.__name__} declares the {capability.backend!r} checkpoint "
|
|
||||||
f"quantization backend for {component_name!r}, but the native encoder "
|
|
||||||
"loader currently supports only the 'diffusion' backend"
|
|
||||||
)
|
|
||||||
|
|
||||||
quant_method = quant_config.get_name()
|
|
||||||
if quant_method not in capability.methods:
|
|
||||||
raise ComponentCheckpointUnsupportedError(
|
|
||||||
f"{model_cls.__name__} does not support {component_name!r} checkpoints "
|
|
||||||
f"quantized with {quant_method!r}; supported methods for the "
|
|
||||||
f"{capability.backend!r} backend: {sorted(capability.methods)}"
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def _resolve_and_configure_encoder_quantization(
|
def _resolve_and_configure_encoder_quantization(
|
||||||
model_config: EncoderConfig,
|
model_config: EncoderConfig,
|
||||||
@@ -244,6 +223,23 @@ def _process_quantized_encoder_weights(
|
|||||||
return processed_layers
|
return processed_layers
|
||||||
|
|
||||||
|
|
||||||
|
def _require_quantized_encoder_layers(
|
||||||
|
model: nn.Module,
|
||||||
|
component_name: str,
|
||||||
|
) -> None:
|
||||||
|
if any(
|
||||||
|
isinstance(module, LinearBase)
|
||||||
|
and module.quant_method is not None
|
||||||
|
and not isinstance(module.quant_method, UnquantizedLinearMethod)
|
||||||
|
for module in model.modules()
|
||||||
|
):
|
||||||
|
return
|
||||||
|
raise ComponentCheckpointUnsupportedError(
|
||||||
|
f"The native {type(model).__name__} implementation does not construct "
|
||||||
|
f"quantized linear layers for {component_name!r}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def _checkpoint_bytes(model_path: str) -> int:
|
def _checkpoint_bytes(model_path: str) -> int:
|
||||||
"""On-disk size of a checkpoint, readable before any weight of it is."""
|
"""On-disk size of a checkpoint, readable before any weight of it is."""
|
||||||
total = 0
|
total = 0
|
||||||
@@ -624,6 +620,9 @@ class TextEncoderLoader(ComponentLoader):
|
|||||||
)
|
)
|
||||||
model.bind_encoder_tp_group(encoder_tp_group)
|
model.bind_encoder_tp_group(encoder_tp_group)
|
||||||
|
|
||||||
|
if quant_config is not None:
|
||||||
|
_require_quantized_encoder_layers(model, component_name)
|
||||||
|
|
||||||
if component_starts_on_cpu and (
|
if component_starts_on_cpu and (
|
||||||
current_platform.is_mps() or _keep_this_checkpoint_mapped(model_path)
|
current_platform.is_mps() or _keep_this_checkpoint_mapped(model_path)
|
||||||
):
|
):
|
||||||
|
|||||||
@@ -2,8 +2,7 @@
|
|||||||
|
|
||||||
# SPDX-License-Identifier: Apache-2.0
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import field
|
||||||
from typing import Literal
|
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
from torch import nn
|
from torch import nn
|
||||||
@@ -155,19 +154,10 @@ def finalize_encoder_folding(
|
|||||||
config.parallel_folding_mode = None
|
config.parallel_folding_mode = None
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
|
||||||
class CheckpointQuantizationCapability:
|
|
||||||
"""Quantized-checkpoint contract implemented by a native encoder."""
|
|
||||||
|
|
||||||
backend: Literal["diffusion", "srt"]
|
|
||||||
methods: frozenset[str]
|
|
||||||
|
|
||||||
|
|
||||||
class EncoderTensorParallelMixin:
|
class EncoderTensorParallelMixin:
|
||||||
"""Keep an encoder on the TP group that was used to build its shards."""
|
"""Keep an encoder on the TP group that was used to build its shards."""
|
||||||
|
|
||||||
_encoder_tp_group: GroupCoordinator | None = None
|
_encoder_tp_group: GroupCoordinator | None = None
|
||||||
checkpoint_quantization_capability: CheckpointQuantizationCapability | None = None
|
|
||||||
# Some encoders own checkpoint quantization end to end because their weight
|
# Some encoders own checkpoint quantization end to end because their weight
|
||||||
# states or sharding contract cannot use the generic loader lifecycle.
|
# states or sharding contract cannot use the generic loader lifecycle.
|
||||||
manages_checkpoint_quantization = False
|
manages_checkpoint_quantization = False
|
||||||
@@ -192,9 +182,6 @@ class TextEncoder(
|
|||||||
# Qwen2_5_VLCausalLMOutputWithPast). Off by default so a new encoder is
|
# Qwen2_5_VLCausalLMOutputWithPast). Off by default so a new encoder is
|
||||||
# replicated rather than silently broken; flip it once dp is verified there.
|
# replicated rather than silently broken; flip it once dp is verified there.
|
||||||
supports_dp_encode = False
|
supports_dp_encode = False
|
||||||
# Quantized checkpoints are opt-in because an encoder must construct
|
|
||||||
# quantized linears and load the checkpoint's auxiliary scale parameters.
|
|
||||||
supported_checkpoint_quantization_methods: frozenset[str] = frozenset()
|
|
||||||
# Some encoders own checkpoint quantization end to end because their weight
|
# Some encoders own checkpoint quantization end to end because their weight
|
||||||
# states or sharding contract cannot use the generic loader lifecycle.
|
# states or sharding contract cannot use the generic loader lifecycle.
|
||||||
manages_checkpoint_quantization = False
|
manages_checkpoint_quantization = False
|
||||||
|
|||||||
@@ -17,10 +17,7 @@ from sglang.multimodal_gen.configs.models.encoders.minimax_h3_qwen3vl import (
|
|||||||
)
|
)
|
||||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||||
from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
|
from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
|
||||||
from sglang.multimodal_gen.runtime.models.encoders.base import (
|
from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder
|
||||||
CheckpointQuantizationCapability,
|
|
||||||
TextEncoder,
|
|
||||||
)
|
|
||||||
from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLModel
|
from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLModel
|
||||||
|
|
||||||
MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120
|
MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120
|
||||||
@@ -49,10 +46,6 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder):
|
|||||||
layer_names = [*TextEncoder.layer_names, "model.visual.blocks"]
|
layer_names = [*TextEncoder.layer_names, "model.visual.blocks"]
|
||||||
|
|
||||||
supports_dp_encode = True
|
supports_dp_encode = True
|
||||||
checkpoint_quantization_capability = CheckpointQuantizationCapability(
|
|
||||||
backend="diffusion",
|
|
||||||
methods=frozenset({"fp8"}),
|
|
||||||
)
|
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def should_materialize_checkpoint_weight(name: str) -> bool:
|
def should_materialize_checkpoint_weight(name: str) -> bool:
|
||||||
|
|||||||
@@ -52,14 +52,6 @@ class TestImageEncoderQuantizationAdmission(unittest.TestCase):
|
|||||||
"/model/image_encoder", self.server_args, "image_encoder", "transformers"
|
"/model/image_encoder", self.server_args, "image_encoder", "transformers"
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_quantized_clip_checkpoint_is_not_silently_enabled(self):
|
|
||||||
config = self._component_config("CLIPVisionModelWithProjection", quantized=True)
|
|
||||||
with self._config_patch(config), self.assertRaisesRegex(
|
|
||||||
ComponentCheckpointUnsupportedError,
|
|
||||||
"CLIPVisionModel.*image_encoder.*no checkpoint quantization capability",
|
|
||||||
):
|
|
||||||
self.loader.load_customized("/model/image_encoder", self.server_args)
|
|
||||||
|
|
||||||
def test_unknown_quantized_architecture_does_not_fall_back(self):
|
def test_unknown_quantized_architecture_does_not_fall_back(self):
|
||||||
config = self._component_config("UnknownVisionModel", quantized=True)
|
config = self._component_config("UnknownVisionModel", quantized=True)
|
||||||
with self._config_patch(config), self.assertRaises(
|
with self._config_patch(config), self.assertRaises(
|
||||||
|
|||||||
@@ -98,3 +98,31 @@ def test_default_position_ids_batch_shape():
|
|||||||
assert torch.equal(layer.position_ids[0], torch.arange(4))
|
assert torch.equal(layer.position_ids[0], torch.arange(4))
|
||||||
assert torch.equal(layer.position_ids[1], torch.arange(4))
|
assert torch.equal(layer.position_ids[1], torch.arange(4))
|
||||||
assert layer.attention_lengths == (4, 4)
|
assert layer.attention_lengths == (4, 4)
|
||||||
|
|
||||||
|
|
||||||
|
def test_fp8_qkv_scale_uses_the_packed_parameter_loader():
|
||||||
|
model = Qwen3ForCausalLM.__new__(Qwen3ForCausalLM)
|
||||||
|
torch.nn.Module.__init__(model)
|
||||||
|
layer = torch.nn.Module()
|
||||||
|
layer.self_attn = torch.nn.Module()
|
||||||
|
layer.self_attn.qkv_proj = torch.nn.Module()
|
||||||
|
scale = torch.nn.Parameter(torch.zeros(3, 1), requires_grad=False)
|
||||||
|
|
||||||
|
def load_scale(param, loaded_scale, shard_id):
|
||||||
|
param.data[{"q": 0, "k": 1, "v": 2}[shard_id]].copy_(loaded_scale)
|
||||||
|
|
||||||
|
scale.weight_loader = load_scale
|
||||||
|
layer.self_attn.qkv_proj.register_parameter("weight_scale_inv", scale)
|
||||||
|
model.layers = torch.nn.ModuleList([layer])
|
||||||
|
model.config = SimpleNamespace(
|
||||||
|
arch_config=SimpleNamespace(
|
||||||
|
stacked_params_mapping=[(".qkv_proj", ".q_proj", "q")]
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
loaded = model.load_weights(
|
||||||
|
[("model.layers.0.self_attn.q_proj.weight_scale_inv", torch.tensor([2.0]))]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert loaded == {"layers.0.self_attn.qkv_proj.weight_scale_inv"}
|
||||||
|
torch.testing.assert_close(scale[:, 0], torch.tensor([2.0, 0.0, 0.0]))
|
||||||
|
|||||||
@@ -16,12 +16,10 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader
|
|||||||
TextEncoderLoader,
|
TextEncoderLoader,
|
||||||
_configure_encoder_quantization,
|
_configure_encoder_quantization,
|
||||||
_process_quantized_encoder_weights,
|
_process_quantized_encoder_weights,
|
||||||
|
_require_quantized_encoder_layers,
|
||||||
_resolve_and_configure_encoder_quantization,
|
_resolve_and_configure_encoder_quantization,
|
||||||
)
|
)
|
||||||
from sglang.multimodal_gen.runtime.models.encoders.base import (
|
from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder
|
||||||
CheckpointQuantizationCapability,
|
|
||||||
TextEncoder,
|
|
||||||
)
|
|
||||||
from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import (
|
from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import (
|
||||||
MiniMaxH3Qwen3VLEncoder,
|
MiniMaxH3Qwen3VLEncoder,
|
||||||
)
|
)
|
||||||
@@ -182,25 +180,25 @@ class TestTextEncoderQuantization(unittest.TestCase):
|
|||||||
self.addCleanup(self.quant_config_patcher.stop)
|
self.addCleanup(self.quant_config_patcher.stop)
|
||||||
self.serialized = serialized
|
self.serialized = serialized
|
||||||
|
|
||||||
def test_serialized_fp8_checkpoint_configures_h3_encoder(self):
|
def test_serialized_checkpoint_configures_native_encoder(self):
|
||||||
model_config = SimpleNamespace(quant_config=None)
|
model_config = SimpleNamespace(quant_config=None)
|
||||||
_configure_encoder_quantization(
|
_configure_encoder_quantization(
|
||||||
model_config,
|
model_config,
|
||||||
MiniMaxH3Qwen3VLEncoder,
|
TextEncoder,
|
||||||
{},
|
{},
|
||||||
"/model/text_encoder",
|
"/model/text_encoder",
|
||||||
"text_encoder",
|
"text_encoder",
|
||||||
)
|
)
|
||||||
self.assertIs(model_config.quant_config, self.serialized)
|
self.assertIs(model_config.quant_config, self.serialized)
|
||||||
|
|
||||||
def test_encoder_class_must_opt_in(self):
|
def test_encoder_must_use_native_loader(self):
|
||||||
model_config = SimpleNamespace(quant_config=None)
|
model_config = SimpleNamespace(quant_config=None)
|
||||||
with self.assertRaisesRegex(
|
with self.assertRaisesRegex(
|
||||||
ComponentCheckpointUnsupportedError, "does not support"
|
ComponentCheckpointUnsupportedError, "requires an in-tree native encoder"
|
||||||
):
|
):
|
||||||
_configure_encoder_quantization(
|
_configure_encoder_quantization(
|
||||||
model_config,
|
model_config,
|
||||||
TextEncoder,
|
nn.Module,
|
||||||
{},
|
{},
|
||||||
"/model/text_encoder",
|
"/model/text_encoder",
|
||||||
"text_encoder",
|
"text_encoder",
|
||||||
@@ -269,28 +267,6 @@ class TestTextEncoderQuantization(unittest.TestCase):
|
|||||||
"text_encoder",
|
"text_encoder",
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_srt_backend_is_not_admitted_without_an_adapter(self):
|
|
||||||
model_config = SimpleNamespace(quant_config=None)
|
|
||||||
capability = CheckpointQuantizationCapability(
|
|
||||||
backend="srt",
|
|
||||||
methods=frozenset({"fp8"}),
|
|
||||||
)
|
|
||||||
with mock.patch.object(
|
|
||||||
MiniMaxH3Qwen3VLEncoder,
|
|
||||||
"checkpoint_quantization_capability",
|
|
||||||
capability,
|
|
||||||
), self.assertRaisesRegex(
|
|
||||||
ComponentCheckpointUnsupportedError,
|
|
||||||
"'srt'.*only the 'diffusion' backend",
|
|
||||||
):
|
|
||||||
_configure_encoder_quantization(
|
|
||||||
model_config,
|
|
||||||
MiniMaxH3Qwen3VLEncoder,
|
|
||||||
{},
|
|
||||||
"/model/text_encoder",
|
|
||||||
"text_encoder",
|
|
||||||
)
|
|
||||||
|
|
||||||
def test_model_managed_quantization_bypasses_generic_lifecycle(self):
|
def test_model_managed_quantization_bypasses_generic_lifecycle(self):
|
||||||
model_config = SimpleNamespace(quant_config=None)
|
model_config = SimpleNamespace(quant_config=None)
|
||||||
with mock.patch.object(
|
with mock.patch.object(
|
||||||
@@ -341,6 +317,13 @@ class _QuantizedEncoder(nn.Module):
|
|||||||
|
|
||||||
|
|
||||||
class TestQuantizedTextEncoderPostprocess(unittest.TestCase):
|
class TestQuantizedTextEncoderPostprocess(unittest.TestCase):
|
||||||
|
def test_rejects_native_encoder_without_quantized_layers(self):
|
||||||
|
with self.assertRaisesRegex(
|
||||||
|
ComponentCheckpointUnsupportedError,
|
||||||
|
"does not construct quantized linear layers",
|
||||||
|
):
|
||||||
|
_require_quantized_encoder_layers(nn.Linear(2, 2), "text_encoder")
|
||||||
|
|
||||||
def test_processes_quantized_layers_without_moving_the_model(self):
|
def test_processes_quantized_layers_without_moving_the_model(self):
|
||||||
quant_method = _RecordingQuantMethod()
|
quant_method = _RecordingQuantMethod()
|
||||||
model = _QuantizedEncoder(quant_method)
|
model = _QuantizedEncoder(quant_method)
|
||||||
|
|||||||
Reference in New Issue
Block a user