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

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

"""Unit tests for DistributedLayerwiseOffloadHook and backend utilities."""

import gc
import json
import weakref
from types import SimpleNamespace
from unittest.mock import Mock

import pytest
import torch
import torch.distributed as dist
from safetensors.torch import save_file
from torch import nn
from torch.distributed.tensor import DeviceMesh, DTensor, Replicate

import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
from tests.diffusion.offloader.helpers import (
    DummyStream,
    _DummyBlock,
    _PlainEncoder,
    _SingleBlockModel,
    _StagedEncoder,
    _StagedVAE,
    patch_offload_runtime,
)
from vllm_omni.diffusion.data import validate_dlo_host_registration_options
from vllm_omni.diffusion.model_loader.host_weight_plan import (
    HostWeightPlan,
    TensorBinding,
    build_checkpoint_mmap_plan,
)
from vllm_omni.diffusion.offloader import plan_resolver
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
from vllm_omni.diffusion.offloader.block_discovery import (
    get_blocks_attr_names,
    get_blocks_from_dit,
    set_blocks_attr_names,
)
from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
    DistributedLayerwiseOffloadBackend,
    DistributedLayerwiseOffloadHook,
    PinnedResidentLayerGroup,
)
from vllm_omni.diffusion.offloader.host_registration import (
    HostRegistrationCleanupError,
    HostRegistrationError,
)
from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
from vllm_omni.diffusion.offloader.offload_plan import (
    OffloadPlan,
    get_offload_plan,
)
from vllm_omni.diffusion.offloader.plan_resolver import ResolvedComponent, resolve_offload_plan
from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
from vllm_omni.host_weight_runtime import MappedHostRegion
from vllm_omni.platforms import current_omni_platform

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


@pytest.fixture
def patched_offload_runtime(monkeypatch):
    patch_offload_runtime(monkeypatch, dist_backend_module.current_omni_platform, synchronize=True)


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)


