Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_checkpoint / swap.py: 95%

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 

17 

18import functools 

19import threading 

20import warnings 

21 

22from collections import defaultdict 

23from contextlib import contextmanager 

24from contextvars import ContextVar 

25from typing import Any, Dict, Iterator, List, Optional, Set 

26 

27from hyper_parallel.platform import get_platform 

28 

29platform = get_platform() 

30 

31# --------------------------------------------------------------------------- 

32# Module-level buffer pools — process-local, no locking needed for single- 

33# stream training. Each GPU process owns its own Python interpreter, so 

34# these dicts are never shared across processes. 

35# 

36# _CPU_PINNED_POOL: a list of available pinned CPU tensors per dtype_key. 

37# Created via alloc_tensor_buffer(pin_memory=True) on the first miss; the 

38# base tensor is returned here after wait_load and reused in the next 

39# launch_offload, avoiding repeated cudaHostAlloc / cudaFreeHost calls. 

40# --------------------------------------------------------------------------- 

41_CPU_PINNED_POOL: Dict[str, List[Any]] = defaultdict(list) 

42# Cap each group-swap staging allocation. 32 MiB keeps DMA chunks large 

43# while avoiding one huge per-dtype staging tensor in large models. 

44_GROUP_SWAP_MAX_BULK_COPY_BYTES = 32 * 1024 * 1024 

45 

46 

47def _get_cpu_pinned_buf(dtype_key: str, total_numel: int, dtype): 

48 """Pop the smallest sufficient pinned buffer from the pool, or allocate. 

49 

50 Best-fit selection minimises wasted pinned memory. When no buffer in the 

51 pool is large enough, an undersized entry is discarded before allocating a 

52 fresh buffer via alloc_tensor_buffer. 

53 

54 Returns the *full* buffer (capacity >= total_numel). Callers must slice 

55 ``buf[:total_numel]`` for the actual copy so the returned reference can be 

56 passed back to :func:`_return_cpu_pinned_buf` without any platform-specific 

57 introspection. 

58 """ 

59 pool = _CPU_PINNED_POOL[dtype_key] 

60 best_i = -1 

61 for i, buf in enumerate(pool): 

62 if buf.numel() >= total_numel: 

63 if best_i == -1 or buf.numel() < pool[best_i].numel(): 

64 best_i = i 

65 if best_i != -1: 

66 return pool.pop(best_i) 

67 # No suitable buffer — discard one stale undersized entry. 

68 if pool: 

69 pool.pop() 

70 return platform.alloc_tensor_buffer(total_numel, dtype, device='cpu', pin_memory=True) 

71 

72 

73def _return_cpu_pinned_buf(buf): 

74 """Return a full pinned CPU buffer to the pool for reuse.""" 

75 if buf is None: 

76 return 

77 _CPU_PINNED_POOL[str(buf.dtype)].append(buf) 

78 

79 

80def _collect_device_storage_ptrs(tensors: Any) -> Set[int]: 

81 """Collect device storage pointers from a nested tensor structure.""" 

82 storage_ptrs = set() 

83 

84 def _collect(x): 

85 if isinstance(x, platform.Tensor) and str(x.device).lower() != "cpu": 

86 storage_ptrs.add(x.untyped_storage().data_ptr()) 

87 return x 

88 

89 platform.tree_map(_collect, tensors) 

90 return storage_ptrs 

91 

92 

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" 

100 

101 def __init__(self, val: Any, funcname: str, group_swap: bool = False) -> 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 if isinstance(val, platform.Tensor) and str(val.device).lower() != 'cpu': 

109 self.ver = val._version 

110 self._state = self.STATE_DEVICE 

111 val_storage = val.untyped_storage() 

112 self.storage_size = val_storage.size() 

113 self.is_slice_tensor = self.storage_size != val.numel() * platform.get_element_size(val) 

114 self.val_cpu = None 

115 else: 

116 self.ver = None 

117 self._state = self.STATE_NON_TENSOR 

118 self.val_cpu = None 

119 self.is_slice_tensor = False 

120 self.storage_size = 0 

121 

122 def dedup_key(self): 

