Blame for library/custom_offloading_utils.py

from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
import ctypes
from dataclasses import dataclass
import gc
import os
import time
from typing import Any, Optional, Union, Callable, Tuple
import torch
import torch.nn as nn


# Keep these functions here for portability, and private to avoid confusion with the ones in device_utils.py
def _clean_memory_on_device(device: torch.device):
    r"""
    Clean memory on the specified device, will be called from training scripts.
    """
    gc.collect()

    # device may "cuda" or "cuda:0", so we need to check the type of device
    if device.type == "cuda":
        torch.cuda.empty_cache()
    if device.type == "xpu":
        torch.xpu.empty_cache()
    if device.type == "mps":
        torch.mps.empty_cache()


def _synchronize_device(device: torch.device):
    if device.type == "cuda":
        torch.cuda.synchronize()
    elif device.type == "xpu":
        torch.xpu.synchronize()
    elif device.type == "mps":
        torch.mps.synchronize()


def swap_weight_devices_cuda(device: torch.device, layer_to_cpu: nn.Module, layer_to_cuda: nn.Module):
    assert layer_to_cpu.__class__ == layer_to_cuda.__class__

    weight_swap_jobs: list[Tuple[nn.Module, nn.Module, torch.Tensor, torch.Tensor]] = []

    # This is not working for all cases (e.g. SD3), so we need to find the corresponding modules
    # for module_to_cpu, module_to_cuda in zip(layer_to_cpu.modules(), layer_to_cuda.modules()):
    #     print(module_to_cpu.__class__, module_to_cuda.__class__)
    #     if hasattr(module_to_cpu, "weight") and module_to_cpu.weight is not None:
    #         weight_swap_jobs.append((module_to_cpu, module_to_cuda, module_to_cpu.weight.data, module_to_cuda.weight.data))

    modules_to_cpu = {k: v for k, v in layer_to_cpu.named_modules()}
    for module_to_cuda_name, module_to_cuda in layer_to_cuda.named_modules():
        if hasattr(module_to_cuda, "weight") and module_to_cuda.weight is not None:
            module_to_cpu = modules_to_cpu.get(module_to_cuda_name, None)
            if module_to_cpu is not None and module_to_cpu.weight.shape == module_to_cuda.weight.shape:
                weight_swap_jobs.append((module_to_cpu, module_to_cuda, module_to_cpu.weight.data, module_to_cuda.weight.data))
            else:
                if module_to_cuda.weight.data.device.type != device.type:
                    # print(
                    #     f"Module {module_to_cuda_name} not found in CPU model or shape mismatch, so not swapping and moving to device"
                    # )
                    module_to_cuda.weight.data = module_to_cuda.weight.data.to(device)

    torch.cuda.current_stream().synchronize()  # this prevents the illegal loss value

    stream = torch.Stream(device="cuda")
    with torch.cuda.stream(stream):
        # cuda to cpu
        for module_to_cpu, module_to_cuda, cuda_data_view, cpu_data_view in weight_swap_jobs:
            cuda_data_view.record_stream(stream)
            module_to_cpu.weight.data = cuda_data_view.data.to("cpu", non_blocking=True)

        stream.synchronize()

        # cpu to cuda
        for module_to_cpu, module_to_cuda, cuda_data_view, cpu_data_view in weight_swap_jobs:
            cuda_data_view.copy_(module_to_cuda.weight.data, non_blocking=True)
            module_to_cuda.weight.data = cuda_data_view

    stream.synchronize()
    torch.cuda.current_stream().synchronize()  # this prevents the illegal loss value


def swap_weight_devices_no_cuda(device: torch.device, layer_to_cpu: nn.Module, layer_to_cuda: nn.Module):
    """
    not tested
    """
    assert layer_to_cpu.__class__ == layer_to_cuda.__class__

    weight_swap_jobs: list[Tuple[nn.Module, nn.Module, torch.Tensor, torch.Tensor]] = []
    for module_to_cpu, module_to_cuda in zip(layer_to_cpu.modules(), layer_to_cuda.modules()):
        if hasattr(module_to_cpu, "weight") and module_to_cpu.weight is not None:
            weight_swap_jobs.append((module_to_cpu, module_to_cuda, module_to_cpu.weight.data, module_to_cuda.weight.data))

    # device to cpu
    for module_to_cpu, module_to_cuda, cuda_data_view, cpu_data_view in weight_swap_jobs:
        module_to_cpu.weight.data = cuda_data_view.data.to("cpu", non_blocking=True)

    _synchronize_device(device)

    # cpu to device
    for module_to_cpu, module_to_cuda, cuda_data_view, cpu_data_view in weight_swap_jobs:
        cuda_data_view.copy_(module_to_cuda.weight.data, non_blocking=True)
        module_to_cuda.weight.data = cuda_data_view

    _synchronize_device(device)


def weighs_to_device(layer: nn.Module, device: torch.device):
    for module in layer.modules():
        if hasattr(module, "weight") and module.weight is not None:
            module.weight.data = module.weight.data.to(device, non_blocking=True)


class Offloader:
    """
    common offloading class
    """

    def __init__(self, num_blocks: int, blocks_to_swap: int, device: torch.device, debug: bool = False):
        self.num_blocks = num_blocks
        self.blocks_to_swap = blocks_to_swap
        self.device = device
        self.debug = debug

        self.thread_pool = ThreadPoolExecutor(max_workers=1)
        self.futures = {}
        self.cuda_available = device.type == "cuda"

    def swap_weight_devices(self, block_to_cpu: nn.Module, block_to_cuda: nn.Module):
        if self.cuda_available:
            swap_weight_devices_cuda(self.device, block_to_cpu, block_to_cuda)
        else:
            swap_weight_devices_no_cuda(self.device, block_to_cpu, block_to_cuda)

    def _submit_move_blocks(self, blocks, block_idx_to_cpu, block_idx_to_cuda):
        def move_blocks(bidx_to_cpu, block_to_cpu, bidx_to_cuda, block_to_cuda):
            if self.debug:
                start_time = time.perf_counter()
                print(f"Move block {bidx_to_cpu} to CPU and block {bidx_to_cuda} to {'CUDA' if self.cuda_available else 'device'}")

            self.swap_weight_devices(block_to_cpu, block_to_cuda)

            if self.debug:
                print(f"Moved blocks {bidx_to_cpu} and {bidx_to_cuda} in {time.perf_counter() - start_time:.2f}s")
            return bidx_to_cpu, bidx_to_cuda  # , event

        block_to_cpu = blocks[block_idx_to_cpu]
        block_to_cuda = blocks[block_idx_to_cuda]

        self.futures[block_idx_to_cuda] = self.thread_pool.submit(
            move_blocks, block_idx_to_cpu, block_to_cpu, block_idx_to_cuda, block_to_cuda
        )

    def _wait_blocks_move(self, block_idx):
        if block_idx not in self.futures:
            return

        if self.debug:
            print(f"Wait for block {block_idx}")
            start_time = time.perf_counter()

        future = self.futures.pop(block_idx)
        _, bidx_to_cuda = future.result()

        assert block_idx == bidx_to_cuda, f"Block index mismatch: {block_idx} != {bidx_to_cuda}"

        if self.debug:
            print(f"Waited for block {block_idx}: {time.perf_counter() - start_time:.2f}s")