class TestDistributedLayerwiseOffloadHook:
    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 = DistributedLayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            dp_group=None,
            dp_size=1,
            rank=0,
            pin_memory=False,
        )
        hook.initialize_hook(current_block)

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

    def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch):
        current_block = nn.Linear(2, 2)
        next_block = nn.Linear(2, 2)
        expected = {name: tensor.detach().clone() for name, tensor in next_block.state_dict().items()}
        hook = DistributedLayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            dp_group=None,
            dp_size=1,
            rank=0,
            pin_memory=False,
        )

        def fail_buffer_allocation():
            raise RuntimeError("injected buffer allocation failure")

        monkeypatch.setattr(hook, "_allocate_device_buffers", fail_buffer_allocation)

        with pytest.raises(RuntimeError, match="buffer allocation failure"):
            hook.initialize_hook(current_block)

        for name, tensor in next_block.state_dict().items():
            torch.testing.assert_close(tensor, expected[name])

    @pytest.mark.parametrize(
        ("kwargs", "message"),
        [
            ({"dp_size": 0, "rank": 0}, "dp_size must be a positive integer"),
            ({"dp_size": 2, "rank": 2, "dp_group": object()}, "rank must satisfy"),
            ({"dp_size": 2, "rank": 0}, "dp_group is required"),
            (
                {"dp_size": 2, "rank": 0, "dp_group": object(), "rank_local_mmap": True},
                "rank_local_mmap requires dp_size=1",
            ),
        ],
    )
    def test_constructor_rejects_invalid_transport_topology(self, kwargs, message):
        options = {"dp_group": None, "dp_size": 1, "rank": 0, **kwargs}

        with pytest.raises(ValueError, match=message):
            DistributedLayerwiseOffloadHook(
                next_block=nn.Linear(2, 2),
                device=torch.device("cpu"),
                pin_memory=False,
                **options,
            )

    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()
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dlo_use_allgather=False,
            ),
            torch.device("cpu"),
        )

        hooks = backend._install_hook_group(blocks, "dit")
        shared_buffers = backend._allocate_shared_buffers(hooks)
        for hook in hooks:
            hook.gpu_buffers = shared_buffers

        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_shard_and_pin_single_rank(self, dist_group, patched_offload_runtime):
        """With dp_size=1, the shard should equal the full weights."""
        current_block = TinyBlock(_make_values(1.0))
        next_block = TinyBlock(_make_values(10.0))

        hook = DistributedLayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            dp_group=None,
            dp_size=1,
            rank=0,
            copy_stream=DummyStream(),
            comm_stream=DummyStream(),
            pin_memory=False,
        )

        hook.initialize_hook(current_block)

        # With dp_size=1, the CPU shard should contain all 4 elements
        assert next_block.weight.dtype in hook.cpu_shards
        shard = hook.cpu_shards[next_block.weight.dtype]
        assert shard.numel() == 4
        assert torch.equal(shard, _make_values(10.0))

    def test_prefetch_preserves_transposed_weight_stride(
        self,
        dist_group,
        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 = DistributedLayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            dp_group=None,
            dp_size=1,
            rank=0,
            copy_stream=DummyStream(),
            comm_stream=DummyStream(),
            pin_memory=False,
        )
        hook.initialize_hook(current_block)
        hook.prefetch_layer(slot=0, non_blocking=False)

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

    def test_allgather_reconstructs_online_fp8_weight_and_scale(
        self,
        patched_offload_runtime,
        monkeypatch,
    ):
        """DLO must reconstruct finalized online-FP8 tensors byte-for-byte."""

        class OnlineFp8Block(nn.Module):
            def __init__(self):
                super().__init__()
                logical_weight = torch.tensor(
                    [
                        [1.0, 2.0, 3.0, 4.0],
                        [5.0, 6.0, 7.0, 8.0],
                        [9.0, 10.0, 11.0, 12.0],
                    ],
                    dtype=torch.float32,
                ).to(torch.float8_e4m3fn)
                self.weight = nn.Parameter(logical_weight.t(), requires_grad=False)
                self.weight_scale = nn.Parameter(torch.tensor([0.25]), requires_grad=False)

        current_block = OnlineFp8Block()
        rank0_block = OnlineFp8Block()
        rank1_block = OnlineFp8Block()
        expected_weight = rank0_block.weight.detach().clone()
        expected_scale = rank0_block.weight_scale.detach().clone()
        expected_stride = rank0_block.weight.stride()

        rank0_hook = DistributedLayerwiseOffloadHook(
            next_block=rank0_block,
            device=torch.device("cpu"),
            dp_group=object(),
            dp_size=2,
            rank=0,
            copy_stream=DummyStream(),
            comm_stream=DummyStream(),
            pin_memory=False,
        )
        rank1_hook = DistributedLayerwiseOffloadHook(
            next_block=rank1_block,
            device=torch.device("cpu"),
            dp_group=object(),
            dp_size=2,
            rank=1,
            copy_stream=DummyStream(),
            comm_stream=DummyStream(),
            pin_memory=False,
        )
        rank0_hook.initialize_hook(current_block)
        rank1_hook.initialize_hook(OnlineFp8Block())
        rank0_hook.gpu_shard_buffers = [
            {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
        ]

        def fake_allgather(output, local_shard, *, group):
            del group
            remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
            shard_size = local_shard.numel()
            output[:shard_size].copy_(local_shard)
            output[shard_size : 2 * shard_size].copy_(remote_shard)

        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
        rank0_hook.prefetch_layer(slot=0, non_blocking=False)

        assert rank0_block.weight.stride() == expected_stride
        torch.testing.assert_close(
            rank0_block.weight.float(),
            expected_weight.float(),
            rtol=0,
            atol=0,
        )
        torch.testing.assert_close(rank0_block.weight_scale, expected_scale)

    def test_allgather_reconstructs_online_int8_weight_and_scale(
        self,
        patched_offload_runtime,
        monkeypatch,
    ):
        """DLO must reconstruct finalized online-INT8 tensors byte-for-byte.

        One block carries both finalized online-INT8 layouts: the NPU method
        leaves a contiguous pre-transposed (K, N) int8 weight with a 1-D fp32
        per-channel scale, while the CUDA method leaves a transposed-view
        weight (stride (1, K)) with a scalar fp32 scale.  The int8 and fp32
        dtype groups must survive the flat-shard AllGather independently.
        """

        class OnlineInt8Block(nn.Module):
            def __init__(self):
                super().__init__()
                logical = torch.arange(1.0, 13.0).reshape(3, 4)
                # NPU online int8: qweight.t().contiguous() -> (K, N), 1-D scale
                self.npu_weight = nn.Parameter(logical.t().contiguous().to(torch.int8), requires_grad=False)
                self.npu_weight_scale = nn.Parameter(torch.tensor([2.0, 3.0, 4.0]), requires_grad=False)
                # CUDA online int8: Parameter(qweight.t().data) -> transposed
                # view (stride (1, K)) over the (N, K) storage, scalar scale
                self.cuda_weight = nn.Parameter((logical + 100).to(torch.int8).t(), requires_grad=False)
                self.cuda_weight_scale = nn.Parameter(torch.tensor([0.5]), requires_grad=False)

        current_block = OnlineInt8Block()
        rank0_block = OnlineInt8Block()
        rank1_block = OnlineInt8Block()
        expected = {name: tensor.detach().clone() for name, tensor in rank0_block.named_parameters()}
        expected_strides = {name: tensor.stride() for name, tensor in rank0_block.named_parameters()}

        def make_hook(block: OnlineInt8Block, rank: int) -> DistributedLayerwiseOffloadHook:
            return DistributedLayerwiseOffloadHook(
                next_block=block,
                device=torch.device("cpu"),
                dp_group=object(),
                dp_size=2,
                rank=rank,
                copy_stream=DummyStream(),
                comm_stream=DummyStream(),
                pin_memory=False,
            )

        rank0_hook = make_hook(rank0_block, 0)
        rank1_hook = make_hook(rank1_block, 1)
        rank0_hook.initialize_hook(current_block)
        rank1_hook.initialize_hook(OnlineInt8Block())
        rank0_hook.gpu_shard_buffers = [
            {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
        ]

        def fake_allgather(output, local_shard, *, group):
            del group
            remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
            shard_size = local_shard.numel()
            output[:shard_size].copy_(local_shard)
            output[shard_size : 2 * shard_size].copy_(remote_shard)

        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
        rank0_hook.prefetch_layer(slot=0, non_blocking=False)

        for name, tensor in rank0_block.named_parameters():
            assert tensor.stride() == expected_strides[name]
            torch.testing.assert_close(
                tensor.float(),
                expected[name].float(),
                rtol=0,
                atol=0,
            )

    def test_allgather_reconstructs_online_mxfp8_weight_and_scale(
        self,
        patched_offload_runtime,
        monkeypatch,
    ):
        """DLO must reconstruct finalized online-MXFP8 tensors byte-for-byte.

        One block carries both finalized online-MXFP8 layouts: the NPU method
        leaves a contiguous pre-transposed (K, N) fp8 weight with a contiguous
        (K_groups/2, N, 2) e8m0 block scale, while the vLLM-kernel method
        (XPU) leaves an (N, K) fp8 weight whose e8m0 scale is a transposed
        view over a contiguous (K/32, N) buffer.  The fp8 and e8m0 dtype
        groups must survive the flat-shard AllGather independently.
        """

        class OnlineMxfp8Block(nn.Module):
            def __init__(self):
                super().__init__()
                logical = torch.arange(1.0, 13.0).reshape(3, 4)
                # Build scales from raw e8m0 bit patterns (1 byte/elem).
                npu_scale_bits = torch.arange(1, 13, dtype=torch.uint8).reshape(2, 3, 2)
                kernel_scale_bits = torch.arange(1, 7, dtype=torch.uint8).reshape(2, 3)
                # NPU online mxfp8: (N, K).t().contiguous() -> (K, N) fp8
                # weight, contiguous (K_groups/2, N, 2) e8m0 scale
                self.npu_weight = nn.Parameter(logical.t().contiguous().to(torch.float8_e4m3fn), requires_grad=False)
                self.npu_weight_scale = nn.Parameter(npu_scale_bits.view(torch.float8_e8m0fnu), requires_grad=False)
                # vLLM-kernel online mxfp8 (XPU): (N, K) fp8 weight, e8m0
                # scale kept as a transposed view (stride (1, K/32)) over the
                # contiguous (K/32, N) storage
                self.kernel_weight = nn.Parameter((logical + 100).to(torch.float8_e4m3fn), requires_grad=False)
                self.kernel_weight_scale = nn.Parameter(
                    kernel_scale_bits.view(torch.float8_e8m0fnu).t(), requires_grad=False
                )

        current_block = OnlineMxfp8Block()
        rank0_block = OnlineMxfp8Block()
        rank1_block = OnlineMxfp8Block()
        expected = {name: tensor.detach().clone() for name, tensor in rank0_block.named_parameters()}
        expected_strides = {name: tensor.stride() for name, tensor in rank0_block.named_parameters()}

        def make_hook(block: OnlineMxfp8Block, rank: int) -> DistributedLayerwiseOffloadHook:
            return DistributedLayerwiseOffloadHook(
                next_block=block,
                device=torch.device("cpu"),
                dp_group=object(),
                dp_size=2,
                rank=rank,
                copy_stream=DummyStream(),
                comm_stream=DummyStream(),
                pin_memory=False,
            )

        rank0_hook = make_hook(rank0_block, 0)
        rank1_hook = make_hook(rank1_block, 1)
        rank0_hook.initialize_hook(current_block)
        rank1_hook.initialize_hook(OnlineMxfp8Block())
        rank0_hook.gpu_shard_buffers = [
            {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
        ]

        def fake_allgather(output, local_shard, *, group):
            del group
            remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
            shard_size = local_shard.numel()
            output[:shard_size].copy_(local_shard)
            output[shard_size : 2 * shard_size].copy_(remote_shard)

        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
        rank0_hook.prefetch_layer(slot=0, non_blocking=False)

        for name, tensor in rank0_block.named_parameters():
            assert tensor.stride() == expected_strides[name]
            # Compare raw bytes in logical order: float8/e8m0 CPU upcasts are
            # version-sensitive, while contiguous().view(uint8) is an exact
            # byte-level comparison for both contiguous and transposed-view
            # parameters.
            torch.testing.assert_close(
                tensor.detach().contiguous().view(torch.uint8),
                expected[name].contiguous().view(torch.uint8),
                rtol=0,
                atol=0,
            )

    def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
        """DTensor wrapper should be preserved through prefetch/offload cycle."""
        current_block = TinyBlock(_make_values(1.0))
        next_block = TinyBlock(_make_values(10.0))

        hook = DistributedLayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            dp_group=None,
            dp_size=1,
            rank=0,
            copy_stream=DummyStream(),
            comm_stream=DummyStream(),
            pin_memory=False,
        )

        hook.initialize_hook(current_block)

        # After init, next_block weights should be placeholders
        assert isinstance(next_block.weight, DTensor)
        assert next_block.weight.to_local().is_meta

        # Prefetch into slot 0
        hook.prefetch_layer(slot=0, non_blocking=False)

        # After prefetch, next_block weights should be materialized
        assert isinstance(next_block.weight, DTensor)
        assert torch.equal(next_block.weight.to_local(), _make_values(10.0))

        # Offload current block
        hook.offload_layer()
        assert isinstance(current_block.weight, DTensor)
        assert current_block.weight.to_local().is_meta
        assert not hook.is_materialized

    def test_double_buffer_slot_swapping(self, dist_group, patched_offload_runtime):
        """Verify slot swapping works correctly after each layer."""
        current_block = TinyBlock(_make_values(1.0))
        next_block = TinyBlock(_make_values(10.0))

        hook = DistributedLayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            dp_group=None,
            dp_size=1,
            rank=0,
            copy_stream=DummyStream(),
            comm_stream=DummyStream(),
            pin_memory=False,
        )

        hook.initialize_hook(current_block)

        assert hook.current_slot == 0

        # Prefetch into slot 1 (next slot)
        hook.prefetch_layer(slot=1, non_blocking=False)

        # Simulate post_forward: offload + swap
        hook.offload_layer()
        hook.current_slot = 1 - hook.current_slot

        assert hook.current_slot == 1

        # Next iteration: prefetch into slot 0
        hook.prefetch_layer(slot=0, non_blocking=False)
        hook.offload_layer()
        hook.current_slot = 1 - hook.current_slot

        assert hook.current_slot == 0

    def test_device_buffers_allocated(self, dist_group, patched_offload_runtime):
        """Verify exactly two device buffers are allocated."""
        current_block = TinyBlock(_make_values(1.0))
        next_block = TinyBlock(_make_values(10.0))

        hook = DistributedLayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            dp_group=None,
            dp_size=1,
            rank=0,
            pin_memory=False,
        )

        hook.initialize_hook(current_block)

        # Two slots should be allocated
        assert hook.gpu_buffers[0] is not None
        assert hook.gpu_buffers[1] is not None

        # Each slot should have one dtype entry (float32)
        assert len(hook.gpu_buffers[0]) == 1
        assert len(hook.gpu_buffers[1]) == 1

        # Buffer size should match total numel
        dtype = next_block.weight.dtype
        assert hook.gpu_buffers[0][dtype].numel() == 4
        assert hook.gpu_buffers[1][dtype].numel() == 4

    def test_sharding_multiple_ranks(self, dist_group, patched_offload_runtime):
        """Verify weight sharding splits correctly across ranks."""
        block = nn.Module()
        block.weight = nn.Parameter(torch.arange(8, dtype=torch.float32))

        next_block = nn.Module()
        next_block.weight = nn.Parameter(torch.arange(100, 108, dtype=torch.float32))

        # Simulate dp_size=4, rank=1
        hook = DistributedLayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            dp_group=object(),
            dp_size=4,
            rank=1,
            pin_memory=False,
        )

        hook.initialize_hook(block)

        shard = hook.cpu_shards[torch.float32]
        # 8 elements / 4 ranks = 2 elements per shard
        assert shard.numel() == 2
        # Rank 1 should have elements [102, 103]
        assert torch.equal(shard, torch.tensor([102.0, 103.0]))

    def test_sharding_with_remainder(self, dist_group, patched_offload_runtime):
        """Verify sharding handles non-even division with padding."""
        block = nn.Module()
        block.weight = nn.Parameter(torch.arange(3, dtype=torch.float32))

        next_block = nn.Module()
        next_block.weight = nn.Parameter(torch.arange(100, 103, dtype=torch.float32))

        # 3 elements, dp_size=2: both ranks get ceil(3/2)=2 elements (padded)
        hook0 = DistributedLayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            dp_group=object(),
            dp_size=2,
            rank=0,
            pin_memory=False,
        )
        # Need fresh next_block for each hook since init replaces tensors
        next_block0 = nn.Module()
        next_block0.weight = nn.Parameter(torch.arange(100, 103, dtype=torch.float32))
        hook0.next_block = next_block0
        hook0.initialize_hook(block)

        shard0 = hook0.cpu_shards[torch.float32]
        assert shard0.numel() == 2
        assert torch.equal(shard0, torch.tensor([100.0, 101.0]))

        next_block1 = nn.Module()
        next_block1.weight = nn.Parameter(torch.arange(100, 103, dtype=torch.float32))
        block1 = nn.Module()
        block1.weight = nn.Parameter(torch.arange(3, dtype=torch.float32))
        hook1 = DistributedLayerwiseOffloadHook(
            next_block=next_block1,
            device=torch.device("cpu"),
            dp_group=object(),
            dp_size=2,
            rank=1,
            pin_memory=False,
        )
        hook1.initialize_hook(block1)

        shard1 = hook1.cpu_shards[torch.float32]
        # Equal-sized shards: rank 1 gets [102, 0] (zero-padded)
        assert shard1.numel() == 2
        assert torch.equal(shard1, torch.tensor([102.0, 0.0]))

    def test_mmap_layout_transform_is_applied_before_sharding(self, dist_group, patched_offload_runtime):
        block = nn.Module()
        block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
        next_block = nn.Module()
        next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))

        def transform(tensor):
            return tensor.flip(0)

        hook = DistributedLayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            dp_group=object(),
            dp_size=2,
            rank=0,
            pin_memory=False,
            tensor_transforms={id(next_block.weight): transform},
        )
        hook.initialize_hook(block)

        assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([3.0, 2.0]))

    def test_transform_is_not_applied_without_a_loader_plan(self, dist_group, patched_offload_runtime):
        block = nn.Module()
        block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
        next_block = nn.Module()
        next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))

        hook = DistributedLayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            dp_group=object(),
            dp_size=2,
            rank=0,
            pin_memory=False,
        )
        hook.initialize_hook(block)

        assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0]))

    def test_rank_local_mmap_retains_source_and_uses_staging(self, dist_group, patched_offload_runtime):
        current_block = nn.Linear(2, 2, bias=False)
        next_block = nn.Linear(2, 2, bias=False)
        next_block.weight.data.copy_(torch.arange(4, dtype=torch.float32).view(2, 2))
        source_storage = next_block.weight.untyped_storage().data_ptr()

        def transform(tensor):
            return tensor.t()

        hook = DistributedLayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            dp_group=None,
            dp_size=1,
            rank=0,
            pin_memory=False,
            rank_local_mmap=True,
            tensor_transforms={id(next_block.weight): transform},
        )
        hook.initialize_hook(current_block)

        source = hook.cpu_sources[torch.float32][0]["tensor"]
        assert source.untyped_storage().data_ptr() == source_storage
        assert hook.cpu_shards == {}
        assert next_block.weight.numel() == 0

        hook.cpu_staging_buffers = [
            {torch.float32: torch.empty(4)},
            {torch.float32: torch.empty(4)},
        ]
        hook.prefetch_layer(slot=0, non_blocking=False)

        assert torch.equal(
            next_block.weight,
            torch.tensor([[0.0, 2.0], [1.0, 3.0]]),
        )
        assert next_block.weight.stride() == (1, 2)
        assert torch.equal(source, torch.arange(4, dtype=torch.float32).view(2, 2))

    def test_rank_local_mmap_staging_is_bounded_by_largest_block(self, patched_offload_runtime):
        hooks = []
        for size in (4, 9):
            current_block = nn.Linear(1, 1, bias=False)
            next_block = nn.Module()
            next_block.weight = nn.Parameter(torch.arange(size, dtype=torch.float32))
            hook = DistributedLayerwiseOffloadHook(
                next_block=next_block,
                device=torch.device("cpu"),
                dp_group=None,
                dp_size=1,
                rank=0,
                pin_memory=False,
                shared_buffers=[None, None],
                rank_local_mmap=True,
            )
            hook.initialize_hook(current_block)
            hooks.append(hook)

        resident_block = nn.Module()
        resident_block.weight = nn.Parameter(torch.arange(12, dtype=torch.float32))
        resident_group = PinnedResidentLayerGroup(
            [resident_block],
            device=torch.device("cpu"),
            copy_stream=DummyStream(),
            pin_memory=False,
            rank_local_mmap=True,
        )
        staging = DistributedLayerwiseOffloadBackend._allocate_shared_cpu_staging_buffers(
            hooks,
            resident_group,
        )

        assert len(staging) == 2
        assert staging[0][torch.float32].numel() == 12
        assert staging[1][torch.float32].numel() == 12


