vllm-omni/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Distributed Layerwise Offload backend with double-buffered H2D.

This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:

* Optionally shards model weights across DP ranks and reconstructs each block
  with AllGather, or streams a complete rank-local block without a collective.
* Can retain compatible checkpoint tensors as node-shared mmap sources instead
  of creating a persistent private host copy in every rank.
* Uses a fixed double-buffer scheme that keeps only two layers' worth of
  weights on each device at any time.
* Pipelines H2D transfers and, when enabled, AllGather communications on
  dedicated streams, overlapping them with computation.
* Is hardware-agnostic, supporting both NVIDIA GPU (CUDA) and Ascend NPU
  (CANN) platforms via vLLM-Omni's platform abstraction layer.
"""

from __future__ import annotations

import threading
import time
import weakref
from collections.abc import Sequence
from itertools import chain
from typing import Any

import torch
import torch.distributed
from torch import nn
from vllm.logger import init_logger

from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
from vllm_omni.diffusion.model_loader.host_weight_plan import (
    HostWeightPlan,
)
from vllm_omni.host_weight_runtime import HostWeightLease
from vllm_omni.platforms import current_omni_platform

from .base import OffloadBackend, OffloadConfig, run_cleanup_steps
from .component_utils import (
    clear_encoder_layerwise_state,
    iter_streamable_dits,
    move_non_block_state_to_device,
    prepare_component,
    prepare_pipeline_components,
    set_encoder_layerwise_state,
)
from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT
from .host_registration import (
    HostRegistration,
    HostRegistrationCleanupError,
    HostRegistrationError,
    register_host_mappings,
)
from .plan_resolver import BlockStack, ResolvedComponent, resolve_offload_plan
from .tensor_utils import (
    clear_block_storage,
    clear_tensor_storage,
    describe_tensor_storage,
    flatten_physical_storage,
    group_named_tensors_by_dtype,
    is_materialized_tensor,
    make_offload_placeholder,
    materialization_probe,
    module_materialization_probe,
    restore_tensor_storage,
    set_tensor_storage,
    tensor_storage_metadata,
)
from .tensor_utils import (
    dtype_size as _dtype_size,
)

logger = init_logger(__name__)

# A backend normally owns both objects below. This process-lifetime safety
# owner prevents HostWeightLease.__del__ from unmapping storage when cleanup
# failure unwinds startup and the backend itself becomes unreachable. A clean
# retry removes the pair before closing the lease.
_ACTIVE_HWR_REGISTRATIONS: list[tuple[HostRegistration, HostWeightLease]] = []
_ACTIVE_HWR_REGISTRATIONS_LOCK = threading.Lock()


def _retain_active_hwr_registration(
    registration: HostRegistration,
    lease: HostWeightLease,
) -> None:
    with _ACTIVE_HWR_REGISTRATIONS_LOCK:
        if not any(
            candidate is registration and candidate_lease is lease
            for candidate, candidate_lease in _ACTIVE_HWR_REGISTRATIONS
        ):
            _ACTIVE_HWR_REGISTRATIONS.append((registration, lease))


def _forget_active_hwr_registration(
    registration: HostRegistration,
    lease: HostWeightLease,
) -> None:
    with _ACTIVE_HWR_REGISTRATIONS_LOCK:
        _ACTIVE_HWR_REGISTRATIONS[:] = [
            (candidate, candidate_lease)
            for candidate, candidate_lease in _ACTIVE_HWR_REGISTRATIONS
            if candidate is not registration or candidate_lease is not lease
        ]


# Threshold (in MB) for deciding whether a non-block DiT submodule should
# use layerwise offload (streaming hooks) or be moved to GPU as a resident
# module.  Submodules larger than this are offloaded to save HBM; smaller
# ones stay resident for lower latency.


class DistributedLayerwiseOffloadHook(ModelHook):
    """Hook for distributed layerwise offloading with fixed double-buffer.

    Each rank stores only a shard of each block's weights on host memory.
    Two device slots alternate: one holds current weights, one holds next weights.
    H2D and AllGather run asynchronously on dedicated streams, overlapped
    with computation.

    Supports both NVIDIA GPU (CUDA) and Ascend NPU (CANN) platforms.
    """

    _HOOK_NAME = "distributed_layerwise_offload"

    def __init__(
        self,
        next_block: nn.Module,
        device: torch.device,
        dp_group: torch.distributed.ProcessGroup | None,
        dp_size: int,
        rank: int,
        copy_stream: Any | None = None,
        comm_stream: Any | None = None,
        pin_memory: bool = True,
        shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None,
        rank_local_mmap: bool = False,
        tensor_transforms: dict[int, Any] | None = None,
        materialization_probe_tensor: torch.Tensor | None = None,
    ):
        assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`"
        if type(dp_size) is not int or dp_size < 1:
            raise ValueError(f"dp_size must be a positive integer, got {dp_size!r}")
        if type(rank) is not int or not 0 <= rank < dp_size:
            raise ValueError(f"rank must satisfy 0 <= rank < dp_size, got rank={rank!r}, dp_size={dp_size}")
        if dp_size > 1 and dp_group is None:
            raise ValueError("dp_group is required when dp_size is greater than one")
        if rank_local_mmap and dp_size != 1:
            raise ValueError("rank_local_mmap requires dp_size=1")

        self.next_block = next_block
        self.device = device
        self.dp_group = dp_group
        self.dp_size = dp_size
        self.rank = rank
        self.pin_memory = pin_memory
        self.rank_local_mmap = rank_local_mmap
        self.registered_mmap = False
        self.tensor_transforms = tensor_transforms or {}
        self._materialization_probe = materialization_probe_tensor

        self.copy_stream = copy_stream or current_omni_platform.Stream()
        self.comm_stream = comm_stream or current_omni_platform.Stream()

        # Double buffers: either shared (from backend) or self-allocated (lazy)
        if shared_buffers is not None:
            self.gpu_buffers: list[dict[torch.dtype, torch.Tensor] | None] = shared_buffers
            self._owns_buffers = False
        else:
            self.gpu_buffers = [None, None]
            self._owns_buffers = True
        self.ready_events: list[Any | None] = [None, None]

        # Sharded host weights for the next block, keyed by dtype
        self.cpu_shards: dict[torch.dtype, torch.Tensor] = {}
        # File-backed source tensors for rank-local mmap.  Unlike cpu_shards,
        # these remain immutable views of the checkpoint and are never pinned
        # or flattened into a model-sized private allocation.
        self.cpu_sources: dict[torch.dtype, list[dict[str, Any]]] = {}
        self.metadata: dict[torch.dtype, list[dict[str, Any]]] = {}

        # Rank-local mmap uses two host staging slots shared by every hook in
        # this worker.  They are assigned by the backend after all block sizes
        # are known, mirroring the shared device-buffer allocation.
        self.cpu_staging_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None]
        self.cpu_staging_events: list[Any | None] = [None, None]

        # Current slot index (0 or 1).  Updated dynamically by the previous
        # hook's prefetch_layer call via _prefetched_slot.  This ensures
        # correct slot tracking for ALL block counts (including odd N).
        self.current_slot = 0
        self._prefetched_slot: int | None = None

        # Backward link to previous hook for fallback (cache-dit skip)
        self._prev_hook: DistributedLayerwiseOffloadHook | None = None

        # Marks the first hook in a shared-buffer group.  When multiple DiT
        # groups share the same 2 GPU buffers, another group may have
        # overwritten this group's slot between forwards.  The first block
        # must sync-prefetch on entry to ensure it loads the correct
        # weights, even if is_materialized sees a non-empty tensor left by
        # the other group.
        self._is_group_first: bool = False

        # Group ID for per-slot contamination tracking.  Shared across all
        # hooks in the same group.  When a hook prefetches into a slot, it
        # stamps the slot with its group ID.  On group entry, the first
        # hook checks whether the slot was last written by its own group
        # (tail hook's async prefetch) — if so, skip the sync-prefetch.
        self._group_id: int = -1
        self._shared_slot_group: list[int] | None = None  # [-1, -1], shared

        # Parameters/buffers of the current and next blocks
        self.block_parameters: dict[str, nn.Parameter] = {}
        self.block_buffers: dict[str, torch.Tensor] = {}
        self.next_block_parameters: dict[str, nn.Parameter] = {}
        self.next_block_buffers: dict[str, torch.Tensor] = {}

        # Per-block synchronization primitive: set after H2D copy completes.
        self._prefetch_done: Any | None = None

        # Shared shard (AllGather input) buffers — assigned by backend.
        self.gpu_shard_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None]

        self._cached_repoint: tuple[tuple[Any, ...], ...] = ()

    # ------------------------------------------------------------------ #
    #  DTensor helpers (shared with LayerwiseOffloadHook)                 #
    # ------------------------------------------------------------------ #

    def initialize_hook(self, module: nn.Module) -> nn.Module:
        module = super().initialize_hook(module)

        self.block_parameters = dict(module.named_parameters())
        self.block_buffers = dict(module.named_buffers())
        if self._materialization_probe is None:
            self._materialization_probe = materialization_probe(self.block_parameters, self.block_buffers)

        self.next_block_parameters = dict(self.next_block.named_parameters())
        self.next_block_buffers = dict(self.next_block.named_buffers())

        if self.rank_local_mmap:
            self.cpu_sources, self.metadata = self._collect_mmap_sources(
                self.next_block_parameters,
                self.next_block_buffers,
                self.tensor_transforms,
            )
        else:
            # Shard next block's weights and store local shard in pinned CPU memory.
            self.cpu_shards, self.metadata = self._shard_and_pin(
                self.next_block_parameters,
                self.next_block_buffers,
                self.dp_size,
                self.rank,
                self.pin_memory,
                self.tensor_transforms,
            )

        # Allocate device buffers only if not using shared buffers from backend
        if self._owns_buffers:
            self._allocate_device_buffers()

        # Cache parameter re-pointing metadata to avoid per-layer dict lookups.
        self._cached_repoint = tuple(
            (
                self.next_block_parameters[m["name"]]
                if m["name"] in self.next_block_parameters
                else self.next_block_buffers[m["name"]],
                dtype,
                m["offset"],
                m["numel"],
                m["shape"],
                m["stride"],
            )
            for dtype, metas in self.metadata.items()
            for m in metas
        )

        # Pre-compute AG output sizes (avoid sum() per layer).
        self._ag_output_sizes: dict[torch.dtype, int] = {}
        if self.dp_size > 1:
            for dtype in self.metadata:
                shard_numel = self.cpu_shards[dtype].numel()
                self._ag_output_sizes[dtype] = shard_numel * self.dp_size

        # Commit the storage mutation only after every allocation and metadata
        # calculation above has succeeded. Until here, hook setup is retryable.
        clear_tensor_storage(chain(self.next_block_parameters.values(), self.next_block_buffers.values()))

        return module

    @staticmethod
    def _collect_mmap_sources(
        params: dict[str, nn.Parameter],
        bufs: dict[str, torch.Tensor],
        tensor_transforms: dict[int, Any] | None = None,
    ) -> tuple[dict[torch.dtype, list[dict[str, Any]]], dict[torch.dtype, list[dict[str, Any]]]]:
        """Retain file-backed tensors and describe their runtime layout.

        The returned sources preserve safetensors mmap storage.  Runtime-layout
        adapters are applied only while packing a bounded staging slot, so no
        full-model anonymous CPU copy is retained by a worker.
        """
        cpu_sources: dict[torch.dtype, list[dict[str, Any]]] = {}
        metadata: dict[torch.dtype, list[dict[str, Any]]] = {}
        offsets: dict[torch.dtype, int] = {}
        for name, target in chain(params.items(), bufs.items()):
            source = target.to_local() if hasattr(target, "to_local") else target
            if source.device.type != "cpu":
                raise ValueError(
                    f"Rank-local mmap storage requires CPU checkpoint views, but {name!r} is on {source.device}."
                )

            dtype = source.dtype
            offset = offsets.get(dtype, 0)
            transform = (tensor_transforms or {}).get(id(target))
            runtime_source = transform(source) if callable(transform) else source
            if runtime_source.dtype != dtype or runtime_source.shape != source.shape:
                raise ValueError(
                    "mmap weight transform changed tensor metadata for "
                    f"{name!r}: expected dtype={dtype}, shape={tuple(source.shape)}, "
                    f"got dtype={runtime_source.dtype}, shape={tuple(runtime_source.shape)}"
                )
            stride = runtime_source.stride()
            storage_numel = (
                0
                if runtime_source.numel() == 0
                else 1 + sum((size - 1) * axis_stride for size, axis_stride in zip(runtime_source.shape, stride))
            )

            cpu_sources.setdefault(dtype, []).append(
                {
                    "name": name,
                    "tensor": source.detach(),
                    "transform": transform,
                }
            )
            metadata.setdefault(dtype, []).append(
                {
                    "name": name,
                    "offset": offset,
                    "numel": storage_numel,
                    "shape": runtime_source.shape,
                    "stride": stride,
                }
            )
            offsets[dtype] = offset + storage_numel

        return cpu_sources, metadata

    @staticmethod
    def _shard_and_pin(
        params: dict[str, nn.Parameter],
        bufs: dict[str, torch.Tensor],
        dp_size: int,
        rank: int,
        pin_memory: bool,
        tensor_transforms: dict[int, Any] | None = None,
    ) -> tuple[dict[torch.dtype, torch.Tensor], dict[torch.dtype, list[dict[str, Any]]]]:
        """Flatten params+buffers by dtype, split into DP shards, store local shard.

        Each rank stores only ``1/dp_size`` of the total weights. The full
        tensor is reconstructed at runtime via AllGather.
        """
        dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {}
        cpu_shards: dict[torch.dtype, torch.Tensor] = {}

        for dtype, named_weights in group_named_tensors_by_dtype(params, bufs).items():
            # Apply loader-declared layout conversions block by block while
            # preserving the physical tensor layout used by the kernels.
            specs = describe_tensor_storage(named_weights, tensor_transforms)
            total_numel = sum(spec.storage_numel for spec in specs)

            # Equal-sized shards (ceil division) for all_gather_into_tensor
            shard_size = (total_numel + dp_size - 1) // dp_size  # ceil
            shard_start = rank * shard_size
            shard_end = min(shard_start + shard_size, total_numel)

            # Allocate ONLY the shard (1/dp_size), zero-padded to ceil.
            # Avoids materialising the full block on CPU.
            shard = torch.zeros(
                shard_size,
                dtype=dtype,
                device="cpu",
                pin_memory=pin_memory,
            )

            current_offset = 0
            for spec in specs:
                # Offsets remain relative to the FULL flattened buffer
                # (needed for correct AllGather reconstruction).
                dtype_metadata.setdefault(dtype, []).append(tensor_storage_metadata(spec, current_offset))

                # Copy ONLY the portion within [shard_start, shard_end)
                overlap_start = max(current_offset, shard_start)
                overlap_end = min(current_offset + spec.storage_numel, shard_end)
                if overlap_start < overlap_end:
                    flat_storage = flatten_physical_storage(spec.value, spec.storage_numel)
                    src_start = overlap_start - current_offset
                    src_end = overlap_end - current_offset
                    dst_start = overlap_start - shard_start
                    dst_end = overlap_end - shard_start
                    shard[dst_start:dst_end].copy_(flat_storage[src_start:src_end])

                current_offset += spec.storage_numel

            cpu_shards[dtype] = shard

        return cpu_shards, dtype_metadata

    def _allocate_device_buffers(self) -> None:
        """Pre-allocate exactly two device buffers (one per slot)."""
        for slot in range(2):
            gpu_weights: dict[torch.dtype, torch.Tensor] = {}
            for dtype, metas in self.metadata.items():
                total_numel = sum(m["numel"] for m in metas)
                # AllGather output = dp_size * shard_size (padded)
                padded = total_numel
                if self.dp_size > 1:
                    shard_sz = (total_numel + self.dp_size - 1) // self.dp_size
                    padded = shard_sz * self.dp_size
                gpu_weights[dtype] = torch.empty(padded, dtype=dtype, device=self.device)
            self.gpu_buffers[slot] = gpu_weights

    @property
    def is_materialized(self) -> bool:
        """Check whether this block's parameters hold real data on device."""
        return self._materialization_probe is None or is_materialized_tensor(self._materialization_probe)

    # ------------------------------------------------------------------ #
    #  Prefetch: H2D + AllGather (overlapped on dedicated streams)      #
    # ------------------------------------------------------------------ #

    @staticmethod
    def _resolve_mmap_source(
        source_info: dict[str, Any],
        meta: dict[str, Any],
        dtype: torch.dtype,
    ) -> torch.Tensor:
        source = source_info["tensor"]
        transform = source_info["transform"]
        if callable(transform):
            source = transform(source)
        if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
            raise ValueError(
                "mmap weight transform changed tensor layout for "
                f"{source_info['name']!r}: expected dtype={dtype}, "
                f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
                f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
                f"stride={source.stride()}"
            )
        return source

    @staticmethod
    def _pack_mmap_sources(
        cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
        metadata: dict[torch.dtype, list[dict[str, Any]]],
        slot_buffers: dict[torch.dtype, torch.Tensor],
    ) -> dict[torch.dtype, torch.Tensor]:
        staged: dict[torch.dtype, torch.Tensor] = {}
        for dtype, metas in metadata.items():
            total_numel = sum(meta["numel"] for meta in metas)
            destination = slot_buffers[dtype][:total_numel]
            sources = cpu_sources[dtype]
            for source_info, meta in zip(sources, metas, strict=True):
                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
                start = meta["offset"]
                physical_storage = destination[start : start + meta["numel"]]
                if source.is_contiguous():
                    physical_storage.copy_(source.flatten())
                else:
                    torch.as_strided(
                        physical_storage,
                        size=source.shape,
                        stride=source.stride(),
                    ).copy_(source)
            staged[dtype] = destination
        return staged

    def _stage_mmap_sources(self, slot: int) -> dict[torch.dtype, torch.Tensor]:
        """Pack this block's mmap views into one bounded host staging slot."""
        previous_copy = self.cpu_staging_events[slot]
        if previous_copy is not None:
            synchronize = getattr(previous_copy, "synchronize", None)
            if callable(synchronize):
                synchronize()
            else:
                # Platform events normally expose synchronize().  Retain a
                # correctness fallback for test and non-CUDA platform shims.
                current_omni_platform.synchronize()
            self.cpu_staging_events[slot] = None

        slot_buffers = self.cpu_staging_buffers[slot]
        if slot_buffers is None:
            raise RuntimeError(f"cpu_staging_buffers[{slot}] was not allocated")
        return self._pack_mmap_sources(self.cpu_sources, self.metadata, slot_buffers)

    @staticmethod
    def _copy_mmap_sources_to_device(
        cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
        metadata: dict[torch.dtype, list[dict[str, Any]]],
        gpu_buffers: dict[torch.dtype, torch.Tensor],
        *,
        non_blocking: bool,
    ) -> None:
        """Copy registered source views directly into flattened device buffers."""
        for dtype, metas in metadata.items():
            destination = gpu_buffers[dtype]
            sources = cpu_sources[dtype]
            for source_info, meta in zip(sources, metas, strict=True):
                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
                start = meta["offset"]
                physical_storage = destination[start : start + meta["numel"]]
                async_copy = non_blocking and source.is_pinned()
                if source.is_contiguous():
                    physical_storage.copy_(source.flatten(), non_blocking=async_copy)
                else:
                    torch.as_strided(
                        physical_storage,
                        size=source.shape,
                        stride=source.stride(),
                    ).copy_(source, non_blocking=async_copy)

    @torch.compiler.disable
    def prefetch_layer(self, slot: int, non_blocking: bool = True) -> None:
        """Prepare next block's weights into the shared device buffer for *slot*.

        Uses the pre-allocated ``self.gpu_buffers[slot]`` instead of
        allocating fresh tensors every layer.  This enforces the fixed
        double-buffer memory bound (exactly 2 blocks on device).
        """
        self.copy_stream.wait_stream(current_omni_platform.current_stream())

        evt = current_omni_platform.Event()
        gpu_weights = self.gpu_buffers[slot]
        assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated"

        if self.dp_size <= 1 or self.dp_group is None:
            if self.rank_local_mmap and self.registered_mmap:
                with current_omni_platform.stream(self.copy_stream):
                    self._copy_mmap_sources_to_device(
                        self.cpu_sources,
                        self.metadata,
                        gpu_weights,
                        non_blocking=non_blocking,
                    )
                    evt.record(self.copy_stream)
            else:
                cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
                with current_omni_platform.stream(self.copy_stream):
                    for dtype, cpu_shard in cpu_weights.items():
                        gw = gpu_weights[dtype]
                        async_copy = non_blocking and cpu_shard.is_pinned()
                        gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
                    evt.record(self.copy_stream)
                if self.rank_local_mmap:
                    # The CPU slot may be overwritten only after this H2D copy has
                    # finished.  The shared event protects reuse by another hook.
                    self.cpu_staging_events[slot] = evt
        else:
            gpu_shards: dict[torch.dtype, torch.Tensor] = {}
            shard_bufs = self.gpu_shard_buffers[slot]
            assert shard_bufs is not None, f"gpu_shard_buffers[{slot}] not allocated"
            with current_omni_platform.stream(self.copy_stream):
                for dtype, cpu_shard in self.cpu_shards.items():
                    gpu_shard = shard_bufs[dtype][: cpu_shard.numel()]
                    gpu_shard.copy_(cpu_shard, non_blocking=non_blocking)
                    gpu_shards[dtype] = gpu_shard

            self.comm_stream.wait_stream(self.copy_stream)
            with current_omni_platform.stream(self.comm_stream):
                for dtype, local_shard in gpu_shards.items():
                    # Slice the shared (max-sized) output buffer down to this
                    # block's actual AllGather output size. The buffers are
                    # sized to the *largest* block across all groups, so for any
                    # smaller block the full buffer would violate the
                    # all_gather_into_tensor contract
                    # (output.numel() == world_size * input.numel()).
                    # Repoint offsets are relative to the block's flattened
                    # buffer, so a prefix slice is safe.
                    gw = gpu_weights[dtype][: self._ag_output_sizes[dtype]]
                    torch.distributed.all_gather_into_tensor(
                        gw,
                        local_shard,
                        group=self.dp_group,
                    )
                evt.record(self.comm_stream)

        self.ready_events[slot] = evt
        self._prefetch_done = evt
        self._prefetched_slot = slot

        # Stamp the slot with this hook's group ID so that group-first
        # hooks can detect whether another group has overwritten the slot.
        if self._shared_slot_group is not None:
            self._shared_slot_group[slot] = self._group_id

        # Re-point using cached metadata (avoids per-layer dict lookups).
        for target, dtype, offset, numel, shape, stride in self._cached_repoint:
            set_tensor_storage(
                target,
                torch.as_strided(
                    gpu_weights[dtype][offset : offset + numel],
                    size=shape,
                    stride=stride,
                ),
            )

    def _wait_for_weights(self, slot: int) -> None:
        """Wait until the slot holding this block is ready for compute.

        The ready event for this slot was set by the *previous* hook's
        prefetch_layer (which prefetched THIS block's weights into the
        shared buffer).  This hook's own ready_events[slot] is None because
        it never prefetched into this slot itself.  Fall back to the
        previous hook's event to ensure the compute stream waits for the
        AllGather (or H2D) to complete before reading the buffer.
        """
        evt = self.ready_events[slot]
        if evt is None and self._prev_hook is not None:
            evt = self._prev_hook.ready_events[slot]
        if evt is not None:
            current_omni_platform.current_stream().wait_event(evt)

    # ------------------------------------------------------------------ #
    #  Offload: free device memory for current block                     #
    # ------------------------------------------------------------------ #

    @torch.compiler.disable
    def offload_layer(self) -> None:
        """Free GPU memory for current block by replacing tensors with placeholders."""
        clear_block_storage(self.block_parameters, self.block_buffers, self._prefetch_done)
        self._prefetch_done = None

    @torch.compiler.disable
    def restore_next_block_to_cpu(self) -> None:
        """Restore hook-owned master weights before removing the hook.

        Every circular hook owns the host backing for its ``next_block``.
        Dropping that hook while the module points at rotating device buffers
        or placeholders would make a later enable shard invalid tensors.
        """

        if self.rank_local_mmap:
            for dtype, metas in self.metadata.items():
                for source_info, meta in zip(self.cpu_sources[dtype], metas, strict=True):
                    source = self._resolve_mmap_source(source_info, meta, dtype)
                    target = (
                        self.next_block_parameters[meta["name"]]
                        if meta["name"] in self.next_block_parameters
                        else self.next_block_buffers[meta["name"]]
                    )
                    restore_tensor_storage(target, source, device="cpu")
            return

        if self.dp_size <= 1:
            for dtype, metas in self.metadata.items():
                flat = self.cpu_shards[dtype]
                for meta in metas:
                    value = torch.as_strided(
                        flat[meta["offset"] : meta["offset"] + meta["numel"]],
                        size=meta["shape"],
                        stride=meta["stride"],
                    )
                    target = (
                        self.next_block_parameters[meta["name"]]
                        if meta["name"] in self.next_block_parameters
                        else self.next_block_buffers[meta["name"]]
                    )
                    restore_tensor_storage(target, value, device="cpu")
            return

        if self.dp_group is None:
            raise RuntimeError("Cannot restore distributed offload weights without the DLO process group")
        self.prefetch_layer(0, non_blocking=False)
        current_omni_platform.synchronize()
        for target in chain(self.next_block_parameters.values(), self.next_block_buffers.values()):
            local = target.to_local() if hasattr(target, "to_local") else target
            restore_tensor_storage(target, local, device="cpu")

    # ------------------------------------------------------------------ #
    #  ModelHook interface                                                #
    # ------------------------------------------------------------------ #

    def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
        # Dynamic slot tracking: read current_slot from the previous hook's
        # _prefetched_slot (the slot where this block's data was actually
        # loaded).  This corrects the initial i%2 assignment for odd block
        # counts, where the circular tail→head prefetch would otherwise
        # collide with the head's read slot.
        if self._prev_hook is not None and self._prev_hook._prefetched_slot is not None:
            self.current_slot = self._prev_hook._prefetched_slot

        # Group-first hook: check whether the shared buffer slot was
        # overwritten by another group since our last forward.  If the
        # slot still contains our own data (from the tail hook's async
        # prefetch), skip the sync-prefetch and just wait for the event.
        if self._is_group_first and self._prev_hook is not None:
            slot_contaminated = True
            if self._shared_slot_group is not None:
                slot_contaminated = self._shared_slot_group[self.current_slot] != self._group_id
            if slot_contaminated:
                # Another group (or no group) wrote to our slot — re-fetch
                self._prev_hook.prefetch_layer(self.current_slot, non_blocking=False)
            # Always wait for data to be ready (handles both sync and async paths)
            self._prev_hook._wait_for_weights(self.current_slot)
        elif not self.is_materialized and self._prev_hook is not None:
            # Previous hook was skipped (e.g. by cache-dit), sync-prefetch
            self._prev_hook.prefetch_layer(self.current_slot, non_blocking=False)
            self._prev_hook._wait_for_weights(self.current_slot)

        # Prefetch next layer into the other slot (overlapped with compute).
        # No explicit get_weights() here — offload_layer() in the previous
        # hook's post_forward already enqueued wait_event for the current
        # layer's H2D/AllGather, which is sufficient to ensure compute_stream
        # waits for weights before reading them.  The redundant wait_event
        # caused cascading NPU stream synchronization stalls (~10ms/layer).
        next_slot = 1 - self.current_slot
        self.prefetch_layer(next_slot, non_blocking=True)

        return args, kwargs

    def post_forward(self, module: nn.Module, output: Any) -> Any:
        self.offload_layer()
        return output


