[diffusion] feat: admit compatible quantized native encoders (#35962)

This commit is contained in:
Mick
2026-08-22 18:51:01 +08:00
committed by GitHub
parent 382343f860
commit 489e605b35
8 changed files with 66 additions and 105 deletions
+1 -1
View File
@@ -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
+1 -22
View File
@@ -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
@@ -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)