chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py

@@ -0,0 +1,558 @@
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config
+from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
+    ModelOptFp8Config,
+)
+from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
+    _ModelOptFp8OffloadAdapter,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers import (
+    component_resident_strategies as component_resident_strategies_mod,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers import (
+    layerwise_offload as layerwise_offload_mod,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
+    ComponentUse,
+    build_component_residency_strategy,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_resident_strategies import (
+    LayerwiseOffloadStrategy,
+    ResidentStrategy,
+    VanillaD2HStrategy,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+    LayerwiseOffloadManager,
+    configure_layerwise_offload_modules,
+    get_layerwise_offload_component_names_for_pipeline,
+    is_layerwise_offloaded_module,
+)
+
+
+class _FakeStream:
+    def wait_stream(self, _stream) -> None:
+        return None
+
+    def wait_event(self, _event) -> None:
+        return None
+
+
+class _FakeEvent:
+    def record(self, _stream) -> None:
+        return None
+
+
+class _FakeDeviceModule:
+    Stream = _FakeStream
+    Event = _FakeEvent
+
+    @staticmethod
+    def is_available() -> bool:
+        return True
+
+    @staticmethod
+    def current_device() -> int:
+        return 0
+
+    @staticmethod
+    def current_stream() -> _FakeStream:
+        return _FakeStream()
+
+    @staticmethod
+    def stream(_stream):
+        return nullcontext()
+
+
+class _DummyBlock(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        base = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+        self.weight = torch.nn.Parameter(base.t())
+        self.bias = torch.nn.Parameter(torch.arange(3, dtype=torch.float32))
+
+
+class _DummyModel(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.blocks = torch.nn.ModuleList([_DummyBlock()])
+
+
+class _NestedDummyModel(torch.nn.Module, LayerwiseOffloadableModuleMixin):
+    layer_names = ["encoder.blocks"]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.encoder = _DummyModel()
+
+
+class _SharedBuffer(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.register_buffer(
+            "cache", torch.arange(12, dtype=torch.float32).reshape(6, 2)
+        )
+
+
+class _SharedBufferLayer(torch.nn.Module):
+    def __init__(self, shared: _SharedBuffer) -> None:
+        super().__init__()
+        self.shared = shared
+        self.weight = torch.nn.Parameter(torch.ones(2, 2, dtype=torch.float32))
+
+
+class _SharedBufferModel(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        shared = _SharedBuffer()
+        self.blocks = torch.nn.ModuleList(
+            [_SharedBufferLayer(shared), _SharedBufferLayer(shared)]
+        )
+
+
+class _OrderedLinearLayer(torch.nn.Module):
+    def __init__(self, scale: float) -> None:
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.eye(2, dtype=torch.float32) * scale)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x @ self.weight
+
+
+class _ReverseLayerwiseModel(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.blocks = torch.nn.ModuleList(
+            [
+                _OrderedLinearLayer(2.0),
+                _OrderedLinearLayer(3.0),
+                _OrderedLinearLayer(5.0),
+            ]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for block in reversed(self.blocks):
+            x = block(x)
+        return x
+
+
+class _NestedEncoderDummyModel(_NestedDummyModel):
+    layerwise_offload_dit_group_enabled = False
+
+
+class _LayerwiseComponent(torch.nn.Module, LayerwiseOffloadableModuleMixin):
+    layer_names = ["blocks"]
+
+    def __init__(self, enabled: bool) -> None:
+        super().__init__()
+        self.blocks = torch.nn.ModuleList([_DummyBlock()])
+        self.layerwise_offload_managers = [SimpleNamespace(enabled=enabled)]
+
+
+def _server_args(**kwargs):
+    defaults = dict(
+        use_fsdp_inference=False,
+        dit_cpu_offload=False,
+        text_encoder_cpu_offload=False,
+        image_encoder_cpu_offload=False,
+        vae_cpu_offload=False,
+        dit_offload_prefetch_size=1,
+        pin_cpu_memory=False,
+    )
+    defaults.update(kwargs)
+    return SimpleNamespace(**defaults)
+
+
+def test_layerwise_offload_preserves_non_contiguous_stride(monkeypatch):
+    monkeypatch.setattr(
+        layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule
+    )
+    monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu")
+
+    model = _DummyModel()
+    original_weight = model.blocks[0].weight.detach().clone()
+    original_stride = model.blocks[0].weight.stride()
+    assert not model.blocks[0].weight.is_contiguous()
+
+    manager = LayerwiseOffloadManager(
+        model=model,
+        layers_attr_str="blocks",
+        num_layers=1,
+        enabled=True,
+        pin_cpu_memory=False,
+        prefetch_size=1,
+    )
+
+    meta = manager._weight_metadata[0]["blocks.0.weight"]
+    assert meta["preserve_strides"] is True
+
+    restored_weight = model.blocks[0].weight.data
+    assert restored_weight.shape == original_weight.shape
+    assert restored_weight.stride() == original_stride
+    assert not restored_weight.is_contiguous()
+    assert torch.equal(restored_weight, original_weight)
+
+    manager.release_layer(0)
+    manager.prefetch_layer(0, non_blocking=False)
+
+    reloaded_weight = model.blocks[0].weight.data
+    assert reloaded_weight.stride() == original_stride
+    assert not reloaded_weight.is_contiguous()
+    assert torch.equal(reloaded_weight, original_weight)
+
+
+def test_layerwise_offload_uses_normal_tensors_under_inference_mode(monkeypatch):
+    monkeypatch.setattr(
+        layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule
+    )
+    monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu")
+
+    model = _DummyModel()
+    manager = LayerwiseOffloadManager(
+        model=model,
+        layers_attr_str="blocks",
+        num_layers=1,
+        enabled=True,
+        pin_cpu_memory=False,
+        prefetch_size=1,
+    )
+
+    with torch.inference_mode():
+        manager.release_layer(0)
+        manager.prefetch_layer(0, non_blocking=False)
+
+    assert model.blocks[0].weight._version >= 0
+    assert model.blocks[0].bias._version >= 0
+
+
+def test_layerwise_offload_keeps_shared_buffers_resident(monkeypatch):
+    monkeypatch.setattr(
+        layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule
+    )
+    monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu")
+
+    model = _SharedBufferModel()
+    original_cache = model.blocks[0].shared.cache.detach().clone()
+
+    manager = LayerwiseOffloadManager(
+        model=model,
+        layers_attr_str="blocks",
+        num_layers=2,
+        enabled=True,
+        pin_cpu_memory=False,
+        prefetch_size=1,
+    )
+
+    assert not any(
+        "cache" in name
+        for metadata in manager._weight_metadata.values()
+        for name in metadata
+    )
+    manager.release_layer(0)
+
+    cache = model.blocks[1].shared.cache
+    assert torch.equal(cache, original_cache)
+    assert torch.equal(cache.index_select(0, torch.tensor([2])), original_cache[2:3])
+
+
+def test_layerwise_offload_loads_current_layer_for_reverse_execution(monkeypatch):
+    monkeypatch.setattr(
+        layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule
+    )
+    monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu")
+
+    model = _ReverseLayerwiseModel()
+    x = torch.ones(1, 2, dtype=torch.float32)
+    expected = model(x)
+
+    LayerwiseOffloadManager(
+        model=model,
+        layers_attr_str="blocks",
+        num_layers=3,
+        enabled=True,
+        pin_cpu_memory=False,
+        prefetch_size=1,
+    )
+
+    assert torch.equal(model(x), expected)
+
+
+def test_modelopt_fp8_adapter_keeps_layerwise_offload_enabled():
+    server_args = SimpleNamespace(
+        dit_cpu_offload=True,
+        dit_layerwise_offload=True,
+    )
+    quant_config = ModelOptFp8Config(is_checkpoint_fp8_serialized=True)
+
+    _ModelOptFp8OffloadAdapter._maybe_disable_incompatible_dit_offload_modes(
+        server_args=server_args,
+        quant_config=quant_config,
+    )
+
+    assert server_args.dit_cpu_offload is False
+    assert server_args.dit_layerwise_offload is True
+
+
+def test_modelopt_fp8_adapter_does_not_change_online_fp8_offload():
+    server_args = SimpleNamespace(
+        dit_cpu_offload=True,
+        dit_layerwise_offload=False,
+        quantization="fp8",
+    )
+
+    _ModelOptFp8OffloadAdapter._maybe_disable_incompatible_dit_offload_modes(
+        server_args=server_args,
+        quant_config=Fp8Config(),
+    )
+
+    assert server_args.dit_cpu_offload is True
+
+
+def test_layerwise_capability_selects_layerwise_strategy_for_any_component():
+    module = _LayerwiseComponent(enabled=True)
+
+    assert is_layerwise_offloaded_module(module)
+    strategy = build_component_residency_strategy(
+        "text_encoder", module, _server_args(text_encoder_cpu_offload=True)
+    )
+
+    assert isinstance(strategy, LayerwiseOffloadStrategy)
+
+
+def test_layerwise_pipeline_selection_uses_dit_group(monkeypatch):
+    monkeypatch.setattr(
+        layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule
+    )
+    monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu")
+    layerwise_module = _NestedDummyModel()
+    modules = {
+        "text_encoder": layerwise_module,
+        "text_encoder_alias": layerwise_module,
+        "scheduler": object(),
+    }
+
+    selected = get_layerwise_offload_component_names_for_pipeline(modules)
+    configured = configure_layerwise_offload_modules(modules, _server_args())
+
+    assert selected == ["text_encoder", "text_encoder_alias"]
+    assert configured == ["text_encoder"]
+    assert is_layerwise_offloaded_module(layerwise_module)
+
+
+def test_layerwise_configuration_filters_by_component_name(monkeypatch):
+    monkeypatch.setattr(
+        layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule
+    )
+    monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu")
+    text_encoder = _NestedEncoderDummyModel()
+    transformer = _NestedDummyModel()
+    vae = _NestedDummyModel()
+    modules = {
+        "custom_encoder_name": text_encoder,
+        "custom_transformer_name": transformer,
+        "custom_vae_name": vae,
+    }
+
+    configured = configure_layerwise_offload_modules(
+        modules, _server_args(), component_names=["custom_encoder_name"]
+    )
+
+    assert configured == ["custom_encoder_name"]
+    assert is_layerwise_offloaded_module(text_encoder)
+    assert not is_layerwise_offloaded_module(transformer)
+    assert not is_layerwise_offloaded_module(vae)
+
+
+def test_layerwise_configuration_default_group_selects_non_dit_defaults(monkeypatch):
+    monkeypatch.setattr(
+        layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule
+    )
+    monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu")
+    text_encoder = _NestedEncoderDummyModel()
+    text_encoder_2 = _NestedEncoderDummyModel()
+    transformer = _NestedDummyModel()
+    image_encoder = _NestedEncoderDummyModel()
+    vae = _NestedEncoderDummyModel()
+    audio_vae = _NestedEncoderDummyModel()
+    vocoder = _NestedEncoderDummyModel()
+    spatial_upsampler = _NestedEncoderDummyModel()
+    condition_image_encoder = _NestedEncoderDummyModel()
+    modules = {
+        "text_encoder": text_encoder,
+        "text_encoder_2": text_encoder_2,
+        "transformer": transformer,
+        "image_encoder": image_encoder,
+        "vae": vae,
+        "audio_vae": audio_vae,
+        "vocoder": vocoder,
+        "spatial_upsampler": spatial_upsampler,
+        "condition_image_encoder": condition_image_encoder,
+    }
+
+    configured = configure_layerwise_offload_modules(
+        modules, _server_args(), component_names=["default"]
+    )
+
+    assert get_layerwise_offload_component_names_for_pipeline(modules, ["default"]) == [
+        "text_encoder",
+        "text_encoder_2",
+        "image_encoder",
+        "vae",
+        "condition_image_encoder",
+    ]
+    assert configured == [
+        "text_encoder",
+        "text_encoder_2",
+        "image_encoder",
+        "vae",
+        "condition_image_encoder",
+    ]
+    assert is_layerwise_offloaded_module(text_encoder)
+    assert is_layerwise_offloaded_module(text_encoder_2)
+    assert not is_layerwise_offloaded_module(transformer)
+    assert is_layerwise_offloaded_module(image_encoder)
+    assert is_layerwise_offloaded_module(vae)
+    assert not is_layerwise_offloaded_module(audio_vae)
+    assert not is_layerwise_offloaded_module(vocoder)
+    assert not is_layerwise_offloaded_module(spatial_upsampler)
+    assert is_layerwise_offloaded_module(condition_image_encoder)
+
+    for component_name, module in (
+        ("audio_vae", audio_vae),
+        ("vocoder", vocoder),
+        ("spatial_upsampler", spatial_upsampler),
+    ):
+        configured = configure_layerwise_offload_modules(
+            modules, _server_args(), component_names=[component_name]
+        )
+        assert configured == [component_name]
+        assert is_layerwise_offloaded_module(module)
+
+
+def test_layerwise_configuration_all_selects_every_capable_component(monkeypatch):
+    monkeypatch.setattr(
+        layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule
+    )
+    monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu")
+    text_encoder = _NestedEncoderDummyModel()
+    transformer = _NestedDummyModel()
+    modules = {
+        "custom_encoder_name": text_encoder,
+        "custom_transformer_name": transformer,
+        "scheduler": object(),
+    }
+
+    configured = configure_layerwise_offload_modules(
+        modules, _server_args(), component_names=["all"]
+    )
+
+    assert configured == ["custom_encoder_name", "custom_transformer_name"]
+    assert is_layerwise_offloaded_module(text_encoder)
+    assert is_layerwise_offloaded_module(transformer)
+
+
+def test_component_cpu_offload_strategy_remains_flag_driven():
+    strategy = build_component_residency_strategy(
+        "text_encoder", _DummyModel(), _server_args(text_encoder_cpu_offload=True)
+    )
+    assert isinstance(strategy, VanillaD2HStrategy)
+
+    strategy = build_component_residency_strategy(
+        "unknown_component", _DummyModel(), _server_args(text_encoder_cpu_offload=True)
+    )
+    assert isinstance(strategy, ResidentStrategy)
+
+
+def test_resident_strategy_prepares_local_device_without_dtype(monkeypatch):
+    calls = []
+
+    def fake_module_to_local_device(module, *, dtype=None):
+        calls.append((module, dtype))
+
+    monkeypatch.setattr(
+        component_resident_strategies_mod,
+        "_module_to_local_device",
+        fake_module_to_local_device,
+    )
+    module = _DummyModel()
+
+    ResidentStrategy().prepare_for_use(
+        module,
+        ComponentUse(stage_name="DenoisingStage", component_name="transformer"),
+        SimpleNamespace(),
+    )
+
+    assert calls == [(module, None)]
+
+
+def test_resident_strategy_keeps_fsdp_managed_module_owned_by_fsdp(monkeypatch):
+    calls = []
+
+    def fake_module_to_local_device(module, *, dtype=None):
+        calls.append((module, dtype))
+
+    monkeypatch.setattr(
+        component_resident_strategies_mod,
+        "_module_to_local_device",
+        fake_module_to_local_device,
+    )
+    module = type("FSDPDummyModel", (_DummyModel,), {})()
+
+    ResidentStrategy().prepare_for_use(
+        module,
+        ComponentUse(stage_name="TextEncodingStage", component_name="text_encoder"),
+        SimpleNamespace(),
+    )
+
+    assert calls == []
+
+
+def test_layerwise_offload_aligns_contiguous_tensor_offsets(monkeypatch):
+    monkeypatch.setattr(
+        layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule
+    )
+    monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu")
+
+    class _AlignedDummyBlock(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weight = torch.nn.Parameter(
+                torch.arange(9, dtype=torch.float32).reshape(3, 3)
+            )
+            self.bias = torch.nn.Parameter(torch.arange(3, dtype=torch.float32))
+
+    class _AlignedDummyModel(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.blocks = torch.nn.ModuleList([_AlignedDummyBlock()])
+
+    model = _AlignedDummyModel()
+    original_weight = model.blocks[0].weight.detach().clone()
+    original_bias = model.blocks[0].bias.detach().clone()
+
+    manager = LayerwiseOffloadManager(
+        model=model,
+        layers_attr_str="blocks",
+        num_layers=1,
+        enabled=True,
+        pin_cpu_memory=False,
+        prefetch_size=1,
+    )
+
+    weight_meta = manager._weight_metadata[0]["blocks.0.weight"]
+    bias_meta = manager._weight_metadata[0]["blocks.0.bias"]
+    assert weight_meta["preserve_strides"] is False
+    assert bias_meta["preserve_strides"] is False
+    assert weight_meta["offset"] == 0
+    assert bias_meta["offset"] % 8 == 0
+
+    restored_weight = model.blocks[0].weight.data
+    restored_bias = model.blocks[0].bias.data
+    assert restored_weight.data_ptr() % 32 == 0
+    assert restored_bias.data_ptr() % 32 == 0
+    assert torch.equal(restored_weight, original_weight)
+    assert torch.equal(restored_bias, original_bias)