class TestPinnedResidentLayerGroup:
    def test_constructor_failure_leaves_every_block_materialized(self, monkeypatch):
        blocks = [nn.Linear(2, 2), nn.Linear(2, 2)]
        expected = [block.weight.detach().clone() for block in blocks]
        shard_and_pin = DistributedLayerwiseOffloadHook._shard_and_pin
        calls = 0

        def fail_second_block(*args, **kwargs):
            nonlocal calls
            calls += 1
            if calls == 2:
                raise RuntimeError("injected resident setup failure")
            return shard_and_pin(*args, **kwargs)

        monkeypatch.setattr(DistributedLayerwiseOffloadHook, "_shard_and_pin", fail_second_block)

        with pytest.raises(RuntimeError, match="resident setup failure"):
            PinnedResidentLayerGroup(
                blocks,
                device=torch.device("cpu"),
                copy_stream=DummyStream(),
                pin_memory=False,
            )

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

    def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
        blocks = [nn.Linear(2, 2), nn.Linear(2, 2)]
        expected = []
        for index, block in enumerate(blocks):
            block.weight.data.fill_(index + 1)
            block.bias.data.fill_(index + 3)
            expected.append((block.weight.detach().clone(), block.bias.detach().clone()))

        group = PinnedResidentLayerGroup(
            blocks,
            device=torch.device("cpu"),
            copy_stream=DummyStream(),
            pin_memory=False,
        )

        assert all(block.weight.numel() == 0 for block in blocks)

        group.load()
        for block, (weight, bias) in zip(blocks, expected):
            assert torch.equal(block.weight, weight)
            assert torch.equal(block.bias, bias)

        # Device-side changes must not overwrite the persistent host master.
        blocks[0].weight.data.zero_()
        group.offload()
        assert all(block.weight.numel() == 0 for block in blocks)

        group.load()
        for block, (weight, bias) in zip(blocks, expected):
            assert torch.equal(block.weight, weight)
            assert torch.equal(block.bias, bias)

        group.offload()

    def test_load_preserves_noncontiguous_weight_stride(self, patched_offload_runtime):
        block = nn.Module()
        weight = torch.arange(6, dtype=torch.float32).reshape(2, 3).t()
        expected = weight.detach().clone()
        expected_stride = weight.stride()
        block.weight = nn.Parameter(weight)
        group = PinnedResidentLayerGroup(
            [block],
            device=torch.device("cpu"),
            copy_stream=DummyStream(),
            pin_memory=False,
        )

        group.load()
        assert torch.equal(block.weight, expected)
        assert block.weight.stride() == expected_stride
        group.offload()

    def test_rank_local_mmap_resident_layers_retain_sources(self, patched_offload_runtime):
        blocks = [nn.Linear(2, 2, bias=False) for _ in range(3)]
        expected = []
        source_storage = []
        for index, block in enumerate(blocks):
            block.weight.data.fill_(index + 1)
            expected.append(block.weight.detach().clone())
            source_storage.append(block.weight.untyped_storage().data_ptr())

        group = PinnedResidentLayerGroup(
            blocks,
            device=torch.device("cpu"),
            copy_stream=DummyStream(),
            pin_memory=False,
            rank_local_mmap=True,
        )

        for index, state in enumerate(group._states):
            source = state["cpu_sources"][torch.float32][0]["tensor"]
            assert source.untyped_storage().data_ptr() == source_storage[index]
            assert state["cpu_shards"] == {}
        assert len(group._cpu_staging_buffers) == 2

        group.load()
        for block, weight in zip(blocks, expected):
            assert torch.equal(block.weight, weight)

        blocks[0].weight.data.zero_()
        group.offload()
        group.load()
        for block, weight in zip(blocks, expected):
            assert torch.equal(block.weight, weight)

        group.offload()
        assert all(block.weight.numel() == 0 for block in blocks)

    def test_shared_allgather_output_is_narrowed_to_current_block(
        self, dist_group, patched_offload_runtime, monkeypatch
    ):
        block = nn.Module()
        block.weight = nn.Parameter(torch.zeros(3, dtype=torch.float32))
        next_block = nn.Module()
        next_block.weight = nn.Parameter(torch.arange(3, dtype=torch.float32))

        hook = DistributedLayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            dp_group=object(),
            dp_size=2,
            rank=0,
            pin_memory=False,
        )
        hook.initialize_hook(block)
        # Simulate unified buffers sized for a larger heterogeneous block.
        hook.gpu_buffers = [
            {torch.float32: torch.empty(12)},
            {torch.float32: torch.empty(12)},
        ]
        hook.gpu_shard_buffers = [
            {torch.float32: torch.empty(6)},
            {torch.float32: torch.empty(6)},
        ]
        output_sizes = []

        def fake_allgather(output, local_shard, *, group):
            del group
            output_sizes.append(output.numel())
            output.zero_()
            output[: local_shard.numel()].copy_(local_shard)

        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)

        hook.prefetch_layer(slot=0, non_blocking=False)

        assert output_sizes == [4]


class TestPinnedModuleStager:
    def test_offload_rebinds_immutable_cpu_master(self, patched_offload_runtime):
        module = nn.Linear(2, 2)
        expected_weight = module.weight.detach().clone()
        expected_bias = module.bias.detach().clone()
        stager = PinnedModuleStager(
            module,
            torch.device("cpu"),
            pin_memory=False,
            copy_stream=DummyStream(),
        )

        stager.load()
        module.weight.data.zero_()
        module.bias.data.zero_()
        stager.offload()

        assert torch.equal(module.weight, expected_weight)
        assert torch.equal(module.bias, expected_bias)
        stager.load()
        assert torch.equal(module.weight, expected_weight)
        assert torch.equal(module.bias, expected_bias)
        stager.offload()


class _HWRPipeline(nn.Module):
    def __init__(self):
        super().__init__()
        self.transformer = _SingleBlockModel(num_blocks=3)


class _FakeHostWeightLease:
    def __init__(self, resolution_id: str, events: list[str] | None = None):
        self.closed = False
        self.provenance = SimpleNamespace(resolution_id=resolution_id)
        self.mapped_regions = (MappedHostRegion("weights.safetensors", 0x1000, 4096),)
        self._events = events

    def close(self):
        if self._events is not None:
            self._events.append("lease")
        self.closed = True


class _FakeLeaseCarrier:
    def __init__(self, lease: _FakeHostWeightLease):
        self.lease = lease
        self.taken = False

    @property
    def closed(self):
        return self.taken or self.lease.closed

    def take(self):
        if self.taken:
            raise RuntimeError("already taken")
        self.taken = True
        return self.lease

    def close(self):
        if not self.taken:
            self.lease.close()


def _fake_hwr_plan(resolution_id: str = "hwr-test"):
    lease = _FakeHostWeightLease(resolution_id)
    carrier = _FakeLeaseCarrier(lease)
    plan = HostWeightPlan(
        backing_kind="host_weight_runtime",
        bindings={},
        lease_carrier=carrier,
    )
    return plan, carrier, lease


def _hwr_backend():
    plan, carrier, lease = _fake_hwr_plan()
    backend = DistributedLayerwiseOffloadBackend(
        OffloadConfig(
            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
            pin_cpu_memory=False,
            dp_size=1,
            dlo_use_allgather=False,
        ),
        torch.device("cpu"),
        host_weight_plan=plan,
    )
    return _HWRPipeline(), backend, carrier, lease


def test_registered_mmap_hook_bypasses_host_staging(patched_offload_runtime):
    current_block = nn.Linear(2, 2, bias=False)
    next_block = nn.Linear(2, 2, bias=False)
    expected = torch.arange(4, dtype=torch.float32).view(2, 2)
    next_block.weight.data.copy_(expected)
    hook = DistributedLayerwiseOffloadHook(
        next_block=next_block,
        device=torch.device("cpu"),
        dp_group=None,
        dp_size=1,
        rank=0,
        pin_memory=False,
        rank_local_mmap=True,
    )
    hook.initialize_hook(current_block)
    hook.registered_mmap = True
    hook._stage_mmap_sources = lambda _slot: pytest.fail("registered mmap must bypass host staging")  # type: ignore[method-assign]

    hook.prefetch_layer(slot=0, non_blocking=True)

    assert torch.equal(next_block.weight, expected)


def test_registered_mmap_resident_group_bypasses_host_staging(patched_offload_runtime):
    block = nn.Linear(2, 2, bias=False)
    expected = torch.arange(4, dtype=torch.float32).view(2, 2)
    block.weight.data.copy_(expected)
    group = PinnedResidentLayerGroup(
        [block],
        device=torch.device("cpu"),
        copy_stream=DummyStream(),
        pin_memory=False,
        rank_local_mmap=True,
    )
    group.registered_mmap = True
    group._cpu_staging_buffers.clear()

    group.load()

    assert torch.equal(block.weight, expected)


def test_hwr_registration_uses_transport_budget(
    monkeypatch: pytest.MonkeyPatch,
    patched_offload_runtime,
):
    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
    plan, _, lease = _fake_hwr_plan()
    backend = DistributedLayerwiseOffloadBackend(
        OffloadConfig(
            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
            pin_cpu_memory=True,
            dp_size=1,
            dlo_use_allgather=False,
            dlo_host_registration_limit_gib=1.5,
        ),
        torch.device("cuda"),
        host_weight_plan=plan,
    )
    backend._host_weight_lease = lease  # type: ignore[assignment]
    calls: list[tuple[object, int | None]] = []

    class Registration:
        total_bytes = 4096
        region_count = 1

        @staticmethod
        def close() -> tuple[str, ...]:
            return ()

    def register(regions, *, device, max_bytes):
        assert device == torch.device("cuda")
        calls.append((regions, max_bytes))
        return Registration()

    monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
    pinned_source = SimpleNamespace(numel=lambda: 1, is_pinned=lambda: True)

    assert backend._try_register_hwr_mmap((pinned_source,))  # type: ignore[arg-type]
    assert calls == [(lease.mapped_regions, int(1.5 * 1024**3))]
    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)]

    backend._release_registered_mmap()

    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == []


def test_hwr_registration_budget_validation_is_transport_scoped():
    assert (
        validate_dlo_host_registration_options(
            limit_gib=1.5,
            enable_dlo=True,
            use_allgather=False,
            hwr_mode="preferred",
        )
        == 1.5
    )
    invalid = (
        dict(limit_gib=-1, enable_dlo=True, use_allgather=False, hwr_mode="preferred"),
        dict(limit_gib=float("inf"), enable_dlo=True, use_allgather=False, hwr_mode="preferred"),
        dict(limit_gib=1, enable_dlo=False, use_allgather=False, hwr_mode="preferred"),
        dict(limit_gib=1, enable_dlo=True, use_allgather=True, hwr_mode="preferred"),
        dict(limit_gib=1, enable_dlo=True, use_allgather=False, hwr_mode="disabled"),
    )
    for options in invalid:
        with pytest.raises(ValueError):
            validate_dlo_host_registration_options(**options)


def test_unregistration_precedes_lease_close_and_retries_failures(
    monkeypatch: pytest.MonkeyPatch,
    patched_offload_runtime,
):
    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
    events: list[str] = []
    responses = iter([("busy",), ()])

    class Registration:
        @staticmethod
        def close() -> tuple[str, ...]:
            events.append("unregister")
            return next(responses)

    backend = DistributedLayerwiseOffloadBackend(
        OffloadConfig(
            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
            pin_cpu_memory=True,
            dp_size=1,
            dlo_use_allgather=False,
        ),
        torch.device("cpu"),
    )
    lease = _FakeHostWeightLease("retry", events)
    backend._host_weight_lease = lease  # type: ignore[assignment]
    backend._host_registration = Registration()
    backend._using_rank_local_mmap = True
    backend._using_registered_mmap = True
    backend.enabled = True
    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)

    with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"):
        backend.disable()
    assert not lease.closed
    assert backend._host_registration is not None
    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)]

    backend.disable()

    assert lease.closed
    assert backend._host_registration is None
    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == []
    assert events == ["unregister", "unregister", "lease"]