123 """Return a stable identity key for duplicate-swap detection.""" 

124 if self._state == self.STATE_NON_TENSOR: 

125 return None 

126 val_storage = self.val.untyped_storage() 

127 return ( 

128 str(self.val.device), 

129 val_storage.data_ptr(), 

130 self.val.storage_offset(), 

131 val_storage.size(), 

132 tuple(self.val.stride()), 

133 ) 

134 

135 def mark_duplicate_swap(self) -> None: 

136 """Mark this wrapper as a duplicate registration in the same swap group.""" 

137 self._duplicate_swap = True 

138 

139 def protect_if_aliases(self, alias_storage_ptrs: Set[int]) -> None: 

140 """Keep tensors that alias externally-owned tensors on device.""" 

141 if self._state == self.STATE_NON_TENSOR: 

142 return 

143 if self.val.untyped_storage().data_ptr() in alias_storage_ptrs: 

144 self._keep_on_device = True 

145 

146 def get_val(self) -> Any: 

147 """Return the underlying tensor value. 

148 

149 Raises RuntimeError if the tensor is not currently in the 'device' state. 

150 Non-tensor values are returned directly regardless of state. 

151 """ 

152 if self._state == self.STATE_NON_TENSOR: 

153 return self.val 

154 if self._state != self.STATE_DEVICE: 

155 raise RuntimeError( 

156 f"Cannot call get_val(): tensor is in '{self._state}' state. " 

157 f"Must be in 'device' state." 

158 ) 

159 return self.val 

160 

161 def resize_device_storage(self): 

162 """Reallocate device memory on compute stream.""" 

163 if self._state == self.STATE_NON_TENSOR or self._duplicate_swap: 

164 return 

165 if self._group_managed: 

166 return 

167 

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) 

174 

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 

181 

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 

188 

189 if self.val_cpu is None: 

190 raise ValueError("val_cpu must not be None during async_load") 

191 with platform.preserve_version_counter(self.val): 

192 if 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 

197 

198 def wait_load(self): 

199 """change state to device after async load is done""" 

200 if self._state == self.STATE_NON_TENSOR or self._keep_on_device or self._duplicate_swap: 

201 return 

202 

203 if self._state == self.STATE_DEVICE: 

204 return # already loaded 

205 if self._state != self.STATE_H2D: 

206 warnings.warn( 

207 f"[SwapTensor.wait_load] Called in invalid state: {self._state}. " 

208 f"Expected 'h2d'. Skipped." 

209 ) 

210 return 

211 self._state = self.STATE_DEVICE 

212 

213 def async_offload(self): 

214 """async offload tensor from device to host""" 

215 if self._state == self.STATE_NON_TENSOR or self._keep_on_device or self._duplicate_swap: 

216 return 

217 if self._group_managed: 

218 return 

219 

220 if self._state != self.STATE_DEVICE: 

221 warnings.warn( 

222 f"[SwapTensor.async_offload] Invalid state: current={self._state}, " 

223 f"expected 'device'. Operation skipped." 

224 ) 

225 return 

226 

227 if self.storage_size != self.val.untyped_storage().size(): 

228 raise RuntimeError( 

229 f"There is a tensor from {self.funcname} cannot be SWAPPED! Its storage has been resized " 

230 f"presize:{self.storage_size}, current size:{self.val.untyped_storage().size()}" 

231 ) 

232 if self.ver != self.val._version: 

233 raise RuntimeError( 

234 f"There is a tensor from {self.funcname} cannot be SWAPPED! In-place modification happened " 

235 f"preversion:{self.ver}, current version:{self.val._version}" 

236 ) 

237 

238 if self.val_cpu is None: 

239 self.val_cpu = platform.empty_like( 

240 self.val, device="cpu", pin_memory=True 

241 ) 

242 if self.is_slice_tensor: 

243 self.val_cpu.copy_(self.val, non_blocking=True) 

244 else: 

245 self.val_cpu.untyped_storage().copy_(self.val.untyped_storage(), non_blocking=True) 

246 self._state = self.STATE_D2H 

247 

248 def wait_offload(self): 

249 """wait offload to host and free device memory""" 