# ---------------------------------------------------------------------- #
#  Module-level helpers                                                   #
# ---------------------------------------------------------------------- #


def apply_distributed_block_hook(
    module: nn.Module,
    next_block: nn.Module,
    device: torch.device,
    dp_group: torch.distributed.ProcessGroup | None,
    dp_size: int,
    rank: int,
    copy_stream: Any | None = None,
    comm_stream: Any | None = None,
    pin_memory: bool = True,
    shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None,
    rank_local_mmap: bool = False,
    tensor_transforms: dict[int, Any] | None = None,
    materialization_probe_tensor: torch.Tensor | None = None,
) -> DistributedLayerwiseOffloadHook:
    """Register a DistributedLayerwiseOffloadHook on *module*."""
    registry = HookRegistry.get_or_create(module)
    hook = DistributedLayerwiseOffloadHook(
        next_block=next_block,
        device=device,
        dp_group=dp_group,
        dp_size=dp_size,
        rank=rank,
        copy_stream=copy_stream,
        comm_stream=comm_stream,
        pin_memory=pin_memory,
        shared_buffers=shared_buffers,
        rank_local_mmap=rank_local_mmap,
        tensor_transforms=tensor_transforms,
        materialization_probe_tensor=materialization_probe_tensor,
    )
    registry.register_hook(DistributedLayerwiseOffloadHook._HOOK_NAME, hook)
    return hook