def test_failed_unregistration_retains_lease_after_backend_is_dropped(
    monkeypatch: pytest.MonkeyPatch,
    patched_offload_runtime,
):
    events: list[str] = []

    class Registration:
        @staticmethod
        def close() -> tuple[str, ...]:
            events.append("unregister")
            return ("busy",)

    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
    backend = DistributedLayerwiseOffloadBackend(
        OffloadConfig(
            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
            pin_cpu_memory=True,
            dp_size=1,
            dlo_use_allgather=False,
        ),
        torch.device("cpu"),
    )
    lease = _FakeHostWeightLease("retained", events)
    registration = Registration()
    backend._host_weight_lease = lease  # type: ignore[assignment]
    backend._host_registration = registration
    backend._using_rank_local_mmap = True
    backend.enabled = True
    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
    lease_ref = weakref.ref(lease)

    with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"):
        backend.disable()
    del backend
    del lease
    gc.collect()

    retained_lease = dist_backend_module._ACTIVE_HWR_REGISTRATIONS[0][1]
    assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(registration, retained_lease)]
    assert lease_ref() is retained_lease
    assert not retained_lease.closed

    dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
    retained_lease.close()


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 _NoAttrsModel(nn.Module):
    def __init__(self, num_blocks: int = 2):
        super().__init__()
        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])


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

    def __init__(self):
        super().__init__()
        self.time_embedder = nn.Linear(2, 2, bias=False)
        self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(2)])
        self.post_load_calls = 0

    def post_load_weights(self) -> None:
        self.time_embedder.to(torch.float32)
        self.post_load_calls += 1


class _MmapPostLoadPipeline(nn.Module):
    def __init__(self):
        super().__init__()
        with torch.device("meta"):
            self.transformer = _MmapPostLoadModel()

    @staticmethod
    def _remap_ckpt_key(key: str) -> str:
        return key


class TestMmapWeightLoading:
    @pytest.mark.parametrize("initial_device", ["cpu", "meta"])
    @pytest.mark.parametrize("use_allgather", [False, True])
    def test_mmap_keeps_undeclared_nested_blocks_streamed(
        self, tmp_path, patched_offload_runtime, monkeypatch, initial_device, use_allgather
    ):
        class Refiner(nn.Module):
            def __init__(self):
                super().__init__()
                self.layers = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])

            def forward(self, x):
                for block in self.layers:
                    x = block(x)
                return x

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

            def __init__(self):
                super().__init__()
                self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
                self.refiner = Refiner()

            def forward(self, x):
                x = self.refiner(x)
                for block in self.blocks:
                    x = block(x)
                return x

        pipeline = nn.Module()
        pipeline.transformer = Transformer()
        weights = {name: param.detach().clone() for name, param in pipeline.named_parameters()}
        x = torch.randn(1, 2)
        with torch.no_grad():
            expected = pipeline.transformer(x)
        checkpoint_file = tmp_path / "model.safetensors"
        save_file(weights, str(checkpoint_file))
        # Exercise the real mmap path with a small model instead of allocating
        # the 1 GiB needed to trigger the production compatibility threshold.
        monkeypatch.setattr(plan_resolver, "_NESTED_OFFLOAD_THRESHOLD_MB", 0)
        if initial_device == "meta":
            pipeline.transformer.to_empty(device="meta")
        plan = HostWeightPlan(
            backing_kind="checkpoint_mmap",
            bindings={name: TensorBinding(checkpoint_key=name, file_path=str(checkpoint_file)) for name in weights},
        )
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dlo_use_allgather=use_allgather,
            ),
            torch.device("cpu"),
            host_weight_plan=plan,
        )

        for _ in range(2):
            try:
                backend.enable(pipeline)
                assert len(backend._all_hook_groups) == 2
                with torch.no_grad():
                    torch.testing.assert_close(pipeline.transformer(x), expected)
            finally:
                backend.disable()
            for name, parameter in pipeline.named_parameters():
                assert not parameter.is_meta
                torch.testing.assert_close(parameter, weights[name])
            for module in pipeline.modules():
                registry = getattr(module, "_hook_registry", None)
                if registry is not None:
                    assert registry.get_hook("distributed_layerwise_offload") is None

    def test_runs_model_post_load_hook(self, tmp_path, patched_offload_runtime):
        pipeline = _MmapPostLoadPipeline()
        weights = {name: torch.ones(param.shape, dtype=torch.bfloat16) for name, param in pipeline.named_parameters()}
        weight_file = tmp_path / "model.safetensors"
        save_file(weights, str(weight_file))
        weight_map = {name: weight_file.name for name in weights}
        (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map}))

        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
            ),
            torch.device("cpu"),
        )
        dits = (ResolvedComponent(path="transformer", module=pipeline.transformer, selected=True),)
        plan = HostWeightPlan(
            backing_kind="checkpoint_mmap",
            bindings={
                name: TensorBinding(
                    checkpoint_key=name,
                    file_path=str(weight_file),
                )
                for name in weights
            },
        )

        backend._load_weights_via_mmap(pipeline, dits, plan)

        assert pipeline.transformer.post_load_calls == 1
        assert pipeline.transformer.time_embedder.weight.dtype == torch.float32
        assert pipeline.transformer.blocks[0].weight.dtype == torch.bfloat16

    def test_allgather_flag_with_group_size_one_uses_rank_local_mmap(
        self,
        tmp_path,
        patched_offload_runtime,
    ):
        class Transformer(nn.Module):
            _layerwise_offload_blocks_attrs = ["blocks"]

            def __init__(self):
                super().__init__()
                self.blocks = nn.ModuleList(
                    [
                        nn.Linear(2, 2, bias=False),
                        nn.Linear(2, 2, bias=False),
                    ]
                )

        class Pipeline(nn.Module):
            def __init__(self):
                super().__init__()
                self.transformer = Transformer()

            @staticmethod
            def _remap_ckpt_key(key):
                return key

        pipeline = Pipeline()
        weights = {name: torch.ones_like(param) for name, param in pipeline.named_parameters()}
        checkpoint_file = tmp_path / "model.safetensors"
        save_file(weights, str(checkpoint_file))
        result = build_checkpoint_mmap_plan(
            pipeline,
            dit_modules=(("transformer", pipeline.transformer),),
            sources=(),
            model_path=str(tmp_path),
            tensor_parallel_size=1,
            use_hsdp=False,
            online_quantization=False,
        )
        assert result.plan is not None
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dp_size=1,
                dlo_use_allgather=True,
            ),
            torch.device("cpu"),
            host_weight_plan=result.plan,
        )

        backend.enable(pipeline)

        assert backend._using_rank_local_mmap
        assert backend.dp_group is None
        assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group)
        backend.disable()

    def test_hwr_carrier_is_taken_and_released_after_bounded_staging(self, patched_offload_runtime):
        """A warm HWR plan uses the existing two-slot rank-local transport."""
        pipeline, backend, carrier, lease = _hwr_backend()
        model_bytes = sum(param.numel() * param.element_size() for param in pipeline.transformer.parameters())

        backend.enable(pipeline)

        assert carrier.taken
        assert backend._using_rank_local_mmap
        assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group)
        staging_bytes = sum(
            buffer.numel() * buffer.element_size()
            for buffer in backend._all_hook_groups[0][0].cpu_staging_buffers[0].values()
        )
        assert staging_bytes * 2 < model_bytes
        assert not lease.closed

        backend.disable()
        assert lease.closed

        backend.enable(pipeline)
        assert backend.host_weight_plan is None
        assert backend.enabled
        backend.disable()

    def test_rank_local_disable_drains_pending_transfers(self, patched_offload_runtime, monkeypatch):
        pipeline = _HWRPipeline()
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dp_size=1,
                dlo_use_allgather=False,
            ),
            torch.device("cpu"),
        )
        backend.enable(pipeline)
        backend._all_hook_groups[0][0].prefetch_layer(0)
        synchronize = Mock()
        monkeypatch.setattr(current_omni_platform, "synchronize", synchronize)

        backend.disable()

        synchronize.assert_called()

    def test_hwr_registration_failure_falls_back_to_bounded_staging(
        self,
        monkeypatch,
        patched_offload_runtime,
    ):
        plan, carrier, lease = _fake_hwr_plan("hwr-registration-fallback")
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=True,
                dp_size=1,
                dlo_use_allgather=False,
            ),
            torch.device("cpu"),
            host_weight_plan=plan,
        )
        pipeline = _HWRPipeline()
        original_staging_allocator = backend._allocate_shared_cpu_staging_buffers

        def fail_registration(*args, **kwargs):
            del args, kwargs
            raise HostRegistrationError("registration unavailable in CPU test")

        def allocate_unpinned_staging(hooks, resident_group=None):
            # Registration selection still observes pin_cpu_memory=True. The
            # CPU-only test avoids asking the host for CUDA-pinned allocations.
            for hook in hooks:
                hook.pin_memory = False
            return original_staging_allocator(hooks, resident_group)

        monkeypatch.setattr(dist_backend_module, "register_host_mappings", fail_registration)
        monkeypatch.setattr(backend, "_allocate_shared_cpu_staging_buffers", allocate_unpinned_staging)

        backend.enable(pipeline)

        assert carrier.taken
        assert backend._using_rank_local_mmap
        assert not backend._using_registered_mmap
        assert backend._host_registration is None
        assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group)
        assert not lease.closed

        backend.disable()
        assert lease.closed

    def test_hwr_backend_failure_drains_partial_setup_before_lease_close(
        self,
        monkeypatch,
        patched_offload_runtime,
    ):
        pipeline, backend, carrier, lease = _hwr_backend()
        monkeypatch.setattr(
            backend,
            "_allocate_shared_buffers",
            lambda hooks: (_ for _ in ()).throw(RuntimeError("device buffer allocation failed")),
        )

        with pytest.raises(RuntimeError, match="device buffer allocation failed"):
            backend.enable(pipeline)

        assert carrier.taken
        assert lease.closed
        assert not backend.enabled
        assert not backend._blocks
        assert not backend._all_hook_groups


class TestGetBlocksFromDit:
    def test_get_blocks_from_dit_single_block_attr(self):
        model = _SingleBlockModel(num_blocks=3)
        attr_names, blocks = 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 = get_blocks_from_dit(model)
        assert set(attr_names) == {"transformer_blocks", "single_transformer_blocks"}
        assert len(blocks) == 5

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

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


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

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

    def test_set_blocks_attr_names(self):
        model = _NoAttrsModel()
        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 TestCrossGroupSharedBuffer:
    """Regression test: when multiple DiT groups share the same 2 GPU
    buffers, the first block of each group must always sync-prefetch on
    entry, because another group may have overwritten the shared slot.
    """

    def test_group_first_hook_forces_sync_prefetch(self, dist_group, patched_offload_runtime):
        """Verify _is_group_first=True forces sync-prefetch in pre_forward
        even when is_materialized returns True."""
        block_a = TinyBlock(_make_values(1.0))
        block_b = TinyBlock(_make_values(10.0))

        hook_a = DistributedLayerwiseOffloadHook(
            next_block=block_b,
            device=torch.device("cpu"),
            dp_group=None,
            dp_size=1,
            rank=0,
            pin_memory=False,
        )
        hook_a.initialize_hook(block_a)

        # Prefetch to make block_a "materialized" (non-empty params)
        hook_a.prefetch_layer(hook_a.current_slot, non_blocking=False)
        assert hook_a.is_materialized

        # Track sync-prefetch calls
        sync_called = [False]
        orig = hook_a.prefetch_layer

        def tracking(slot, non_blocking=True):
            if not non_blocking:
                sync_called[0] = True
            orig(slot, non_blocking=non_blocking)

        hook_a.prefetch_layer = tracking

        # Case 1: _is_group_first=False, materialized → no sync-prefetch
        hook_a._is_group_first = False
        sync_called[0] = False
        hook_a.pre_forward(block_a)
        assert not sync_called[0], "Materialized block should not sync-prefetch without _is_group_first"

        # Case 2: _is_group_first=True, materialized, slot contaminated
        # (no _shared_slot_group set → defaults to contaminated) → MUST sync-prefetch
        hook_a._is_group_first = True
        hook_a._prev_hook = hook_a
        sync_called[0] = False
        hook_a.pre_forward(block_a)
        assert sync_called[0], "_is_group_first must force sync-prefetch when slot is contaminated"

        # Case 3: _is_group_first=True, slot NOT contaminated (same group owns slot)
        # → skip sync-prefetch, just wait for async event
        hook_a._shared_slot_group = [hook_a._group_id, -1]
        sync_called[0] = False
        hook_a.pre_forward(block_a)
        assert not sync_called[0], "Should skip sync-prefetch when slot_group matches (non-contaminated)"