250 if self._state == self.STATE_NON_TENSOR or self._keep_on_device or self._duplicate_swap: 

251 return 

252 

253 if self._state == self.STATE_HOST: 

254 return 

255 if self._state != self.STATE_D2H: 

256 warnings.warn( 

257 f"[SwapTensor.wait_offload] Called in invalid state: {self._state}. " 

258 f"Expected 'd2h'. Skipped." 

259 ) 

260 return 

261 storage = self.val.untyped_storage() 

262 if storage.size() != 0: 

263 storage.resize_(0) 

264 self._state = self.STATE_HOST 

265 

266 @property 

267 def state(self) -> str: 

268 """Return the current swap state of this tensor (device, host, d2h, h2d, or non_tensor).""" 

269 return self._state 

270 

271 def __repr__(self): 

272 if self._state == self.STATE_NON_TENSOR: 

273 return f"<SwapTensor state=non_tensor, val_type={type(self.val).__name__}>" 

274 return ( 

275 f"<SwapTensor state={self._state}, duplicate={self._duplicate_swap}, " 

276 f"device_val={'exists' if self.val is not None else 'None'}>" 

277 ) 

278 

279 

280class Storage: 

281 """Manage a collection of tensors for swapping operations. 

282 

283 Supports dict-like access: ``storage[key].append(item)``, ``storage.clear()``, 

284 ``for batch in storage.values(): ...``. 

285 """ 

286 

287 def __init__(self): 

288 self._data: Dict[Any, List[Any]] = defaultdict(list) 

289 

290 def __getitem__(self, key: Any) -> List[Any]: 

291 return self._data[key] 

292 

293 def values(self): 

294 """Return an iterable view of all stored lists.""" 

295 return self._data.values() 

296 

297 def clear(self): 

298 """Remove all entries from the storage.""" 

299 self._data.clear() 

300 

301 def iter_swap_tensors(self): 

302 """Iterate all SwapTensor objects stored in this storage.""" 

303 collected = [] 

304 

305 def _collect(x): 

306 if isinstance(x, SwapTensor): 

307 collected.append(x) 

308 return x 

309 

310 for storage_list in self.values(): 

311 for item in storage_list: 

312 platform.tree_map(_collect, item) 

313 return collected 

314 

315 def mark_duplicate_swaps(self, seen_keys) -> int: 

316 """Mark tensors already registered in the same swap group as duplicates.""" 

317 duplicate_count = 0 

318 for swap_tensor in self.iter_swap_tensors(): 

319 dedup_key = swap_tensor.dedup_key() 

320 if dedup_key is None: 

321 continue 

322 if dedup_key in seen_keys: 

323 swap_tensor.mark_duplicate_swap() 

324 duplicate_count += 1 

325 continue 

326 seen_keys.add(dedup_key) 

327 return duplicate_count 

328 

329 def protect_alias_storage_ptrs(self, alias_storage_ptrs: Set[int]): 

330 """Avoid offloading swap entries that alias externally-owned storage.""" 

331 if not alias_storage_ptrs: 

332 return 

333 

334 def _protect_tensor(x): 

335 if isinstance(x, SwapTensor): 

336 x.protect_if_aliases(alias_storage_ptrs) 

337 return x 

338 

339 for storage_list in self.values(): 

340 for item in storage_list: 

341 platform.tree_map(_protect_tensor, item) 

342 

343 def launch_load(self): 

344 """launch async load for all tensors in swap storage""" 

345 def _async_load(x): 

346 if isinstance(x, SwapTensor): 

347 x.async_load() 

348 return x 

349 

350 for storage_list in self.values(): 

351 for item in storage_list: 

352 platform.tree_map(_async_load, item) 

353 

354 def resize_device_storage(self): 

355 """Resize device storage for all swap tensors (runs on compute stream).""" 

356 def _resize(x): 

357 if isinstance(x, SwapTensor): 

358 x.resize_device_storage() 

359 return x 

360 for storage_list in self.values(): 

361 for item in storage_list: 

362 platform.tree_map(_resize, item) 

363 

364 def wait_load(self): 

365 """wait load for all tensors in swap storage""" 

366 def _wait_load(x): 