# Gradient tensors
_grad_t = Union[tuple[torch.Tensor, ...], torch.Tensor]


class ModelOffloader(Offloader):
    """
    supports forward offloading
    """

    def __init__(
        self,
        blocks: Union[list[nn.Module], nn.ModuleList],
        blocks_to_swap: int,
        device: torch.device,
        supports_backward: bool = True,
        debug: bool = False,
    ):
        super().__init__(len(blocks), blocks_to_swap, device, debug)

        self.supports_backward = supports_backward
        self.forward_only = not supports_backward  # forward only offloading: can be changed to True for inference

        if self.supports_backward:
            # register backward hooks
            self.remove_handles = []
            for i, block in enumerate(blocks):
                hook = self.create_backward_hook(blocks, i)
                if hook is not None:
                    handle = block.register_full_backward_hook(hook)
                    self.remove_handles.append(handle)

    def set_forward_only(self, forward_only: bool):
        # switching must wait for all pending transfers
        for block_idx in list(self.futures.keys()):
            self._wait_blocks_move(block_idx)
        self.forward_only = forward_only

    def __del__(self):
        if self.supports_backward:
            for handle in self.remove_handles:
                handle.remove()

    def create_backward_hook(
        self, blocks: Union[list[nn.Module], nn.ModuleList], block_index: int
    ) -> Optional[Callable[[nn.Module, _grad_t, _grad_t], Union[None, _grad_t]]]:
        # -1 for 0-based index
        num_blocks_propagated = self.num_blocks - block_index - 1
        swapping = num_blocks_propagated > 0 and num_blocks_propagated <= self.blocks_to_swap
        waiting = block_index > 0 and block_index <= self.blocks_to_swap

        if not swapping and not waiting:
            return None

        # create  hook
        block_idx_to_cpu = self.num_blocks - num_blocks_propagated
        block_idx_to_cuda = self.blocks_to_swap - num_blocks_propagated
        block_idx_to_wait = block_index - 1

        def backward_hook(module: nn.Module, grad_input: _grad_t, grad_output: _grad_t):
            if self.debug:
                print(f"Backward hook for block {block_index}")

            if swapping:
                self._submit_move_blocks(blocks, block_idx_to_cpu, block_idx_to_cuda)
            if waiting:
                self._wait_blocks_move(block_idx_to_wait)
            return None

        return backward_hook

    def prepare_block_devices_before_forward(self, blocks: Union[list[nn.Module], nn.ModuleList]):
        if self.blocks_to_swap is None or self.blocks_to_swap == 0:
            return

        if self.debug:
            print(f"Prepare block devices before forward")

        # wait for all pending transfers
        for block_idx in list(self.futures.keys()):
            self._wait_blocks_move(block_idx)

        for b in blocks[0 : self.num_blocks - self.blocks_to_swap]:
            b.to(self.device)
            weighs_to_device(b, self.device)  # make sure weights are on device

        for b in blocks[self.num_blocks - self.blocks_to_swap :]:
            b.to(self.device)  # move block to device first. this makes sure that buffers (non weights) are on the device
            weighs_to_device(b, torch.device("cpu"))  # make sure weights are on cpu

        _synchronize_device(self.device)
        _clean_memory_on_device(self.device)

    def wait_for_block(self, block_idx: int):
        if self.blocks_to_swap is None or self.blocks_to_swap == 0:
            return
        self._wait_blocks_move(block_idx)

    def submit_move_blocks_forward(self, blocks: Union[list[nn.Module], nn.ModuleList], block_idx: int):
        # check if blocks_to_swap is enabled
        if self.blocks_to_swap is None or self.blocks_to_swap == 0:
            return

        # if backward is enabled, we do not swap blocks in forward pass more than blocks_to_swap, because it should be on GPU
        if not self.forward_only and block_idx >= self.blocks_to_swap:
            return

        block_idx_to_cpu = block_idx
        block_idx_to_cuda = self.num_blocks - self.blocks_to_swap + block_idx
        # this works for forward-only offloading. move upstream blocks to cuda
        block_idx_to_cuda = block_idx_to_cuda % self.num_blocks
        self._submit_move_blocks(blocks, block_idx_to_cpu, block_idx_to_cuda)


# endregion


# region H2D-only block swap (ported from Musubi Tuner)


def default_swap_tensor_selector(block: nn.Module) -> list[tuple[nn.Module, str]]:
    """Default streaming selection rule: the ``weight`` of every ``*Linear`` module.

    An offloader can be given a different selector to also stream tensors that a quantization scheme hangs off
    the same modules (per-group scales etc.); each entry is ``(module, attribute_name)`` where the attribute is a
    parameter or a registered buffer of the module.
    """
    jobs = []
    for _, module in block.named_modules():
        if hasattr(module, "weight") and module.weight is not None and module.__class__.__name__.endswith("Linear"):
            jobs.append((module, "weight"))
    return jobs


