Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / swap.py: 88%
671 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Swap tensor and swap manager implementation for activation checkpointing"""
16# pylint: disable=W0212
18import functools
19import threading
20import warnings
22from collections import defaultdict
23from contextlib import contextmanager
24from contextvars import ContextVar
25from typing import Any, Dict, Iterable, Iterator, List, Optional, Set
27from hyper_parallel.core.dtensor.dtensor import DTensor
28from . import _backend
30# ---------------------------------------------------------------------------
31# Module-level buffer pools — process-local, no locking needed for single-
32# stream training. Each GPU process owns its own Python interpreter, so
33# these dicts are never shared across processes.
34#
35# _CPU_PINNED_POOL: a list of available pinned CPU tensors per dtype_key.
36# Created via alloc_tensor_buffer(pin_memory=True) on the first miss; the
37# base tensor is returned here after wait_load and reused in the next
38# launch_offload, avoiding repeated cudaHostAlloc / cudaFreeHost calls.
39# ---------------------------------------------------------------------------
40_CPU_PINNED_POOL: Dict[str, List[Any]] = defaultdict(list)
41# Cap each group-swap staging allocation. 32 MiB keeps DMA chunks large
42# while avoiding one huge per-dtype staging tensor in large models.
43_GROUP_SWAP_MAX_BULK_COPY_BYTES = 32 * 1024 * 1024
46def _get_cpu_pinned_buf(dtype_key: str, total_numel: int, dtype):
47 """Pop the smallest sufficient pinned buffer from the pool, or allocate.
49 Best-fit selection minimises wasted pinned memory. When no buffer in the
50 pool is large enough, an undersized entry is discarded before allocating a
51 fresh buffer via alloc_tensor_buffer.
53 Returns the *full* buffer (capacity >= total_numel). Callers must slice
54 ``buf[:total_numel]`` for the actual copy so the returned reference can be
55 passed back to :func:`_return_cpu_pinned_buf` without any platform-specific
56 introspection.
57 """
58 pool = _CPU_PINNED_POOL[dtype_key]
59 best_i = -1
60 for i, buf in enumerate(pool):
61 if buf.numel() >= total_numel:
62 if best_i == -1 or buf.numel() < pool[best_i].numel():
63 best_i = i
64 if best_i != -1:
65 return pool.pop(best_i)
66 # No suitable buffer — discard one stale undersized entry.
67 if pool:
68 pool.pop()
69 return _backend.alloc_tensor_buffer(total_numel, dtype, device='cpu', pin_memory=True)
72def _return_cpu_pinned_buf(buf):
73 """Return a full pinned CPU buffer to the pool for reuse."""
74 if buf is None:
75 return
76 _CPU_PINNED_POOL[str(buf.dtype)].append(buf)
79def _collect_device_storage_ptrs(tensors: Any) -> Set[int]:
80 """Collect device storage pointers from a nested tensor structure."""
81 storage_ptrs = set()
83 def _collect(x):
84 local_tensor = x.to_local() if isinstance(x, DTensor) else x
85 if isinstance(local_tensor, _backend.Tensor) and str(local_tensor.device).lower() != "cpu":
86 storage_ptrs.add(local_tensor.untyped_storage().data_ptr())
87 return x
89 _backend.tree_map(_collect, tensors)
90 return storage_ptrs
93class SwapTensor:
94 """A tensor that can be swapped between device and host memory asynchronously."""
95 STATE_DEVICE = "device"
96 STATE_HOST = "host"
97 STATE_D2H = "d2h"
98 STATE_H2D = "h2d"
99 STATE_NON_TENSOR = "non_tensor"
101 def __init__(self, val: Any, funcname: str, group_swap: bool = False, cpu_pool=None) -> None:
102 self.val = val
103 self.funcname = funcname
104 self._keep_on_device = False
105 self._duplicate_swap = False
106 self._group_managed = False # True when this tensor is handled by SwapGroup bulk copy
107 self.group_swap = group_swap # opt-in for group copy fusion (MUST_SWAP tensors only)
108 self.cpu_pool = cpu_pool
109 self._cpu_pool_buffer = None
110 if isinstance(val, _backend.Tensor) and str(val.device).lower() != 'cpu':
111 self.ver = val._version
112 self._state = self.STATE_DEVICE
113 val_storage = val.untyped_storage()
114 self.storage_size = val_storage.size()
115 self.is_slice_tensor = self.storage_size != val.numel() * _backend.get_element_size(val)
116 self.val_cpu = None
117 else:
118 self.ver = None
119 self._state = self.STATE_NON_TENSOR
120 self.val_cpu = None
121 self.is_slice_tensor = False
122 self.storage_size = 0
124 def dedup_key(self):
125 """Return a stable identity key for duplicate-swap detection."""
126 if self._state == self.STATE_NON_TENSOR:
127 return None
128 val_storage = self.val.untyped_storage()
129 return (
130 str(self.val.device),
131 val_storage.data_ptr(),
132 self.val.storage_offset(),
133 val_storage.size(),
134 tuple(self.val.stride()),
135 )
137 def mark_duplicate_swap(self) -> None:
138 """Mark this wrapper as a duplicate registration in the same swap group."""
139 self._duplicate_swap = True
141 def protect_if_aliases(self, alias_storage_ptrs: Set[int]) -> None:
142 """Keep tensors that alias externally-owned tensors on device."""
143 if self._state == self.STATE_NON_TENSOR:
144 return
145 if self.val.untyped_storage().data_ptr() in alias_storage_ptrs:
146 self._keep_on_device = True
148 def get_val(self) -> Any:
149 """Return the underlying tensor value.
151 Raises RuntimeError if the tensor is not currently in the 'device' state.
152 Non-tensor values are returned directly regardless of state.
153 """
154 if self._state == self.STATE_NON_TENSOR:
155 return self.val
156 if self._state != self.STATE_DEVICE:
157 raise RuntimeError(
158 f"Cannot call get_val(): tensor is in '{self._state}' state. "
159 f"Must be in 'device' state."
160 )
161 return self.val
163 def resize_device_storage(self):
164 """Reallocate device memory on compute stream."""
165 if self._state == self.STATE_NON_TENSOR or self._duplicate_swap:
166 return
168 if self._state != self.STATE_HOST:
169 return
170 storage = self.val.untyped_storage()
171 if storage.size() == self.storage_size:
172 return
173 storage.resize_(self.storage_size)
175 def async_load(self):
176 """async load tensor from host to device"""
177 if self._state == self.STATE_NON_TENSOR or self._keep_on_device or self._duplicate_swap:
178 return
179 if self._group_managed:
180 return
182 if self._state != self.STATE_HOST:
183 warnings.warn(
184 f"[SwapTensor.async_load] Invalid state: current={self._state}, "
185 f"expected 'host'. Operation skipped."
186 )
187 return
189 if self.val_cpu is None:
190 raise ValueError("val_cpu must not be None during async_load")
191 with _backend.preserve_version_counter(self.val):
192 if self.cpu_pool is not None or self.is_slice_tensor:
193 self.val.data.copy_(self.val_cpu, non_blocking=True)
194 else:
195 self.val.untyped_storage().copy_(self.val_cpu.untyped_storage(), non_blocking=True)
196 self._state = self.STATE_H2D
198 def async_group_load(self, source):
199 """Copy a packed device-buffer slice back into the original storage."""
200 if self._state == self.STATE_NON_TENSOR or self._keep_on_device or self._duplicate_swap:
201 return
202 if not self._group_managed:
203 return
204 if self._state != self.STATE_HOST:
205 warnings.warn(
206 f"[SwapTensor.async_group_load] Invalid state: current={self._state}, "
207 f"expected 'host'. Operation skipped."
208 )
209 return
210 if self.val.untyped_storage().size() != self.storage_size:
211 raise RuntimeError(
212 f"Cannot load grouped tensor from {self.funcname}: device storage was not restored. "
213 f"expected size:{self.storage_size}, current size:{self.val.untyped_storage().size()}"
214 )
215 with _backend.preserve_version_counter(self.val):
216 self.val.copy_(source.reshape(self.val.shape), non_blocking=True)
217 self._state = self.STATE_H2D
219 def release_cpu_buffer(self, event=None):
220 """Release an explicitly pooled host buffer exactly once."""
221 if self.cpu_pool is None or self._cpu_pool_buffer is None:
222 return
223 release_tensor = self.val_cpu if self.val_cpu is not None else self._cpu_pool_buffer
224 self.cpu_pool.release(release_tensor, event=event)
225 self._cpu_pool_buffer = None
226 self.val_cpu = None
228 def wait_load(self, release_event=None):
229 """change state to device after async load is done"""
230 if self._state == self.STATE_NON_TENSOR or self._keep_on_device or self._duplicate_swap:
231 return
233 if self._state == self.STATE_DEVICE:
234 return # already loaded
235 if self._state != self.STATE_H2D:
236 warnings.warn(
237 f"[SwapTensor.wait_load] Called in invalid state: {self._state}. "
238 f"Expected 'h2d'. Skipped."
239 )
240 return
241 self._state = self.STATE_DEVICE
242 self.release_cpu_buffer(release_event)
243 if self.cpu_pool is None:
244 self.val_cpu = None
246 def _copy_to_host(self):
247 """Copy the device tensor to its host buffer, releasing the pooled buffer on failure."""
248 try:
249 if self.cpu_pool is not None or self.is_slice_tensor:
250 self.val_cpu.copy_(self.val, non_blocking=True)
251 else:
252 self.val_cpu.untyped_storage().copy_(self.val.untyped_storage(), non_blocking=True)
253 except Exception as exc:
254 if self.cpu_pool is not None and self._cpu_pool_buffer is not None:
255 release_event = _backend.new_event()
256 release_event.record(_backend.get_current_stream())
257 self.release_cpu_buffer(release_event)
258 self.val_cpu = None
259 copy_mode = "tensor" if self.cpu_pool is not None or self.is_slice_tensor else "storage"
260 raise RuntimeError(
261 "Failed to offload activation tensor from device to CPU: "
262 f"source={self.funcname!r}, shape={tuple(self.val.shape)}, dtype={self.val.dtype}, "
263 f"device={self.val.device}, copy_mode={copy_mode}, "
264 f"cpu_pool={'enabled' if self.cpu_pool is not None else 'disabled'}. "
265 f"Original error: {exc}"
266 ) from exc
268 def async_offload(self):
269 """async offload tensor from device to host"""
270 if self._state == self.STATE_NON_TENSOR or self._keep_on_device or self._duplicate_swap:
271 return
272 if self._group_managed:
273 return
275 if self._state != self.STATE_DEVICE:
276 warnings.warn(
277 f"[SwapTensor.async_offload] Invalid state: current={self._state}, "
278 f"expected 'device'. Operation skipped."
279 )
280 return
282 if self.storage_size != self.val.untyped_storage().size():
283 raise RuntimeError(
284 f"There is a tensor from {self.funcname} cannot be SWAPPED! Its storage has been resized "
285 f"presize:{self.storage_size}, current size:{self.val.untyped_storage().size()}"
286 )
287 if self.ver != self.val._version:
288 raise RuntimeError(
289 f"There is a tensor from {self.funcname} cannot be SWAPPED! In-place modification happened "
290 f"preversion:{self.ver}, current version:{self.val._version}"
291 )
293 if self.val_cpu is None:
294 if self.cpu_pool is None:
295 self.val_cpu = _backend.empty_like(
296 self.val, device="cpu", pin_memory=True
297 )
298 else:
299 logical_bytes = self.val.numel() * _backend.get_element_size(self.val)
300 self._cpu_pool_buffer = self.cpu_pool.acquire(logical_bytes)
301 try:
302 self.val_cpu = self._cpu_pool_buffer.view(self.val.dtype).reshape(self.val.shape)
303 except Exception as exc:
304 pool_buffer = self._cpu_pool_buffer
305 self.release_cpu_buffer()
306 raise RuntimeError(
307 "Failed to create a typed CPU-pool view for activation tensor "
308 f"from {self.funcname!r}: shape={tuple(self.val.shape)}, dtype={self.val.dtype}, "
309 f"logical_bytes={logical_bytes}, pool_buffer_shape={tuple(pool_buffer.shape)}, "
310 f"pool_buffer_dtype={pool_buffer.dtype}. Original error: {exc}"
311 ) from exc
313 self._copy_to_host()
314 self._state = self.STATE_D2H
316 def wait_offload(self):
317 """wait offload to host and free device memory"""
318 if self._state == self.STATE_NON_TENSOR or self._keep_on_device or self._duplicate_swap:
319 return
321 if self._state == self.STATE_HOST:
322 return
323 if self._state != self.STATE_D2H:
324 warnings.warn(
325 f"[SwapTensor.wait_offload] Called in invalid state: {self._state}. "
326 f"Expected 'd2h'. Skipped."
327 )
328 return
329 storage = self.val.untyped_storage()
330 if storage.size() != 0:
331 storage.resize_(0)
332 self._state = self.STATE_HOST
334 @property
335 def state(self) -> str:
336 """Return the current swap state of this tensor (device, host, d2h, h2d, or non_tensor)."""
337 return self._state
339 def __repr__(self):
340 if self._state == self.STATE_NON_TENSOR:
341 return f"<SwapTensor state=non_tensor, val_type={type(self.val).__name__}>"
342 return (
343 f"<SwapTensor state={self._state}, duplicate={self._duplicate_swap}, "
344 f"device_val={'exists' if self.val is not None else 'None'}>"
345 )
348class Storage:
349 """Manage a collection of tensors for swapping operations.
351 Supports dict-like access: ``storage[key].append(item)``, ``storage.clear()``,
352 ``for batch in storage.values(): ...``.
353 """
355 def __init__(self):
356 self._data: Dict[Any, List[Any]] = defaultdict(list)
358 def __getitem__(self, key: Any) -> List[Any]:
359 return self._data[key]
361 def values(self):
362 """Return an iterable view of all stored lists."""
363 return self._data.values()
365 def clear(self):
366 """Remove all entries from the storage."""
367 self._data.clear()
369 def iter_swap_tensors(self):
370 """Iterate all SwapTensor objects stored in this storage."""
371 collected = []
373 def _collect(x):
374 if isinstance(x, SwapTensor):
375 collected.append(x)
376 return x
378 for storage_list in self.values():
379 for item in storage_list:
380 _backend.tree_map(_collect, item)
381 return collected
383 def mark_duplicate_swaps(self, seen_keys) -> int:
384 """Mark tensors already registered in the same swap group as duplicates."""
385 duplicate_count = 0
386 for swap_tensor in self.iter_swap_tensors():
387 dedup_key = swap_tensor.dedup_key()
388 if dedup_key is None:
389 continue
390 if dedup_key in seen_keys:
391 swap_tensor.mark_duplicate_swap()
392 duplicate_count += 1
393 continue
394 seen_keys.add(dedup_key)
395 return duplicate_count
397 def protect_alias_storage_ptrs(self, alias_storage_ptrs: Set[int]):
398 """Avoid offloading swap entries that alias externally-owned storage."""
399 if not alias_storage_ptrs:
400 return
402 def _protect_tensor(x):
403 if isinstance(x, SwapTensor):
404 x.protect_if_aliases(alias_storage_ptrs)
405 return x
407 for storage_list in self.values():
408 for item in storage_list:
409 _backend.tree_map(_protect_tensor, item)
411 def launch_load(self):
412 """launch async load for all tensors in swap storage"""
413 def _async_load(x):
414 if isinstance(x, SwapTensor):
415 x.async_load()
416 return x
418 for storage_list in self.values():
419 for item in storage_list:
420 _backend.tree_map(_async_load, item)
422 def resize_device_storage(self):
423 """Resize device storage for all swap tensors (runs on compute stream)."""
424 def _resize(x):
425 if isinstance(x, SwapTensor):
426 x.resize_device_storage()
427 return x
428 for storage_list in self.values():
429 for item in storage_list:
430 _backend.tree_map(_resize, item)
432 def wait_load(self, release_event=None):
433 """wait load for all tensors in swap storage"""
434 def _wait_load(x):
435 if isinstance(x, SwapTensor):
436 x.wait_load(release_event=release_event)
437 return x
439 for storage_list in self.values():
440 for item in storage_list:
441 _backend.tree_map(_wait_load, item)
442 self.clear()
444 def release_cpu_buffers(self, event=None):
445 """Release all explicitly pooled host buffers held by this storage."""
446 def _release(x):
447 if isinstance(x, SwapTensor):
448 x.release_cpu_buffer(event=event)
449 return x
451 for storage_list in self.values():
452 for item in storage_list:
453 _backend.tree_map(_release, item)
455 def wait_offload(self):
456 """wait offload for all tensors in swap storage"""
457 def _wait_offload(x):
458 if isinstance(x, SwapTensor):
459 x.wait_offload()
460 return x
462 for storage_list in self.values():
463 for item in storage_list:
464 _backend.tree_map(_wait_offload, item)
466 def launch_offload(self):
467 """launch async offload for all tensors in swap storage"""
468 def _async_offload(x):
470 if isinstance(x, SwapTensor):
471 x.async_offload()
472 return x
474 for storage_list in self.values():
475 for item in storage_list:
476 _backend.tree_map(_async_offload, item)
479class SwapGroup:
480 """Manager for a group of storages to coordinate swap operations.
482 Non-slice tensors within the group are packed into bounded contiguous device
483 buffers before D2H transfer, and loaded back from bounded H2D buffers.
484 Each packed slice is copied back into its tensor's original device storage so
485 autograd views that alias that storage remain valid during backward.
487 Slice tensors (storage larger than logical data) fall back to the original
488 per-tensor copy path.
489 """
491 def __init__(self, group_name: str):
492 self.group_name = group_name
493 self.is_last_group: bool = False
494 self._storages: List[Storage] = []
495 self._load_event: Optional[Any] = None
496 self._offload_event: Optional[Any] = None
497 # Group-level contiguous buffers for non-slice tensors.
498 self._packed_tensor_info: List = [] # [(SwapTensor, bucket_key, element_offset), ...]
499 self._packed_buckets: Dict[str, Dict[str, Any]] = {}
500 self._group_cpu_buf = None # pinned CPU bufs; live offload→load
501 self._group_device_buf = None # temp device bufs; cleared after each phase
502 # Persistent dedup set accumulated across add() calls; avoids O(N²) rebuild.
503 # mark_duplicate_swaps mutates it in-place, so new keys are added automatically.
504 # Reset at wait_load() so stale data_ptrs don't leak into the next iteration.
505 self._seen_dedup_keys: set = set()
506 # Per-bucket SwapTensor lists built in _collect_packable_tensors and consumed
507 # in launch_offload, eliminating a redundant pass over _packed_tensor_info.
508 self._packed_by_bucket: Dict[str, List] = {}
510 def add(self, storage):
511 """Add a storage to the swap group."""
512 duplicate_count = storage.mark_duplicate_swaps(self._seen_dedup_keys)
513 if duplicate_count > 0:
514 warnings.warn(
515 f"SwapGroup '{self.group_name}' skipped {duplicate_count} duplicate tensor swap registration(s)."
516 )
517 self._storages.append(storage)
519 def protect_alias_tensors(self, tensors: Any):
520 """Protect externally-owned tensors from premature offload."""
521 alias_storage_ptrs = _collect_device_storage_ptrs(tensors)
522 if not alias_storage_ptrs:
523 return
524 for storage in self._storages:
525 storage.protect_alias_storage_ptrs(alias_storage_ptrs)
527 @staticmethod
528 def _bucket_tensor(candidate_buckets, dtype_key, x):
529 """Append ``x`` to its dtype bucket, opening a new bucket when full."""
530 dtype_buckets = candidate_buckets.setdefault(dtype_key, [])
531 if (not dtype_buckets or
532 dtype_buckets[-1]["total_bytes"] + x.storage_size > _GROUP_SWAP_MAX_BULK_COPY_BYTES):
533 dtype_buckets.append({
534 "bucket_key": f"{dtype_key}#{len(dtype_buckets)}",
535 "dtype": x.val.dtype,
536 "dtype_key": str(x.val.dtype),
537 "device": x.val.device,
538 "tensors": [],
539 "total_bytes": 0,
540 "total_numel": 0,
541 "cpu_pool": x.cpu_pool,
542 })
543 bucket = dtype_buckets[-1]
544 bucket["tensors"].append(x)
545 bucket["total_bytes"] += x.storage_size
546 bucket["total_numel"] += x.val.numel()
547 return x
549 @staticmethod
550 def _finalize_packed_buckets(candidate_buckets, packed_info, packed_buckets, packed_by_bucket) -> int:
551 """Turn multi-tensor candidate buckets into packed buckets; return owned bytes."""
552 total_bytes = 0
553 for dtype_bucket_list in candidate_buckets.values():
554 for candidate_bucket in dtype_bucket_list:
555 tensors = candidate_bucket["tensors"]
556 if len(tensors) < 2:
557 continue
558 bucket_key = candidate_bucket["bucket_key"]
559 packed_buckets[bucket_key] = {
560 "dtype": candidate_bucket["dtype"],
561 "dtype_key": candidate_bucket["dtype_key"],
562 "device": candidate_bucket["device"],
563 "total_numel": candidate_bucket["total_numel"],
564 "total_bytes": candidate_bucket["total_bytes"],
565 "cpu_pool": candidate_bucket["cpu_pool"],
566 }
567 element_offset = 0
568 for tensor in tensors:
569 tensor._group_managed = True
570 tensor._state = SwapTensor.STATE_D2H
571 packed_info.append((tensor, bucket_key, element_offset))
572 element_offset += tensor.val.numel()
573 packed_by_bucket[bucket_key] = tensors
574 total_bytes += candidate_bucket["total_bytes"]
575 return total_bytes
577 def _collect_packable_tensors(self) -> int:
578 """Identify tensors eligible for group packing and mark them for bulk copy.
580 A tensor is eligible only when it is contiguous, not a slice tensor,
581 not a duplicate, not sharing storage with another live swap tensor, and
582 has ``group_swap=True``. Dtype buckets are split before their staging
583 allocation would exceed ``_GROUP_SWAP_MAX_BULK_COPY_BYTES``. A packed
584 bucket with fewer than two tensors is left on the original per-tensor
585 path because it has no batch-copy benefit. Non-contiguous
586 tensors are excluded because the packing step copies storage-order
587 bytes while restore uses the original stride; those tensors fall back to
588 the per-tensor copy path.
589 Shared-storage tensors also fall back together because group packing
590 frees the original storage after packing, which would invalidate any
591 non-packed aliases such as transpose views before their own offload.
593 Side effects: marks each eligible tensor with ``_group_managed=True``
594 and ``_state=STATE_D2H``, and populates ``_packed_tensor_info`` /
595 ``_packed_buckets``.
597 Returns:
598 Total byte count of all packable tensors.
599 """
600 candidate_buckets: Dict[str, List[Dict[str, Any]]] = {}
601 packed_info: List = []
602 packed_buckets: Dict[str, Dict[str, Any]] = {}
603 packed_by_bucket: Dict[str, List] = {}
605 def _try_pack(x):
606 if not isinstance(x, SwapTensor):
607 return x
608 no_pack = (not x.group_swap or x._state != SwapTensor.STATE_DEVICE or x._keep_on_device
609 or x.is_slice_tensor or x._duplicate_swap or x.storage_size >= _GROUP_SWAP_MAX_BULK_COPY_BYTES
610 or not x.val.is_contiguous())
611 if no_pack:
612 return x
613 if x.storage_size != x.val.untyped_storage().size():
614 raise RuntimeError(
615 f"There is a tensor from {x.funcname} cannot be SWAPPED! Its storage has been resized "
616 f"presize:{x.storage_size}, current size:{x.val.untyped_storage().size()}"
617 )
618 if x.ver != x.val._version:
619 raise RuntimeError(
620 f"There is a tensor from {x.funcname} cannot be SWAPPED! In-place modification happened "
621 f"preversion:{x.ver}, current version:{x.val._version}"
622 )
623 dtype_key = (str(x.val.dtype), id(x.cpu_pool))
624 return self._bucket_tensor(candidate_buckets, dtype_key, x)
626 for storage in self._storages:
627 for storage_list in storage.values():
628 for item in storage_list:
629 _backend.tree_map(_try_pack, item)
631 total_bytes = self._finalize_packed_buckets(
632 candidate_buckets, packed_info, packed_buckets, packed_by_bucket
633 )
635 self._packed_tensor_info = packed_info
636 self._packed_buckets = packed_buckets
637 self._packed_by_bucket = packed_by_bucket
638 return total_bytes
640 def _pack_device_buffers(self):
641 """Concatenate each packed bucket's tensors into one contiguous device buffer."""
642 group_device_bufs = {}
643 for bucket_key, swap_tensors in self._packed_by_bucket.items():
644 group_device_bufs[bucket_key] = _backend.cat(
645 [st.val.reshape(-1) for st in swap_tensors], dim=0
646 )
647 return group_device_bufs
649 @staticmethod
650 def _protect_device_bufs(group_device_bufs, copy_stream):
651 """Keep staging device buffers reserved for the copy stream.
653 A staging buffer is otherwise referenced only by a local, so a failure
654 part-way through the copy loop can drop the last reference while the
655 copy stream is still reading it -- and the caching allocator may then
656 hand that memory to the compute stream.
657 """
658 for buf in group_device_bufs.values():
659 device = getattr(buf, "device", None)
660 if device is not None and device.type not in ("cpu", "meta"):
661 buf.record_stream(copy_stream)
663 def _acquire_bucket_cpu_buf(self, bucket_key, bucket):
664 """Acquire the pinned CPU buffer that receives one bucket's D2H copy."""
665 numel = bucket["total_numel"]
666 if bucket["cpu_pool"] is None:
667 return _get_cpu_pinned_buf(bucket["dtype_key"], numel, bucket["dtype"])
668 raw_buf = bucket["cpu_pool"].acquire(bucket["total_bytes"])
669 try:
670 return raw_buf.view(bucket["dtype"])
671 except Exception as exc:
672 bucket["cpu_pool"].release(raw_buf)
673 raise RuntimeError(
674 "Failed to create a typed CPU-pool view for packed activation bucket: "
675 f"group={self.group_name!r}, bucket={bucket_key!r}, "
676 f"requested_dtype={bucket['dtype']}, total_numel={numel}, "
677 f"total_bytes={bucket['total_bytes']}, raw_buffer_shape={tuple(raw_buf.shape)}, "
678 f"raw_buffer_dtype={raw_buf.dtype}. Original error: {exc}"
679 ) from exc
681 def _release_offloaded_bucket_bufs(self, group_cpu_bufs, copy_stream):
682 """Give back every buffer already filled before a failed D2H loop."""
683 release_event = _backend.new_event()
684 release_event.record(copy_stream)
685 for bucket_key, cpu_buf in group_cpu_bufs.items():
686 cpu_pool = self._packed_buckets[bucket_key]["cpu_pool"]
687 if cpu_pool is not None:
688 cpu_pool.release(cpu_buf, event=release_event)
689 else:
690 _return_cpu_pinned_buf(cpu_buf)
692 def _offload_buckets_d2h(self, group_device_bufs, copy_stream):
693 """One-shot D2H per packed bucket."""
694 group_cpu_bufs = {}
695 active_bucket_key = None
696 try:
697 for bucket_key, bucket in self._packed_buckets.items():
698 active_bucket_key = bucket_key
699 cpu_buf = self._acquire_bucket_cpu_buf(bucket_key, bucket)
700 group_cpu_bufs[bucket_key] = cpu_buf
701 cpu_buf[:bucket["total_numel"]].copy_(
702 group_device_bufs[bucket_key], non_blocking=True
703 )
704 except Exception as exc:
705 self._release_offloaded_bucket_bufs(group_cpu_bufs, copy_stream)
706 failed_bucket = self._packed_buckets.get(active_bucket_key, {})
707 raise RuntimeError(
708 "Failed to offload packed activation bucket from device to CPU: "
709 f"group={self.group_name!r}, bucket={active_bucket_key!r}, "
710 f"dtype={failed_bucket.get('dtype', 'unknown')}, "
711 f"total_numel={failed_bucket.get('total_numel', 'unknown')}, "
712 f"total_bytes={failed_bucket.get('total_bytes', 'unknown')}, "
713 f"cpu_pool={'enabled' if failed_bucket.get('cpu_pool') is not None else 'disabled'}. "
714 f"Original error: {exc}"
715 ) from exc
716 return group_cpu_bufs
718 def launch_offload(self, copy_stream):
719 """Launch async offload for all storages in the group.
721 Non-slice tensors are first packed into bounded contiguous device
722 buffers, then transferred to pinned CPU memory. Slice tensors are
723 offloaded individually via the existing per-tensor path.
724 """
725 total_bytes = self._collect_packable_tensors()
726 with _backend.no_grad():
727 group_device_bufs = self._pack_device_buffers() if total_bytes > 0 else {}
729 compute_event = _backend.new_event()
730 compute_event.record(_backend.get_current_stream())
731 self._offload_event = _backend.new_event()
732 stream_context = _backend.get_stream_context()
733 with _backend.no_grad(), stream_context(copy_stream):
734 compute_event.wait(copy_stream)
736 if total_bytes > 0:
737 # Publish and protect the staging buffers before the first async
738 # D2H: if the copy loop fails part-way, the stack unwinding must
739 # not drop the last reference to memory the copy stream is still
740 # reading.
741 self._group_device_buf = group_device_bufs
742 self._protect_device_bufs(group_device_bufs, copy_stream)
743 self._group_cpu_buf = self._offload_buckets_d2h(group_device_bufs, copy_stream)
745 # Slice tensors use the existing per-tensor path.
746 # Group-managed tensors are already STATE_D2H so async_offload is a no-op.
747 for storage in self._storages:
748 storage.launch_offload()
749 self._offload_event.record(copy_stream)
751 def wait_offload(self):
752 """Wait for offload to complete for all storages in the group."""
753 if self._offload_event is None:
754 raise RuntimeError(
755 f"SwapGroup '{self.group_name}' wait_offload() called before launch_offload()."
756 )
757 compute_stream = _backend.get_current_stream()
758 stream_context = _backend.get_stream_context()
759 with _backend.no_grad(), stream_context(compute_stream):
760 self._offload_event.wait(compute_stream)
761 self._offload_event = None
762 for storage in self._storages:
763 storage.wait_offload()
764 # Release the temporary device packing buffer; _group_cpu_buf persists until launch_load.
765 self._group_device_buf = None
767 def launch_load(self, copy_stream):
768 """Prepare storage and launch async load for all storages in the group.
770 Non-slice tensors are loaded from pinned CPU memory into bounded
771 contiguous device buffers, then copied device-to-device into their
772 original storages. Slice tensors use the existing per-tensor path.
773 """
774 # Restore original storages before scheduling copies. Keeping the same
775 # storage object is required for autograd-saved views of packed tensors.
776 with _backend.no_grad():
777 for storage in self._storages:
778 storage.resize_device_storage()
780 compute_event = _backend.new_event()
781 compute_event.record(_backend.get_current_stream())
782 self._load_event = _backend.new_event()
783 stream_context = _backend.get_stream_context()
784 with _backend.no_grad(), stream_context(copy_stream):
785 compute_event.wait(copy_stream)
787 if self._packed_tensor_info and self._group_cpu_buf is not None:
788 group_device_bufs = {}
789 for bucket_key, bucket in self._packed_buckets.items():
790 cpu_buf = self._group_cpu_buf.get(bucket_key)
791 if cpu_buf is None:
792 continue
793 numel = bucket["total_numel"]
794 group_device_bufs[bucket_key] = _backend.alloc_tensor_buffer(
795 numel, bucket["dtype"], bucket["device"]
796 )
797 # Publish and protect the staging buffers before the first async
798 # H2D, for the same reason as launch_offload.
799 self._group_device_buf = group_device_bufs
800 self._protect_device_bufs(group_device_bufs, copy_stream)
802 for bucket_key, bucket in self._packed_buckets.items():
803 group_device_buf = group_device_bufs.get(bucket_key)
804 if group_device_buf is None:
805 continue
806 # One-shot H2D per packed bucket.
807 group_device_buf.copy_(
808 self._group_cpu_buf[bucket_key][:bucket["total_numel"]],
809 non_blocking=True,
810 )
812 # Unpack with D2D copies into the original storages. Rebinding
813 # st.val with set_() would leave existing aliases on freed storage.
814 for st, bucket_key, element_offset in self._packed_tensor_info:
815 group_device_buf = group_device_bufs.get(bucket_key)
816 if group_device_buf is None:
817 continue
818 source = group_device_buf[element_offset:element_offset + st.val.numel()]
819 st.async_group_load(source)
821 # Slice tensors use the existing per-tensor path.
822 # Group-managed tensors skip async_load via _group_managed flag.
823 for storage in self._storages:
824 storage.launch_load() # Only copy, no resize
825 self._load_event.record(copy_stream)
827 def release_cpu_buffers(self, event=None):
828 """Release staging buffers immediately or defer until ``event`` completes."""
829 if self._group_cpu_buf is None:
830 return
831 for bucket_key, buf in self._group_cpu_buf.items():
832 bucket = self._packed_buckets.get(bucket_key)
833 if bucket is not None and bucket["cpu_pool"] is not None:
834 bucket["cpu_pool"].release(buf, event=event)
835 else:
836 _return_cpu_pinned_buf(buf)
837 self._group_cpu_buf = None
839 def wait_load(self):
840 """Wait for grouped H2D and D2D loads to complete."""
841 if self._load_event is None:
842 raise RuntimeError(
843 f"SwapGroup '{self.group_name}' wait_load() called before launch_load()."
844 )
845 compute_stream = _backend.get_current_stream()
846 load_event = self._load_event
847 stream_context = _backend.get_stream_context()
848 with _backend.no_grad(), stream_context(compute_stream):
849 load_event.wait(compute_stream)
850 for storage in self._storages:
851 storage.wait_load(release_event=load_event)
852 self._load_event = None
853 self._storages.clear()
854 # Keep explicit-pool buffers pending until the copy event completes;
855 # legacy buffers retain their existing immediate-reuse behavior.
856 self.release_cpu_buffers(event=load_event)
857 # Device buffer: the pool holds the staging reference; just drop
858 # the local reference. Tensors aliasing _group_device_buf's
859 # storage keep it alive via their own storage references until
860 # they are consumed in backward.
861 self._group_device_buf = None
862 self._packed_tensor_info = []
863 self._packed_buckets = {}
864 self._packed_by_bucket = {}
865 self._seen_dedup_keys = set()
868class SwapManager:
869 """Singleton manager for swap groups and their operations."""
870 _instance: Optional["SwapManager"] = None
871 _lock = threading.Lock()
872 _FORWARD_PREFETCH_HOOK_HANDLE_ATTRS = (
873 "_swap_forward_pre_hook_handle",
874 "_swap_forward_hook_handle",
875 "_swap_backward_pre_hook_handle",
876 "_swap_backward_hook_handle",
877 )
879 def __init__(self) -> None:
880 """Initialize process-local swap groups once for the singleton."""
881 if hasattr(self, '_groups'):
882 return
883 self._groups: Dict[str, SwapGroup] = {}
884 self._current_group_name: ContextVar[str] = ContextVar(
885 "swap_current_group_name", default=""
886 )
887 self._layer_count: int = 0
888 self._copy_stream: Optional[Any] = None
889 self._missing_group_warned: bool = False
891 def __new__(cls):
892 if cls._instance is None:
893 with cls._lock:
894 if cls._instance is None:
895 cls._instance = super().__new__(cls)
896 return cls._instance
898 def add_storage(self, group_name: str, storage: Storage) -> None:
899 """Add a storage to a specified swap group."""
900 self.ensure_group(group_name)
901 self._groups[group_name].add(storage)
903 def ensure_group(self, group_name: str) -> None:
904 """Create the swap group if it does not exist yet."""
905 if group_name not in self._groups:
906 self._groups[group_name] = SwapGroup(group_name)
908 def launch_offload(self, group_name: str, copy_stream=None):
909 """Launch async offload for a specified swap group."""
910 group = self._groups.get(group_name)
911 if group is None:
912 raise RuntimeError(f"Group {group_name} does not exist.")
913 if copy_stream is None:
914 copy_stream = self._get_copy_stream()
915 group.launch_offload(copy_stream)
917 def protect_alias_tensors(self, group_name: str, tensors: Any):
918 """Keep tensors that alias externally-owned tensors on device."""
919 group = self._groups.get(group_name)
920 if group is None:
921 raise RuntimeError(f"Group {group_name} does not exist.")
922 group.protect_alias_tensors(tensors)
924 def wait_offload(self, group_name: str):
925 """Wait for offload to complete for a specified swap group."""
926 group = self._groups.get(group_name)
927 if group is None:
928 raise RuntimeError(f"Group {group_name} does not exist.")
929 group.wait_offload()
931 def launch_load(self, group_name: str, copy_stream=None):
932 """Launch async load for a specified swap group."""
933 group = self._groups.get(group_name)
934 if group is None:
935 raise RuntimeError(f"Group {group_name} does not exist.")
936 if copy_stream is None:
937 copy_stream = self._get_copy_stream()
938 group.launch_load(copy_stream)
940 def wait_load(self, group_name: str):
941 """Wait for load to complete for a specified swap group."""
942 group = self._groups.get(group_name)
943 if group is None:
944 raise RuntimeError(f"Group {group_name} does not exist.")
945 group.wait_load()
947 def release_group_storage(self, group_name: str) -> None:
948 """Release storage references held by the swap group."""
949 group = self._groups.get(group_name)
950 if group is not None:
951 group._storages.clear()
953 def abort_group(self, group_name: str) -> None:
954 """Synchronize in-flight transfers and remove a failed run's group."""
955 group = self._groups.pop(group_name, None)
956 if group is None:
957 return
958 for event in (group._offload_event, group._load_event):
959 if event is not None:
960 event.synchronize()
961 for storage in group._storages:
962 storage.release_cpu_buffers()
963 group.release_cpu_buffers()
964 group._storages.clear()
966 def get_current_group_name(self) -> str:
967 """Return the name of the currently active swap group."""
968 return self._current_group_name.get()
970 def set_current_group_name(self, group_name: str) -> None:
971 """Set the name of the currently active swap group."""
972 self._current_group_name.set(group_name)
974 def warn_missing_group_once(self) -> None:
975 """Warn once per process when ``MUST_SWAP`` runs outside any swap group.
977 Activation swap only serves pipeline chunks the schedule marked as
978 swappable; every other region keeps its activations on device. A single
979 warning keeps that fallback discoverable without repeating for every
980 dispatched operator.
981 """
982 if self._missing_group_warned:
983 return
984 self._missing_group_warned = True
985 warnings.warn(
986 "Activation swap requested MUST_SWAP outside an active swap group; "
987 "the activations were kept on device instead of being offloaded.",
988 UserWarning,
989 stacklevel=2,
990 )
992 def active_group_count(self) -> int:
993 """Return the number of live swap groups for lifecycle diagnostics."""
994 return len(self._groups)
996 @contextmanager
997 def group_context(self, group_name: str) -> Iterator[None]:
998 """Activate a swap group within the current execution context."""
999 token = self._current_group_name.set(group_name)
1000 try:
1001 yield
1002 finally:
1003 self._current_group_name.reset(token)
1005 def is_last_group(self, group_name: Optional[str] = None) -> bool:
1006 """Return whether the specified swap group is the terminal group in the chain."""
1007 group_name = self.get_current_group_name() if group_name is None else group_name
1008 group = self._groups.get(group_name)
1009 if group is None:
1010 return False
1011 return group.is_last_group
1013 def unregister_forward_prefetch_hooks(self, module: Any) -> int:
1014 """Remove hooks installed by :meth:`set_forward_prefetch_layer`.
1016 Args:
1017 module: Module whose layer-level swap hooks should be removed.
1019 Returns:
1020 Number of removed hook handles.
1021 """
1022 removed_count = 0
1023 for attr_name in self._FORWARD_PREFETCH_HOOK_HANDLE_ATTRS:
1024 if not hasattr(module, attr_name):
1025 continue
1026 handle = getattr(module, attr_name)
1027 if handle is not None:
1028 handle.remove()
1029 delattr(module, attr_name)
1030 removed_count += 1
1031 return removed_count
1033 def unregister_forward_prefetch_layer(self, module: Any) -> int:
1034 """Tear down prefetch wiring registered on ``module`` and its group.
1036 Removes the four forward/backward swap hooks installed on ``module``
1037 by :meth:`set_forward_prefetch_layer` and drops its swap group from the
1038 process-wide singleton. Mirrors the reverse of :meth:`set_forward_prefetch_layer`
1039 so models destroyed in a long-lived process do not leave stale ``SwapGroup``
1040 entries or hook handles behind (see the swap-inputs path in
1041 ``distributed/activation_checkpoint.py``, which installs these hooks but
1042 previously had no teardown).
1044 Args:
1045 module: A layer that participated in a prefetch chain.
1047 Returns:
1048 Number of removed hook handles.
1049 """
1050 removed_count = self.unregister_forward_prefetch_hooks(module)
1051 group_name = getattr(module, "_swap_group_name", None)
1052 if group_name is not None:
1053 self.abort_group(group_name)
1054 delattr(module, "_swap_group_name")
1055 if hasattr(module, "_swap_group_order"):
1056 delattr(module, "_swap_group_order")
1057 # Also drop the swap state so a module re-registered through
1058 # set_forward_prefetch_layer does not inherit a stale "pre_backward"
1059 # flag that would short-circuit the new forward pre-hook.
1060 if hasattr(module, "_swap_state"):
1061 delattr(module, "_swap_state")
1062 return removed_count
1064 def set_forward_prefetch_layer(self, first_layer, second_layer):
1065 """
1066 Configure prefetching and offloading order between two consecutive layers.
1068 Usage:
1069 for i in range(len(model.layers) - 1):
1070 set_forward_prefetch_layer(model.layers[i], model.layers[i + 1])
1072 Ensures idempotency: safe to call multiple times on the same layer pair.
1073 """
1074 if first_layer is second_layer:
1075 warnings.warn(
1076 "set_forward_prefetch_layer: "
1077 "Prefetching between identical layers has no effect.",
1078 UserWarning,
1079 stacklevel=2,
1080 )
1082 def _ensure_group_name(module):
1083 """Assign a unique swap group name to the module if not already assigned."""
1084 if not hasattr(module, "_swap_group_name"):
1085 name = f"swap_group_{self._layer_count}"
1086 self._layer_count += 1
1087 module._swap_group_name = name
1088 module._swap_group_order = {"prev": None, "next": None}
1089 return module._swap_group_name
1090 first_name = _ensure_group_name(first_layer)
1091 second_name = _ensure_group_name(second_layer)
1093 if first_name not in self._groups:
1094 self._groups[first_name] = SwapGroup(first_name)
1095 if second_name not in self._groups:
1096 self._groups[second_name] = SwapGroup(second_name)
1098 if first_layer._swap_group_order["next"] is None:
1099 first_layer._swap_group_order["next"] = second_name
1100 if second_layer._swap_group_order["prev"] is None:
1101 second_layer._swap_group_order["prev"] = first_name
1103 self._groups[first_name].is_last_group = first_layer._swap_group_order["next"] is None
1104 self._groups[second_name].is_last_group = second_layer._swap_group_order["next"] is None
1106 def _forward_pre_hook(group_name, module, _): # pylint: disable=W0613
1107 if getattr(module, "_swap_state", None) == "pre_backward":
1108 return
1109 SwapManager().set_current_group_name(group_name)
1111 def _forward_hook(group_name, module, args, output): # pylint: disable=W0613
1112 """
1113 Forward post-hook executed immediately after forward computation
1114 of the current layer finishes.
1116 Execution timeline (example with 3 layers, forward order: L0 → L1 → L2):
1118 Time →
1119 Forward Compute Stream:
1120 | Fwd L0 | post(L0) | Fwd L1 | post(L1) | Fwd L2 |
1122 Copy Stream (offload):
1123 | Offload L0 | - | Offload L1 |
1124 ↑ ↑
1125 offload at post(L0) offload at post(L1)
1127 Swap rules:
1128 1. After forward computation of the current layer completes:
1129 - If a next layer exists, asynchronously offload the activations
1130 of the current layer (launch_offload).
1132 Example:
1133 - At post-forward of L0, offload activations of L0.
1134 - At post-forward of L1, offload activations of L1.
1136 2. To limit device memory peak:
1137 - If a previous layer exists, wait until its offload operation
1138 has completed (wait_offload).
1140 Notes:
1141 - Offload operations are issued on the copy stream to overlap data transfer
1142 with forward computation of subsequent layers.
1143 - If the module is already in 'pre_backward' state, this hook is skipped
1144 to avoid triggering offload during backward phase.
1145 """
1146 if getattr(module, "_swap_state", None) == "pre_backward":
1147 return
1148 next_name = module._swap_group_order.get('next', None)
1149 if next_name:
1150 SwapManager().protect_alias_tensors(group_name, output)
1151 SwapManager().launch_offload(group_name)
1152 prev_name = module._swap_group_order.get('prev', None)
1153 if prev_name:
1154 SwapManager().wait_offload(prev_name)
1156 def _backward_pre_hook(group_name, module, grad_input): # pylint: disable=W0613
1157 """
1158 Pre-backward hook executed immediately before backward computation
1159 of the current layer starts.
1161 Execution timeline (example with 3 layers, backward order: L2 → L1 → L0):
1163 Time →
1164 Backward Compute Stream:
1165 | pre(L2) | Grad L2 | pre(L1) | Grad L1 | pre(L0) | Grad L0 |
1167 Copy Stream (load):
1168 | Load L1 | - | Load L0 |
1169 ↑ ↑
1170 prefetch at pre(L2) prefetch at pre(L1)
1172 Swap rules:
1173 1. At the beginning of backward for the current layer:
1174 - If a previous layer exists in backward order, asynchronously
1175 prefetch its activations (launch_load).
1177 Example:
1178 - At pre-backward of L2, prefetch activations of L1.
1179 - At pre-backward of L1, prefetch activations of L0.
1181 2. Before starting backward computation of the current layer:
1182 - Ensure that the activations of the current layer have already
1183 been loaded back to device memory (wait_load).
1185 Notes:
1186 - Load operations are issued on the copy stream to overlap data transfer
1187 with backward computation of the current layer.
1188 - The swap state is marked as 'pre_backward' to prevent forward hooks
1189 from issuing offload operations during backward phase.
1190 """
1191 module._swap_state = "pre_backward"
1192 prev_name = module._swap_group_order.get('prev', None)
1193 if prev_name:
1194 SwapManager().launch_load(prev_name)
1196 next_name = module._swap_group_order.get('next', None)
1197 if next_name:
1198 SwapManager().wait_load(group_name)
1199 SwapManager().release_group_storage(group_name)
1201 def _backward_hook(group_name, module, grad_input, grad_output): # pylint: disable=W0613
1202 module._swap_state = "backward"
1204 def _register_hooks_once(module, group_name):
1205 hooks = [
1206 ("_swap_forward_pre_hook_handle",
1207 lambda h: _backend.register_forward_pre_hook(module, h, prepend=True),
1208 functools.partial(_forward_pre_hook, group_name)),
1210 ("_swap_forward_hook_handle",
1211 module.register_forward_hook,
1212 functools.partial(_forward_hook, group_name)),
1214 ("_swap_backward_pre_hook_handle",
1215 lambda h: _backend.register_full_backward_pre_hook(module, h, prepend=True),
1216 functools.partial(_backward_pre_hook, group_name)),
1218 ("_swap_backward_hook_handle",
1219 lambda h: _backend.register_full_backward_hook(module, h),
1220 functools.partial(_backward_hook, group_name)),
1221 ]
1223 for attr_name, register_func, hook in hooks:
1224 if not hasattr(module, attr_name):
1225 handle = register_func(hook)
1226 setattr(module, attr_name, handle)
1227 # Register for both layers
1228 _register_hooks_once(first_layer, first_name)
1229 _register_hooks_once(second_layer, second_name)
1231 def _get_copy_stream(self):
1232 """Return a singleton copy stream, created on first access."""
1233 if self._copy_stream is None:
1234 self._copy_stream = _backend.new_stream()
1235 return self._copy_stream
1238def _teardown_wired_swap_layers(wired_modules: Iterable[Any]) -> None:
1239 """Release swap groups and prefetch hooks for an iterable of wired layers.
1241 Teardown for layers already registered through
1242 :meth:`SwapManager.set_forward_prefetch_layer` (the swap-inputs paths in
1243 ``distributed/activation_checkpoint.py`` and ``distributed/attention_swap.py``).
1244 Those paths register the modules in the process-wide singleton and install
1245 hook handles, but never otherwise tear them down, so a long-lived process
1246 that rebuilds or discards models leaks ``SwapGroup`` entries and hook
1247 handles. This is called from each path's life-cycle finalizer with the wired
1248 modules held directly, because the model container is unreachable at that
1249 point and its module tree can no longer be traversed. Idempotent: a module
1250 with no ``_swap_group_name`` (already released) is a no-op.
1252 Args:
1253 wired_modules: Layers that participated in a swap prefetch chain.
1254 """
1255 manager = SwapManager()
1256 for module in wired_modules:
1257 manager.unregister_forward_prefetch_layer(module)