class TestAllGatherOutputSlicing:
    """Regression test: with heterogeneous block sizes (e.g. MiniMax H3's
    token_refiner with blocks of ~1231MB and ~239MB), the shared AllGather
    output buffers are sized to the *largest* block across all groups.
    prefetch_layer must slice the output buffer down to each block's actual
    AllGather output size, otherwise output.numel() != dp_size * input.numel()
    and all_gather_into_tensor raises during enable()'s first prefetch.
    """

    def test_prefetch_slices_shared_output_buffer_to_block_size(self, dist_group, patched_offload_runtime, monkeypatch):
        dp_size = 2

        # Heterogeneous block sizes: 8 vs 32 elements.
        small_block = nn.Module()
        small_block.weight = nn.Parameter(torch.arange(8, dtype=torch.float32))
        large_block = nn.Module()
        large_block.weight = nn.Parameter(torch.arange(100, 132, dtype=torch.float32))

        def make_hook(next_block: nn.Module) -> DistributedLayerwiseOffloadHook:
            hook = DistributedLayerwiseOffloadHook(
                next_block=next_block,
                device=torch.device("cpu"),
                dp_group=dist.group.WORLD,  # non-None -> AllGather path (mocked below)
                dp_size=dp_size,
                rank=0,
                pin_memory=False,
            )
            current_block = nn.Module()
            current_block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
            hook.initialize_hook(current_block)
            return hook

        hook_small = make_hook(small_block)
        hook_large = make_hook(large_block)
        hooks = [hook_small, hook_large]

        # Assign max-sized shared buffers exactly as backend.enable() does.
        shared_buffers = DistributedLayerwiseOffloadBackend._allocate_shared_buffers(hooks)
        shared_shard_buffers = DistributedLayerwiseOffloadBackend._allocate_shared_shard_buffers(hooks)
        for hook in hooks:
            hook.gpu_buffers = shared_buffers
            hook.gpu_shard_buffers = shared_shard_buffers

        # Shared output buffer must be sized to the largest block (32).
        assert shared_buffers[0][torch.float32].numel() == 32

        # Fake AllGather that enforces the real torch size contract and
        # emulates reconstruction of the full block from all ranks' shards.
        full_values = {
            8: torch.arange(8, dtype=torch.float32),
            32: torch.arange(100, 132, dtype=torch.float32),
        }
        ag_output_numels: list[int] = []

        def fake_all_gather(output, input, group=None):
            assert output.numel() == dp_size * input.numel(), (
                f"AllGather contract violated: output.numel()={output.numel()} "
                f"!= {dp_size} * input.numel()={input.numel()}"
            )
            ag_output_numels.append(output.numel())
            output.copy_(full_values[output.numel()])

        monkeypatch.setattr(dist, "all_gather_into_tensor", fake_all_gather)

        # Small block: output must be sliced to dp_size * shard = 8, not 32.
        hook_small.prefetch_layer(slot=0, non_blocking=False)
        assert ag_output_numels[-1] == 8
        assert torch.equal(small_block.weight, full_values[8])

        # Large block: uses the full max-sized buffer.
        hook_large.prefetch_layer(slot=1, non_blocking=False)
        assert ag_output_numels[-1] == 32
        assert torch.equal(large_block.weight, full_values[32])


class TestOffloadPlan:
    """Test declarative OffloadPlan metadata."""

    def test_get_offload_plan_returns_none_when_not_declared(self):
        """Pipelines without _offload_plan should return None."""
        model = _SingleBlockModel(num_blocks=3)
        assert get_offload_plan(model) is None

    def test_get_offload_plan_returns_declared_plan(self):
        """Pipelines with _offload_plan should return it."""
        plan = OffloadPlan(
            block_attrs={"transformer": ("gen_layers",)},
            offload_submodules={"context_encoder": "layers"},
            resident_dit_paths=frozenset({"transformer"}),
        )

        class PipelineWithPlan(nn.Module):
            _offload_plan = plan

        model = PipelineWithPlan()
        result = get_offload_plan(model)
        assert result is plan
        assert result.block_attrs == {"transformer": ("gen_layers",)}
        assert result.offload_submodules == {"context_encoder": "layers"}
        assert result.resident_dit_paths == frozenset({"transformer"})

    def test_offload_plan_defaults_to_empty(self):
        """OffloadPlan with no arguments should have empty dicts."""
        plan = OffloadPlan()
        assert plan.block_attrs == {}
        assert plan.offload_submodules == {}
        assert plan.resident_dit_paths == frozenset()
        assert plan.encoder_dlo_weight_replication == frozenset()

    def test_offload_plan_is_frozen(self):
        """OffloadPlan should be immutable (frozen=True)."""
        plan = OffloadPlan()
        with pytest.raises(Exception):
            plan.block_attrs = {"x": ("y",)}  # type: ignore

    def test_all_resident_dit_still_prepares_planned_submodules(self, patched_offload_runtime):
        class TokenRefiner(nn.Module):
            def __init__(self):
                super().__init__()
                self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])

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

            def __init__(self):
                super().__init__()
                self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
                self.token_refiner = TokenRefiner()

        class Pipeline(nn.Module):
            _offload_plan = OffloadPlan(
                offload_submodules={"token_refiner": "blocks"},
                resident_dit_paths=frozenset({"transformer"}),
            )

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

        pipeline = Pipeline()
        expected_parameters = {name: parameter.detach().clone() for name, parameter in pipeline.named_parameters()}
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dlo_use_allgather=False,
                dlo_resident_layers=2,
            ),
            torch.device("cpu"),
        )

        backend.enable(pipeline)

        assert len(backend._resident_blocks) == 2
        assert len(backend._all_hook_groups) == 1
        assert len(backend._all_hook_groups[0]) == 2
        assert pipeline._dlo_residency_controller is backend
        assert backend.enabled

        backend.disable()

        assert pipeline._dlo_residency_controller is None

        for name, parameter in pipeline.named_parameters():
            torch.testing.assert_close(parameter, expected_parameters[name])

        backend.enable(pipeline)
        assert len(backend._resident_blocks) == 2
        backend.disable()
        for name, parameter in pipeline.named_parameters():
            torch.testing.assert_close(parameter, expected_parameters[name])

    @pytest.mark.parametrize("shared_owner", ["parent", "encoder", "sibling"])
    def test_nested_duplicate_ownership_is_rejected_before_any_hook(self, patched_offload_runtime, shared_owner):
        class Pipeline(nn.Module):
            _offload_plan = OffloadPlan(
                offload_submodules={"refiner": "blocks", "other_refiner": "blocks"},
                encoder_block_attrs={"text_encoder": ("encoder.block",)},
            )

            def __init__(self):
                super().__init__()
                self.transformer = _SingleBlockModel(num_blocks=2)
                self.transformer.refiner = _SingleBlockModel(num_blocks=2)
                self.transformer.other_refiner = _SingleBlockModel(num_blocks=2)
                self.text_encoder = _PlainEncoder()
                owners = {
                    "parent": self.transformer.blocks,
                    "encoder": self.text_encoder.encoder.block,
                    "sibling": self.transformer.other_refiner.blocks,
                }
                self.transformer.refiner.blocks = owners[shared_owner]

        pipeline = Pipeline()
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False),
            torch.device("cpu"),
        )
        storage = {name: parameter.data_ptr() for name, parameter in pipeline.named_parameters()}
        try:
            with pytest.raises(ValueError, match="claimed by both"):
                backend.enable(pipeline)
            assert not backend.enabled
            assert not backend._all_hook_groups
            assert {name: parameter.data_ptr() for name, parameter in pipeline.named_parameters()} == storage
            assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False)
            for module in pipeline.modules():
                assert getattr(module, "_hook_registry", None) is None
        finally:
            backend.disable()

    def test_unstageable_submodule_is_rejected_before_any_hook(self, patched_offload_runtime):
        """Plan-dependent rejection happens before earlier components are hooked."""

        class Pipeline(nn.Module):
            _offload_plan = OffloadPlan(
                offload_submodules={"refiner": "missing"},
                encoder_component_types={"text_encoder": "text_encoder"},
                encoder_block_attrs={"text_encoder": ("encoder.block",)},
            )

            def __init__(self):
                super().__init__()
                self.transformer = _SingleBlockModel(num_blocks=2)
                # Declared for submodule offload but without a block list or the
                # load_to_device/offload_to_cpu lifecycle.
                self.transformer.refiner = nn.Linear(2, 2)
                self.text_encoder = _PlainEncoder()

        pipeline = Pipeline()
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False),
            torch.device("cpu"),
        )

        # Storage identity proves the encoder was never staged and rolled back.
        encoder_storage = [block.weight.data_ptr() for block in pipeline.text_encoder.encoder.block]

        with pytest.raises(ValueError, match="must implement load_to_device"):
            backend.enable(pipeline)

        assert not backend.enabled
        assert [block.weight.data_ptr() for block in pipeline.text_encoder.encoder.block] == encoder_storage
        assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False)
        for block in pipeline.transformer.blocks:
            assert getattr(block, "_hook_registry", None) is None

    def test_undeclared_submodule_block_scan_warns(self, patched_offload_runtime, monkeypatch):
        """The size-triggered attribute scan still works and is deprecated."""

        class Refiner(nn.Module):
            def __init__(self):
                super().__init__()
                self.layers = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])

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

            def __init__(self):
                super().__init__()
                self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
                self.refiner = Refiner()

        class Pipeline(nn.Module):
            def __init__(self):
                super().__init__()
                self.transformer = Transformer()

        warnings: list[str] = []

        class _Recorder:
            def warning(self, message, *args):
                warnings.append(message % args if args else message)

            def __getattr__(self, _name):
                return lambda *args, **kwargs: None

        monkeypatch.setattr(plan_resolver, "logger", _Recorder())
        monkeypatch.setattr(plan_resolver, "_NESTED_OFFLOAD_THRESHOLD_MB", 0)
        plan_resolver._warn_nested_block_scan.cache_clear()
        plan_resolver._warn_legacy_discovery.cache_clear()
        pipeline = Pipeline()
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False),
            torch.device("cpu"),
        )

        backend.enable(pipeline)

        assert len(backend._all_hook_groups) == 2
        assert sum("OffloadPlan.offload_submodules" in message for message in warnings) == 1

        backend.disable()


