[diffusion] refactor: reuse plain state-dict loading without per-model classes (#38127)

This commit is contained in:
Mick
2026-09-06 18:39:21 +08:00
committed by GitHub
parent 97c6978369
commit 938dc5621d
5 changed files with 262 additions and 7 deletions
@@ -359,6 +359,54 @@ class MyModelPipeline(LoRAPipeline, ComposedPipelineBase):
EntryClass = [MyModelPipeline]
```
#### Reuse component loaders
Most components should keep the default loader for their role (transformer,
text encoder, VAE, scheduler, or tokenizer). For an auxiliary module that accepts
`model_cls(**config)` and loads an unchanged state dict, select
`PlainStateDictComponentLoader` instead of adding a model-specific loader class:
```python
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
PlainStateDictComponentLoader,
)
class MyModelPipeline(ComposedPipelineBase):
_required_config_modules = ["transformer", "queryformer", "text_projection"]
component_loaders = {
"queryformer": PlainStateDictComponentLoader,
"text_projection": PlainStateDictComponentLoader,
}
# Wire the stages as in the examples above.
```
The mapping uses exact pipeline component names, including aliases declared in
`_extra_config_module_map`. It is local to this pipeline; omitted components keep
their existing dispatch. An explicitly selected loader fails on errors instead
of falling back to a different implementation.
Register each module through `EntryClass` in its model file, or
`ModelRegistry.register_model` for an out-of-tree package. The shared loader:
- Resolves the registered class from `config.json`'s `_class_name`, falling back
to the architecture in `model_index.json`, and passes non-metadata config
fields to its constructor.
- Loads a single safetensors file or an indexed sharded checkpoint using the
shared checkpoint selector, with strict key and shape checks.
- Honors exact overrides such as `--component-weights-paths.queryformer PATH`
and `--component-precisions.queryformer bf16`. Precision defaults to the
pipeline's `dit_precision`.
- Leaves eval mode and final CPU/GPU placement to the shared loading lifecycle.
It does not implement TP/FSDP weight sharding, quantization, or direct GPU
loading. Explicit quantization and direct-GPU-loading overrides are rejected.
Keep a specialized loader when the component needs a config-object constructor,
weight remapping or fusion, sharding, or a non-standard checkpoint format. Such
a loader can also be selected through `component_loaders`; registering a model
alone does not opt it into the plain state-dict protocol.
### 4. Last-Resort Before-Denoising Stage
A `BeforeDenoisingStage` is not a catch-all replacement for the native stages.