@dataclass
class BlockSwapConfig:
    """
    Construction policy for a block-swap offloader, assembled once by the training/inference script and
    passed through each architecture's ``enable_block_swap`` to ``create_offloader``.

    It holds everything the offloader constructor needs that is *not* architecture-specific. The
    architecture-specific arguments (the block-type label, the block list, and its counts) are supplied by
    ``enable_block_swap`` per block list. Adding a new offloader knob (or a whole new offloader type) means
    adding a field here and a branch in ``create_offloader`` -- the per-architecture ``enable_block_swap``
    signatures stay ``(blocks_to_swap, config)`` and never change.
    """

    device: torch.device
    supports_backward: bool
    use_pinned_memory: bool = False
    h2d_only: bool = False  # frozen-base (LoRA / LoHa / LoKr) only: H2D-only streaming, no device->host copy
    ring_size: int = 2  # (h2d_only) number of GPU ring buffers for streamed blocks; 2 = double buffering
    debug: bool = False

    @classmethod
    def from_args(cls, args, device: torch.device, supports_backward: bool) -> "BlockSwapConfig":
        """Build from a parsed-args namespace, tolerating scripts whose parser lacks the optional knobs."""
        h2d_only = getattr(args, "block_swap_h2d_only", False)

        # H2D-only streams frozen weights into a reused GPU ring buffer and loads them with an in-place copy_,
        # which bumps the autograd version of the saved weight view. Gradient checkpointing re-reads the weight
        # at recompute time (a short, safe window), so the version stays consistent; without it the backward
        # pass fails with an opaque "modified by an inplace operation ... version N expected M" error that does
        # not point at the offloader. Fail early with an actionable message instead. (Inference / forward-only
        # has no backward, so it is unaffected.)
        if h2d_only and supports_backward and not getattr(args, "gradient_checkpointing", False):
            raise ValueError(
                "--block_swap_h2d_only requires --gradient_checkpointing for training. H2D-only block swap streams"
                " frozen weights through a reused GPU ring buffer, which advances the autograd version of weights"
                " saved for backward; gradient checkpointing re-reads them at recompute time and avoids this."
                " / --block_swap_h2d_only は学習時に --gradient_checkpointing が必須です（リングバッファの上書きで"
                "backward 用に保存された重みの version が進むため。gradient checkpointing は再計算時に読み直すので回避できます）。"
            )

        ring_size = getattr(args, "block_swap_ring_size", 2)
        if ring_size < 1:
            raise ValueError("--block_swap_ring_size must be >= 1")

        return cls(
            device=device,
            supports_backward=supports_backward,
            use_pinned_memory=getattr(args, "use_pinned_memory_for_block_swap", False),
            h2d_only=h2d_only,
            ring_size=ring_size,
        )


def create_offloader(block_type: str, blocks: list[nn.Module], num_blocks: int, blocks_to_swap: int, config: BlockSwapConfig):
    """
    Create the block-swap offloader for one block list. This is the single place that selects the offloader
    implementation from ``config``; ``enable_block_swap`` only supplies the per-block-list arguments, so a new
    offloader type plugs in here without touching any architecture.
    """
    if config.h2d_only:
        # H2D-only streaming for frozen-base (LoRA) training: keep a CPU master, copy Host->Device only.
        return LoRAStreamOffloader(
            block_type,
            blocks,
            num_blocks,
            blocks_to_swap,
            config.supports_backward,
            config.device,
            ring_size=config.ring_size,
            use_pinned_memory=config.use_pinned_memory,
            debug=config.debug,
        )
    # classic exchange offloader. `use_pinned_memory` is not used by this implementation.
    return ModelOffloader(blocks, blocks_to_swap, config.device, supports_backward=config.supports_backward, debug=config.debug)


class _DirectCopier:
    """
    Transfer engine for ``LoRAStreamOffloader`` when the masters are *pinned*: copies a flat pinned host
    buffer straight into a flat device buffer on a private copy stream (async Host->Device only).

    This is the seam that isolates all CUDA stream/event primitives from the residency logic. The residency
    side only ever calls ``submit`` / ``wait``; ``_StagedCopier`` (pageable master + worker thread) and a
    future non-CUDA device shim plug in behind the same interface without the residency code changing. Keep
    every ``torch.cuda.*`` transfer primitive inside copiers like this one.
    """

    def __init__(self, device: torch.device, debug: bool = False):
        self.device = device
        self.debug = debug
        self.copy_stream = torch.cuda.Stream(device=device)
        self._events = {}  # key -> cuda.Event recording H2D completion
        self._xfer = {}  # key -> (ev_a, ev_b) transfer-timing events (debug only)

    def submit(self, key, dst_flat: torch.Tensor, src_flat: torch.Tensor, gate_event=None):
        """Begin the H2D copy of ``src_flat`` into ``dst_flat``, gated behind ``gate_event`` (compute done with dst)."""
        ev_a = ev_b = None
        with torch.cuda.stream(self.copy_stream):
            if gate_event is not None:
                self.copy_stream.wait_event(gate_event)  # do not overwrite a buffer still used by compute
            if self.debug:
                ev_a = torch.cuda.Event(enable_timing=True)
                ev_b = torch.cuda.Event(enable_timing=True)
                ev_a.record(self.copy_stream)
            dst_flat.copy_(src_flat, non_blocking=True)  # src is pinned -> async H2D
            if self.debug:
                ev_b.record(self.copy_stream)
                self._xfer[key] = (ev_a, ev_b)
        self._events[key] = self.copy_stream.record_event()

    def wait(self, key):
        """Return the H2D-completion event for ``key`` (None if nothing is in flight for it)."""
        return self._events.get(key)

    def pop_xfer_timing(self, key):
        """Pop the (start, end) transfer-timing events recorded for ``key`` (debug only)."""
        return self._xfer.pop(key, None)

    def sync(self):
        self.copy_stream.synchronize()

    def reset(self):
        self._events.clear()
        self._xfer.clear()