def remove_distributed_block_hook(module: nn.Module) -> None:
    """Remove the distributed layerwise offload hook from *module*."""
    registry: HookRegistry | None = getattr(module, "_hook_registry", None)
    if registry is not None:
        registry.remove_hook(DistributedLayerwiseOffloadHook._HOOK_NAME)
        logger.debug("Removed distributed offload hook from %s", module.__class__.__name__)


class PinnedResidentLayerGroup:
    """Keep selected layers available for stage-scoped device residency.


    TODO(offload): Extract this alongside PinnedModuleStager after the
    distributed shard-and-pin operation becomes a shared storage primitive.
    It currently remains here because it depends on DLO's local-shard layout.
    Unlike ``module.to(device)``/``module.to("cpu")``, this group retains a
    pinned CPU master copy (or mmap source plus bounded staging) and never
    copies generated device weights back to host. Entering the denoise stage
    performs one asynchronous H2D pass;
    leaving it only restores zero-sized placeholders and releases the device
    buffers.  This lets the following VAE stage reuse the same HBM.

    The resident path is intentionally local-shard only.  With tensor
    parallelism, the regular model loader has already produced each rank's TP
    shard, so no DP AllGather is required or desirable here.
    """

    def __init__(
        self,
        blocks: list[nn.Module],
        device: torch.device,
        copy_stream: Any,
        pin_memory: bool,
        rank_local_mmap: bool = False,
        defer_staging: bool = False,
        tensor_transforms: dict[int, Any] | None = None,
    ) -> None:
        self.device = device
        self.copy_stream = copy_stream
        self.loaded = False
        self.rank_local_mmap = rank_local_mmap
        self.registered_mmap = False
        self.pin_memory = pin_memory
        self._states: list[dict[str, Any]] = []
        self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []

        for block in blocks:
            params = dict(block.named_parameters())
            bufs = dict(block.named_buffers())
            targets: dict[str, torch.Tensor] = {**params, **bufs}
            if rank_local_mmap:
                cpu_sources, metadata = DistributedLayerwiseOffloadHook._collect_mmap_sources(
                    params,
                    bufs,
                    tensor_transforms,
                )
                cpu_shards = {}
            else:
                cpu_shards, metadata = DistributedLayerwiseOffloadHook._shard_and_pin(
                    params,
                    bufs,
                    dp_size=1,
                    rank=0,
                    pin_memory=pin_memory,
                    tensor_transforms=tensor_transforms,
                )
                cpu_sources = {}
            self._states.append(
                {
                    "targets": targets,
                    "cpu_shards": cpu_shards,
                    "cpu_sources": cpu_sources,
                    "metadata": metadata,
                }
            )

        self._cpu_staging_buffers: list[dict[torch.dtype, torch.Tensor]] = []
        self._cpu_staging_events: list[Any | None] = [None, None]
        if rank_local_mmap and not defer_staging:
            max_sizes: dict[torch.dtype, int] = {}
            for state in self._states:
                for dtype, metas in state["metadata"].items():
                    total = sum(meta["numel"] for meta in metas)
                    max_sizes[dtype] = max(max_sizes.get(dtype, 0), total)
            for _ in range(2):
                buffers = {}
                for dtype, total in max_sizes.items():
                    buffer = torch.empty(
                        total,
                        dtype=dtype,
                        device="cpu",
                        pin_memory=pin_memory,
                    )
                    buffers[dtype] = buffer
                self._cpu_staging_buffers.append(buffers)

        # All host masters and optional staging buffers are now ready. Clear
        # model storage as one final constructor commit so an earlier failure
        # leaves every resident block intact.
        clear_tensor_storage(target for state in self._states for target in state["targets"].values())

    def load(self) -> None:
        if self.loaded:
            return

        # Allocate on the compute stream so the caching allocator can reuse
        # blocks released by the encoder stage.  Allocating on the copy stream
        # would create a separate stream-local pool and inflate peak HBM.
        gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
        for state in self._states:
            block_buffers: dict[torch.dtype, torch.Tensor] = {}
            for dtype, metas in state["metadata"].items():
                total = sum(meta["numel"] for meta in metas)
                block_buffers[dtype] = torch.empty(total, dtype=dtype, device=self.device)
            gpu_buffers.append(block_buffers)

        self.copy_stream.wait_stream(current_omni_platform.current_stream())
        ready = current_omni_platform.Event()
        with current_omni_platform.stream(self.copy_stream):
            for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)):
                if self.rank_local_mmap and self.registered_mmap:
                    DistributedLayerwiseOffloadHook._copy_mmap_sources_to_device(
                        state["cpu_sources"],
                        state["metadata"],
                        block_buffers,
                        non_blocking=True,
                    )
                    continue
                if self.rank_local_mmap:
                    slot = index % 2
                    previous_copy = self._cpu_staging_events[slot]
                    if previous_copy is not None:
                        synchronize = getattr(previous_copy, "synchronize", None)
                        if callable(synchronize):
                            synchronize()
                        else:
                            current_omni_platform.synchronize()
                    cpu_weights = DistributedLayerwiseOffloadHook._pack_mmap_sources(
                        state["cpu_sources"],
                        state["metadata"],
                        self._cpu_staging_buffers[slot],
                    )
                else:
                    cpu_weights = state["cpu_shards"]

                for dtype, cpu_weight in cpu_weights.items():
                    block_buffers[dtype].copy_(
                        cpu_weight,
                        non_blocking=cpu_weight.is_pinned(),
                    )
                if self.rank_local_mmap:
                    slot_ready = current_omni_platform.Event()
                    slot_ready.record(self.copy_stream)
                    self._cpu_staging_events[slot] = slot_ready
            ready.record(self.copy_stream)

        for state, block_buffers in zip(self._states, gpu_buffers):
            targets = state["targets"]
            for dtype, metas in state["metadata"].items():
                gpu_buffer = block_buffers[dtype]
                for meta in metas:
                    set_tensor_storage(
                        targets[meta["name"]],
                        torch.as_strided(
                            gpu_buffer[meta["offset"] : meta["offset"] + meta["numel"]],
                            size=meta["shape"],
                            stride=meta["stride"],
                        ),
                    )

        current_omni_platform.current_stream().wait_event(ready)
        self._gpu_buffers = gpu_buffers
        self.loaded = True

    def offload(self) -> None:
        if not self.loaded:
            return

        # Denoise kernels consume these buffers on the current compute stream.
        # Synchronize once at the stage boundary; no D2H copy is necessary
        # because the pinned CPU master weights were never overwritten.
        current_omni_platform.synchronize()
        for state in self._states:
            for target in state["targets"].values():
                set_tensor_storage(target, make_offload_placeholder(target))
        self._gpu_buffers.clear()
        self.loaded = False

    def restore_to_cpu(self) -> None:
        """Materialize the persistent host masters back into the module.

        Stage-scoped ``offload()`` deliberately leaves placeholders in the
        module while this group owns the CPU backing. ``disable()`` discards
        the group, so it must first restore ordinary CPU tensors to make a
        later enable cycle safe.
        """
        self.offload()

        for state in self._states:
            targets = state["targets"]
            for dtype, metas in state["metadata"].items():
                if self.rank_local_mmap:
                    for source_info, meta in zip(state["cpu_sources"][dtype], metas, strict=True):
                        source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
                        restore_tensor_storage(targets[meta["name"]], source, device="cpu")
                    continue

                flat = state["cpu_shards"][dtype]
                for meta in metas:
                    source = torch.as_strided(
                        flat[meta["offset"] : meta["offset"] + meta["numel"]],
                        size=meta["shape"],
                        stride=meta["stride"],
                    )
                    restore_tensor_storage(targets[meta["name"]], source, device="cpu")