367 if isinstance(x, SwapTensor): 

368 x.wait_load() 

369 return x 

370 

371 for storage_list in self.values(): 

372 for item in storage_list: 

373 platform.tree_map(_wait_load, item) 

374 self.clear() 

375 

376 def wait_offload(self): 

377 """wait offload for all tensors in swap storage""" 

378 def _wait_offload(x): 

379 if isinstance(x, SwapTensor): 

380 x.wait_offload() 

381 return x 

382 

383 for storage_list in self.values(): 

384 for item in storage_list: 

385 platform.tree_map(_wait_offload, item) 

386 

387 def launch_offload(self): 

388 """launch async offload for all tensors in swap storage""" 

389 def _async_offload(x): 

390 

391 if isinstance(x, SwapTensor): 

392 x.async_offload() 

393 return x 

394 

395 for storage_list in self.values(): 

396 for item in storage_list: 

397 platform.tree_map(_async_offload, item) 

398 

399 

400class SwapGroup: 

401 """Manager for a group of storages to coordinate swap operations. 

402 

403 Non-slice tensors within the group are packed into bounded contiguous device 

404 buffers before D2H transfer, and loaded back from bounded H2D buffers. 

405 Each tensor then aliases its slice of the relevant buffer via 

406 ``Tensor.set_()``, avoiding per-tensor memory fragmentation. 

407 

408 Slice tensors (storage larger than logical data) fall back to the original 

409 per-tensor copy path. 

410 """ 

411 

412 def __init__(self, group_name: str): 

413 self.group_name = group_name 

414 self.is_last_group: bool = False 

415 self._storages: List[Storage] = [] 

416 self._load_event: Optional[Any] = None 

417 self._offload_event: Optional[Any] = None 

418 # Group-level contiguous buffers for non-slice tensors. 

419 self._packed_tensor_info: List = [] # [(SwapTensor, bucket_key, element_offset), ...] 

420 self._packed_buckets: Dict[str, Dict[str, Any]] = {} 

421 self._group_cpu_buf = None # pinned CPU bufs; live offload→load 

422 self._group_device_buf = None # temp device bufs; cleared after each phase 

423 # Persistent dedup set accumulated across add() calls; avoids O(N²) rebuild. 

424 # mark_duplicate_swaps mutates it in-place, so new keys are added automatically. 

425 # Reset at wait_load() so stale data_ptrs don't leak into the next iteration. 

426 self._seen_dedup_keys: set = set() 

427 # Per-bucket SwapTensor lists built in _collect_packable_tensors and consumed 

428 # in launch_offload, eliminating a redundant pass over _packed_tensor_info. 

429 self._packed_by_bucket: Dict[str, List] = {} 

430 

431 def add(self, storage): 

432 """Add a storage to the swap group.""" 

433 duplicate_count = storage.mark_duplicate_swaps(self._seen_dedup_keys) 

434 if duplicate_count > 0: 

435 warnings.warn( 

436 f"SwapGroup '{self.group_name}' skipped {duplicate_count} duplicate tensor swap registration(s)." 

437 ) 

438 self._storages.append(storage) 

439 

440 def protect_alias_tensors(self, tensors: Any): 

441 """Protect externally-owned tensors from premature offload.""" 

442 alias_storage_ptrs = _collect_device_storage_ptrs(tensors) 

443 if not alias_storage_ptrs: 

444 return 

445 for storage in self._storages: 

446 storage.protect_alias_storage_ptrs(alias_storage_ptrs) 

447 

448 def _collect_packable_tensors(self) -> int: 

449 """Identify tensors eligible for group packing and mark them for bulk copy. 

450 

451 A tensor is eligible only when it is contiguous, not a slice tensor, 

452 not a duplicate, not sharing storage with another live swap tensor, and 

453 has ``group_swap=True``. Dtype buckets are split before their staging 

454 allocation would exceed ``_GROUP_SWAP_MAX_BULK_COPY_BYTES``. A packed 

455 bucket with fewer than two tensors is left on the original per-tensor 

456 path because it has no batch-copy benefit. Non-contiguous 

457 tensors are excluded because the packing step copies storage-order 

458 bytes while restore uses the original stride; those tensors fall back to 

459 the per-tensor copy path. 

460 Shared-storage tensors also fall back together because group packing 

461 frees the original storage after packing, which would invalidate any 

462 non-packed aliases such as transpose views before their own offload. 

463 

464 Side effects: marks each eligible tensor with ``_group_managed=True`` 

465 and ``_state=STATE_D2H``, and populates ``_packed_tensor_info`` / 

466 ``_packed_buckets``. 

467 

468 Returns: 

469 Total byte count of all packable tensors. 

470 """ 