class TestMmapValidation:
    """Tests for loader preflight and backend plan realization."""

    def test_component_source_prefix_builds_model_namespace(self, tmp_path):
        source_root = tmp_path / "partition"
        checkpoint_dir = source_root / "transformer"
        checkpoint_dir.mkdir(parents=True)
        checkpoint_file = checkpoint_dir / "model.safetensors"
        save_file(
            {"blocks.0.weight": torch.arange(4, dtype=torch.float32).reshape(1, 4)},
            str(checkpoint_file),
        )

        class Pipeline(nn.Module):
            def __init__(self):
                super().__init__()
                self.transformer = nn.Module()
                self.transformer.blocks = nn.ModuleList([nn.Linear(4, 1, bias=False)])
                self.weights_sources = [
                    SimpleNamespace(
                        model_or_path=str(source_root),
                        subfolder="transformer",
                        revision=None,
                        prefix="transformer.",
                    )
                ]

        pipeline = Pipeline()
        result = build_checkpoint_mmap_plan(
            pipeline,
            dit_modules=(("transformer", pipeline.transformer),),
            sources=tuple(pipeline.weights_sources),
            model_path=None,
            tensor_parallel_size=1,
            use_hsdp=False,
            online_quantization=False,
        )

        assert result.fallback_reason is None
        assert result.plan is not None
        assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding(
            checkpoint_key="blocks.0.weight",
            file_path=str(checkpoint_file),
        )
        assert result.plan.planned_source_prefixes == frozenset({"transformer."})

    def test_non_dedicated_component_source_falls_back_before_discovery(self):
        pipeline = nn.Module()
        pipeline.transformer = nn.Linear(2, 2, bias=False)
        sources = (
            SimpleNamespace(
                model_or_path="unused",
                subfolder=None,
                revision=None,
                prefix="",
            ),
        )

        result = build_checkpoint_mmap_plan(
            pipeline,
            dit_modules=(("transformer", pipeline.transformer),),
            sources=sources,
            model_path=None,
            tensor_parallel_size=1,
            use_hsdp=False,
            online_quantization=False,
        )

        assert result.plan is None
        assert "dedicated component weight source" in result.fallback_reason

    def test_preflight_falls_back_before_mutation_for_missing_model_tensors(self, tmp_path):
        source_root = tmp_path / "partition"
        checkpoint_dir = source_root / "transformer"
        checkpoint_dir.mkdir(parents=True)
        save_file(
            {"blocks.0.weight": torch.ones((2, 2), dtype=torch.float32)},
            str(checkpoint_dir / "model.safetensors"),
        )

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

            def __init__(self):
                super().__init__()
                self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False), nn.Linear(2, 2, bias=False)])

        class Pipeline(nn.Module):
            def __init__(self):
                super().__init__()
                self.transformer = Transformer()
                self.weights_sources = [
                    SimpleNamespace(
                        model_or_path=str(source_root),
                        subfolder="transformer",
                        revision=None,
                        prefix="transformer.",
                    )
                ]

        pipeline = Pipeline()
        original_weight = pipeline.transformer.blocks[0].weight
        result = build_checkpoint_mmap_plan(
            pipeline,
            dit_modules=(("transformer", pipeline.transformer),),
            sources=tuple(pipeline.weights_sources),
            model_path=None,
            tensor_parallel_size=1,
            use_hsdp=False,
            online_quantization=False,
        )

        assert result.plan is None
        assert "1 required DiT tensors" in result.fallback_reason
        assert pipeline.transformer.blocks[0].weight is original_weight
        assert not original_weight.is_meta

    def test_tp_falls_back_before_checkpoint_discovery(self):
        class SimpleModel(nn.Module):
            _layerwise_offload_blocks_attrs = ["blocks"]

            def __init__(self):
                super().__init__()
                self.blocks = nn.ModuleList([nn.Linear(4, 4, bias=False)])

        class SimplePipeline(nn.Module):
            def __init__(self):
                super().__init__()
                self.transformer = SimpleModel()

            @staticmethod
            def _remap_ckpt_key(key):
                return key

        pipeline = SimplePipeline()
        result = build_checkpoint_mmap_plan(
            pipeline,
            dit_modules=(("transformer", pipeline.transformer),),
            sources=(),
            model_path="unused",
            tensor_parallel_size=2,
            use_hsdp=False,
            online_quantization=False,
        )

        assert result.plan is None
        assert result.fallback_reason == "TP=2 requires the ordinary loader"

    @pytest.mark.parametrize(
        ("use_hsdp", "online_quantization", "expected_reason"),
        [
            (True, False, "HSDP requires the ordinary loader"),
            (False, True, "online quantization requires the ordinary loader"),
        ],
    )
    def test_runtime_layouts_that_need_loader_callbacks_fall_back(
        self,
        use_hsdp,
        online_quantization,
        expected_reason,
    ):
        pipeline = nn.Module()
        pipeline.transformer = nn.Linear(2, 2, bias=False)

        result = build_checkpoint_mmap_plan(
            pipeline,
            dit_modules=(("transformer", pipeline.transformer),),
            sources=(),
            model_path="unused",
            tensor_parallel_size=1,
            use_hsdp=use_hsdp,
            online_quantization=online_quantization,
        )

        assert result.plan is None
        assert result.fallback_reason == expected_reason

    @pytest.mark.parametrize(
        ("checkpoint_tensor", "expected_reason"),
        [
            (torch.ones((3, 2), dtype=torch.float32), "shape mismatch"),
            (torch.ones((2, 2), dtype=torch.bfloat16), "dtype mismatch"),
        ],
    )
    def test_source_metadata_mismatch_falls_back(self, tmp_path, checkpoint_tensor, expected_reason):
        checkpoint_file = tmp_path / "model.safetensors"
        save_file({"transformer.weight": checkpoint_tensor}, str(checkpoint_file))

        class Pipeline(nn.Module):
            def __init__(self):
                super().__init__()
                self.transformer = nn.Linear(2, 2, bias=False)

            @staticmethod
            def _remap_ckpt_key(key):
                return key

        pipeline = Pipeline()

        result = build_checkpoint_mmap_plan(
            pipeline,
            dit_modules=(("transformer", pipeline.transformer),),
            sources=(),
            model_path=str(tmp_path),
            tensor_parallel_size=1,
            use_hsdp=False,
            online_quantization=False,
        )

        assert result.plan is None
        assert expected_reason in result.fallback_reason

    def test_non_persistent_buffers_do_not_need_checkpoint_bindings(self, tmp_path):
        class Transformer(nn.Module):
            def __init__(self):
                super().__init__()
                self.weight = nn.Parameter(torch.ones(2, 2))
                self.register_buffer("derived", torch.ones(2), persistent=False)

        checkpoint_file = tmp_path / "model.safetensors"
        save_file({"transformer.weight": torch.ones(2, 2)}, str(checkpoint_file))

        class Pipeline(nn.Module):
            def __init__(self):
                super().__init__()
                self.transformer = Transformer()

            @staticmethod
            def _remap_ckpt_key(key):
                return key

        pipeline = Pipeline()

        result = build_checkpoint_mmap_plan(
            pipeline,
            dit_modules=(("transformer", pipeline.transformer),),
            sources=(),
            model_path=str(tmp_path),
            tensor_parallel_size=1,
            use_hsdp=False,
            online_quantization=False,
        )

        assert result.plan is not None
        assert set(result.plan.bindings) == {"transformer.weight"}

    def test_unadapted_custom_weight_loader_falls_back_at_tp1(self, tmp_path):
        def custom_weight_loader(param, weight):
            param.data.copy_(weight)

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

            def __init__(self):
                super().__init__()
                self.linear = nn.Linear(4, 4, bias=False)
                # Simulate QKVParallelLinear: custom weight_loader even at TP=1
                self.linear.weight.weight_loader = custom_weight_loader
                self.blocks = nn.ModuleList([nn.Linear(4, 4, bias=False)])

        class PipelineWithLoader(nn.Module):
            def __init__(self):
                super().__init__()
                self.transformer = ModelWithCustomLoader()

            @staticmethod
            def _remap_ckpt_key(key):
                return key

        pipeline = PipelineWithLoader()
        weights = {name: torch.ones(p.shape, dtype=p.dtype) for name, p in pipeline.named_parameters() if not p.is_meta}
        save_file(weights, str(tmp_path / "model.safetensors"))
        (tmp_path / "model.safetensors.index.json").write_text(
            json.dumps({"weight_map": {k: "model.safetensors" for k in weights}})
        )

        result = build_checkpoint_mmap_plan(
            pipeline,
            dit_modules=(("transformer", pipeline.transformer),),
            sources=(),
            model_path=str(tmp_path),
            tensor_parallel_size=1,
            use_hsdp=False,
            online_quantization=False,
        )

        assert result.plan is None
        assert "custom weight_loader" in result.fallback_reason

    def test_cosmos3_adapter_preserves_existing_tp1_direct_mmap_contract(self, tmp_path):
        from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
            Cosmos3OmniDiffusersPipeline,
        )
        from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
            Cosmos3VFMTransformer,
        )

        transformer = object.__new__(Cosmos3VFMTransformer)
        nn.Module.__init__(transformer)
        transformer.proj_in = nn.Linear(2, 2, bias=False)
        transformer.proj_in.weight.weight_loader = lambda *_args: None
        pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
        nn.Module.__init__(pipeline)
        pipeline.transformer = transformer

        source_root = tmp_path / "partition"
        checkpoint_dir = source_root / "transformer"
        checkpoint_dir.mkdir(parents=True)
        checkpoint_file = checkpoint_dir / "model.safetensors"
        save_file({"proj_in.weight": torch.ones(2, 2)}, str(checkpoint_file))
        sources = (
            SimpleNamespace(
                model_or_path=str(source_root),
                subfolder="transformer",
                revision=None,
                prefix="transformer.",
            ),
        )

        result = build_checkpoint_mmap_plan(
            pipeline,
            dit_modules=(("transformer", transformer),),
            sources=sources,
            model_path=None,
            tensor_parallel_size=1,
            use_hsdp=False,
            online_quantization=False,
        )

        assert result.plan is not None
        assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding(
            checkpoint_key="proj_in.weight",
            file_path=str(checkpoint_file),
        )

    def test_validate_loaded_weights_called(self, tmp_path, patched_offload_runtime):
        """validate_loaded_weights should be called after mmap loading."""

        class ModelWithValidate(nn.Module):
            _layerwise_offload_blocks_attrs = ["blocks"]
            validate_called = False

            def __init__(self):
                super().__init__()
                self.blocks = nn.ModuleList([nn.Linear(4, 4, bias=False) for _ in range(2)])

            def post_load_weights(self):
                pass

            def validate_loaded_weights(self, loaded):
                self.validate_called = True
                assert len(loaded) > 0, "Should have loaded weights"

        class PipelineWithValidate(nn.Module):
            def __init__(self):
                super().__init__()
                self.transformer = ModelWithValidate()

            @staticmethod
            def _remap_ckpt_key(key):
                return key

        pipeline = PipelineWithValidate()
        # Save on meta to trigger mmap path
        with torch.device("meta"):
            pipeline.transformer = ModelWithValidate()

        weights = {
            name: torch.ones(param.shape, dtype=param.dtype if not param.is_meta else torch.float32)
            for name, param in pipeline.named_parameters()
        }
        save_file(weights, str(tmp_path / "model.safetensors"))
        (tmp_path / "model.safetensors.index.json").write_text(
            json.dumps({"weight_map": {k: "model.safetensors" for k in weights}})
        )

        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
            ),
            torch.device("cpu"),
        )
        dits = (ResolvedComponent(path="transformer", module=pipeline.transformer, selected=True),)
        result = build_checkpoint_mmap_plan(
            pipeline,
            dit_modules=(("transformer", pipeline.transformer),),
            sources=(),
            model_path=str(tmp_path),
            tensor_parallel_size=1,
            use_hsdp=False,
            online_quantization=False,
        )

        assert result.plan is not None
        backend._load_weights_via_mmap(pipeline, dits, result.plan)
        assert pipeline.transformer.validate_called, "validate_loaded_weights should be called"