# ---------------------------------------------------------------------- #
#  Backend                                                                #
# ---------------------------------------------------------------------- #


class DistributedLayerwiseOffloadBackend(OffloadBackend):
    """Distributed layer-wise (block-level) offloading backend.

    Supports both GPU (CUDA) and NPU (CANN) platforms.
    Device type is determined by the device passed to enable().

    Each rank stores only a shard of each block's weights on host memory.
    Two device slots alternate: one holds current weights, one holds next
    weights. H2D and AllGather run asynchronously on dedicated streams,
    overlapped with computation.
    """

    def __init__(
        self,
        config: OffloadConfig,
        device: torch.device,
        host_weight_plan: HostWeightPlan | None = None,
    ):
        super().__init__(config, device)

        self.copy_stream = current_omni_platform.Stream()
        self.comm_stream = current_omni_platform.Stream()
        self.dp_group: torch.distributed.ProcessGroup | None = None
        self.dp_size = config.dp_size
        self.rank = 0
        self._blocks: list[list[nn.Module]] = []
        self._all_hook_groups: list[list[DistributedLayerwiseOffloadHook]] = []
        self._resident_blocks: list[nn.Module] = []
        self._resident_layer_group: PinnedResidentLayerGroup | None = None
        self._residency_pipeline_ref: weakref.ReferenceType[nn.Module] | None = None
        self._encoder_modules: list[nn.Module] = []
        self._staged_components: list[nn.Module] = []
        self._using_mmap = False
        self._using_rank_local_mmap = False
        self._using_registered_mmap = False
        self.host_weight_plan = host_weight_plan
        self._host_weight_lease: HostWeightLease | None = None
        self._host_registration: HostRegistration | None = None
        self._mmap_transforms_by_tensor_id: dict[int, Any] = {}
        self._poisoned_reason: str | None = None

    def load_resident_layers(self) -> None:
        """Load the model-declared leading blocks for the denoise stage."""
        if self._resident_layer_group is not None:
            self._resident_layer_group.load()

    def offload_resident_layers(self) -> None:
        """Release leading blocks before VAE decode to bound peak HBM."""
        if self._resident_layer_group is not None:
            self._resident_layer_group.offload()

    def _clear_residency_controller(self) -> None:
        pipeline_ref = self._residency_pipeline_ref
        self._residency_pipeline_ref = None
        pipeline = None if pipeline_ref is None else pipeline_ref()
        if pipeline is not None and getattr(pipeline, "_dlo_residency_controller", None) is self:
            pipeline._dlo_residency_controller = None

    def _rank_local_source_tensors(
        self,
        hooks: list[DistributedLayerwiseOffloadHook],
    ) -> tuple[torch.Tensor, ...]:
        """Return the exact CPU sources that direct H2D would consume."""
        tensors: list[torch.Tensor] = []
        seen: set[int] = set()

        def collect(cpu_sources: dict[torch.dtype, list[dict[str, Any]]]) -> None:
            for sources in cpu_sources.values():
                for source in sources:
                    tensor = source["tensor"]
                    if id(tensor) not in seen:
                        seen.add(id(tensor))
                        tensors.append(tensor)

        for hook in hooks:
            collect(hook.cpu_sources)
        if self._resident_layer_group is not None:
            for state in self._resident_layer_group._states:
                collect(state["cpu_sources"])
        return tuple(tensors)

    def _try_register_hwr_mmap(self, source_tensors: tuple[torch.Tensor, ...]) -> bool:
        """Register the complete final-layout lease under pinned-memory policy."""
        lease = self._host_weight_lease
        if lease is None:
            return False
        if not self.config.pin_cpu_memory:
            logger.info("HWR mmap registration disabled by pin_cpu_memory=False; using bounded host staging")
            return False
        if not lease.mapped_regions or not source_tensors:
            logger.warning("HWR mmap registration found no mapped sources; using bounded host staging")
            return False

        limit_gib = self.config.dlo_host_registration_limit_gib
        max_bytes = int(limit_gib * 1024**3) if limit_gib > 0 else None
        started = time.perf_counter()
        try:
            registration = register_host_mappings(
                lease.mapped_regions,
                device=self.device,
                max_bytes=max_bytes,
            )
            try:
                unpinned = [tensor for tensor in source_tensors if tensor.numel() and not tensor.is_pinned()]
            except Exception as exc:
                errors = registration.close()
                if errors:
                    self._host_registration = registration
                    _retain_active_hwr_registration(registration, lease)
                    raise HostRegistrationCleanupError(
                        "CUDA registration succeeded but pinned-source verification failed, "
                        f"and rollback failed: {errors[:3]}"
                    ) from exc
                raise HostRegistrationError(f"cannot verify registered HWR sources: {exc}") from exc
            if unpinned:
                errors = registration.close()
                if errors:
                    self._host_registration = registration
                    _retain_active_hwr_registration(registration, lease)
                    raise HostRegistrationCleanupError(
                        "CUDA registration succeeded but PyTorch rejected mapped sources, "
                        f"and rollback failed: {errors[:3]}"
                    )
                raise HostRegistrationError(
                    "CUDA registration succeeded but PyTorch did not recognize "
                    f"{len(unpinned)} mapped source(s) as pinned"
                )
        except HostRegistrationCleanupError as exc:
            # Falling back could close a lease while the platform still owns
            # one of its mappings. Fail startup and retain ownership for retry.
            active_registration = exc.active_registration
            if active_registration is not None:
                self._host_registration = active_registration
                _retain_active_hwr_registration(active_registration, lease)
            logger.exception("HWR mmap registration rollback failed")
            raise
        except HostRegistrationError as exc:
            logger.warning("HWR registered direct H2D unavailable (%s); using bounded host staging", exc)
            return False

        self._host_registration = registration
        _retain_active_hwr_registration(registration, lease)
        logger.info(
            "Registered %.2f GiB of HWR mmap in %d range(s) for direct H2D in %.3f s",
            registration.total_bytes / 1024**3,
            registration.region_count,
            time.perf_counter() - started,
        )
        return True

    def _configure_hwr_transfer(self, hooks: list[DistributedLayerwiseOffloadHook]) -> None:
        """Select registered direct H2D or bounded staging once per backend."""
        plan = self.host_weight_plan
        if (
            (not hooks and self._resident_layer_group is None)
            or not self._using_rank_local_mmap
            or plan is None
            or plan.backing_kind != "host_weight_runtime"
        ):
            return

        source_tensors = self._rank_local_source_tensors(hooks)
        self._using_registered_mmap = self._try_register_hwr_mmap(source_tensors)
        for hook in hooks:
            hook.registered_mmap = self._using_registered_mmap
        if self._resident_layer_group is not None:
            self._resident_layer_group.registered_mmap = self._using_registered_mmap
            if self._using_registered_mmap:
                self._resident_layer_group._cpu_staging_buffers.clear()

    def _release_registered_mmap(self) -> None:
        """Release every platform registration before closing the HWR lease."""
        registration = self._host_registration
        if registration is None:
            return
        errors = registration.close()
        if errors:
            lease = self._host_weight_lease
            if lease is not None and not lease.closed:
                _retain_active_hwr_registration(registration, lease)
            logger.error("HWR mmap unregistration failed; retaining lease mappings for retry: %s", errors[:3])
            raise HostRegistrationCleanupError(f"failed to unregister {len(errors)} HWR mmap range(s)")
        lease = self._host_weight_lease
        if lease is not None:
            _forget_active_hwr_registration(registration, lease)
        self._host_registration = None
        logger.info("Unregistered HWR mmap ranges")

    def _load_weights_via_mmap(
        self,
        pipeline: nn.Module,
        dits: tuple[ResolvedComponent, ...],
        plan: HostWeightPlan,
    ) -> None:
        """Load DiT checkpoint tensors as file-backed safetensors views.

        When the transformer is created on meta device, this method
        replaces meta params with mmap views of the checkpoint files.
        The views point to OS page cache shared across ranks. AllGather mode
        copies only the rank's 1/dp_size shard to a persistent private buffer;
        rank-local mode retains the views and packs one block at a time into
        bounded staging storage.

        Non-DiT modules (VAE, encoders) are NOT affected — they were
        created on CPU with real weights via from_pretrained.
        """
        from safetensors import safe_open

        logger.info(
            "Loading DiT weights via mmap (meta -> shared page cache): %d tensors",
            len(plan.bindings),
        )

        # --- Convert DiT modules to meta device ---
        # The transformer was created normally (with random weights and
        # correct non-persistent buffers from __init__).  We convert it
        # to meta device to release the random weights (they're useless),
        # then load real weights via mmap.
        #
        # Non-persistent buffers (e.g. RoPE inv_freq, timestep freqs) are
        # NOT in the checkpoint — they are computed from formulas in
        # __init__.  Save them before meta conversion and restore after
        # mmap loading, so we don't need model-specific buffer rebuild code.
        #
        # The loader proved topology, coverage, source metadata, and adapter
        # compatibility before it skipped ordinary weight materialization.
        # This method only realizes that exact plan; it does not select or
        # rediscover a checkpoint layout.

        #
        # Important: when DiT modules are nested (e.g. transformer contains
        # transformer.language_model), to_empty("meta") on the parent
        # converts the child's buffers too.  So we must save ALL buffers
        # from ALL DiT modules BEFORE any to_empty call.
        saved_buffers: dict[int, dict[str, torch.Tensor]] = {}
        for dit_module in (component.module for component in dits):
            bufs: dict[str, torch.Tensor] = {}
            for name, buf in dit_module.named_buffers():
                # Check if this buffer is non-persistent on its OWNING module
                owner = dit_module
                parts = name.split(".")
                for part in parts[:-1]:
                    owner = getattr(owner, part)
                buf_name = parts[-1]
                is_non_persistent = buf_name in owner._non_persistent_buffers_set
                # Save if non-persistent OR already meta (shouldn't happen
                # before to_empty, but be safe)
                if is_non_persistent:
                    bufs[name] = buf.detach().clone()
            saved_buffers[id(dit_module)] = bufs

        # Now convert all DiT modules to meta (after saving all buffers).
        # Skip modules that are already meta (happens when a parent DiT
        # module contains a child DiT module — to_empty on the parent
        # already converted the child).
        for dit_module in (component.module for component in dits):
            if any(p.is_meta for p in dit_module.parameters()):
                logger.info(
                    "%s already on meta device (skipping to_empty, %d buffers saved)",
                    dit_module.__class__.__name__,
                    len(saved_buffers.get(id(dit_module), {})),
                )
                continue
            dit_module.to_empty(device="meta")
            logger.info(
                "Converted %s to meta device (%d non-persistent buffers saved)",
                dit_module.__class__.__name__,
                len(saved_buffers.get(id(dit_module), {})),
            )

        # Cache open file handles
        file_cache: dict[str, Any] = {}
        loaded_names: set[str] = set()

        # Realize the loader's exact runtime-name bindings.  In particular,
        # do not reconstruct names from DLO block discovery: storage planning
        # and transfer topology are separate contracts.
        for runtime_name, binding in plan.bindings.items():
            parent_path, _, leaf_name = runtime_name.rpartition(".")
            try:
                parent = pipeline.get_submodule(parent_path)
            except AttributeError as exc:
                raise RuntimeError(f"Host-weight plan target module {parent_path!r} no longer exists") from exc

            target = parent._parameters.get(leaf_name)
            is_parameter = target is not None
            if target is None:
                target = parent._buffers.get(leaf_name)
            if target is None:
                raise RuntimeError(f"Host-weight plan target tensor {runtime_name!r} no longer exists")

            if binding.file_path not in file_cache:
                file_cache[binding.file_path] = safe_open(
                    binding.file_path,
                    framework="pt",
                    device="cpu",
                )
            tensor = file_cache[binding.file_path].get_tensor(binding.checkpoint_key)
            if is_parameter:
                replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad)
                parent._parameters[leaf_name] = replacement
            else:
                replacement = tensor
                parent._buffers[leaf_name] = replacement
            loaded_names.add(runtime_name)

        logger.info("Realized %d loader-planned tensors as mmap views", len(loaded_names))

        # Keep file handles open — _shard_and_pin will read from the mmap views.
        # They will be released after _shard_and_pin completes (when params are
        # replaced with offload placeholders).
        self._mmap_file_cache = file_cache
        # The regular loader runs model-specific post-load transforms after
        # assigning checkpoint tensors. The mmap path bypasses that loader, so
        # preserve the same lifecycle for transforms such as Cosmos3's fp32
        # timestep embedder.
        for dit_name, dit_module in ((component.path, component.module) for component in dits):
            # Restore non-persistent buffers before post-load hooks and strict
            # validation. They are constructor-derived and intentionally have
            # no checkpoint binding.
            bufs = saved_buffers.get(id(dit_module), {})
            for name, buf in bufs.items():
                parent_path, _, leaf_name = name.rpartition(".")
                parent = dit_module.get_submodule(parent_path)
                restore_device = torch.device("cpu") if self._using_rank_local_mmap else self.device
                parent._buffers[leaf_name] = buf.to(restore_device)

            post_load_weights = getattr(dit_module, "post_load_weights", None)
            if callable(post_load_weights):
                post_load_weights()

            # Call validate_loaded_weights if the model defines it.
            # This preserves the sound-weight and action-weight validation
            # that AutoWeightsLoader.load_weights() triggers in the regular
            # load path (e.g. Cosmos3 checks for missing audio/action weights).
            validate = getattr(dit_module, "validate_loaded_weights", None)
            if callable(validate):
                local_loaded_names = {
                    name.removeprefix(f"{dit_name}.") for name in loaded_names if name.startswith(f"{dit_name}.")
                }
                validate(local_loaded_names)

        # Post-load hooks may rebind parameters (for example while casting a
        # submodule), so associate deferred bounded transforms with the final
        # runtime tensor objects rather than the initial mmap replacements.
        self._mmap_transforms_by_tensor_id.clear()
        for runtime_name, binding in plan.bindings.items():
            if binding.transform is None:
                continue
            parent_path, _, leaf_name = runtime_name.rpartition(".")
            parent = pipeline.get_submodule(parent_path)
            target = parent._parameters.get(leaf_name)
            if target is None:
                target = parent._buffers.get(leaf_name)
            if target is None:
                raise RuntimeError(
                    f"Host-weight transform target {runtime_name!r} no longer exists after post-load processing"
                )
            self._mmap_transforms_by_tensor_id[id(target)] = binding.transform

        remaining_meta: list[str] = []
        for dit_name, dit_module in ((component.path, component.module) for component in dits):
            remaining_meta.extend(
                f"{dit_name}.{name}" for name, tensor in dit_module.named_parameters() if tensor.is_meta
            )
            for name, tensor in dit_module.named_buffers():
                parent_path, _, leaf_name = name.rpartition(".")
                owner = dit_module.get_submodule(parent_path)
                if leaf_name not in owner._non_persistent_buffers_set and tensor.is_meta:
                    remaining_meta.append(f"{dit_name}.{name}")
        if remaining_meta:
            raise RuntimeError(
                "The prevalidated host-weight plan left "
                f"{len(remaining_meta)} DiT tensors on the meta device "
                f"(first 5: {remaining_meta[:5]})."
            )

    def _init_dp_group(self) -> None:
        """Reuse the process group initialized by parallel_state.

        When DP > 1, uses the DP group (handles strided groups with SP/CFG/PP).
        When DP = 1 but SP > 1 (and dp_size was set to sp_size in OffloadConfig),
        uses the SP group so weights are sharded across SP ranks.
        """
        if self.dp_size <= 1:
            logger.info("Distributed layerwise offload: dp_size=1, running without AllGather")
            self.dp_group = None
            return

        if not torch.distributed.is_initialized():
            raise RuntimeError(
                "torch.distributed is not initialized. "
                "Distributed layerwise offload with dp_size > 1 requires "
                "an initialized process group."
            )

        from vllm_omni.diffusion.distributed.parallel_state import (
            get_data_parallel_world_size,
            get_dp_group,
        )

        # Determine which parallel group to use for sharding.
        # When data_parallel_size > 1, use the DP group.
        # When data_parallel_size == 1 but dp_size > 1 (set from sp_size in
        # OffloadConfig), use the SP group for weight sharding.
        dp_world = get_data_parallel_world_size()
        if dp_world > 1:
            coord = get_dp_group()
        else:
            from vllm_omni.diffusion.distributed.parallel_state import get_sp_group

            coord = get_sp_group()
            logger.info(
                "Distributed layerwise offload: DP=1, using SP group (world_size=%d) for weight sharding",
                coord.world_size,
            )

        self.dp_group = coord.device_group
        self.rank = coord.rank_in_group
        self.dp_size = coord.world_size

        logger.info(
            "Distributed layerwise offload: dp_size=%d, rank_in_group=%d, global_rank=%d, group_ranks=%s",
            self.dp_size,
            self.rank,
            coord.rank,
            coord.ranks,
        )

    def _component_transport(
        self,
        component: str,
    ) -> tuple[torch.distributed.ProcessGroup | None, int, int]:
        if self.config.uses_allgather(component):
            return self.dp_group, self.dp_size, self.rank
        return None, 1, 0

    def _has_multirank_allgather(self) -> bool:
        components = self.config.components or frozenset({DIT_COMPONENT})
        return self.dp_size > 1 and any(self.config.uses_allgather(component) for component in components)

    def _install_hook_group(
        self,
        blocks: Sequence[nn.Module] | nn.ModuleList,
        component: str,
        *,
        use_dit_mmap: bool = False,
    ) -> list[DistributedLayerwiseOffloadHook]:
        """Install one circular block ring with the shared slot protocol."""
        block_list = list(blocks)
        if len(block_list) <= 1:
            raise ValueError("A distributed layerwise hook group requires at least two blocks")

        group, group_size, group_rank = self._component_transport(component)
        hooks: list[DistributedLayerwiseOffloadHook] = []
        self._all_hook_groups.append(hooks)
        self._blocks.append(block_list)
        probes = {id(block): module_materialization_probe(block) for block in block_list}
        for block, next_block in zip(
            chain((block_list[-1],), block_list[:-1]),
            block_list,
            strict=True,
        ):
            hooks.append(
                apply_distributed_block_hook(
                    block,
                    next_block,
                    self.device,
                    group,
                    group_size,
                    group_rank,
                    self.copy_stream,
                    self.comm_stream,
                    self.config.pin_cpu_memory,
                    shared_buffers=[None, None],
                    rank_local_mmap=self._using_rank_local_mmap if use_dit_mmap else False,
                    tensor_transforms=self._mmap_transforms_by_tensor_id if use_dit_mmap else None,
                    materialization_probe_tensor=probes[id(block)],
                )
            )

        # hooks = [last -> first, block0 -> block1, ...]. Alternating slots
        # keep a prefetch from overwriting the current block for any ring size.
        for index, hook in enumerate(hooks):
            hook._prev_hook = hooks[index - 1]
            hook.current_slot = index % 2
        hooks[1]._is_group_first = True
        return hooks

    def _try_layerwise_offload_encoder(self, component: ResolvedComponent) -> bool:
        """Apply DLO hooks to a text encoder's resolved block stacks.

        Rank-local transfer is always safe because it preserves the tensors
        produced by the model loader. Multi-rank AllGather is enabled only
        when the model declares that those tensors are replicated across the
        selected DLO group; the resolver rejects every other layout before
        this method installs anything.
        """
        block_groups = [stack.blocks for stack in component.stacks]
        if not block_groups:
            return False

        module = component.module
        encoder_hooks: list[DistributedLayerwiseOffloadHook] = []
        for blocks in block_groups:
            encoder_hooks.extend(self._install_hook_group(blocks, TEXT_ENCODER_COMPONENT))
        # Track the module before placement, which may fail, so outer rollback
        # also removes its partially-installed block hooks and marker state.
        self._encoder_modules.append(module)
        if not component.on_demand:
            move_non_block_state_to_device(module, block_groups, self.device)
        set_encoder_layerwise_state(
            module,
            encoder_hooks,
            block_groups,
        )
        logger.info(
            "Enabled %s DLO transfer for text encoder %s (%d blocks across %d stacks, group_size=%d)",
            self.config.transfer_for(TEXT_ENCODER_COMPONENT).value,
            component.path,
            sum(len(blocks) for blocks in block_groups),
            len(block_groups),
            self._component_transport(TEXT_ENCODER_COMPONENT)[1],
        )
        return True

    def _offload_nested_submodule(self, component: ResolvedComponent) -> None:
        """Stream one resolved DiT submodule through its own hook ring."""
        stack = component.stacks[0]
        logger.info(
            "Distributed layerwise offload for submodule '%s.%s' (%d blocks, %.0f MB total, group_size=%d)",
            component.path,
            stack.attrs[0],
            len(stack.blocks),
            sum(p.nelement() * p.element_size() for p in component.module.parameters()) / 1048576,
            self._component_transport(DIT_COMPONENT)[1],
        )
        # Keep the submodule's own non-block state resident by block identity,
        # so an attribute aliasing a streamed block is not placed with it.
        move_non_block_state_to_device(component.module, (stack.blocks,), self.device)
        self._install_hook_group(stack.blocks, DIT_COMPONENT, use_dit_mmap=True)

    def _prepare_dit_non_block_modules(
        self,
        component: ResolvedComponent,
        stack: BlockStack,
    ) -> None:
        """Place or hook the DiT parts that are outside its repeated blocks.

        This must run even when every repeated block is resident.  Otherwise
        an all-resident stage skips placement for modules such as H3's token
        refiner and enters the forward pass with CPU or meta tensors.
        """
        dit_module = component.module
        nested = {child.path.rsplit(".", 1)[-1]: child for child in component.children}
        for name, module in dit_module.named_children():
            if name in stack.attrs:
                logger.debug("Skipped blocks module %s", name)
                continue

            child = nested.get(name)
            if child is not None:
                if not child.selected:
                    logger.info("Submodule '%s' is already a DiT module, skipping layerwise offload", name)
                elif child.stacks:
                    self._offload_nested_submodule(child)
                else:
                    prepare_component(
                        child.module,
                        child.path,
                        device=self.device,
                        stage_on_demand=True,
                        blockwise=False,
                        staged_components=self._staged_components,
                    )
                continue

            try:
                module.to(self.device)
            except (NotImplementedError, RuntimeError):
                # Non-persistent buffers such as RoPE frequencies do not
                # exist in the checkpoint and must be reconstructed.
                has_meta_buffer = any(getattr(buffer, "is_meta", False) for buffer in module.buffers(recurse=True))
                if not has_meta_buffer:
                    raise
                saved_params = {
                    param_name: param.data.clone()
                    for param_name, param in module.named_parameters()
                    if not getattr(param, "is_meta", False)
                }
                module.to_empty(device=self.device)
                for submodule in module.modules():
                    if hasattr(submodule, "reset_parameters"):
                        submodule.reset_parameters()
                for param_name, param in module.named_parameters():
                    if param_name in saved_params:
                        param.data.copy_(saved_params[param_name])
                module.to(self.device)

        for param in dit_module._parameters.values():
            if param is not None and not getattr(param, "is_meta", False):
                param.data = param.data.to(self.device, non_blocking=True)
        for buffer in dit_module._buffers.values():
            if buffer is not None:
                buffer.data = buffer.data.to(self.device, non_blocking=True)

    def enable(self, pipeline: nn.Module) -> None:
        """Enable DLO and make partial startup failures transactional."""
        if self._poisoned_reason is not None:
            raise RuntimeError(self._poisoned_reason)

        # A rank-local rollback can reconstruct ordinary tensors from its host
        # masters. A multi-rank AllGather rollback is unsafe: another rank may
        # never enter the matching restoration collective after startup fails.
        restore_allgather_weights = not self._has_multirank_allgather()
        try:
            self._enable(pipeline)
        except BaseException:
            try:
                self._disable(restore_allgather_weights=restore_allgather_weights)
            except BaseException:
                logger.exception("DistributedLayerwiseOffloadBackend cleanup failed while handling an enable failure")
            raise

    def _enable(self, pipeline: nn.Module) -> None:
        if self.enabled:
            logger.warning("DistributedLayerwiseOffloadBackend already enabled")
            return

        # Initialize DP group (if not already done by early init)
        if self.dp_group is None and self._has_multirank_allgather():
            self._init_dp_group()

        # One resolver call validates the whole topology; an explicit selection
        # has already failed by the time any of the checks below run.
        resolved = resolve_offload_plan(pipeline, self.config)
        if not resolved.dits and self.config.offloads(DIT_COMPONENT):
            if self.host_weight_plan is not None:
                raise RuntimeError(
                    "DLO received a loader-owned host-weight plan, but no DiT modules were discovered to consume it"
                )
            logger.warning("No DiT/transformer modules found for selected DiT offload")

        # Storage selection belongs to the loader.  DLO consumes the exact
        # prevalidated plan that caused the loader to skip materialization;
        # without a plan, all weights must already come from the ordinary
        # loader.  The transfer protocol is selected independently below.
        host_weight_plan = self.host_weight_plan
        self._using_mmap = host_weight_plan is not None
        # A one-rank AllGather transport is rank-local in practice. Preserve
        # the mmap source as the host master instead of eagerly closing it.
        self._using_rank_local_mmap = self._using_mmap and self._component_transport(DIT_COMPONENT)[1] <= 1
        if host_weight_plan is not None:
            if host_weight_plan.backing_kind == "host_weight_runtime":
                carrier = host_weight_plan.lease_carrier
                if carrier is None:
                    raise RuntimeError("DLO received a Host Weight Runtime plan without a lease carrier")
                self._host_weight_lease = carrier.take()
                if self._host_weight_lease.closed:
                    raise RuntimeError("DLO received a closed Host Weight Runtime lease")
                # The final-layout restorer has already rebound the model to
                # immutable host tensors.  Treat those tensors as mmap-like
                # sources; transport setup below selects registered direct
                # H2D or the bounded two-slot staging fallback.
                self._using_rank_local_mmap = True
                logger.info(
                    "DLO consuming final-layout Host Weight Runtime lease %s",
                    self._host_weight_lease.provenance.resolution_id,
                )
            elif host_weight_plan.backing_kind == "checkpoint_mmap":
                self._load_weights_via_mmap(
                    pipeline,
                    resolved.dits,
                    host_weight_plan,
                )
            else:
                raise ValueError(f"Unsupported DLO host-weight backing: {host_weight_plan.backing_kind}")
            if self._using_rank_local_mmap:
                logger.info(
                    "DLO rank-local host storage enabled: source pages are "
                    "node-shared; transfer setup will select registered direct H2D or bounded host staging"
                )
        else:
            remaining_meta = [
                name
                for component in resolved.dits
                for name, tensor in chain(
                    component.module.named_parameters(),
                    component.module.named_buffers(),
                )
                if getattr(tensor, "is_meta", False)
            ]
            if remaining_meta:
                raise RuntimeError(
                    f"DLO received meta tensors without a loader-owned host-weight plan (first 5: {remaining_meta[:5]})"
                )
            logger.info("DLO is using host tensors materialized by the ordinary loader")

        # Apply each selected encoder transfer while keeping explicit VAEs and
        # unselected components resident.
        prepare_pipeline_components(
            resolved,
            device=self.device,
            staged_components=self._staged_components,
            enable_encoder_blocks=self._try_layerwise_offload_encoder,
        )

        if any(component.selected for component in resolved.dits):
            logger.info(
                "Applying distributed layer-wise offloading on %s",
                [component.path for component in resolved.dits],
            )

        # Apply hooks for each DiT module
        for component, stack in iter_streamable_dits(resolved, self.device):
            streaming = list(stack.streaming)
            if stack.resident_head:
                self._resident_blocks.extend(stack.resident)
                logger.info(
                    "Keeping %d leading blocks resident on %s; streaming %d tail blocks",
                    stack.resident_head,
                    component.path,
                    len(streaming),
                )

            self._prepare_dit_non_block_modules(component, stack)

            if not streaming:
                logger.info("All blocks for %s are resident; no streaming hooks required", component.path)
                continue

            self._install_hook_group(streaming, DIT_COMPONENT, use_dit_mmap=True)
        if self._resident_blocks:
            self._resident_layer_group = PinnedResidentLayerGroup(
                self._resident_blocks,
                self.device,
                self.copy_stream,
                self.config.pin_cpu_memory,
                rank_local_mmap=self._using_rank_local_mmap,
                defer_staging=bool(self._all_hook_groups),
                tensor_transforms=self._mmap_transforms_by_tensor_id,
            )
            pipeline._dlo_residency_controller = self
            self._residency_pipeline_ref = weakref.ref(pipeline)

        all_hooks = [hook for group in self._all_hook_groups for hook in group]
        self._configure_hwr_transfer(all_hooks)

        if not self._all_hook_groups:
            self.enabled = bool(self._resident_blocks or self._encoder_modules or self._staged_components)
            if self._using_mmap and not self.enabled:
                self._release_mmap_handles()
            return

        # Unified allocation: 2 shared output buffers + 2 shared shard buffers
        # sized to the max block across ALL module groups (gen_layers +
        # language_model).  Groups execute sequentially, so 2 buffers suffice.
        unified_buffers = self._allocate_shared_buffers(all_hooks)
        allgather_hooks = [hook for hook in all_hooks if hook.dp_size > 1]
        mmap_hooks = [hook for hook in all_hooks if hook.rank_local_mmap]
        unified_shard_buffers = self._allocate_shared_shard_buffers(allgather_hooks) if allgather_hooks else None
        unified_cpu_staging = None
        cpu_staging_events = None
        if self._using_rank_local_mmap and not self._using_registered_mmap:
            unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
                mmap_hooks,
                self._resident_layer_group,
            )
            cpu_staging_events = [None, None]
            if self._resident_layer_group is not None:
                # Resident and streamed layers execute in the same stage and
                # reuse the same host slots. Events serialize slot reuse.
                self._resident_layer_group._cpu_staging_buffers = [
                    buffers for buffers in unified_cpu_staging if buffers is not None
                ]
                self._resident_layer_group._cpu_staging_events = cpu_staging_events

        # Shared slot-group tracker: _shared_slot_group[slot] = group_id
        # that last wrote to that slot.  Group-first hooks use this to
        # skip sync-prefetch when the slot still contains their own data.
        shared_slot_group = [-1, -1]

        for group_idx, group in enumerate(self._all_hook_groups):
            for hook in group:
                hook.gpu_buffers = unified_buffers
                hook._owns_buffers = False
                if unified_shard_buffers is not None and hook.dp_size > 1:
                    hook.gpu_shard_buffers = unified_shard_buffers
                if hook.rank_local_mmap and unified_cpu_staging is not None and cpu_staging_events is not None:
                    hook.cpu_staging_buffers = unified_cpu_staging
                    hook.cpu_staging_events = cpu_staging_events
                hook._group_id = group_idx
                hook._shared_slot_group = shared_slot_group

        # Defer every group's first prefetch until its first forward. DP ranks
        # may intentionally own different rank-local groups (for example only
        # rank 0 owns the text encoder) while sharing DiT AllGather groups.
        # Entering a collective for each rank's locally first group here would
        # give the ranks different collective orders and deadlock startup. The
        # group-first hook treats the initially unowned slot as contaminated
        # and performs the required synchronous prefetch on first use.

        total_blocks = sum(len(b) for b in self._blocks)
        transfer_summary = ", ".join(
            f"{component}: {self.config.transfer_for(component).value}"
            for component in sorted(self.config.components or {DIT_COMPONENT})
        )
        logger.info(
            f"Distributed layer-wise offloading enabled on {total_blocks} blocks "
            f"across {len(self._all_hook_groups)} group(s), "
            f"transfers={{{transfer_summary}}}, "
            f"unified shared_buffers=2"
        )

        self.enabled = True

        if self._using_mmap and not self._using_rank_local_mmap:
            # AllGather mode copied each rank's persistent shard, so the source
            # mappings are no longer needed. Rank-local mode retains them as
            # the node-shared host master until disable().
            self._release_mmap_handles()

        self._cleanup_after_loading()

    def _release_mmap_handles(self) -> None:
        """Release source handles and the transport-owned HWR lease."""
        if self._host_registration is not None:
            raise HostRegistrationCleanupError("cannot close HWR mappings while host registration is still active")
        self._mmap_transforms_by_tensor_id.clear()
        if hasattr(self, "_mmap_file_cache"):
            self._mmap_file_cache.clear()
            del self._mmap_file_cache
            logger.info("Released safetensors mmap file handles")
        lease = self._host_weight_lease
        self._host_weight_lease = None
        if lease is not None and not lease.closed:
            lease.close()
            logger.info("Released Host Weight Runtime lease %s", lease.provenance.resolution_id)
        if self.host_weight_plan is not None and self.host_weight_plan.lease_carrier is not None:
            self.host_weight_plan.lease_carrier.close()

    def _cleanup_after_loading(self) -> None:
        """Synchronize and release freed device/CPU memory after sharding."""
        current_omni_platform.synchronize()
        current_omni_platform.empty_cache()
        import ctypes as _ctypes
        import gc as _gc

        _gc.collect()
        try:
            _ctypes.CDLL("libc.so.6").malloc_trim(0)
        except Exception:
            pass

    def _disable(self, *, restore_allgather_weights: bool) -> None:
        has_open_lease = self._host_weight_lease is not None and not self._host_weight_lease.closed
        has_registration = self._host_registration is not None
        has_carrier = (
            self.host_weight_plan is not None
            and self.host_weight_plan.lease_carrier is not None
            and not self.host_weight_plan.lease_carrier.closed
        )
        has_partial_hooks = bool(
            self._blocks
            or self._all_hook_groups
            or self._encoder_modules
            or self._staged_components
            or self._resident_layer_group is not None
            or self._residency_pipeline_ref is not None
        )
        if (
            not self.enabled
            and not hasattr(self, "_mmap_file_cache")
            and not has_open_lease
            and not has_registration
            and not has_carrier
            and not has_partial_hooks
        ):
            return

        self._clear_residency_controller()

        # A hook can leave the circular tail prefetch queued after the final
        # forward. Drain every transport before releasing hook-owned host or
        # device buffers, including the ordinary rank-local path.
        sync_error = run_cleanup_steps([("synchronizing pending DLO transfers", current_omni_platform.synchronize)])

        unique_hooks: list[DistributedLayerwiseOffloadHook] = []
        seen_blocks: set[int] = set()
        for hook in chain.from_iterable(self._all_hook_groups):
            block_id = id(hook.next_block)
            if block_id not in seen_blocks:
                seen_blocks.add(block_id)
                unique_hooks.append(hook)

        allgather_hooks = [hook for hook in unique_hooks if hook.dp_size > 1]
        rank_local_hooks = [hook for hook in unique_hooks if hook.dp_size <= 1]
        skipped_allgather = bool(allgather_hooks) and not restore_allgather_weights
        if skipped_allgather:
            # Startup rollback cannot safely enter a collective that a failed
            # peer may never reach, and shutdown has no later user. Those blocks
            # cannot be reconstructed, so make accidental reuse explicit
            # instead of accepting zero weights.
            self._poisoned_reason = (
                "Distributed layerwise offload skipped AllGather weight restoration; "
                "recreate the backend and reload the pipeline before retrying"
            )

        collective_error = None
        if restore_allgather_weights:
            # Run collective-bearing restores before any rank-local operation:
            # a local failure must never keep this rank out of a later AllGather.
            collective_error = run_cleanup_steps(
                ("restoring an AllGather block", hook.restore_next_block_to_cpu) for hook in allgather_hooks
            )
        rank_local_error = run_cleanup_steps(
            ("restoring a rank-local block", hook.restore_next_block_to_cpu) for hook in rank_local_hooks
        )
        removal_error = run_cleanup_steps(
            (
                "removing a distributed block hook",
                lambda block=block: remove_distributed_block_hook(block),
            )
            for blocks in self._blocks
            for block in blocks
        )
        encoder_error = run_cleanup_steps(
            (
                "clearing distributed encoder state",
                lambda module=module: clear_encoder_layerwise_state(module),
            )
            for module in self._encoder_modules
        )

        resident_error = None
        if self._resident_layer_group is not None:
            # Resident layers are always rank-local (the public parser rejects
            # resident DiT layers combined with AllGather).
            resident_error = run_cleanup_steps(
                [("restoring resident DLO blocks", self._resident_layer_group.restore_to_cpu)]
            )

        lifecycle_error = next(
            (
                error
                for error in (
                    sync_error,
                    collective_error,
                    rank_local_error,
                    removal_error,
                    encoder_error,
                    resident_error,
                )
                if error is not None
            ),
            None,
        )

        # Unregistration is independent once streams have drained, but the
        # lease/file mappings must stay alive if restoration or hook removal
        # needs a retry.
        registration_error = None
        if sync_error is None:
            registration_error = run_cleanup_steps(
                [("releasing registered HWR mappings", self._release_registered_mmap)]
            )
        cleanup_error = lifecycle_error or registration_error
        if cleanup_error is not None:
            raise cleanup_error

        release_error = run_cleanup_steps([("releasing DLO mmap handles", self._release_mmap_handles)])

        self._blocks.clear()
        self._all_hook_groups.clear()
        self._resident_blocks.clear()
        self._resident_layer_group = None
        self._encoder_modules.clear()
        self._staged_components.clear()
        # Loader plans are single-use (HWR carriers in particular). A clean
        # disable can rebuild from restored CPU tensors; a poisoned backend is
        # guarded above and must be recreated.
        self.host_weight_plan = None
        self._using_mmap = False
        self._using_rank_local_mmap = False
        self._using_registered_mmap = False
        self.enabled = False
        logger.info("Distributed layer-wise offloading disabled")
        if release_error is not None:
            raise release_error

    def disable(self) -> None:
        self._disable(restore_allgather_weights=True)

    def shutdown(self) -> None:
        # Restoring AllGather blocks rebuilds the full DiT in every rank's host
        # memory, which can outlast the executor's shutdown grace period.
        self._disable(restore_allgather_weights=False)

    @staticmethod
    def _allocate_shared_buffers(
        hooks: list[DistributedLayerwiseOffloadHook],
    ) -> list[dict[torch.dtype, torch.Tensor] | None]:
        """Allocate exactly 2 shared device buffers sized to the largest block.

        All hooks share these 2 buffers. At any time, slot 0 holds the current
        layer's weights and slot 1 holds the next layer's weights (or vice
        versa). This ensures only 2 layers' worth of weights reside on device,
        regardless of the total number of blocks.
        """
        max_sizes: dict[torch.dtype, int] = {}
        for hook in hooks:
            dp = hook.dp_size
            for dtype, metas in hook.metadata.items():
                total = sum(m["numel"] for m in metas)
                # AllGather output = dp * ceil(total/dp) (padded for equal shards)
                if dp > 1:
                    total = ((total + dp - 1) // dp) * dp
                if dtype not in max_sizes or total > max_sizes[dtype]:
                    max_sizes[dtype] = total

        device = hooks[0].device
        shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None]
        for slot in range(2):
            gpu_weights: dict[torch.dtype, torch.Tensor] = {}
            for dtype, total_numel in max_sizes.items():
                gpu_weights[dtype] = torch.empty(total_numel, dtype=dtype, device=device)
            shared_buffers[slot] = gpu_weights

        logger.info(
            "Allocated 2 shared device buffers (max block size: %s)",
            {str(k): f"{v * _dtype_size(k) / 1024 / 1024:.1f}MB" for k, v in max_sizes.items()},
        )
        return shared_buffers

    @staticmethod
    def _allocate_shared_cpu_staging_buffers(
        hooks: list[DistributedLayerwiseOffloadHook],
        resident_group: PinnedResidentLayerGroup | None = None,
    ) -> list[dict[torch.dtype, torch.Tensor] | None]:
        """Allocate two bounded host slots for rank-local mmap -> device copies."""
        max_sizes: dict[torch.dtype, int] = {}
        for hook in hooks:
            for dtype, metas in hook.metadata.items():
                total = sum(meta["numel"] for meta in metas)
                max_sizes[dtype] = max(max_sizes.get(dtype, 0), total)
        if resident_group is not None:
            for state in resident_group._states:
                for dtype, metas in state["metadata"].items():
                    total = sum(meta["numel"] for meta in metas)
                    max_sizes[dtype] = max(max_sizes.get(dtype, 0), total)

        pin_memory = hooks[0].pin_memory if hooks else bool(resident_group and resident_group.pin_memory)
        shared_staging: list[dict[torch.dtype, torch.Tensor] | None] = [None, None]
        for slot in range(2):
            buffers: dict[torch.dtype, torch.Tensor] = {}
            for dtype, total_numel in max_sizes.items():
                buffer = torch.empty(
                    total_numel,
                    dtype=dtype,
                    device="cpu",
                    pin_memory=pin_memory,
                )
                buffers[dtype] = buffer
            shared_staging[slot] = buffers

        logger.info(
            "Allocated 2 shared host staging buffers for rank-local mmap (max block size: %s, pinned=%s)",
            {str(k): f"{v * _dtype_size(k) / 1024 / 1024:.1f}MB" for k, v in max_sizes.items()},
            pin_memory,
        )
        return shared_staging

    @staticmethod
    def _allocate_shared_shard_buffers(
        hooks: list[DistributedLayerwiseOffloadHook],
    ) -> list[dict[torch.dtype, torch.Tensor] | None]:
        """Allocate 2 shared shard (AllGather input) buffers sized to the
        largest per-rank shard.

        All hooks share these 2 input buffers — the same device address is
        reused for every layer's AllGather so HCCL can reuse its internal
        communication buffers (preventing rank-dependent HBM imbalance).
        Cost: 2 × max_shard_size (~184 MB) instead of 2 × N_hooks × shard.
        """
        max_shard_sizes: dict[torch.dtype, int] = {}
        for hook in hooks:
            for dtype, shard in hook.cpu_shards.items():
                numel = shard.numel()
                if dtype not in max_shard_sizes or numel > max_shard_sizes[dtype]:
                    max_shard_sizes[dtype] = numel

        device = hooks[0].device
        shared_shard_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None]
        for slot in range(2):
            shard_bufs: dict[torch.dtype, torch.Tensor] = {}
            for dtype, numel in max_shard_sizes.items():
                shard_bufs[dtype] = torch.empty(numel, dtype=dtype, device=device)
            shared_shard_buffers[slot] = shard_bufs

        logger.info(
            "Allocated 2 shared shard buffers (max shard size: %s)",
            {str(k): f"{v * _dtype_size(k) / 1024 / 1024:.1f}MB" for k, v in max_shard_sizes.items()},
        )
        return shared_shard_buffers