471 candidate_buckets: Dict[str, List[Dict[str, Any]]] = {} 

472 packed_info: List = [] 

473 packed_buckets: Dict[str, Dict[str, Any]] = {} 

474 packed_by_bucket: Dict[str, List] = {} 

475 total_bytes = 0 

476 

477 def _try_pack(x): 

478 if not isinstance(x, SwapTensor): 

479 return x 

480 no_pack = (not x.group_swap or x._state != SwapTensor.STATE_DEVICE or x._keep_on_device 

481 or x.is_slice_tensor or x._duplicate_swap or x.storage_size >= _GROUP_SWAP_MAX_BULK_COPY_BYTES 

482 or not x.val.is_contiguous()) 

483 if no_pack: 

484 return x 

485 if x.storage_size != x.val.untyped_storage().size(): 

486 raise RuntimeError( 

487 f"There is a tensor from {x.funcname} cannot be SWAPPED! Its storage has been resized " 

488 f"presize:{x.storage_size}, current size:{x.val.untyped_storage().size()}" 

489 ) 

490 if x.ver != x.val._version: 

491 raise RuntimeError( 

492 f"There is a tensor from {x.funcname} cannot be SWAPPED! In-place modification happened " 

493 f"preversion:{x.ver}, current version:{x.val._version}" 

494 ) 

495 dtype_key = str(x.val.dtype) 

496 dtype_buckets = candidate_buckets.setdefault(dtype_key, []) 

497 if (not dtype_buckets or 

498 dtype_buckets[-1]["total_bytes"] + x.storage_size > _GROUP_SWAP_MAX_BULK_COPY_BYTES): 

499 dtype_buckets.append({ 

500 "bucket_key": f"{dtype_key}#{len(dtype_buckets)}", 

501 "dtype": x.val.dtype, 

502 "dtype_key": dtype_key, 

503 "device": x.val.device, 

504 "tensors": [], 

505 "total_bytes": 0, 

506 "total_numel": 0, 

507 }) 

508 bucket = dtype_buckets[-1] 

509 bucket["tensors"].append(x) 

510 bucket["total_bytes"] += x.storage_size 

511 bucket["total_numel"] += x.val.numel() 

512 return x 

513 

514 for storage in self._storages: 

515 for storage_list in storage.values(): 

516 for item in storage_list: 

517 platform.tree_map(_try_pack, item) 

518 

519 for dtype_bucket_list in candidate_buckets.values(): 

520 for candidate_bucket in dtype_bucket_list: 

521 tensors = candidate_bucket["tensors"] 

522 if len(tensors) < 2: 

523 continue 

524 bucket_key = candidate_bucket["bucket_key"] 

525 packed_buckets[bucket_key] = { 

526 "dtype": candidate_bucket["dtype"], 

527 "dtype_key": candidate_bucket["dtype_key"], 

528 "device": candidate_bucket["device"], 

529 "total_numel": candidate_bucket["total_numel"], 

530 } 

531 element_offset = 0 

532 for tensor in tensors: 

533 tensor._group_managed = True 

534 tensor._state = SwapTensor.STATE_D2H 

535 packed_info.append((tensor, bucket_key, element_offset)) 

536 element_offset += tensor.val.numel() 

537 packed_by_bucket[bucket_key] = tensors 

538 total_bytes += candidate_bucket["total_bytes"] 

539 

540 self._packed_tensor_info = packed_info 

541 self._packed_buckets = packed_buckets 

542 self._packed_by_bucket = packed_by_bucket 

543 return total_bytes 

544 

545 def launch_offload(self, copy_stream): 