class TestConfigValidation:
    """Tests for configuration validation in OffloadConfig / execute_request."""

    @pytest.mark.parametrize(
        ("dp", "sp", "tp", "use_allgather", "expected_dlo_group"),
        [
            (1, 1, 1, True, 1),
            (2, 1, 1, True, 2),
            (1, 4, 1, True, 4),
            (2, 2, 2, True, 2),
            (2, 2, 2, False, 2),
            (1, 4, 2, False, 4),
        ],
    )
    def test_dlo_group_selection_covers_dp_sp_tp_matrix(
        self,
        dp,
        sp,
        tp,
        use_allgather,
        expected_dlo_group,
    ):
        """DLO sharding follows DP, falls back to SP, and never uses TP."""

        config = OffloadConfig.from_od_config(
            SimpleNamespace(
                enable_cpu_offload=False,
                enable_layerwise_offload=False,
                enable_distributed_layerwise_offload=True,
                dlo_use_allgather=use_allgather,
                dlo_resident_layers=0,
                pin_cpu_memory=False,
                parallel_config=SimpleNamespace(
                    data_parallel_size=dp,
                    use_hsdp=False,
                    hsdp_shard_size=-1,
                    hsdp_replicate_size=1,
                    sequence_parallel_size=sp,
                    tensor_parallel_size=tp,
                ),
            )
        )
        assert config.dp_size == expected_dlo_group
        assert config.dlo_use_allgather is use_allgather

    def test_loader_plan_cannot_be_silently_dropped_on_unsupported_platform(self, monkeypatch):
        import vllm_omni.diffusion.offloader as offloader_module

        od_config = SimpleNamespace(
            enable_cpu_offload=False,
            enable_layerwise_offload=False,
            enable_distributed_layerwise_offload=True,
            dlo_use_allgather=False,
            dlo_resident_layers=0,
            pin_cpu_memory=False,
            model="unused",
            parallel_config=SimpleNamespace(
                use_hsdp=False,
                data_parallel_size=1,
                sequence_parallel_size=1,
            ),
        )
        plan = HostWeightPlan(
            backing_kind="checkpoint_mmap",
            bindings={},
        )
        monkeypatch.setattr(
            offloader_module.current_omni_platform,
            "supports_cpu_offload",
            lambda: False,
        )

        with pytest.raises(RuntimeError, match="skipped ordinary weight materialization"):
            offloader_module.get_offload_backend(
                od_config,
                device=torch.device("cpu"),
                host_weight_plan=plan,
            )

    def test_startup_recovery_stays_inside_offloader_boundary(self, monkeypatch):
        import vllm_omni.diffusion.offloader as offloader_module

        class FailingBackend:
            def __init__(self):
                self.disabled = False

            def enable(self, pipeline):
                del pipeline
                raise RuntimeError("initial prefetch failed")

            def disable(self):
                self.disabled = True

        warm_pipeline = nn.Module()
        canonical_pipeline = nn.Module()
        plan, carrier, _lease = _fake_hwr_plan("hwr-recovery")
        failing_backend = FailingBackend()
        calls = []

        def fake_backend(od_config, device, host_weight_plan):
            del od_config, device
            calls.append(host_weight_plan)
            return failing_backend if host_weight_plan is plan else None

        attach_offload_startup_state(
            warm_pipeline,
            OffloadStartupState(
                host_weight_plan=plan,
                fresh_model_loader=lambda: canonical_pipeline,
                allow_fresh_retry=True,
            ),
        )
        monkeypatch.setattr(offloader_module, "get_offload_backend", fake_backend)

        recovered, backend = offloader_module.enable_offload_backend(
            SimpleNamespace(),
            warm_pipeline,
            device=torch.device("cpu"),
        )

        assert recovered is canonical_pipeline
        assert backend is None
        assert failing_backend.disabled
        assert carrier.closed
        assert calls == [plan, None]

    def test_hsdp_with_allgather_rejected(self):
        """HSDP + DLO + AllGather should raise ValueError (double sharding)."""
        from vllm_omni.diffusion.offloader.base import OffloadConfig

        class FakePC:
            data_parallel_size = 1
            use_hsdp = True
            hsdp_shard_size = 2
            hsdp_replicate_size = 1
            sequence_parallel_size = 1
            tensor_parallel_size = 1
            cfg_parallel_size = 1
            pipeline_parallel_size = 1
            ulysses_degree = 1
            ring_degree = 1
            allgather_degree = 1
            enable_expert_parallel = False

        class FakeODConfig:
            enable_cpu_offload = False
            enable_layerwise_offload = False
            enable_distributed_layerwise_offload = True
            dlo_use_allgather = True
            pin_cpu_memory = True
            parallel_config = FakePC()
            model = "/fake/path"

        with pytest.raises(ValueError, match="incompatible with HSDP"):
            OffloadConfig.from_od_config(FakeODConfig())

    def test_hsdp_without_allgather_allowed(self):
        """HSDP + DLO + no-AllGather uses standard-loader rank-local weights."""
        from vllm_omni.diffusion.offloader.base import OffloadConfig

        class FakePC:
            data_parallel_size = 1
            use_hsdp = True
            hsdp_shard_size = 2
            hsdp_replicate_size = 1
            sequence_parallel_size = 1
            tensor_parallel_size = 1
            cfg_parallel_size = 1
            pipeline_parallel_size = 1
            ulysses_degree = 1
            ring_degree = 1
            allgather_degree = 1
            enable_expert_parallel = False

        class FakeODConfig:
            enable_cpu_offload = False
            enable_layerwise_offload = False
            enable_distributed_layerwise_offload = True
            dlo_use_allgather = False  # no AllGather → should be allowed
            dlo_resident_layers = 20
            pin_cpu_memory = True
            parallel_config = FakePC()
            model = "/fake/path"

        config = OffloadConfig.from_od_config(FakeODConfig())
        assert config.dp_size == 2  # topology is retained; the hook selects group size 1
        assert config.dlo_resident_layers == 20

    def test_resident_layers_with_allgather_rejected(self):
        class FakePC:
            data_parallel_size = 2
            use_hsdp = False
            sequence_parallel_size = 1

        class FakeODConfig:
            enable_cpu_offload = False
            enable_layerwise_offload = False
            enable_distributed_layerwise_offload = True
            dlo_use_allgather = True
            dlo_resident_layers = 20
            pin_cpu_memory = True
            parallel_config = FakePC()
            model = "/fake/path"

        with pytest.raises(ValueError, match="requires the DiT DLO transfer to be rank-local"):
            OffloadConfig.from_od_config(FakeODConfig())


class TestDynamicSlotTracking:
    """Tests for dynamic slot tracking with odd and even block counts.

    With i%2 initialization alone, odd block counts cause the circular
    tail→head prefetch to collide with the head's read slot (both map
    to slot 0).  Dynamic slot tracking via _prefetched_slot corrects
    this by reading the actual slot from the previous hook's prefetch.
    """

    @staticmethod
    def _trace_slots(num_blocks: int, num_iterations: int = 1):
        """Trace slot assignments through forward passes (pure logic, no hooks).

        Returns a list of (read_slot, write_slot) per block per iteration.
        """
        # Initial i%2 assignment (same as enable())
        current_slots = [i % 2 for i in range(num_blocks)]
        prefetched_slots: list[int | None] = [None] * num_blocks

        # Sync-prefetch: tail hook prefetches into slot 0
        prefetched_slots[-1] = 0

        results = []
        for _ in range(num_iterations):
            for i in range(num_blocks):
                prev_idx = (i - 1) % num_blocks

                # Dynamic slot tracking: read from prev's prefetched slot
                prefetched_slot = prefetched_slots[prev_idx]
                if prefetched_slot is not None:
                    current_slots[i] = prefetched_slot

                read_slot = current_slots[i]

                # Prefetch into other slot
                next_slot = 1 - read_slot
                current_slots[i] = next_slot
                prefetched_slots[i] = next_slot

                results.append((read_slot, next_slot))

        return results

    def test_odd_block_count_no_collision(self):
        """3 blocks (odd): no read/write slot collision."""
        results = self._trace_slots(3)
        for i, (read, write) in enumerate(results):
            assert read != write, f"Block {i % 3} (step {i}): read={read} == write={write}"

    def test_even_block_count_no_collision(self):
        """4 blocks (even): no read/write slot collision."""
        results = self._trace_slots(4)
        for i, (read, write) in enumerate(results):
            assert read != write, f"Block {i % 4} (step {i}): read={read} == write={write}"

    def test_large_odd_block_count(self):
        """7 blocks (odd): no collision across all blocks."""
        results = self._trace_slots(7)
        for i, (read, write) in enumerate(results):
            assert read != write, f"Block {i % 7} (step {i}): read={read} == write={write}"

    def test_5_blocks_two_iterations(self):
        """5 blocks (odd), 2 iterations: no collision in either pass."""
        results = self._trace_slots(5, num_iterations=2)
        n = 5
        for iteration in range(2):
            for i in range(n):
                idx = iteration * n + i
                read, write = results[idx]
                assert read != write, f"Iter {iteration} block {i}: read={read} == write={write}"

    def test_3_blocks_three_iterations(self):
        """3 blocks (odd), 3 iterations: no collision across all passes."""
        results = self._trace_slots(3, num_iterations=3)
        n = 3
        for iteration in range(3):
            for i in range(n):
                idx = iteration * n + i
                read, write = results[idx]
                assert read != write, f"Iter {iteration} block {i}: read={read} == write={write}"

    def test_72_blocks_cosmos3(self):
        """72 blocks (Cosmos3, even): no collision."""
        results = self._trace_slots(72)
        for i, (read, write) in enumerate(results):
            assert read != write, f"Block {i % 72} (step {i}): read={read} == write={write}"

    def test_slots_alternate_within_iteration(self):
        """Within one iteration, consecutive blocks should use alternating slots."""
        results = self._trace_slots(6)
        for i in range(6):
            read, write = results[i]
            # Each block reads from one slot and writes to the other
            assert {read, write} == {0, 1}, f"Block {i}: slots {read},{write} don't cover both"

    def test_prefetched_slot_set_after_prefetch(self, dist_group, patched_offload_runtime):
        """prefetch_layer should set _prefetched_slot to the slot argument."""
        block = TinyBlock(_make_values(1.0))
        next_block = TinyBlock(_make_values(2.0))

        hook = DistributedLayerwiseOffloadHook(
            next_block=next_block,
            device=torch.device("cpu"),
            dp_group=None,
            dp_size=1,
            rank=0,
            copy_stream=DummyStream(),
            comm_stream=DummyStream(),
            pin_memory=False,
        )
        hook.initialize_hook(block)

        assert hook._prefetched_slot is None  # initially None

        hook.prefetch_layer(slot=0, non_blocking=False)
        assert hook._prefetched_slot == 0

        hook.prefetch_layer(slot=1, non_blocking=False)
        assert hook._prefetched_slot == 1

    def test_dynamic_slot_overrides_initial_assignment(self, dist_group, patched_offload_runtime):
        """When _prev_hook._prefetched_slot is set, pre_forward should use
        that value instead of the initial i%2 assignment."""
        block_a = TinyBlock(_make_values(1.0))
        block_b = TinyBlock(_make_values(2.0))

        hook_a = DistributedLayerwiseOffloadHook(
            next_block=block_b,
            device=torch.device("cpu"),
            dp_group=None,
            dp_size=1,
            rank=0,
            copy_stream=DummyStream(),
            comm_stream=DummyStream(),
            pin_memory=False,
        )
        hook_a.initialize_hook(block_a)

        hook_b = DistributedLayerwiseOffloadHook(
            next_block=block_a,
            device=torch.device("cpu"),
            dp_group=None,
            dp_size=1,
            rank=0,
            copy_stream=DummyStream(),
            comm_stream=DummyStream(),
            pin_memory=False,
        )
        hook_b.initialize_hook(block_b)

        # Wire: hook_b._prev_hook = hook_a
        hook_b._prev_hook = hook_a

        # hook_a prefetched into slot 1
        hook_a.prefetch_layer(slot=1, non_blocking=False)
        assert hook_a._prefetched_slot == 1

        # hook_b's initial slot is 0 (from i%2)
        hook_b.current_slot = 0

        # pre_forward should override to slot 1
        hook_b.pre_forward(block_b)
        assert hook_b.current_slot == 1, "pre_forward should read _prev_hook._prefetched_slot=1, not keep initial 0"