class _StagedCopier:
    """
    Transfer engine for ``LoRAStreamOffloader`` when the masters are kept in *pageable* host memory: each
    H2D is staged through a small pool of pinned buffers filled by a background worker thread. Same
    ``submit`` / ``wait`` interface as ``_DirectCopier``, so the residency code is identical.

    Why: pinning every streaming master (what ``_DirectCopier`` needs) can cost tens of GB of non-pageable
    "shared GPU memory" (notably on Windows). Here the masters stay pageable and only ``num_staging`` pinned
    buffers exist. A pageable->device copy is, inside the CUDA driver, already a synchronous serial
    pageable->pinned->device transfer; doing it ourselves on a worker thread lets the pageable->pinned
    memcpy overlap the previous block's pinned->device DMA and -- crucially -- keeps it off the main thread
    so compute kernel launches are never blocked. That is the whole speed difference versus the old
    synchronous pageable path this replaces.

    Hazards handled:
      - a staging buffer is reused only after the H2D that consumed it has completed (one consume-event per
        buffer; the worker waits on it before overwriting);
      - ``wait`` blocks the caller until the worker has actually enqueued the H2D and recorded its
        completion event, so the residency layer can ``current_stream().wait_event`` it as usual.
    """

    def __init__(self, device: torch.device, num_staging: int = 2, debug: bool = False):
        self.device = device
        self.num_staging = max(1, num_staging)
        self.debug = debug
        self.copy_stream = torch.cuda.Stream(device=device)
        self._pool = ThreadPoolExecutor(max_workers=1)
        self._futures = {}  # key -> Future -> (load_event, (ev_a, ev_b))
        self._xfer = {}  # key -> (ev_a, ev_b) transfer-timing events (debug only)
        self._staging = None  # list[pinned uint8 tensor] (lazily sized to the flat block)
        self._staging_free = None  # list[cuda.Event | None]: the H2D that last consumed each buffer
        self._rr = 0  # round-robin staging index (touched only on the worker thread)

    def _ensure_staging(self, nbytes: int):
        if self._staging is None:
            self._staging = [torch.empty(nbytes, dtype=torch.uint8).pin_memory(device=self.device) for _ in range(self.num_staging)]
            self._staging_free = [None] * self.num_staging

    def _run(self, dst_flat: torch.Tensor, src_flat: torch.Tensor, gate_event):
        # runs on the worker thread: pageable->pinned memcpy, then enqueue the async pinned->device H2D
        idx = self._rr
        self._rr = (idx + 1) % self.num_staging
        staging = self._staging[idx]

        free_ev = self._staging_free[idx]
        if free_ev is not None:
            free_ev.synchronize()  # CPU(worker) waits until the prior H2D released this buffer -- off main thread

        # pageable -> pinned host memcpy. Not `staging.copy_(src_flat)`: torch parallelizes a large CPU copy over all
        # intra-op threads and the OpenMP workers keep spinning afterwards, so the repeated per-block copies kept
        # nearly every core busy (46 of 64 cores on average in a MOSS-Image run). A single-threaded memmove
        # (about 18 GB/s here) is still faster than the PCIe DMA that consumes the buffer, releases the GIL and
        # does not touch the OpenMP pool.
        ctypes.memmove(staging.data_ptr(), src_flat.data_ptr(), src_flat.numel())

        ev_a = ev_b = None
        with torch.cuda.stream(self.copy_stream):
            if gate_event is not None:
                self.copy_stream.wait_event(gate_event)  # do not overwrite a dst still used by compute
            if self.debug:
                ev_a = torch.cuda.Event(enable_timing=True)
                ev_b = torch.cuda.Event(enable_timing=True)
                ev_a.record(self.copy_stream)
            dst_flat.copy_(staging, non_blocking=True)  # pinned -> device, async
            if self.debug:
                ev_b.record(self.copy_stream)
            done = self.copy_stream.record_event()  # H2D complete == this staging buffer is free again
        self._staging_free[idx] = done
        return done, (ev_a, ev_b)

    def submit(self, key, dst_flat: torch.Tensor, src_flat: torch.Tensor, gate_event=None):
        self._ensure_staging(src_flat.numel())  # flat uint8 -> numel == byte size
        self._futures[key] = self._pool.submit(self._run, dst_flat, src_flat, gate_event)

    def wait(self, key):
        fut = self._futures.get(key)
        if fut is None:
            return None
        load_event, xfer = fut.result()  # block until the worker enqueued the H2D (tiny if it ran ahead)
        if self.debug and xfer[0] is not None:
            self._xfer[key] = xfer
        return load_event

    def pop_xfer_timing(self, key):
        return self._xfer.pop(key, None)

    def _flush(self):
        for fut in list(self._futures.values()):
            fut.result()

    def sync(self):
        self._flush()
        self.copy_stream.synchronize()

    def reset(self):
        self._flush()
        self.copy_stream.synchronize()  # ensure all in-flight H2D are done before clearing staging state
        self._futures.clear()
        self._xfer.clear()
        if self._staging is not None:
            self._staging_free = [None] * self.num_staging

    def __del__(self):
        pool = getattr(self, "_pool", None)
        if pool is not None:
            pool.shutdown(wait=False)


