File size: 6,049 Bytes
f2c0505 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | from __future__ import annotations
import torch
class Animate2KVCacheCPUOffloader:
"""Keep Wan-Animate-2 reference K/V caches in CPU memory.
The upstream source stores one immutable pre-RoPE K and V tensor per
transformer layer. `forward_gen` only reads those tensors. This allows us
to preserve one CPU master copy after `forward_ref`, load only the
current layer's K/V to CUDA immediately before its generation block, and
restore the dictionary entry to the CPU master immediately after the block.
This does not quantize or otherwise alter K/V; it is a placement change only.
"""
def __init__(
self, model: torch.nn.Module, device: str | torch.device = "cuda", *, pin_memory: bool = False
):
self.model = model
self.device = torch.device(device)
self.pin_memory = bool(pin_memory)
self.handles = []
self._masters: dict[tuple[int, int, str], torch.Tensor] = {}
self._current_cache_pair: tuple[int, int] | None = None
self.ref_offload_bytes = 0
self.gen_load_bytes = 0
self.gen_loads = 0
self.layers_offloaded = set()
self.peak_layer_bytes = 0
@staticmethod
def _nbytes(t: torch.Tensor) -> int:
return int(t.numel() * t.element_size())
def _cpu_copy(self, t: torch.Tensor) -> torch.Tensor:
if t.device.type == "cpu" and (not self.pin_memory or t.is_pinned()):
return t
# Keeping an entire 40-layer 480p/720p KV cache page-locked can itself
# exhaust host pinned memory. Default to ordinary CPU memory; an
# explicit cpu-pinned mode is available on hosts with ample RAM.
out = torch.empty_like(t, device="cpu", pin_memory=self.pin_memory)
out.copy_(t, non_blocking=False)
return out
def _begin_cache_pair(self, k_cache: dict, v_cache: dict, index: int):
pair = (id(k_cache), id(v_cache))
# Every segment creates fresh dicts. When layer 0 of a new reference
# pass arrives, release the previous segment's pinned masters.
if int(index) == 0 and self._current_cache_pair != pair:
self._masters.clear()
self._current_cache_pair = pair
def _pre(self, module, args, kwargs):
if kwargs.get("method") != "forward_gen" or len(args) < 4:
return
_x, index, k_cache, v_cache = args[:4]
index = int(index)
pair = (id(k_cache), id(v_cache))
for kind, cache in (("k", k_cache), ("v", v_cache)):
key = (pair[0] if kind == "k" else pair[1], index, kind)
master = self._masters.get(key)
if master is None:
# This indicates the caller bypassed the reference pass or the
# cache dictionaries changed unexpectedly. Do not silently
# fall back to GPU residency.
raise RuntimeError(
f"missing CPU {kind.upper()} master for Animate-2 layer {index}; "
"forward_ref must complete before forward_gen"
)
gpu = master.to(self.device, non_blocking=self.pin_memory)
cache[index] = gpu
self.gen_load_bytes += self._nbytes(master)
self.gen_loads += 1
def _post(self, module, args, kwargs, output):
method = kwargs.get("method")
if method not in {"forward_ref", "forward_gen"} or len(args) < 4:
return output
_x, index, k_cache, v_cache = args[:4]
index = int(index)
pair = (id(k_cache), id(v_cache))
if method == "forward_ref":
self._begin_cache_pair(k_cache, v_cache, index)
layer_bytes = 0
for kind, cache in (("k", k_cache), ("v", v_cache)):
t = cache.get(index)
if t is None:
raise RuntimeError(f"Animate-2 reference pass did not populate {kind}_cache[{index}]")
master = self._cpu_copy(t.detach())
key = (pair[0] if kind == "k" else pair[1], index, kind)
self._masters[key] = master
cache[index] = master
nb = self._nbytes(master)
self.ref_offload_bytes += nb
layer_bytes += nb
self.layers_offloaded.add(index)
self.peak_layer_bytes = max(self.peak_layer_bytes, layer_bytes)
else:
# K/V are immutable in forward_gen. Restore the already-existing
# pinned master instead of copying the CUDA tensor back to host.
for kind, cache in (("k", k_cache), ("v", v_cache)):
key = (pair[0] if kind == "k" else pair[1], index, kind)
master = self._masters.get(key)
if master is None:
raise RuntimeError(f"lost CPU {kind.upper()} master for layer {index}")
cache[index] = master
return output
def install(self):
if self.device.type != "cuda":
raise ValueError("Animate2KVCacheCPUOffloader requires a CUDA device")
for block in self.model.blocks:
self.handles.append(block.register_forward_pre_hook(self._pre, with_kwargs=True))
self.handles.append(block.register_forward_hook(self._post, with_kwargs=True))
return self
def remove(self):
for handle in self.handles:
handle.remove()
self.handles.clear()
self._masters.clear()
def report(self) -> dict:
return {
"mode": "pinned_cpu_per_layer_stream" if self.pin_memory else "cpu_per_layer_stream",
"pinned_host_master": self.pin_memory,
"layers_offloaded": len(self.layers_offloaded),
"reference_offload_gib": self.ref_offload_bytes / 2**30,
"generation_h2d_gib": self.gen_load_bytes / 2**30,
"generation_layer_loads": int(self.gen_loads),
"peak_single_layer_cache_gib": self.peak_layer_bytes / 2**30,
"full_cache_gpu_residency": False,
}
|