class _DistributedComponentPipeline(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(num_blocks=2)
        self.text_encoder = _StagedEncoder()
        self.vae = _StagedVAE()


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

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


class TestDistributedComponentSelection:
    def test_on_demand_only_encoder_cannot_claim_allgather(self, patched_offload_runtime):
        class Pipeline(nn.Module):
            _offload_plan = OffloadPlan(
                encoder_component_types={"text_encoder": "text_encoder"},
                on_demand_component_paths=frozenset({"text_encoder"}),
                encoder_dlo_weight_replication=frozenset({"text_encoder"}),
            )

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

        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dp_size=2,
                components=frozenset({"text_encoder"}),
                dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
            ),
            torch.device("cpu"),
        )
        backend.dp_group = object()

        with pytest.raises(ValueError, match="cannot use AllGather without.*streamable block plan"):
            backend.enable(Pipeline())

    def test_explicit_single_block_dit_is_rejected(self, patched_offload_runtime):
        pipeline = nn.Module()
        pipeline.transformer = _SingleBlockModel(num_blocks=1)
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                components=frozenset({"dit"}),
            ),
            torch.device("cpu"),
        )

        with pytest.raises(ValueError, match="leaves only one streaming block"):
            backend.enable(pipeline)

    def test_enable_defers_collectives_until_first_forward(self, patched_offload_runtime, monkeypatch):
        pipeline = _GenericDistributedEncoderPipeline()
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dp_size=2,
                components=frozenset({"dit", "text_encoder"}),
                dlo_transfers={"dit": "allgather", "text_encoder": "rank-local"},
            ),
            torch.device("cpu"),
        )
        backend.dp_group = object()

        def reject_startup_collective(*_args, **_kwargs):
            raise AssertionError("enable() entered a rank-asymmetric collective")

        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", reject_startup_collective)

        backend.enable(pipeline)

        assert all(hook._prefetched_slot is None for group in backend._all_hook_groups for hook in group)
        encoder_hook = pipeline.text_encoder.encoder.block[0]._hook_registry.get_hook("distributed_layerwise_offload")
        dit_hook = pipeline.transformer.blocks[0]._hook_registry.get_hook("distributed_layerwise_offload")
        assert encoder_hook.gpu_shard_buffers == [None, None]
        assert all(buffer is not None for buffer in dit_hook.gpu_shard_buffers)

        monkeypatch.setattr(
            torch.distributed,
            "all_gather_into_tensor",
            lambda output, local, group: output.copy_(local.repeat(2)),
        )
        backend.disable()

    def test_partial_dit_enable_failure_restores_weights_and_hooks(
        self,
        patched_offload_runtime,
        monkeypatch,
    ):
        pipeline = _DistributedComponentPipeline()
        expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks]
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dlo_use_allgather=False,
                components=frozenset({"dit"}),
            ),
            torch.device("cpu"),
        )
        original_apply = dist_backend_module.apply_distributed_block_hook
        calls = 0

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

        monkeypatch.setattr(dist_backend_module, "apply_distributed_block_hook", fail_second_hook)

        with pytest.raises(RuntimeError, match="injected DLO 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("distributed_layerwise_offload") is None
            torch.testing.assert_close(block.weight, expected)

    def test_multirank_enable_failure_cleanup_skips_restore_collective(self, monkeypatch, mocker):
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dp_size=2,
                dlo_use_allgather=True,
                components=frozenset({"dit"}),
            ),
            torch.device("cpu"),
        )
        monkeypatch.setattr(backend, "_enable", mocker.Mock(side_effect=RuntimeError("injected startup failure")))
        cleanup = mocker.Mock()
        monkeypatch.setattr(backend, "_disable", cleanup)

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

        cleanup.assert_called_once_with(restore_allgather_weights=False)

    def test_multirank_failure_restores_only_rank_local_hooks(self, patched_offload_runtime):
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dp_size=2,
                components=frozenset({"dit", "text_encoder"}),
                dlo_transfers={"dit": "allgather", "text_encoder": "rank-local"},
            ),
            torch.device("cpu"),
        )
        allgather_hook = Mock(dp_size=2, next_block=nn.Linear(1, 1))
        rank_local_hook = Mock(dp_size=1, next_block=nn.Linear(1, 1))
        backend._all_hook_groups = [[allgather_hook], [rank_local_hook]]

        backend._disable(restore_allgather_weights=False)

        allgather_hook.restore_next_block_to_cpu.assert_not_called()
        rank_local_hook.restore_next_block_to_cpu.assert_called_once_with()

        with pytest.raises(RuntimeError, match="recreate the backend and reload the pipeline"):
            backend.enable(nn.Module())

    def test_worker_shutdown_skips_allgather_restore(self, patched_offload_runtime, mocker):
        # Restoring AllGather blocks rebuilds the full DiT on every rank's host,
        # which outlived the executor's shutdown grace period (issue 8081).
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dp_size=2,
                components=frozenset({"dit", "text_encoder"}),
                dlo_transfers={"dit": "allgather", "text_encoder": "rank-local"},
            ),
            torch.device("cpu"),
        )
        allgather_hook = Mock(dp_size=2, next_block=nn.Linear(1, 1))
        rank_local_hook = Mock(dp_size=1, next_block=nn.Linear(1, 1))
        backend._all_hook_groups = [[allgather_hook], [rank_local_hook]]
        backend.enabled = True
        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env")
        worker = DiffusionWorker.__new__(DiffusionWorker)
        worker.model_runner = SimpleNamespace(offload_backend=backend, kv_transfer_manager=None)

        worker.shutdown()

        allgather_hook.restore_next_block_to_cpu.assert_not_called()
        rank_local_hook.restore_next_block_to_cpu.assert_called_once_with()
        assert not backend.enabled
        assert not backend._all_hook_groups

    def test_disable_runs_collectives_before_best_effort_local_cleanup(
        self,
        patched_offload_runtime,
        monkeypatch,
    ):
        events: list[str] = []
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dp_size=2,
            ),
            torch.device("cpu"),
        )
        rank_local_hook = Mock(dp_size=1, next_block=nn.Linear(1, 1))
        allgather_hook = Mock(dp_size=2, next_block=nn.Linear(1, 1))
        allgather_hook.restore_next_block_to_cpu.side_effect = lambda: events.append("allgather")

        def fail_rank_local_restore():
            events.append("rank-local")
            raise RuntimeError("injected rank-local restore failure")

        rank_local_hook.restore_next_block_to_cpu.side_effect = fail_rank_local_restore
        backend._all_hook_groups = [[rank_local_hook], [allgather_hook]]
        backend._blocks = [[nn.Module()]]
        backend.enabled = True
        monkeypatch.setattr(
            dist_backend_module,
            "remove_distributed_block_hook",
            lambda _block: events.append("remove"),
        )
        monkeypatch.setattr(backend, "_release_registered_mmap", lambda: events.append("unregister"))
        release_handles = Mock()
        monkeypatch.setattr(backend, "_release_mmap_handles", release_handles)

        with pytest.raises(RuntimeError, match="injected rank-local restore failure"):
            backend.disable()

        assert events == ["allgather", "rank-local", "remove", "unregister"]
        release_handles.assert_not_called()
        assert backend.enabled
        assert backend._all_hook_groups

    def test_all_streams_encoder_and_keeps_vae_resident(self, patched_offload_runtime, monkeypatch):
        pipeline = _DistributedComponentPipeline()
        move_non_block_state = Mock()
        monkeypatch.setattr(dist_backend_module, "move_non_block_state_to_device", move_non_block_state)
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dlo_use_allgather=False,
                components=frozenset({"dit", "text_encoder"}),
            ),
            torch.device("cpu"),
        )

        backend.enable(pipeline)

        assert pipeline.text_encoder._omni_layerwise_enabled
        assert len(pipeline.text_encoder._omni_layerwise_hooks) == 4
        # The pipeline-owned stage lifecycle releases the encoder's non-block
        # state after DLO has installed blockwise hooks.
        assert pipeline.text_encoder.offload_calls == 1
        assert pipeline.vae.offload_calls == 0
        assert pipeline.vae.to_calls == 1
        move_non_block_state.assert_not_called()

        backend.disable()

        assert not pipeline.text_encoder._omni_layerwise_enabled
        encoder_blocks = [
            *pipeline.text_encoder.vision.blocks,
            *pipeline.text_encoder.text_model.layers,
        ]
        assert all(block._hook_registry.get_hook("distributed_layerwise_offload") is None for block in encoder_blocks)

    def test_dit_and_encoder_choose_transfer_independently(self, patched_offload_runtime, monkeypatch):
        pipeline = _GenericDistributedEncoderPipeline()
        expected_parameters = {name: parameter.detach().clone() for name, parameter in pipeline.named_parameters()}
        full_weights_by_rank_zero_shard = {
            tuple(parameter.flatten()[: (parameter.numel() + 1) // 2].tolist()): parameter.flatten()
            for parameter in expected_parameters.values()
            if parameter.numel() > 4
        }
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dp_size=2,
                components=frozenset({"dit", "text_encoder"}),
                dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
            ),
            torch.device("cpu"),
        )
        backend.dp_group = object()

        def fake_allgather(output, local, group):
            del group
            output.copy_(full_weights_by_rank_zero_shard[tuple(local.tolist())])

        monkeypatch.setattr(
            torch.distributed,
            "all_gather_into_tensor",
            fake_allgather,
        )

        backend.enable(pipeline)

        encoder_hook = pipeline.text_encoder.encoder.block[0]._hook_registry.get_hook("distributed_layerwise_offload")
        dit_hook = pipeline.transformer.blocks[0]._hook_registry.get_hook("distributed_layerwise_offload")
        assert encoder_hook.dp_size == 2
        assert dit_hook.dp_size == 1
        assert all(buffer is not None for buffer in encoder_hook.gpu_shard_buffers)
        assert dit_hook.gpu_shard_buffers == [None, None]

        backend.disable()
        actual_parameters = dict(pipeline.named_parameters())
        for name, expected in expected_parameters.items():
            torch.testing.assert_close(actual_parameters[name], expected)

    def test_encoder_allgather_rejects_undeclared_replication(self, patched_offload_runtime):
        pipeline = _DistributedComponentPipeline()
        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
                dp_size=2,
                components=frozenset({"text_encoder"}),
                dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
            ),
            torch.device("cpu"),
        )
        backend.dp_group = object()

        with pytest.raises(ValueError, match="not declared replicated"):
            backend.enable(pipeline)

    def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self):
        """Every rank must reject an unsafe encoder group, including stub ranks."""

        class StubEncoderPipeline(nn.Module):
            _offload_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("missing.blocks",)})

            def __init__(self):
                super().__init__()
                self.text_encoder = nn.Module()

        config = OffloadConfig(
            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
            pin_cpu_memory=False,
            dp_size=2,
            components=frozenset({"text_encoder"}),
            dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
        )

        with pytest.raises(ValueError, match="not declared replicated"):
            resolve_offload_plan(StubEncoderPipeline(), config)