546 """Launch async offload for all storages in the group. 

547 

548 Non-slice tensors are first packed into bounded contiguous device 

549 buffers, then transferred to pinned CPU memory. Slice tensors are 

550 offloaded individually via the existing per-tensor path. 

551 """ 

552 total_bytes = self._collect_packable_tensors() 

553 with platform.no_grad(): 

554 if total_bytes > 0: 

555 group_device_bufs = {} 

556 group_cpu_bufs = {} 

557 for bucket_key, swap_tensors in self._packed_by_bucket.items(): 

558 group_device_bufs[bucket_key] = platform.cat( 

559 [st.val.reshape(-1) for st in swap_tensors], dim=0 

560 ) 

561 

562 compute_event = platform.new_event() 

563 compute_event.record(platform.get_current_stream()) 

564 self._offload_event = platform.new_event() 

565 stream_context = platform.get_stream_context() 

566 with platform.no_grad(), stream_context(copy_stream): 

567 compute_event.wait(copy_stream) 

568 

569 if total_bytes > 0: 

570 # One-shot D2H per packed bucket. MindSpore requires tensor/storage dtype consistency. 

571 for bucket_key, bucket in self._packed_buckets.items(): 

572 dtype_key = bucket["dtype_key"] 

573 numel = bucket["total_numel"] 

574 cpu_buf = _get_cpu_pinned_buf(dtype_key, numel, bucket["dtype"]) 

575 group_cpu_bufs[bucket_key] = cpu_buf 

576 cpu_buf[:numel].copy_(group_device_bufs[bucket_key], non_blocking=True) 

577 self._group_device_buf = group_device_bufs 

578 self._group_cpu_buf = group_cpu_bufs 

579 

580 # Slice tensors use the existing per-tensor path. 

581 # Group-managed tensors are already STATE_D2H so async_offload is a no-op. 

582 for storage in self._storages: 

583 storage.launch_offload() 

584 self._offload_event.record(copy_stream) 

585 

586 def wait_offload(self): 

587 """Wait for offload to complete for all storages in the group.""" 

588 if self._offload_event is None: 

589 raise RuntimeError( 

590 f"SwapGroup '{self.group_name}' wait_offload() called before launch_offload()." 

591 ) 

592 compute_stream = platform.get_current_stream() 

593 stream_context = platform.get_stream_context() 

594 with platform.no_grad(), stream_context(compute_stream): 

595 self._offload_event.wait(compute_stream) 

596 self._offload_event = None 

597 for storage in self._storages: 

598 storage.wait_offload() 

599 # Release the temporary device packing buffer; _group_cpu_buf persists until launch_load. 

600 self._group_device_buf = None 

601 

602 def launch_load(self, copy_stream): 

603 """Prepare storage and launch async load for all storages in the group. 

604 

605 Non-slice tensors are loaded from pinned CPU memory into bounded 

606 contiguous device buffers. Tensors will alias their slice of the 

607 relevant buffer after ``wait_load``. Slice tensors use the existing 

608 per-tensor path. 

609 """ 

610 # Resize device storage for slice tensors only. 

611 # Group-managed tensors skip resize_device_storage via _group_managed flag. 

612 with platform.no_grad(): 

613 for storage in self._storages: 

614 storage.resize_device_storage() 

615 

616 compute_event = platform.new_event() 

617 compute_event.record(platform.get_current_stream()) 

618 self._load_event = platform.new_event() 

619 stream_context = platform.get_stream_context() 

620 with platform.no_grad(), stream_context(copy_stream): 

621 compute_event.wait(copy_stream) 

622 

623 if self._packed_tensor_info and self._group_cpu_buf is not None: 

624 group_device_bufs = {} 

625 for bucket_key, bucket in self._packed_buckets.items(): 

626 cpu_buf = self._group_cpu_buf.get(bucket_key) 

627 if cpu_buf is None: 

628 continue 

629 numel = bucket["total_numel"] 

630 group_device_bufs[bucket_key] = platform.alloc_tensor_buffer( 

631 numel, bucket["dtype"], bucket["device"] 

632 ) 

633 # One-shot H2D per packed bucket. 

