vllm-omni/tests/diffusion/offloader/test_layerwise_backend.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""Unit tests for LayerwiseOffloadHook and LayerWiseOffloadBackend utilities."""

from multiprocessing.reduction import ForkingPickler
from types import SimpleNamespace
from unittest.mock import Mock

import pytest
import torch
from torch import nn
from torch.distributed.tensor import DeviceMesh, DTensor, Replicate

import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module
import vllm_omni.diffusion.offloader.tensor_utils as tensor_utils_module
from tests.diffusion.offloader.helpers import (
    DummyStream,
    _DummyBlock,
    _PlainEncoder,
    _SingleBlockModel,
    _StagedEncoder,
    _StagedVAE,
    patch_offload_runtime,
)
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
from vllm_omni.diffusion.offloader.layerwise_backend import (
    LayerWiseOffloadBackend,
    LayerwiseOffloadHook,
)
from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan

pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]


@pytest.fixture
def patched_offload_runtime(monkeypatch):
    patch_offload_runtime(monkeypatch, layerwise_backend_module.current_omni_platform)


class TinyBlock(nn.Module):
    def __init__(self, values: torch.Tensor):
        super().__init__()
        mesh = DeviceMesh("cpu", [0])
        dtensor = DTensor.from_local(values, mesh, [Replicate()])
        self.weight = nn.Parameter(dtensor)


def _make_values(start: float) -> torch.Tensor:
    return torch.arange(start, start + 4, dtype=torch.float32)


def test_clear_tensor_storage_rolls_back_partial_commit(monkeypatch):
    first = nn.Parameter(torch.tensor([1.0, 2.0]))
    second = nn.Parameter(torch.tensor([3.0, 4.0]))
    expected_first = first.detach().clone()
    expected_second = second.detach().clone()
    second_data_ptr = second.data_ptr()
    original_set_storage = tensor_utils_module.set_tensor_storage
    calls = 0

    def fail_second_commit(target, value):
        nonlocal calls
        calls += 1
        if calls == 2:
            raise RuntimeError("injected storage clear failure")
        original_set_storage(target, value)

    monkeypatch.setattr(tensor_utils_module, "set_tensor_storage", fail_second_commit)

    with pytest.raises(RuntimeError, match="injected storage clear failure"):
        tensor_utils_module.clear_tensor_storage([first, second])

    assert calls == 3  # first commit, failed second commit, first rollback
    torch.testing.assert_close(first, expected_first)
    torch.testing.assert_close(second, expected_second)
    assert second.data_ptr() == second_data_ptr


class TestLayerwiseOffloadHook:
    def test_buffer_only_block_reports_offloaded_state(self, patched_offload_runtime):
        current_block = nn.Module()
        current_block.register_buffer("state", torch.ones(2))
        next_block = nn.Module()
        next_block.register_buffer("state", torch.ones(2))
        hook = LayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            stream=DummyStream(),
            pin_memory=False,
        )
        hook.initialize_hook(current_block)

        assert hook.is_materialized
        hook.offload_layer()
        assert not hook.is_materialized

    def test_ring_probes_are_captured_before_any_block_is_cleared(self, patched_offload_runtime):
        blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(3)])
        expected_middle = blocks[1].weight.detach().clone()

        hooks = layerwise_backend_module._install_layerwise_hook_group(
            blocks,
            torch.device("cpu"),
            DummyStream(),
            pin_memory=False,
        )

        assert all(not hook.is_materialized for hook in hooks)
        hooks[2].pre_forward(blocks[1])
        torch.testing.assert_close(blocks[1].weight, expected_middle)

    def test_dtensor_wrapper_is_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
        current_block = TinyBlock(_make_values(1.0))
        next_block = TinyBlock(_make_values(10.0))

        hook = LayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            stream=DummyStream(),
            pin_memory=False,
        )

        hook.initialize_hook(current_block)

        assert isinstance(next_block.weight, DTensor)
        assert next_block.weight.to_local().is_meta
        assert next_block.weight.to_local().shape == torch.Size([4])
        assert hook.dtype_metadata[next_block.weight.dtype][0]["shape"] == torch.Size([4])

        hook.prefetch_layer(non_blocking=False)
        assert isinstance(next_block.weight, DTensor)
        assert torch.equal(next_block.weight.to_local(), _make_values(10.0))
        assert next_block.weight.to_local().shape == torch.Size([4])

        hook.offload_layer()
        assert isinstance(current_block.weight, DTensor)
        assert current_block.weight.to_local().is_meta
        assert current_block.weight.to_local().shape == torch.Size([4])
        assert not hook.is_materialized

    def test_prefetch_preserves_transposed_weight_stride(self, patched_offload_runtime):
        """Online-FP8 Cutlass weights must retain their transposed layout."""

        class StridedBlock(nn.Module):
            def __init__(self, start: float):
                super().__init__()
                base = torch.arange(start, start + 12).reshape(3, 4)
                self.weight = nn.Parameter(base.t(), requires_grad=False)

        current_block = StridedBlock(1.0)
        next_block = StridedBlock(20.0)
        expected = next_block.weight.detach().clone()
        expected_stride = next_block.weight.stride()
        assert expected_stride == (1, 4)

        hook = LayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            stream=DummyStream(),
            pin_memory=False,
        )
        hook.initialize_hook(current_block)
        hook.prefetch_layer(non_blocking=False)

        assert next_block.weight.stride() == expected_stride
        assert torch.equal(next_block.weight, expected)


class _MultiBlockModel(nn.Module):
    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]

    def __init__(self, num_transformer: int = 2, num_single: int = 2):
        super().__init__()
        self.transformer_blocks = nn.ModuleList([_DummyBlock() for _ in range(num_transformer)])
        self.single_transformer_blocks = nn.ModuleList([_DummyBlock() for _ in range(num_single)])


class _EmptyBlocksModel(nn.Module):
    _layerwise_offload_blocks_attrs = ["blocks"]

    def __init__(self):
        super().__init__()
        self.blocks = nn.ModuleList([])


class _InvalidAttrModel(nn.Module):
    _layerwise_offload_blocks_attrs = ["nonexistent_blocks", "blocks"]

    def __init__(self, num_blocks: int = 2):
        super().__init__()
        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])


class _DeprecatedSingleAttrModel(nn.Module):
    _layerwise_offload_blocks_attr = "blocks"

    def __init__(self, num_blocks: int = 2):
        super().__init__()
        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])


class _NoAttrsModel(nn.Module):
    def __init__(self, num_blocks: int = 2):
        super().__init__()
        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])


class TestGetBlocksFromDit:
    def test_get_blocks_from_dit_single_block_attr(self):
        model = _SingleBlockModel(num_blocks=3)
        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
        assert attr_names == ["blocks"]
        assert len(blocks) == 3
        assert all(isinstance(b, _DummyBlock) for b in blocks)

    def test_get_blocks_from_dit_multi_block_attrs(self):
        model = _MultiBlockModel(num_transformer=2, num_single=3)
        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
        assert set(attr_names) == {"transformer_blocks", "single_transformer_blocks"}
        assert len(blocks) == 5
        assert all(isinstance(b, _DummyBlock) for b in blocks)

    def test_get_blocks_from_dit_empty_blocks(self):
        model = _EmptyBlocksModel()
        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
        assert attr_names == []
        assert blocks == []

    def test_get_blocks_from_dit_invalid_attr_name(self):
        model = _InvalidAttrModel(num_blocks=2)
        with pytest.raises(
            AttributeError,
            match="Attribute 'nonexistent_blocks' declared in _layerwise_offload_blocks_attrs does not exist",
        ):
            LayerWiseOffloadBackend.get_blocks_from_dit(model)

    def test_get_blocks_from_dit_no_attrs_defined(self):
        model = _NoAttrsModel(num_blocks=3)
        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
        assert attr_names == []
        assert blocks == []

    def test_get_blocks_from_dit_deprecated_single_attr(self):
        model = _DeprecatedSingleAttrModel(num_blocks=2)
        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
        assert attr_names == ["blocks"]
        assert len(blocks) == 2


class TestGetBlocksAttrNames:
    def test_get_blocks_attr_names_new_format(self):
        model = _MultiBlockModel()
        attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model)
        assert attrs == ["transformer_blocks", "single_transformer_blocks"]

    def test_get_blocks_attr_names_no_attrs(self):
        model = _NoAttrsModel()
        attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model)
        assert attrs == []

    def test_set_blocks_attr_names(self):
        model = _NoAttrsModel()
        LayerWiseOffloadBackend.set_blocks_attr_names(model, ["new_blocks"])
        assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs")
        assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"]


class _ComponentPipeline(nn.Module):
    _offload_plan = OffloadPlan(
        encoder_component_types={"text_encoder": "text_encoder"},
        encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")},
        on_demand_component_paths=frozenset({"text_encoder", "vae"}),
    )

    def __init__(self):
        super().__init__()
        self.transformer = _SingleBlockModel()
        self.text_encoder = _StagedEncoder()
        self.vae = _StagedVAE()


class _LegacyComponentPipeline(nn.Module):
    """Pipeline with DiT block metadata but no auxiliary OffloadPlan."""

    def __init__(self):
        super().__init__()
        self.transformer = _SingleBlockModel()
        self.text_encoder = _StagedEncoder()
        self.vae = _StagedVAE()


class _LegacyEncoderOnlyPipeline(nn.Module):
    _offload_plan = _ComponentPipeline._offload_plan

    def __init__(self):
        super().__init__()
        self.text_encoder = _StagedEncoder()
        self.vae = _StagedVAE()


class _PlannedDitPipeline(nn.Module):
    _offload_plan = OffloadPlan(block_attrs={"transformer": ("blocks",)})

    def __init__(self):
        super().__init__()
        self.transformer = _NoAttrsModel()


class _GenericEncoderPipeline(nn.Module):
    _offload_plan = OffloadPlan(
        encoder_component_types={"text_encoder": "text_encoder"},
        encoder_block_attrs={"text_encoder": ("encoder.block",)},
    )

    def __init__(self):
        super().__init__()
        self.transformer = _SingleBlockModel()
        self.text_encoder = _PlainEncoder()


def _layer_backend(components: set[str] | None = None) -> LayerWiseOffloadBackend:
    options = {}
    if components is not None:
        options = {"components": frozenset(components)}
    return LayerWiseOffloadBackend(
        OffloadConfig(
            strategy=OffloadStrategy.LAYER_WISE,
            pin_cpu_memory=False,
            **options,
        ),
        torch.device("cpu"),
    )


class TestLayerwiseComponentSelection:
    def test_plan_block_attrs_drive_dit_discovery(self, patched_offload_runtime):
        pipeline = _PlannedDitPipeline()
        backend = _layer_backend({"dit"})

        backend.enable(pipeline)

        assert hasattr(pipeline.transformer.blocks[0], "_hook_registry")
        backend.disable()

    def test_legacy_missing_dit_preserves_noop(self, patched_offload_runtime):
        pipeline = _LegacyEncoderOnlyPipeline()
        backend = _layer_backend()

        backend.enable(pipeline)

        assert not backend.enabled
        assert pipeline.text_encoder.offload_calls == 0
        assert pipeline.vae.offload_calls == 0

    def test_encoder_only_streams_planned_blocks(self, patched_offload_runtime, monkeypatch):
        pipeline = _ComponentPipeline()
        backend = _layer_backend({"text_encoder"})
        move_non_block_state = Mock()
        monkeypatch.setattr(layerwise_backend_module, "move_non_block_state_to_device", move_non_block_state)

        backend.enable(pipeline)

        assert pipeline.text_encoder._omni_layerwise_enabled
        assert len(pipeline.text_encoder._omni_layerwise_hooks) == 4
        assert pipeline.text_encoder.offload_calls == 1
        assert pipeline.vae.offload_calls == 0
        assert pipeline.vae.to_calls == 1
        assert not hasattr(pipeline.transformer.blocks[0], "_hook_registry")
        assert backend.enabled
        move_non_block_state.assert_not_called()

        backend.disable()
        assert not pipeline.text_encoder._omni_layerwise_enabled

    def test_single_gpu_dit_only_keeps_encoder_and_vae_resident(self, patched_offload_runtime):
        pipeline = _ComponentPipeline()
        expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks]
        backend = _layer_backend({"dit"})

        backend.enable(pipeline)

        assert not hasattr(pipeline.text_encoder, "_omni_layerwise_enabled")
        assert pipeline.text_encoder.to_calls == 1
        assert pipeline.text_encoder.offload_calls == 0
        assert pipeline.vae.to_calls == 1
        assert pipeline.vae.offload_calls == 0
        assert hasattr(pipeline.transformer.blocks[0], "_hook_registry")

        backend.disable()

        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
            torch.testing.assert_close(block.weight, expected)

        backend.enable(pipeline)
        backend.disable()
        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
            torch.testing.assert_close(block.weight, expected)

    def test_partial_dit_enable_failure_restores_weights_and_hooks(self, patched_offload_runtime, monkeypatch):
        pipeline = _ComponentPipeline()
        expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks]
        backend = _layer_backend({"dit"})
        original_apply = layerwise_backend_module.apply_block_hook
        calls = 0

        def fail_second_hook(*args, **kwargs):
            nonlocal calls
            calls += 1
            if calls == 2:
                raise RuntimeError("injected hook failure")
            return original_apply(*args, **kwargs)

        monkeypatch.setattr(layerwise_backend_module, "apply_block_hook", fail_second_hook)

        with pytest.raises(RuntimeError, match="injected hook failure"):
            backend.enable(pipeline)

        assert not backend.enabled
        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
            registry = getattr(block, "_hook_registry", None)
            assert registry is None or registry.get_hook("layerwise_offload") is None
            torch.testing.assert_close(block.weight, expected)

    def test_later_component_failure_removes_installed_encoder_hooks(
        self,
        patched_offload_runtime,
        monkeypatch,
    ):
        pipeline = _ComponentPipeline()
        blocks = [*pipeline.text_encoder.vision.blocks, *pipeline.text_encoder.text_model.layers]
        expected = [block.weight.detach().clone() for block in blocks]
        backend = _layer_backend({"text_encoder"})

        def fail_vae_placement(*_args, **_kwargs):
            raise RuntimeError("injected VAE placement failure")

        monkeypatch.setattr(pipeline.vae, "to", fail_vae_placement)

        with pytest.raises(RuntimeError, match="injected VAE placement failure"):
            backend.enable(pipeline)

        assert not backend.enabled
        assert not pipeline.text_encoder._omni_layerwise_enabled
        assert not backend._encoder_modules
        for block, original in zip(blocks, expected, strict=True):
            assert block._hook_registry.get_hook("layerwise_offload") is None
            torch.testing.assert_close(block.weight, original)

    def test_encoder_only_requires_streamable_offload_plan(self, patched_offload_runtime):
        pipeline = nn.Module()
        pipeline.transformer = _SingleBlockModel()
        pipeline.text_encoder = _StagedEncoder()
        backend = _layer_backend({"text_encoder"})

        with pytest.raises(ValueError, match="Selected text encoder 'text_encoder' requires"):
            backend.enable(pipeline)

    def test_explicit_text_encoder_selection_requires_a_matching_module(self, patched_offload_runtime):
        pipeline = nn.Module()
        pipeline.transformer = _SingleBlockModel(num_blocks=2)
        backend = _layer_backend({"dit", "text_encoder"})

        with pytest.raises(ValueError, match="No text encoder modules found"):
            backend.enable(pipeline)

        assert not hasattr(pipeline.transformer.blocks[0], "_hook_registry")

    def test_every_selected_encoder_requires_its_own_plan(self, patched_offload_runtime):
        class Pipeline(nn.Module):
            _offload_plan = OffloadPlan(
                encoder_component_types={
                    "text_encoder": "text_encoder",
                    "text_encoder_2": "text_encoder",
                },
                encoder_block_attrs={"text_encoder": ("encoder.block",)},
            )

            def __init__(self):
                super().__init__()
                self.transformer = _SingleBlockModel()
                self.text_encoder = _PlainEncoder()
                self.text_encoder_2 = _PlainEncoder()

        pipeline = Pipeline()
        backend = _layer_backend({"text_encoder"})

        with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"):
            backend.enable(pipeline)

        # The plan resolver rejects the topology before any encoder is hooked.
        assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False)

    def test_default_selection_preserves_unplanned_auxiliaries(self, patched_offload_runtime):
        pipeline = _LegacyComponentPipeline()
        backend = _layer_backend()

        backend.enable(pipeline)

        assert hasattr(pipeline.transformer.blocks[0], "_hook_registry")
        assert not hasattr(pipeline.text_encoder, "_omni_layerwise_enabled")
        assert pipeline.text_encoder.to_calls == 1
        assert pipeline.text_encoder.offload_calls == 0
        assert pipeline.vae.to_calls == 1
        assert pipeline.vae.offload_calls == 0

        backend.disable()

    def test_disable_failure_keeps_host_masters_for_retry(self, patched_offload_runtime, monkeypatch):
        pipeline = nn.Module()
        pipeline.transformer = _SingleBlockModel(num_blocks=3)
        expected = {name: tensor.detach().clone() for name, tensor in pipeline.state_dict().items()}
        backend = _layer_backend()
        backend.enable(pipeline)
        failing_hook = backend._dit_hooks[0]
        restore = failing_hook.restore_next_block
        attempts = 0

        def fail_once():
            nonlocal attempts
            attempts += 1
            if attempts == 1:
                raise RuntimeError("injected restore failure")
            restore()

        monkeypatch.setattr(failing_hook, "restore_next_block", fail_once)

        with pytest.raises(RuntimeError, match="Failed to fully disable"):
            backend.disable()

        assert backend.enabled
        assert backend._dit_hooks

        backend.disable()

        assert not backend.enabled
        for name, tensor in pipeline.state_dict().items():
            torch.testing.assert_close(tensor, expected[name])

    def test_encoder_disable_failure_keeps_host_masters_for_retry(
        self,
        patched_offload_runtime,
        monkeypatch,
    ):
        pipeline = _ComponentPipeline()
        encoder = pipeline.text_encoder
        expected = {name: tensor.detach().clone() for name, tensor in encoder.state_dict().items()}
        backend = _layer_backend({"text_encoder"})
        backend.enable(pipeline)
        failing_hook = encoder._omni_layerwise_hooks[0]
        restore = failing_hook.restore_next_block
        attempts = 0

        def fail_once():
            nonlocal attempts
            attempts += 1
            if attempts == 1:
                raise RuntimeError("injected encoder restore failure")
            restore()

        monkeypatch.setattr(failing_hook, "restore_next_block", fail_once)

        with pytest.raises(RuntimeError, match="Failed to fully disable"):
            backend.disable()

        assert backend.enabled
        assert backend._encoder_modules == [encoder]
        assert encoder._omni_layerwise_enabled
        assert encoder._omni_layerwise_hooks

        backend.disable()

        assert not backend.enabled
        assert not encoder._omni_layerwise_enabled
        for name, tensor in encoder.state_dict().items():
            torch.testing.assert_close(tensor, expected[name])

    def test_legacy_selector_preserves_planned_encoder_and_vae_lifecycle(self, patched_offload_runtime):
        pipeline = _ComponentPipeline()
        backend = _layer_backend()

        backend.enable(pipeline)

        assert pipeline.text_encoder._omni_layerwise_enabled
        assert pipeline.text_encoder.offload_calls == 1
        assert pipeline.vae.offload_calls == 1

        backend.disable()

    def test_standard_encoder_needs_only_declared_block_paths(self, patched_offload_runtime):
        pipeline = _GenericEncoderPipeline()
        backend = _layer_backend({"text_encoder"})

        backend.enable(pipeline)

        blocks = pipeline.text_encoder.encoder.block
        assert pipeline.text_encoder._omni_layerwise_enabled
        assert all(block._hook_registry.get_hook("layerwise_offload") is not None for block in blocks)
        assert pipeline.text_encoder.final_norm.weight.numel() == 4

        backend.disable()

        assert not pipeline.text_encoder._omni_layerwise_enabled
        assert all(block._hook_registry.get_hook("layerwise_offload") is None for block in blocks)
        assert all(block.weight.numel() == 100 for block in blocks)


def _offload_od_config(**overrides):
    values = {
        "diffusion_offload_config": None,
        "enable_cpu_offload": False,
        "enable_layerwise_offload": False,
        "enable_distributed_layerwise_offload": False,
        "dlo_use_allgather": True,
        "dlo_resident_layers": 0,
        "pin_cpu_memory": True,
        "parallel_config": None,
        "model": "/fake/model",
    }
    values.update(overrides)
    return SimpleNamespace(**values)


def _resolve_offload_config(public=None, **overrides) -> OffloadConfig:
    return OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=public, **overrides))


class TestLayerwiseComponentConfig:
    def test_layer_mode_selects_components_without_backend_jargon(self):
        config = _resolve_offload_config(
            {
                "mode": "layer",
                "components": ["dit", "text_encoder"],
            }
        )

        assert config.strategy is OffloadStrategy.LAYER_WISE
        assert config.components == frozenset({"dit", "text_encoder"})
        assert not config.uses_allgather("dit")
        assert not config.uses_allgather("text_encoder")
        plan = OffloadPlan(encoder_component_types={"mllm": "text_encoder"})
        assert config.offloads_encoder("mllm", plan)

    @pytest.mark.parametrize("component", ["image_encoder", "vae", "scheduler", "text-encoder"])
    def test_unknown_or_noncanonical_component_is_rejected(self, component):
        with pytest.raises(ValueError, match="Unknown diffusion offload component"):
            OffloadConfig.from_od_config(
                _offload_od_config(diffusion_offload_config={"mode": "layer", "components": [component]})
            )

    def test_components_is_selection_only_not_an_options_mapping(self):
        with pytest.raises(TypeError, match="components must be a non-empty list"):
            OffloadConfig.from_od_config(
                _offload_od_config(
                    diffusion_offload_config={
                        "mode": "layer",
                        "components": {"dit": {}},
                    }
                )
            )

    def test_layer_options_requires_component_name_keys(self):
        with pytest.raises(TypeError, match="layer_options keys must be strings"):
            OffloadConfig.from_od_config(
                _offload_od_config(
                    diffusion_offload_config={
                        "mode": "layer",
                        "components": ["dit"],
                        "layer_options": {0: {}},
                    }
                )
            )

    @pytest.mark.parametrize(
        ("config", "message"),
        [
            ({0: "layer"}, "diffusion_offload_config keys must be strings"),
            (
                {
                    "mode": "layer",
                    "components": ["dit"],
                    "layer_options": {"dit": {0: "allgather"}},
                },
                "layer_options\\['dit'\\] keys must be strings",
            ),
        ],
    )
    def test_all_public_mapping_levels_require_string_keys(self, config, message):
        with pytest.raises(TypeError, match=message):
            OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=config))

    @pytest.mark.parametrize(
        "config,match",
        [
            ({"components": ["dit"]}, "requires 'mode'"),
            ({"mode": "layer"}, "requires 'components'"),
            ({"mode": "layer", "components": []}, "must not be empty"),
            ({"mode": "layer", "components": [" dit "]}, "Unknown diffusion offload component"),
            ({"mode": "layerwise", "components": ["dit"]}, "Unknown diffusion offload mode"),
            (
                {
                    "mode": "layer",
                    "components": ["dit"],
                    "layer_options": {"dit": {"weight_transfer": "rank_local"}},
                },
                "Unknown offload transfer",
            ),
            (
                {
                    "mode": "layer",
                    "components": ["dit"],
                    "layer_options": {"dit": {"prefetch": 2}},
                },
                "Unknown diffusion offload setting",
            ),
            (
                {
                    "mode": "layer",
                    "components": ["text_encoder"],
                    "layer_options": {"text_encoder": {"resident_layers": 1}},
                },
                "supports only the 'dit'",
            ),
            (
                {
                    "mode": "layer",
                    "components": ["dit"],
                    "layer_options": {"text_encoder": {}},
                },
                "requires selecting the same component",
            ),
        ],
    )
    def test_schema_rejects_ambiguous_or_unsupported_values(self, config, match):
        with pytest.raises(ValueError, match=match):
            OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=config))

    def test_module_mode_accepts_component_selection(self):
        config = _resolve_offload_config(
            {
                "mode": "module",
                "components": ["dit", "text_encoder"],
            }
        )

        assert config.strategy is OffloadStrategy.MODEL_LEVEL
        assert config.components == frozenset({"dit", "text_encoder"})

    @pytest.mark.parametrize("setting", [{"weight_transfer": "rank-local"}, {"resident_layers": 1}])
    def test_module_mode_rejects_layer_settings(self, setting):
        with pytest.raises(ValueError, match="layer_options requires mode='layer'"):
            OffloadConfig.from_od_config(
                _offload_od_config(
                    diffusion_offload_config={
                        "mode": "module",
                        "components": ["dit"],
                        "layer_options": {"dit": setting},
                    }
                )
            )

    def test_each_component_can_choose_its_transfer(self):
        config = _resolve_offload_config(
            {
                "mode": "layer",
                "components": ["dit", "text_encoder"],
                "layer_options": {
                    "dit": {"weight_transfer": "rank-local"},
                    "text_encoder": {"weight_transfer": "allgather"},
                },
            }
        )

        assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
        assert not config.uses_allgather("dit")
        assert config.uses_allgather("text_encoder")
        assert config.dlo_use_allgather is False

    def test_dit_residency_selects_capable_backend(self):
        config = _resolve_offload_config(
            {
                "mode": "layer",
                "components": ["dit"],
                "layer_options": {"dit": {"resident_layers": 20}},
            }
        )

        assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
        assert config.dlo_resident_layers == 20
        assert not config.uses_allgather("dit")

    def test_dit_residency_rejects_allgather(self):
        with pytest.raises(ValueError, match="resident_layers requires dit.weight_transfer='rank-local'"):
            OffloadConfig.from_od_config(
                _offload_od_config(
                    diffusion_offload_config={
                        "mode": "layer",
                        "components": ["dit"],
                        "layer_options": {"dit": {"weight_transfer": "allgather", "resident_layers": 1}},
                    }
                )
            )

    def test_host_weight_runtime_remains_separate_from_compact_config(self):
        with pytest.raises(ValueError, match="cannot be combined with Host Weight Runtime"):
            OffloadConfig.from_od_config(
                _offload_od_config(
                    diffusion_offload_config={"mode": "layer", "components": ["dit"]},
                    host_weight_runtime_mode="preferred",
                )
            )

    def test_pin_memory_override_is_scoped_to_offload_config(self):
        config = _resolve_offload_config(
            {
                "mode": "layer",
                "components": ["dit"],
                "pin_memory": False,
            },
            pin_cpu_memory=True,
        )

        assert config.pin_cpu_memory is False

    def test_legacy_omitted_selector_preserves_dit_only_behavior(self):
        config = _resolve_offload_config(enable_layerwise_offload=True)

        assert config.strategy is OffloadStrategy.LAYER_WISE
        assert config.components is None

    def test_offload_resolution_is_cached_for_runtime_accessors(self, monkeypatch):
        import vllm_omni.diffusion.offloader.config as config_module

        od_config = _offload_od_config(
            diffusion_offload_config={
                "mode": "layer",
                "components": ["dit", "text_encoder"],
            }
        )
        parse_calls = 0
        original_parse = config_module.parse_diffusion_offload_config

        def track_parse(value):
            nonlocal parse_calls
            parse_calls += 1
            return original_parse(value)

        monkeypatch.setattr(config_module, "parse_diffusion_offload_config", track_parse)

        assert config_module.resolve_offload_strategy(od_config) is OffloadStrategy.LAYER_WISE
        assert config_module.selected_offload_components(od_config) == {"dit", "text_encoder"}
        assert not config_module.component_uses_allgather(od_config, "dit")
        assert not config_module.any_selected_component_uses_allgather(od_config)
        assert parse_calls == 1

    def test_cached_policy_is_deeply_immutable_and_pickle_safe(self):
        import vllm_omni.diffusion.offloader.config as config_module

        od_config = _offload_od_config(
            diffusion_offload_config={
                "mode": "layer",
                "components": ["dit"],
                "layer_options": {"dit": {"weight_transfer": "rank-local"}},
            }
        )
        resolved = config_module.resolve_offload(od_config)

        with pytest.raises(TypeError):
            resolved.transfers["dit"] = config_module.DLOTransfer.ALLGATHER  # type: ignore[index]
        assert resolved.public is not None
        with pytest.raises(TypeError):
            resolved.public.layer_options["dit"] = config_module.LayerOffloadOptions()  # type: ignore[index]

        restored = ForkingPickler.loads(ForkingPickler.dumps(resolved))
        assert not restored.uses_allgather("dit")

    def test_dp_text_encoder_allgather_rejects_rank_local_prompt_cache(self):
        config = _offload_od_config(
            diffusion_offload_config={
                "mode": "layer",
                "components": ["text_encoder"],
                "layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
            },
            enable_prompt_embed_cache=True,
            parallel_config=SimpleNamespace(
                data_parallel_size=2,
                sequence_parallel_size=1,
                use_hsdp=False,
            ),
        )

        with pytest.raises(ValueError, match="rank-local cache hits would skip different encoder collectives"):
            OffloadConfig.from_od_config(config)

    def test_dp_allgather_rejects_rank_local_dit_cache_decisions(self):
        config = _offload_od_config(
            diffusion_offload_config={
                "mode": "layer",
                "components": ["dit"],
                "layer_options": {"dit": {"weight_transfer": "allgather"}},
            },
            cache_backend="tea_cache",
            parallel_config=SimpleNamespace(
                data_parallel_size=2,
                sequence_parallel_size=1,
                use_hsdp=False,
            ),
        )

        with pytest.raises(ValueError, match="rank-local cache decisions can skip different weight collectives"):
            OffloadConfig.from_od_config(config)

    def test_internal_transfer_map_requires_every_component(self):
        with pytest.raises(ValueError, match="missing: text_encoder"):
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                dlo_transfers={"dit": "rank-local"},
            )

    def test_legacy_distributed_settings_still_work(self):
        config = _resolve_offload_config(
            enable_distributed_layerwise_offload=True,
            dlo_use_allgather=False,
            dlo_resident_layers=3,
        )

        assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
        assert config.dlo_resident_layers == 3
        assert not config.uses_allgather("dit")
        assert not config.uses_allgather("text_encoder")

    def test_legacy_allgather_remains_dit_only(self):
        config = _resolve_offload_config(
            enable_distributed_layerwise_offload=True,
            dlo_use_allgather=True,
        )

        assert config.uses_allgather("dit")
        assert not config.uses_allgather("text_encoder")

    @pytest.mark.parametrize(
        ("legacy_flags", "expected_strategy"),
        [
            (
                {"enable_cpu_offload": True, "enable_layerwise_offload": True},
                OffloadStrategy.LAYER_WISE,
            ),
            (
                {
                    "enable_cpu_offload": True,
                    "enable_layerwise_offload": True,
                    "enable_distributed_layerwise_offload": True,
                },
                OffloadStrategy.DISTRIBUTED_LAYER_WISE,
            ),
        ],
    )
    def test_legacy_aliases_preserve_existing_priority(self, legacy_flags, expected_strategy):
        config = _resolve_offload_config(**legacy_flags)

        assert config.strategy is expected_strategy

    def test_compact_config_cannot_mix_with_legacy_alias(self):
        with pytest.raises(ValueError, match="cannot be combined"):
            OffloadConfig.from_od_config(
                _offload_od_config(
                    diffusion_offload_config={"mode": "layer", "components": ["dit"]},
                    enable_cpu_offload=True,
                )
            )

    @pytest.mark.parametrize(
        "legacy_options",
        [
            {"dlo_use_allgather": False},
            {"dlo_resident_layers": 12},
        ],
    )
    def test_compact_config_cannot_mix_with_legacy_dlo_options(self, legacy_options):
        with pytest.raises(ValueError, match="cannot be combined with legacy DLO option"):
            OffloadConfig.from_od_config(
                _offload_od_config(
                    diffusion_offload_config={"mode": "layer", "components": ["dit"]},
                    **legacy_options,
                )
            )

    def test_legacy_allgather_resident_layers_fail_during_config_resolution(self):
        with pytest.raises(ValueError, match="requires the DiT DLO transfer to be rank-local"):
            OffloadConfig.from_od_config(
                _offload_od_config(
                    enable_distributed_layerwise_offload=True,
                    dlo_use_allgather=True,
                    dlo_resident_layers=12,
                )
            )