| 1 | """Layerwise CPU weight streaming for the Echo 1.5 DiT. |
| 2 | |
| 3 | The manager keeps an immutable CPU master copy of every Transformer block and |
| 4 | only materializes the blocks needed by the current forward on CUDA. Copies use |
| 5 | a dedicated stream so the next block can overlap with current-block compute. |
| 6 | Immutable inference weights never travel from CUDA back to CPU. |
| 7 | """ |
| 8 | |
| 9 | from __future__ import annotations |
| 10 | |
| 11 | from dataclasses import dataclass |
| 12 | from typing import Iterable |
| 13 | |
| 14 | import torch |
| 15 | |
| 16 | |
| 17 | @dataclass(frozen=True) |
| 18 | class LayerwiseOffloadReport: |
| 19 | enabled: bool |
| 20 | block_count: int |
| 21 | resident_blocks: int |
| 22 | prefetch_blocks: int |
| 23 | cpu_weight_bytes: int |
| 24 | pinned_weight_bytes: int |
| 25 | |
| 26 | |
| 27 | @dataclass |
| 28 | class _TensorRecord: |
| 29 | parameter: torch.nn.Parameter |
| 30 | cpu_tensor: torch.Tensor |
| 31 | |
| 32 | |
| 33 | class DiTLayerwiseOffload: |
| 34 | """Stream ``transformer_blocks`` between CPU and one CUDA device. |
| 35 | |
| 36 | The first ``resident_blocks`` stay on CUDA for the denoising stage. All |
| 37 | remaining blocks are prefetched in execution order and released after |
| 38 | their forward hook. This is inference-only: weights must stay immutable. |
| 39 | """ |
| 40 | |
| 41 | def __init__( |
| 42 | self, |
| 43 | generator: torch.nn.Module, |
| 44 | *, |
| 45 | execution_device: torch.device | str, |
| 46 | resident_blocks: int = 0, |
| 47 | prefetch_blocks: int = 1, |
| 48 | pin_memory: bool = True, |
| 49 | ) -> None: |
| 50 | self.generator = generator |
| 51 | self.execution_device = torch.device(execution_device) |
| 52 | if self.execution_device.type != "cuda": |
| 53 | raise ValueError("DiT layerwise offload requires a CUDA execution device") |
| 54 | self.blocks = list(generator.model.velocity_model.transformer_blocks) |
| 55 | if not 0 <= resident_blocks <= len(self.blocks): |
| 56 | raise ValueError( |
| 57 | f"resident_blocks must be between 0 and {len(self.blocks)}, " |
| 58 | f"got {resident_blocks}" |
| 59 | ) |
| 60 | if prefetch_blocks < 1: |
| 61 | raise ValueError("prefetch_blocks must be at least 1") |
| 62 | |
| 63 | self.resident_blocks = resident_blocks |
| 64 | self.prefetch_blocks = prefetch_blocks |
| 65 | self.pin_memory = pin_memory |
| 66 | self._records: list[list[_TensorRecord]] = [] |
| 67 | self._events: dict[int, torch.cuda.Event] = {} |
| 68 | self._loaded: set[int] = set() |
| 69 | self._handles: list[torch.utils.hooks.RemovableHandle] = [] |
| 70 | self._copy_stream: torch.cuda.Stream | None = None |
| 71 | self._gpu_placeholders: dict[torch.dtype, torch.Tensor] = {} |
| 72 | self._active = False |
| 73 | self._cpu_weight_bytes = 0 |
| 74 | self._pinned_weight_bytes = 0 |
| 75 | |
| 76 | self._capture_cpu_weights() |
| 77 | self._register_hooks() |
| 78 | |
| 79 | @property |
| 80 | def report(self) -> LayerwiseOffloadReport: |
| 81 | return LayerwiseOffloadReport( |
| 82 | enabled=True, |
| 83 | block_count=len(self.blocks), |
| 84 | resident_blocks=self.resident_blocks, |
| 85 | prefetch_blocks=self.prefetch_blocks, |
| 86 | cpu_weight_bytes=self._cpu_weight_bytes, |
| 87 | pinned_weight_bytes=self._pinned_weight_bytes, |
| 88 | ) |
| 89 | |
| 90 | @staticmethod |
| 91 | def _unique_parameters(module: torch.nn.Module) -> Iterable[torch.nn.Parameter]: |
| 92 | seen: set[int] = set() |
| 93 | for parameter in module.parameters(): |
| 94 | if id(parameter) not in seen: |
| 95 | seen.add(id(parameter)) |
| 96 | yield parameter |
| 97 | |
| 98 | def _cpu_copy(self, value: torch.Tensor) -> torch.Tensor: |
| 99 | source = value.detach() |
| 100 | if source.device.type != "cpu": |
| 101 | source = source.to("cpu") |
| 102 | result = torch.empty_strided( |
| 103 | source.size(), |
| 104 | source.stride(), |
| 105 | dtype=source.dtype, |
| 106 | device="cpu", |
| 107 | pin_memory=self.pin_memory, |
| 108 | ) |
| 109 | result.copy_(source) |
| 110 | return result |
| 111 | |
| 112 | def _capture_cpu_weights(self) -> None: |
| 113 | """Capture CPU masters, then leave one-element registered placeholders.""" |
| 114 | |
| 115 | if any( |
| 116 | parameter.device.type != "cpu" |
| 117 | for block in self.blocks |
| 118 | for parameter in block.parameters() |
| 119 | ): |
| 120 | raise ValueError("attach DiT layerwise offload while Transformer blocks are on CPU") |
| 121 | |
| 122 | for block in self.blocks: |
| 123 | block_records: list[_TensorRecord] = [] |
| 124 | for parameter in self._unique_parameters(block): |
| 125 | try: |
| 126 | cpu_tensor = self._cpu_copy(parameter) |
| 127 | except RuntimeError as error: |
| 128 | if not self.pin_memory: |
| 129 | raise |
| 130 | # Pinned memory is an acceleration, not a correctness requirement. |
| 131 | if "pin memory" not in str(error).lower() and "cuda" not in str(error).lower(): |
| 132 | raise |
| 133 | source = parameter.detach().cpu() |
| 134 | cpu_tensor = torch.empty_strided( |
| 135 | source.size(), source.stride(), dtype=source.dtype, device="cpu" |
| 136 | ) |
| 137 | cpu_tensor.copy_(source) |
| 138 | size_bytes = cpu_tensor.numel() * cpu_tensor.element_size() |
| 139 | self._cpu_weight_bytes += size_bytes |
| 140 | if cpu_tensor.is_pinned(): |
| 141 | self._pinned_weight_bytes += size_bytes |
| 142 | block_records.append(_TensorRecord(parameter=parameter, cpu_tensor=cpu_tensor)) |
| 143 | parameter.data = torch.empty(1, dtype=parameter.dtype, device="cpu") |
| 144 | self._records.append(block_records) |
| 145 | |
| 146 | def _register_hooks(self) -> None: |
| 147 | for block_index, block in enumerate(self.blocks): |
| 148 | self._handles.append( |
| 149 | block.register_forward_pre_hook( |
| 150 | lambda _module, _args, index=block_index: self._before_block(index) |
| 151 | ) |
| 152 | ) |
| 153 | self._handles.append( |
| 154 | block.register_forward_hook( |
| 155 | lambda _module, _args, output, index=block_index: self._after_block( |
| 156 | index, output |
| 157 | ) |
| 158 | ) |
| 159 | ) |
| 160 | |
| 161 | def _materialize(self, block_index: int) -> None: |
| 162 | if block_index in self._loaded: |
| 163 | return |
| 164 | if self._copy_stream is None: |
| 165 | raise RuntimeError("layerwise offload is not active") |
| 166 | with torch.cuda.device(self.execution_device), torch.cuda.stream(self._copy_stream): |
| 167 | for record in self._records[block_index]: |
| 168 | source = record.cpu_tensor |
| 169 | target = torch.empty_strided( |
| 170 | source.size(), |
| 171 | source.stride(), |
| 172 | dtype=source.dtype, |
| 173 | device=self.execution_device, |
| 174 | ) |
| 175 | target.copy_(source, non_blocking=source.is_pinned()) |
| 176 | record.parameter.data = target |
| 177 | event = torch.cuda.Event() |
| 178 | event.record(self._copy_stream) |
| 179 | self._events[block_index] = event |
| 180 | self._loaded.add(block_index) |
| 181 | |
| 182 | def _wait_for(self, block_index: int) -> None: |
| 183 | event = self._events.get(block_index) |
| 184 | if event is None: |
| 185 | raise RuntimeError(f"Transformer block {block_index} was not prefetched") |
| 186 | current = torch.cuda.current_stream(self.execution_device) |
| 187 | current.wait_event(event) |
| 188 | for record in self._records[block_index]: |
| 189 | record.parameter.data.record_stream(current) |
| 190 | |
| 191 | def _next_streamed(self, block_index: int) -> Iterable[int]: |
| 192 | stop = min(len(self.blocks), block_index + self.prefetch_blocks + 1) |
| 193 | for candidate in range(block_index + 1, stop): |
| 194 | if candidate < self.resident_blocks or candidate in self._loaded: |
| 195 | continue |
| 196 | yield candidate |
| 197 | |
| 198 | def _before_block(self, block_index: int) -> None: |
| 199 | if not self._active: |
| 200 | return |
| 201 | self._materialize(block_index) |
| 202 | self._wait_for(block_index) |
| 203 | for candidate in self._next_streamed(block_index): |
| 204 | self._materialize(candidate) |
| 205 | |
| 206 | def _release(self, block_index: int) -> None: |
| 207 | if block_index not in self._loaded: |
| 208 | return |
| 209 | for record in self._records[block_index]: |
| 210 | placeholder = self._gpu_placeholders.setdefault( |
| 211 | record.parameter.dtype, |
| 212 | torch.empty(1, dtype=record.parameter.dtype, device=self.execution_device), |
| 213 | ) |
| 214 | record.parameter.data = placeholder |
| 215 | self._events.pop(block_index, None) |
| 216 | self._loaded.remove(block_index) |
| 217 | |
| 218 | def _after_block(self, block_index: int, output): |
| 219 | if self._active and block_index >= self.resident_blocks: |
| 220 | self._release(block_index) |
| 221 | return output |
| 222 | |
| 223 | def activate(self) -> None: |
| 224 | if self._active: |
| 225 | return |
| 226 | with torch.cuda.device(self.execution_device): |
| 227 | self._copy_stream = torch.cuda.Stream(device=self.execution_device) |
| 228 | # Only one-element block placeholders move here; non-block weights move normally. |
| 229 | self.generator.to(self.execution_device) |
| 230 | self._active = True |
| 231 | for block_index in range(self.resident_blocks): |
| 232 | self._materialize(block_index) |
| 233 | if self.resident_blocks == 0: |
| 234 | for block_index in range(min(len(self.blocks), self.prefetch_blocks)): |
| 235 | self._materialize(block_index) |
| 236 | for block_index in range(self.resident_blocks): |
| 237 | self._wait_for(block_index) |
| 238 | |
| 239 | def deactivate(self) -> None: |
| 240 | if not self._active: |
| 241 | return |
| 242 | torch.cuda.synchronize(self.execution_device) |
| 243 | for block_index in list(self._loaded): |
| 244 | self._release(block_index) |
| 245 | self._active = False |
| 246 | self._copy_stream = None |
| 247 | self.generator.to("cpu") |
| 248 | self._gpu_placeholders.clear() |
| 249 | |
| 250 | def close(self) -> None: |
| 251 | self.deactivate() |
| 252 | for handle in self._handles: |
| 253 | handle.remove() |
| 254 | self._handles.clear() |
| 255 |