class LoRAStreamOffloader:
    """
    H2D-only offloader for training where the base weights are frozen (e.g. LoRA / LoHa / LoKr).

    The classic block swap (``ModelOffloader``) *exchanges* a block: it copies one block back to the
    CPU (Device->Host) while copying the next one in (Host->Device). When the base weights never change,
    the CPU already holds an identical copy, so the D2H half is pure overhead. This offloader keeps a
    permanent (pinned) master copy of every streamed weight on the CPU and only ever copies Host->Device,
    removing the D2H transfer and the CPU-side staging memcpy entirely.

    Layout (N = num_blocks, S = blocks_to_swap, B = ring_size):
      - ``N - S`` blocks stay resident on the GPU permanently.
      - ``S`` "streaming" blocks are placed at even intervals; their ``Linear`` weights live on the CPU
        (pinned) as masters and are streamed into a small ring of ``B`` GPU buffers (default 2: one being
        computed on, one being prefetched). Non-Linear params / buffers of streaming blocks stay resident.

    Mechanics:
      - Because the streamed weights are frozen (no ``.grad`` to preserve), we swap the whole ``module.weight``
        *reference* between its CPU master Parameter and a preallocated GPU ring Parameter, instead of the
        ``.data`` storage trick used by ``ModelOffloader``. No per-step ``cudaMalloc``.
      - What is streamed is decided by ``swap_tensor_selector`` (default: every ``*Linear`` ``weight``). A
        selector may return several tensors per module -- e.g. a quantized weight together with its scale
        buffers -- as ``(module, attribute_name)`` pairs; parameters are rebound as ``nn.Parameter`` ring
        views, registered buffers as plain tensor views.
      - All swap weights of a block are packed into a single contiguous (byte) buffer -- one pinned CPU
        buffer per streaming block and one GPU buffer per ring slot -- and the individual weights are
        dtype/shape views into it. Loading a block is therefore a single H2D ``copy_`` instead of one
        per Linear, which removes the per-tensor launch/driver overhead on PCIe/CPU-bound systems.
      - Eviction is free: repoint ``module.weight`` back to the CPU master (no transfer). A ring slot is
        overwritten only when its block is consumed, so the last ``B`` streamed blocks of a pass stay
        resident. On the reverse pass they are exactly the first blocks needed, so the forward<->backward
        turn-around needs no reload.
      - Direction is implicit in the callback: ``submit_move_blocks_forward`` prefetches the block ``B``
        streaming-slots *ahead* (larger index); the backward hook prefetches ``B`` slots *behind* (smaller
        index). No explicit scheduler / no ``prepare_*_before_backward``.
      - ``wait_for_block`` self-heals: if a needed streaming block is not resident in its slot (cold start,
        or an inference pass boundary), it is loaded on demand.

    Requirement for the model's forward loop: the block's activation input(s) must be passed *positionally*
    (``block(x, ...)``, not ``block(x=x, ...)``). The backward hooks are full backward hooks, which PyTorch fires
    after the block's backward only when a positional input requires grad; otherwise they fire when the gradients of
    the outputs become available, before the block's backward, and the slot would be evicted too early.
    """

    def __init__(
        self,
        block_type: str,
        blocks: list[nn.Module],
        num_blocks: int,
        blocks_to_swap: int,
        supports_backward: bool,
        device: torch.device,
        ring_size: int = 2,
        use_pinned_memory: bool = True,
        debug: bool = False,
        swap_tensor_selector=None,
    ):
        self.block_type = block_type
        self._blocks = blocks
        self.num_blocks = num_blocks  # N
        self.blocks_to_swap = blocks_to_swap  # requested S
        self.device = device
        self.use_pinned_memory = use_pinned_memory
        self.supports_backward = supports_backward
        self.forward_only = not supports_backward
        self.swap_tensor_selector = swap_tensor_selector if swap_tensor_selector is not None else default_swap_tensor_selector

        if not debug:
            debug = os.getenv("MUSUBI_TUNER_OFFLOADER_DEBUG", "0") == "1"
        self.debug = debug
        self.debug_interval = int(os.getenv("MUSUBI_TUNER_OFFLOADER_DEBUG_INTERVAL", "10"))  # steps between prints

        assert device.type == "cuda", "LoRAStreamOffloader currently supports CUDA only"

        # ---- streaming placement: S evenly spaced block indices (midpoint formula -> distinct for S <= N) ----
        stream_idx = sorted({((2 * i + 1) * num_blocks) // (2 * blocks_to_swap) for i in range(blocks_to_swap)})
        self.stream_idx = stream_idx
        self.S = len(stream_idx)  # actual streaming count (>=1; dedup is a no-op unless S is close to N)
        self.rank = {b: k for k, b in enumerate(stream_idx)}  # block_idx -> position in stream_idx
        self.is_stream = [b in self.rank for b in range(num_blocks)]
        self.B = min(ring_size, self.S)  # clamp ring to number of streaming blocks

        # ---- finetuning guard: H2D-only must never lose weight updates ----
        self._job_cache = {}  # block_idx -> [(module, attr_name, is_param) per swap tensor]
        for b in stream_idx:
            for m, name, is_param in self._jobs(b):
                assert not (is_param and getattr(m, name).requires_grad), (
                    "LoRAStreamOffloader requires frozen base weights (LoRA / no full fine-tune). "
                    f"Found a trainable swap tensor {name} in block {b}."
                )

        # ---- transfer engine: owns the copy stream and all H2D primitives ----
        # pinned masters -> direct async H2D; pageable masters -> stage through a worker thread + pinned pool
        # (low "shared GPU memory" footprint, e.g. on Windows). Both expose the same submit/wait interface.
        if use_pinned_memory:
            self.copier = _DirectCopier(device, debug=self.debug)
        else:
            self.copier = _StagedCopier(device, num_staging=self.B, debug=self.debug)

        # ---- runtime state (GPU buffers allocated lazily in prepare_block_devices_before_forward) ----
        self.cpu_master = {}  # block_idx -> [CPU (pinned) Parameter / buffer tensor per swap tensor] (views into cpu_flat)
        self.cpu_flat = {}  # block_idx -> flat (pinned) uint8 CPU tensor backing the masters
        self.ring_param = None  # [slot] -> [GPU nn.Parameter / buffer tensor per swap tensor] (views into ring_flat)
        self.ring_flat = None  # [slot] -> flat uint8 GPU tensor backing the ring params
        self._layout = None  # ([byte offset per swap tensor], total bytes) shared by all streaming blocks
        self.in_slot = [None] * self.B  # slot -> block_idx currently bound to this slot (or None)
        self.free_event = [None] * self.B  # slot -> cuda.Event recording when compute finished using the slot

        self._wait_ctx = "fwd"  # context tag for wait/load timing ("fwd" while in the forward loop, "bwd" in hooks)
        if self.debug:
            self._dbg_reset()

        # ---- backward hooks: only where there is work (prefetch this block, or wait the previous one) ----
        if supports_backward:
            self.remove_handles = []
            for i, block in enumerate(blocks):
                hook = self._create_backward_hook(i)
                if hook is not None:
                    self.remove_handles.append(block.register_full_backward_hook(hook))

        print(
            f"LoRAStreamOffloader[{block_type}]: H2D-only block swap. "
            f"{self.S} streaming / {num_blocks} blocks, ring={self.B}, pinned={use_pinned_memory}. "
            f"streaming indices: {stream_idx}"
        )

    # ------------------------------------------------------------------ helpers

    def _jobs(self, block_idx: int) -> list[tuple[nn.Module, str, bool]]:
        """Swap jobs of a block: ``(module, attr_name, is_param)`` per streamed tensor, cached."""
        cached = self._job_cache.get(block_idx)
        if cached is None:
            cached = []
            for m, name in self.swap_tensor_selector(self._blocks[block_idx]):
                if name in m._parameters:
                    is_param = True
                elif name in m._buffers:
                    is_param = False
                else:
                    raise ValueError(
                        f"swap tensor selector returned {name!r}, which is neither a parameter nor a"
                        f" registered buffer of {m.__class__.__name__} (block {block_idx})"
                    )
                cached.append((m, name, is_param))
            self._job_cache[block_idx] = cached
        return cached

    def _bind(self, block_idx: int, tensors: list[torch.Tensor]):
        # parameters arrive as nn.Parameter, buffers as plain tensors; setattr keeps them in the
        # module's _parameters / _buffers dicts respectively (buffer names are already registered)
        for (m, name, _is_param), t in zip(self._jobs(block_idx), tensors):
            setattr(m, name, t)

    @staticmethod
    def _compute_layout(weights: list[torch.Tensor]) -> tuple[list[int], int]:
        """Byte offsets (aligned so ``view(dtype)`` is valid for any element size) and total size of the flat buffer."""
        align = 256
        offsets = []
        total = 0
        for w in weights:
            total = (total + align - 1) // align * align
            offsets.append(total)
            total += w.numel() * w.element_size()
        return offsets, total

    @staticmethod
    def _flat_views(flat: torch.Tensor, weights: list[torch.Tensor], layout: tuple[list[int], int]) -> list[torch.Tensor]:
        """dtype/shape views into a flat uint8 buffer, one per swap tensor, following ``layout``."""
        offsets, _ = layout
        return [flat[off : off + w.numel() * w.element_size()].view(w.dtype).view(w.shape) for off, w in zip(offsets, weights)]

    def _load(self, rank: int, slot: int, ctx: str = "fwd"):
        """Stream the streaming block at ``rank`` into ring ``slot`` (H2D only), repointing weights."""
        blk = self.stream_idx[rank]

        # already resident (e.g. ring survived a re-prepare): the slot content is still identical to the
        # immutable CPU master, so just rebind the weights to the ring -- no transfer needed
        if self.in_slot[slot] == blk:
            self._bind(blk, self.ring_param[slot])
            return

        # evict the slot's current owner: repoint its weights back to the CPU master (no transfer)
        prev = self.in_slot[slot]
        if prev is not None:
            self._bind(prev, self.cpu_master[prev])

        # single coalesced H2D copy of the whole block, gated behind the slot's free_event (compute done
        # with it). The ring params are views into ring_flat, so the copy bumps their autograd version
        # counter: safe under gradient checkpointing (weights are re-read at recompute time, and a slot is
        # only overwritten after its free_event); without checkpointing autograd raises a version error
        # instead of silently using stale-saved weights -- which the old per-tensor `.data.copy_` hid.
        self.copier.submit(blk, self.ring_flat[slot], self.cpu_flat[blk], self.free_event[slot])

        self._bind(blk, self.ring_param[slot])  # reference swap (no .data trick, no cudaMalloc)
        self.in_slot[slot] = blk

        if self.debug:
            c = self._cur[ctx]
            c["loads"] += 1
            xfer = self.copier.pop_xfer_timing(blk)
            if xfer is not None:
                c["xfer_ev"].append(xfer)

    # ------------------------------------------------------------------ public interface (mirrors ModelOffloader)

    def set_forward_only(self, forward_only: bool):
        self.copier.sync()
        self.forward_only = forward_only

    def __del__(self):
        if getattr(self, "supports_backward", False):
            for handle in getattr(self, "remove_handles", []):
                handle.remove()

    def prepare_block_devices_before_forward(self, blocks: list[nn.Module]):
        if self.S == 0:
            return

        if self.debug:
            print(f"[{self.block_type}] Prepare block devices before forward (H2D-only)")

        first_time = not self.cpu_master  # device placement is only needed on the very first call
        cpu_device = torch.device("cpu")
        for i, b in enumerate(blocks):
            if not self.is_stream[i]:
                if first_time:
                    b.to(self.device)  # resident block: buffers + Linear weights on GPU permanently
                    weighs_to_device(b, self.device)
                continue

            if not first_time:
                # re-prepare (e.g. around in-training sampling): the non-swap parts are already on the device
                # and the CPU masters have stayed on the CPU throughout, so just repoint weights to the masters.
                # IMPORTANT: do NOT call b.to(device) here -- it would drag the CPU masters onto the GPU and
                # they would never be released (every swap block would end up resident).
                self._bind(i, self.cpu_master[i])
                continue

            jobs = self._jobs(i)
            tensors = [getattr(m, name).data for m, name, _ in jobs]
            if self._layout is None:
                self._layout = self._compute_layout(tensors)  # first streaming block defines the shared layout
            flat = torch.empty(self._layout[1], dtype=torch.uint8, device=cpu_device)
            if self.use_pinned_memory:
                flat = flat.pin_memory(device=self.device)
            views = self._flat_views(flat, tensors, self._layout)
            master = []

            if all(t.device.type == "cpu" for t in tensors):
                # CPU-direct: build the flat masters host-to-host, detach the swap tensors from the
                # module so the block's device move cannot drag them, then restore them as the masters.
                # Avoids the H2D+D2H round trip of every streamed tensor that the general path below
                # pays when the model starts on the CPU.
                for (m, name, is_param), t, view in zip(jobs, tensors, views):
                    view.copy_(t)  # host-to-host memcpy into the flat master
                    placeholder = torch.empty(0, dtype=t.dtype, device=cpu_device)
                    if is_param:
                        m._parameters[name].data = placeholder
                    else:
                        m._buffers[name] = placeholder
                b.to(self.device)  # non-swap params/buffers to the device; the placeholders are empty
                for (m, name, is_param), view in zip(jobs, views):
                    if is_param:
                        p = m._parameters[name]
                        p.data = view
                        master.append(p)  # keep the original Parameter object as the persistent master
                    else:
                        m._buffers[name] = view
                        master.append(view)
            else:
                # swap tensors already live on an accelerator: move the whole block to the device
                # (buffers, norms, bias), then pull the swap tensors back into the flat master (one-time D2H)
                b.to(self.device)
                for (m, name, is_param), view in zip(jobs, views):
                    t = getattr(m, name)
                    view.copy_(t.data)
                    if is_param:
                        t.data = view
                        master.append(t)  # keep the original Parameter object as the persistent master
                    else:
                        m._buffers[name] = view
                        master.append(view)

            self.cpu_flat[i] = flat
            self.cpu_master[i] = master

        # validate homogeneous streaming blocks (shared ring template)
        template = self.cpu_master[self.stream_idx[0]]
        for b in self.stream_idx[1:]:
            assert len(self.cpu_master[b]) == len(template), f"block {b} has a different number of swap tensors"
            for p, t in zip(self.cpu_master[b], template):
                assert p.data.shape == t.data.shape and p.data.dtype == t.data.dtype, (
                    f"block {b} swap-tensor shape/dtype differs from the streaming template"
                )

        # preallocate the GPU ring (once); copies happen into these flat buffers, never reallocated
        if self.ring_param is None:
            template_jobs = self._jobs(self.stream_idx[0])
            template_weights = [p.data for p in template]
            self.ring_flat = [torch.empty(self._layout[1], dtype=torch.uint8, device=self.device) for _ in range(self.B)]
            self.ring_param = [
                [
                    nn.Parameter(view, requires_grad=False) if is_param else view
                    for (_m, _name, is_param), view in zip(template_jobs, self._flat_views(flat, template_weights, self._layout))
                ]
                for flat in self.ring_flat
            ]

        # reset transfer state and preload the first B streaming blocks. in_slot deliberately survives:
        # after a completed backward pass the ring already holds ranks 0..B-1, so these loads become
        # no-op rebinds (see _load) instead of redundant cold-start transfers.
        self.free_event = [None] * self.B
        self.copier.reset()
        for k in range(self.B):
            self._load(k, k)

        _synchronize_device(self.device)
        _clean_memory_on_device(self.device)

    def wait_for_block(self, block_idx: int):
        if self.S == 0 or not self.is_stream[block_idx]:
            return
        ctx = self._wait_ctx
        if self.debug and ctx == "fwd" and block_idx == self.stream_idx[0]:
            self._dbg_boundary()  # a forward pass is starting -> finalize the previous step's stats
        j = self.rank[block_idx]
        slot = j % self.B
        if self.in_slot[slot] != block_idx:
            # self-healing: not resident (cold start / inference pass boundary / same-direction restart)
            if self.debug:
                self._cur[ctx]["self_heal"] += 1
            self._load(j, slot, ctx)
        ev = self.copier.wait(block_idx)
        if ev is None:
            return
        if not self.debug:
            torch.cuda.current_stream().wait_event(ev)  # compute waits for H2D (GPU-side, does not block CPU)
            return
        # debug: measure the actual GPU stall (idle time the compute stream spends on wait_event)
        s = torch.cuda.current_stream()
        c = self._cur[ctx]
        c["waits"] += 1
        if not ev.query():
            c["not_ready"] += 1  # H2D not finished when compute reached it -> a real stall (cheap, no sync)
        ev_a = torch.cuda.Event(enable_timing=True)
        ev_b = torch.cuda.Event(enable_timing=True)
        ev_a.record(s)
        s.wait_event(ev)
        ev_b.record(s)
        c["stall_ev"].append((ev_a, ev_b))
        # staged copier records its transfer-timing events on the worker thread, so they only become
        # available once wait() has joined it; the direct copier already popped them in _load (returns None)
        xfer = self.copier.pop_xfer_timing(block_idx)
        if xfer is not None:
            c["xfer_ev"].append(xfer)

    def submit_move_blocks_forward(self, blocks: list[nn.Module], block_idx: int):
        # consumed block_idx in the forward pass: free its slot and prefetch the block B streaming-slots ahead
        if self.S == 0 or not self.is_stream[block_idx]:
            return
        j = self.rank[block_idx]
        self.free_event[j % self.B] = torch.cuda.current_stream().record_event()
        if j + self.B < self.S:
            self._load(j + self.B, (j + self.B) % self.B, "fwd")  # same slot as rank j; evicts rank j
        elif self.forward_only and j == self.S - 1 and self.S > self.B:
            # inference: no backward pass will consume the retained tail, and the next forward pass starts
            # from rank 0 again. Wrap around: all tail slots are free once the last streaming block is
            # consumed, so preload the first B ranks now. The H2D overlaps with the remaining (resident)
            # blocks and the output layers, instead of stalling at the start of the next pass (self-heal).
            for k in range(self.B):
                self._load(k, k, "fwd")

    def _create_backward_hook(self, block_index: int):
        prefetch = self.is_stream[block_index]
        wait_prev = block_index - 1 >= 0 and self.is_stream[block_index - 1]
        if not prefetch and not wait_prev:
            return None

        def backward_hook(module, grad_input, grad_output):
            if not any(g is not None for g in grad_input):
                # No positional input of this block requires grad, so PyTorch fired this hook when the gradients of the
                # *outputs* became available, i.e. before the block's own backward (and its checkpoint recompute) ran.
                # This is the case for the first block that takes part in the backward pass (e.g. when the LoRA is
                # applied only to later blocks): no block runs after it, so there is nothing to free, prefetch or wait
                # for, and evicting it here would break its recompute. (The same happens when a forward loop passes all
                # tensors as keyword arguments, which it must not do; see the model forward loops.)
                return None
            if prefetch:
                # consumed block_index in backward: free its slot and prefetch B streaming-slots behind
                j = self.rank[block_index]
                self.free_event[j % self.B] = torch.cuda.current_stream().record_event()
                if j - self.B >= 0:
                    self._load(j - self.B, (j - self.B) % self.B, "bwd")  # same slot as rank j; evicts rank j
            if wait_prev:
                self._wait_ctx = "bwd"
                self.wait_for_block(block_index - 1)  # ensure the next block to run is ready
                self._wait_ctx = "fwd"
            return None

        return backward_hook

    # ------------------------------------------------------------------ debug timing

    def _dbg_reset(self):
        self._dbg_empty = lambda: {"stall_ev": [], "xfer_ev": [], "waits": 0, "not_ready": 0, "self_heal": 0, "loads": 0}
        self._cur = {"fwd": self._dbg_empty(), "bwd": self._dbg_empty()}
        self._roll = {"fwd": defaultdict(float), "bwd": defaultdict(float)}
        self._step = 0

    def _dbg_boundary(self):
        # finalize the events accumulated for the step that just ended, then print every debug_interval steps
        if not any(self._cur[c]["waits"] or self._cur[c]["loads"] for c in ("fwd", "bwd")):
            return  # very first forward: nothing finished yet
        torch.cuda.synchronize()  # ensure timing events are complete (debug only)
        for ctx in ("fwd", "bwd"):
            c, r = self._cur[ctx], self._roll[ctx]
            r["stall_ms"] += sum(a.elapsed_time(b) for a, b in c["stall_ev"])
            r["xfer_ms"] += sum(a.elapsed_time(b) for a, b in c["xfer_ev"])
            for k in ("waits", "not_ready", "self_heal", "loads"):
                r[k] += c[k]
            self._cur[ctx] = self._dbg_empty()
        self._step += 1
        if self._step % self.debug_interval == 0:
            self._dbg_print()

    def _dbg_print(self):
        n = self.debug_interval
        print(f"[{self.block_type}] H2D-only timing (avg over {n} steps):")
        for ctx in ("fwd", "bwd"):
            r = self._roll[ctx]
            print(
                f"  {ctx}: stall {r['stall_ms'] / n:6.2f} ms/step, H2D {r['xfer_ms'] / n:6.2f} ms/step"
                f" | waits {r['waits'] / n:.1f}, not-ready {r['not_ready'] / n:.1f},"
                f" self-heal {r['self_heal'] / n:.1f}, loads {r['loads'] / n:.1f}"
            )
            self._roll[ctx] = defaultdict(float)


# endregion

# region cpu offload utils


def to_device(x: Any, device: torch.device) -> Any:
    if isinstance(x, torch.Tensor):
        return x.to(device)
    elif isinstance(x, list):
        return [to_device(elem, device) for elem in x]
    elif isinstance(x, tuple):
        return tuple(to_device(elem, device) for elem in x)
    elif isinstance(x, dict):
        return {k: to_device(v, device) for k, v in x.items()}
    else:
        return x


def to_cpu(x: Any) -> Any:
    """
    Recursively moves torch.Tensor objects (and containers thereof) to CPU.

    Args:
        x: A torch.Tensor, or a (possibly nested) list, tuple, or dict containing tensors.

    Returns:
        The same structure as x, with all torch.Tensor objects moved to CPU.
        Non-tensor objects are returned unchanged.
    """
    if isinstance(x, torch.Tensor):
        return x.cpu()
    elif isinstance(x, list):
        return [to_cpu(elem) for elem in x]
    elif isinstance(x, tuple):
        return tuple(to_cpu(elem) for elem in x)
    elif isinstance(x, dict):
        return {k: to_cpu(v) for k, v in x.items()}
    else:
        return x


# Unsloth Offloaded Gradient Checkpointing
# Based on Unsloth Zoo by Daniel Han-Chen & the Unsloth team
try:
    from deepspeed.runtime.activation_checkpointing.checkpointing import detach_variable
except ImportError:

    def detach_variable(inputs, device=None):
        """Detach tensors from computation graph, optionally moving to a device.

        Reimplementation of deepspeed.runtime.activation_checkpointing.checkpointing.detach_variable
        for environments without DeepSpeed installed.
        """
        if isinstance(inputs, tuple):
            out = []
            for inp in inputs:
                if not isinstance(inp, torch.Tensor):
                    out.append(inp)
                    continue
                requires_grad = inp.requires_grad
                if device is not None:
                    x = inp.to(device=device)
                else:
                    x = inp
                x = x.detach()
                x.requires_grad = requires_grad
                out.append(x)
            return tuple(out)
        else:
            raise RuntimeError(
                "Only tuple of tensors is supported. Got Unsupported input type: ",
                type(inputs).__name__,
            )


class UnslothOffloadedGradientCheckpointer(torch.autograd.Function):
    """Saves VRAM by offloading activations to CPU RAM using non-blocking transfers.

    Compared to standard cpu_offload_checkpointing which uses blocking transfers,
    this uses non_blocking=True to hide CPU<->GPU transfer latency behind compute.
    """

    @staticmethod
    @torch.amp.custom_fwd(device_type="cuda")
    def forward(ctx, forward_function, hidden_states, *args):
        # Remember the original device for backward pass (multi-GPU support)
        ctx.input_device = hidden_states.device
        saved_hidden_states = hidden_states.to("cpu", non_blocking=True)
        with torch.no_grad():
            output = forward_function(hidden_states, *args)
        ctx.save_for_backward(saved_hidden_states)
        ctx.forward_function = forward_function
        # NOTE: args stored directly on ctx (not via save_for_backward) because
        # the training loop holds references to these tensors, preventing GC.
        # Using save_for_backward for all args would add complexity for no benefit.
        ctx.args = args
        return output

    @staticmethod
    @torch.amp.custom_bwd(device_type="cuda")
    def backward(ctx, *grads):
        (hidden_states,) = ctx.saved_tensors
        hidden_states = hidden_states.to(ctx.input_device, non_blocking=True).detach()
        hidden_states.requires_grad_(True)
        args = detach_variable(ctx.args)
        inputs = (hidden_states,) + args
        with torch.enable_grad():
            outputs = ctx.forward_function(*inputs)

        output_tensors = []
        grad_tensors = []
        for out, grad in zip(
            outputs if isinstance(outputs, tuple) else (outputs,), grads if isinstance(grads, tuple) else (grads,)
        ):
            if isinstance(out, torch.Tensor) and out.requires_grad:
                output_tensors.append(out)
                grad_tensors.append(grad)
        torch.autograd.backward(output_tensors, grad_tensors)
        return (None,) + tuple(inp.grad if isinstance(inp, torch.Tensor) else None for inp in inputs)


@torch._disable_dynamo
def unsloth_checkpoint(function, *args):
    """Wrapper for UnslothOffloadedGradientCheckpointer."""
    return UnslothOffloadedGradientCheckpointer.apply(function, *args)


def create_cpu_offloading_wrapper(func: Callable, device: torch.device) -> Callable:
    """
    Create a wrapper function that offloads inputs to CPU before calling the original function
    and moves outputs back to the specified device.

    Args:
        func: The original function to wrap.
        device: The device to move outputs back to.

    Returns:
        A wrapped function that offloads inputs to CPU and moves outputs back to the specified device.
    """

    def wrapper(orig_func: Callable) -> Callable:
        def custom_forward(*inputs):
            nonlocal device, orig_func
            cuda_inputs = to_device(inputs, device)
            outputs = orig_func(*cuda_inputs)
            return to_cpu(outputs)

        return custom_forward

    return wrapper(func)


# endregion