634 group_device_bufs[bucket_key].copy_(cpu_buf[:numel], non_blocking=True) 

635 self._group_device_buf = group_device_bufs 

636 # Mirror async_load's STATE_H2D transition: H2D is in flight. 

637 for st, _, _ in self._packed_tensor_info: 

638 st._state = SwapTensor.STATE_H2D 

639 

640 # Slice tensors use the existing per-tensor path. 

641 # Group-managed tensors skip async_load via _group_managed flag. 

642 for storage in self._storages: 

643 storage.launch_load() # Only copy, no resize 

644 self._load_event.record(copy_stream) 

645 

646 def wait_load(self): 

647 """Wait for load to complete for all storages in the group. 

648 

649 After the H2D transfer completes, each group-managed tensor is made to 

650 alias its slice of the contiguous device buffer via ``Tensor.set_()``. 

651 The buffer stays alive through the tensors' own storage references after 

652 ``_group_device_buf`` is cleared here. 

653 """ 

654 if self._load_event is None: 

655 raise RuntimeError( 

656 f"SwapGroup '{self.group_name}' wait_load() called before launch_load()." 

657 ) 

658 compute_stream = platform.get_current_stream() 

659 stream_context = platform.get_stream_context() 

660 with platform.no_grad(), stream_context(compute_stream): 

661 self._load_event.wait(compute_stream) 

662 self._load_event = None 

663 # Restore group-managed tensors: alias into the contiguous device buffer. 

664 if self._group_device_buf is not None: 

665 prev_key = None 

666 group_storage = None 

667 for st, bucket_key, element_offset in self._packed_tensor_info: 

668 if bucket_key != prev_key: 

669 group_device_buf = self._group_device_buf.get(bucket_key) 

670 group_storage = group_device_buf.untyped_storage() if group_device_buf is not None else None 

671 prev_key = bucket_key 

672 if group_storage is None: 

673 continue 

674 with platform.preserve_version_counter(st.val): 

675 st.val.set_(group_storage, element_offset, st.val.shape, st.val.stride()) 

676 st._state = SwapTensor.STATE_DEVICE 

677 for storage in self._storages: 

678 storage.wait_load() 

679 self._storages.clear() 

680 # Return CPU pinned buffers to the pool. By the time wait_load 

681 # returns, _load_event has fired on the compute stream, which 

682 # means the copy stream's H2D transfer has completed and the CPU 

683 # buffer is no longer being read by the DMA engine. The next 

684 # launch_offload (start of the following iteration) will pop these 

685 # buffers from the pool, well after the current H2D is done. 

686 if self._group_cpu_buf is not None: 

687 for buf in self._group_cpu_buf.values(): 

688 _return_cpu_pinned_buf(buf) 

689 self._group_cpu_buf = None 

690 # Device buffer: the pool holds the staging reference; just drop 

691 # the local reference. Tensors aliasing _group_device_buf's 

692 # storage keep it alive via their own storage references until 

693 # they are consumed in backward. 

694 self._group_device_buf = None 

695 self._packed_tensor_info = [] 

696 self._packed_buckets = {} 

697 self._packed_by_bucket = {} 

698 self._seen_dedup_keys = set() 

699 

700 

701class SwapManager: 

702 """Singleton manager for swap groups and their operations.""" 

703 _instance: Optional["SwapManager"] = None 

704 _lock = threading.Lock() 

705 

706 def __init__(self) -> None: 

707 """Initialize process-local swap groups once for the singleton.""" 

708 if hasattr(self, '_groups'): 

709 return 

710 self._groups: Dict[str, SwapGroup] = {} 

711 self._current_group_name: ContextVar[str] = ContextVar( 

712 "swap_current_group_name", default="" 

713 ) 

714 self._layer_count: int = 0 

715 self._copy_stream: Optional[Any] = None 

716 

717 def __new__(cls): 

718 if cls._instance is None: 

719 with cls._lock: 

720 if cls._instance is None: 

721 cls._instance = super().__new__(cls) 

722 return cls._instance 

723 

724 def add_storage(self, group_name: str, storage: Storage) -> None: 

