diff --git a/docs/docs/sglang-diffusion/support_new_models.mdx b/docs/docs/sglang-diffusion/support_new_models.mdx index 1ad630beb..1dabf8da5 100644 --- a/docs/docs/sglang-diffusion/support_new_models.mdx +++ b/docs/docs/sglang-diffusion/support_new_models.mdx @@ -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. diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py index 646e6e6ca..e6009b67e 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py @@ -8,7 +8,7 @@ import pkgutil import traceback from abc import ABC from collections.abc import Callable, Iterator -from typing import Any, Type +from typing import Any import torch import transformers @@ -328,6 +328,7 @@ class ComponentLoader(ABC): *, component_attn_backend: Any = None, component_attn_name: str | None = None, + allow_native_fallback: bool = True, ) -> tuple[AutoModel, float]: """ Template method that standardizes logging around the core load implementation. @@ -399,7 +400,7 @@ class ComponentLoader(ABC): ): raise except Exception as e: - if require_backend_selection: + if require_backend_selection or not allow_native_fallback: raise native_loader_required = isinstance(e, NativeComponentLoaderRequired) if native_loader_required and component_weight_override is not None: @@ -649,6 +650,8 @@ class ComponentLoader(ABC): component_type: str, transformers_or_diffusers: str, component_architecture: str | None = None, + *, + loader_cls: type["ComponentLoader"] | None = None, ) -> "ComponentLoader": """ Factory method to create a component loader for a specific component type. @@ -667,10 +670,9 @@ class ComponentLoader(ABC): transformers_or_diffusers, loader_type ) - if loader_type in component_name_to_loader_cls: - loader_cls: Type[ComponentLoader] = component_name_to_loader_cls[ - loader_type - ] + if loader_cls is None: + loader_cls = component_name_to_loader_cls.get(loader_type) + if loader_cls is not None: expected_library = loader_cls.expected_library # Assert that the library matches what's expected for this component type assert transformers_or_diffusers == expected_library, ( @@ -1013,6 +1015,7 @@ class PipelineComponentLoader: component_attn_backend: Any = None, component_attn_name: str | None = None, component_type: str | None = None, + loader_cls: type[ComponentLoader] | None = None, ): """ Load a pipeline component. @@ -1023,6 +1026,7 @@ class PipelineComponentLoader: transformers_or_diffusers: Whether the component is from transformers or diffusers component_architecture: the class name of the module component_type: structural config slot when it differs from the exact key + loader_cls: explicit pipeline-local loader, with no native fallback """ # Get the appropriate loader for this component type @@ -1030,6 +1034,7 @@ class PipelineComponentLoader: component_type or component_name, transformers_or_diffusers, component_architecture, + loader_cls=loader_cls, ) try: @@ -1040,6 +1045,7 @@ class PipelineComponentLoader: transformers_or_diffusers, component_attn_backend=component_attn_backend, component_attn_name=component_attn_name, + allow_native_fallback=loader_cls is None, ) except Exception: logger.error( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index 8d601c781..5e2545b99 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -9,7 +9,7 @@ This module defines the base class for pipelines that are composed of multiple s import os from abc import ABC, abstractmethod -from typing import Any, Callable, Iterator, Literal, cast +from typing import Any, Callable, ClassVar, Iterator, Literal, cast import torch from tqdm import tqdm @@ -19,6 +19,7 @@ from sglang.multimodal_gen.runtime.disaggregation.roles import ( filter_modules_for_role, ) from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( + ComponentLoader, PipelineComponentLoader, ) from sglang.multimodal_gen.runtime.loader.utils import _normalize_component_type @@ -76,6 +77,8 @@ class ComposedPipelineBase(ABC): _required_config_modules: list[str] = [] _unfiltered_required_config_modules: tuple[str, ...] = () _extra_config_module_map: dict[str, str] = {} + # Exact module keys; unspecified components retain the default loader dispatch. + component_loaders: ClassVar[dict[str, type[ComponentLoader]]] = {} server_args: ServerArgs | None = None modules: dict[str, Any] = {} executor: PipelineExecutor | None = None @@ -616,6 +619,7 @@ class ComposedPipelineBase(ABC): module, memory_usage = PipelineComponentLoader.load_component( component_name=module_name, component_type=load_module_name, + loader_cls=self.component_loaders.get(module_name), component_model_path=component_model_path, transformers_or_diffusers=transformers_or_diffusers, server_args=server_args, diff --git a/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py b/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py index a70472e57..ca031421c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py @@ -106,6 +106,7 @@ class TestComponentLoaderIdentity(unittest.TestCase): load_component.assert_called_once_with( component_name="auxiliary_head", component_type="duration_head_2", + loader_cls=None, component_model_path="/model/duration_head_2", transformers_or_diffusers="ltx2", server_args=server_args, diff --git a/python/sglang/multimodal_gen/test/unit/test_plain_component_loader.py b/python/sglang/multimodal_gen/test/unit/test_plain_component_loader.py new file mode 100644 index 000000000..4974a8e4e --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_plain_component_loader.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 + +import json +from unittest.mock import patch + +import pytest +import torch +from safetensors.torch import save_file +from torch import nn + +from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType +from sglang.multimodal_gen.runtime.loader.component_loaders.adapter_loader import ( + AdapterLoader, +) +from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( + ComponentCheckpointUnsupportedError, + ComponentLoader, + GenericComponentLoader, + PipelineComponentLoader, + PlainStateDictComponentLoader, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( + COMPONENT_OFFLOAD, + RESIDENT, +) +from sglang.multimodal_gen.runtime.models.registry import ModelRegistry +from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.server_args import ServerArgs + + +class _ProjectionPipeline(ComposedPipelineBase): + _required_config_modules = ["projection"] + component_loaders = {"projection": PlainStateDictComponentLoader} + + def initialize_pipeline(self, server_args): + pass + + def create_pipeline_stages(self, server_args): + pass + + +@pytest.fixture +def checkpoint(tmp_path): + component = tmp_path / "projection" + component.mkdir() + config = { + "_class_name": "PlainLoaderTestProjection", + "_diffusers_version": "0", + "_name_or_path": "unused", + "in_features": 4, + "out_features": 3, + } + (component / "config.json").write_text(json.dumps(config)) + (tmp_path / "model_index.json").write_text( + json.dumps( + { + "_class_name": "ProjectionPipeline", + "_diffusers_version": "0", + "projection": ["diffusers", "PlainLoaderTestProjection"], + "scheduler": None, + } + ) + ) + weights = {"weight": torch.arange(12).reshape(3, 4).float(), "bias": torch.ones(3)} + save_file(weights, component / "model.safetensors") + with patch.dict(ModelRegistry.registered_models): + ModelRegistry.register_model("PlainLoaderTestProjection", nn.Linear) + yield component, weights + + +@pytest.fixture +def server_args(): + return ServerArgs( + model_path="x", + component_residency={"projection": COMPONENT_OFFLOAD}, + ) + + +def _load(component, server_args, **kwargs): + with patch.object(current_platform, "get_available_gpu_memory", return_value=16.0): + model, _ = PipelineComponentLoader.load_component( + "projection", + str(component), + "diffusers", + server_args, + loader_cls=PlainStateDictComponentLoader, + **kwargs, + ) + return model + + +@pytest.mark.parametrize("sharded", [False, True]) +@pytest.mark.parametrize("precision", [None, "fp32"]) +def test_pipeline_loads_strict_state_dict(checkpoint, server_args, sharded, precision): + component, weights = checkpoint + if sharded: + for key, weight in weights.items(): + save_file({key: weight}, component / f"{key}.safetensors") + (component / "diffusion_pytorch_model.safetensors.index.json").write_text( + json.dumps({"weight_map": {key: f"{key}.safetensors" for key in weights}}) + ) + # The index must win over this unrelated full checkpoint. + save_file({"wrong": torch.zeros(1)}, component / "model.safetensors") + if precision is not None: + server_args.component_precisions["projection"] = precision + + pipeline = object.__new__(_ProjectionPipeline) + pipeline.model_path = str(component.parent) + pipeline.server_args = server_args + pipeline._disagg_role = RoleType.MONOLITHIC + pipeline.memory_usages = {} + with patch.object(current_platform, "get_available_gpu_memory", return_value=16.0): + model = pipeline.load_modules(server_args)["projection"] + + expected_dtype = torch.float32 if precision else torch.bfloat16 + assert not model.training + assert model.weight.device.type == "cpu" + assert model.weight.dtype == expected_dtype + for key, tensor in model.state_dict().items(): + torch.testing.assert_close(tensor, weights[key].to(expected_dtype)) + assert server_args.model_paths == {"projection": str(component)} + + +def test_weight_override_and_architecture_fallback(checkpoint, server_args, tmp_path): + component, weights = checkpoint + config_path = component / "config.json" + config = json.loads(config_path.read_text()) + config.pop("_class_name") + config_path.write_text(json.dumps(config)) + override = tmp_path / "override.safetensors" + replacement = {key: value + 1 for key, value in weights.items()} + save_file(replacement, override) + server_args.component_weights_paths["projection"] = str(override) + model = _load( + component, server_args, component_architecture="PlainLoaderTestProjection" + ) + torch.testing.assert_close(model.weight, replacement["weight"].bfloat16()) + + +@pytest.mark.parametrize( + "failure", ["missing", "unexpected", "shape", "config", "quantized"] +) +def test_explicit_loader_never_falls_back(checkpoint, server_args, failure): + component, weights = checkpoint + if failure == "missing": + weights.pop("bias") + elif failure == "unexpected": + weights["extra"] = torch.zeros(1) + elif failure == "shape": + weights["bias"] = torch.zeros(4) + else: + config_path = component / "config.json" + config = json.loads(config_path.read_text()) + if failure == "config": + config["unsupported_argument"] = True + else: + config["quantization_config"] = {"quant_method": "fp8"} + config_path.write_text(json.dumps(config)) + save_file(weights, component / "model.safetensors") + with patch.object(ComponentLoader, "load_native") as native: + with pytest.raises( + (RuntimeError, TypeError, ComponentCheckpointUnsupportedError) + ): + _load(component, server_args) + native.assert_not_called() + + +def test_explicit_selection_does_not_change_other_pipelines(): + selected = ComponentLoader.for_component_type( + "duration_head_2", "ltx2", loader_cls=PlainStateDictComponentLoader + ) + assert type(selected) is PlainStateDictComponentLoader + assert selected.component_type == "duration_head_2" + assert isinstance( + ComponentLoader.for_component_type("duration_head_2", "ltx2"), AdapterLoader + ) + assert isinstance( + ComponentLoader.for_component_type("projection", "diffusers"), + GenericComponentLoader, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cuda_residency_and_forward(checkpoint, server_args): + component, weights = checkpoint + server_args.component_residency["projection"] = RESIDENT + model = _load(component, server_args) + assert model.weight.device.type == "cuda" + inputs = torch.ones(2, 4, device=model.weight.device, dtype=torch.bfloat16) + expected = nn.functional.linear( + inputs, weights["weight"].to(inputs), weights["bias"].to(inputs) + ) + torch.testing.assert_close(model(inputs), expected, rtol=0, atol=0)