725 """Add a storage to a specified swap group.""" 

726 self.ensure_group(group_name) 

727 self._groups[group_name].add(storage) 

728 

729 def ensure_group(self, group_name: str) -> None: 

730 """Create the swap group if it does not exist yet.""" 

731 if group_name not in self._groups: 

732 self._groups[group_name] = SwapGroup(group_name) 

733 

734 def launch_offload(self, group_name: str, copy_stream=None): 

735 """Launch async offload for a specified swap group.""" 

736 group = self._groups.get(group_name) 

737 if group is None: 

738 raise RuntimeError(f"Group {group_name} does not exist.") 

739 if copy_stream is None: 

740 copy_stream = self._get_copy_stream() 

741 group.launch_offload(copy_stream) 

742 

743 def protect_alias_tensors(self, group_name: str, tensors: Any): 

744 """Keep tensors that alias externally-owned tensors on device.""" 

745 group = self._groups.get(group_name) 

746 if group is None: 

747 raise RuntimeError(f"Group {group_name} does not exist.") 

748 group.protect_alias_tensors(tensors) 

749 

750 def wait_offload(self, group_name: str): 

751 """Wait for offload to complete for a specified swap group.""" 

752 group = self._groups.get(group_name) 

753 if group is None: 

754 raise RuntimeError(f"Group {group_name} does not exist.") 

755 group.wait_offload() 

756 

757 def launch_load(self, group_name: str, copy_stream=None): 

758 """Launch async load for a specified swap group.""" 

759 group = self._groups.get(group_name) 

760 if group is None: 

761 raise RuntimeError(f"Group {group_name} does not exist.") 

762 if copy_stream is None: 

763 copy_stream = self._get_copy_stream() 

764 group.launch_load(copy_stream) 

765 

766 def wait_load(self, group_name: str): 

767 """Wait for load to complete for a specified swap group.""" 

768 group = self._groups.get(group_name) 

769 if group is None: 

770 raise RuntimeError(f"Group {group_name} does not exist.") 

771 group.wait_load() 

772 

773 def release_group_storage(self, group_name: str) -> None: 

774 """Release storage references held by the swap group.""" 

775 group = self._groups.get(group_name) 

776 if group is not None: 

777 group._storages.clear() 

778 

779 def abort_group(self, group_name: str) -> None: 

780 """Synchronize in-flight transfers and remove a failed run's group.""" 

781 group = self._groups.pop(group_name, None) 

782 if group is None: 

783 return 

784 for event in (group._offload_event, group._load_event): 

785 if event is not None: 

786 event.synchronize() 

787 group._storages.clear() 

788 

789 def get_current_group_name(self) -> str: 

790 """Return the name of the currently active swap group.""" 

791 return self._current_group_name.get() 

792 

793 def set_current_group_name(self, group_name: str) -> None: 

794 """Set the name of the currently active swap group.""" 

795 self._current_group_name.set(group_name) 

796 

797 def active_group_count(self) -> int: 

798 """Return the number of live swap groups for lifecycle diagnostics.""" 

799 return len(self._groups) 

800 

801 @contextmanager 

802 def group_context(self, group_name: str) -> Iterator[None]: 

803 """Activate a swap group within the current execution context.""" 

804 token = self._current_group_name.set(group_name) 

805 try: 

806 yield 

807 finally: 

808 self._current_group_name.reset(token) 

809 

810 def is_last_group(self, group_name: Optional[str] = None) -> bool: 

811 """Return whether the specified swap group is the terminal group in the chain.""" 

812 group_name = self.get_current_group_name() if group_name is None else group_name 

813 group = self._groups.get(group_name) 

814 if group is None: 

815 return False 

816 return group.is_last_group 

817 

818 def set_forward_prefetch_layer(self, first_layer, second_layer): 

819 """ 

820 Configure prefetching and offloading order between two consecutive layers. 

821 

822 Usage: 

823 for i in range(len(model.layers) - 1): 

824 set_forward_prefetch_layer(model.layers[i], model.layers[i + 1]) 

825 

826 Ensures idempotency: safe to call multiple times on the same layer pair. 

827 """ 

828 if first_layer is second_layer: