Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / swap_optimizer_base.py: 96%
1318 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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"""Generic optimizer-state swap runtime, slots and adapter base."""
16# pylint: disable=protected-access
18from __future__ import annotations
20import contextlib
21import copy
22import itertools
23import logging
24from dataclasses import dataclass, field
25from typing import Any, Callable, Dict, Iterable, List, Optional, Protocol, Sequence, Tuple
27# Optimizer runtime is Torch-only; the legacy checker still flags core imports.
28import torch # pylint: disable=forbidden-backend-import
29from torch.optim import optimizer as torch_optimizer # pylint: disable=forbidden-backend-import
31logger = logging.getLogger(__name__)
33# Logical optimizer-state keys every swap family in this codebase understands.
34# This is the vocabulary ``SwapOptimizerConfig.state_keys`` may name; each
35# adapter still validates that a key belongs to its own optimizer family.
36KNOWN_STATE_KEYS = (
37 "exp_avg",
38 "exp_avg_sq",
39 "max_exp_avg_sq",
40 "momentum_buffer",
41)
42_PACKED_ALIGNMENT_BYTES = 512
43_DEVICE_TYPE = "npu"
44_MISSING_CHECKPOINT_STATE = object()
47def _restore_dtensor_metadata(template: Any, local_tensor: Any) -> Any:
48 """Wrap a local tensor with the logical DTensor layout of ``template``."""
49 to_local = getattr(template, "to_local", None)
50 if not callable(to_local):
51 return local_tensor
53 layout = getattr(template, "layout", None)
54 if hasattr(layout, "tensor_stride"):
55 logical_stride = layout.tensor_stride
56 else:
57 logical_stride = tuple(template.stride())
58 if logical_stride is None:
59 logical_stride = [1] * len(template.shape)
60 for index in range(len(logical_stride) - 2, -1, -1):
61 logical_stride[index] = logical_stride[index + 1] * max(int(template.shape[index + 1]), 1)
63 return type(template).from_local(
64 local_tensor,
65 device_mesh=template.device_mesh,
66 placements=template.placements,
67 run_check=False,
68 shape=template.shape,
69 stride=tuple(logical_stride),
70 )
73@dataclass
74class SwapSlot:
75 """One logical optimizer state tensor that may be swapped."""
77 name: str
78 tensor: Any
79 cpu_tensor: Optional[Any] = None
80 storage_nbytes: int = 0
81 swappable: bool = True
82 state: str = "device"
83 event: Optional[Any] = None
84 shape: tuple[int, ...] = ()
85 dtype: Optional[Any] = None
86 device: Optional[Any] = None
87 numel: int = 0
88 host_offset: int = 0
89 packed: bool = False
90 logical_tensor: Optional[Any] = None
92 def bind_tensor(self, tensor: Any) -> None:
93 """Bind the slot to a host or staging tensor view."""
94 device = getattr(tensor, "device", None)
95 device_type = getattr(device, "type", None)
96 if self.logical_tensor is not None:
97 # DTensor exposes no public setter for replacing its local shard.
98 # Rebind the shell on host too, so it cannot retain an arena view.
99 setattr(self.logical_tensor, "_local_tensor", tensor)
100 self.logical_tensor.data = tensor
101 if self.logical_tensor is not None and device_type != "cpu":
102 self.tensor = self.logical_tensor
103 else:
104 self.tensor = tensor
106 @property
107 def checkpoint_tensor(self) -> Any:
108 """Return the CPU tensor when present, otherwise the live tensor."""
109 return self.cpu_tensor if self.cpu_tensor is not None else self.tensor
112class SwapUpdateUnit(Protocol):
113 """Minimal contract the pipeline runtime requires of an update unit.
115 The runtime only ever reads a unit's swappable slots, so an update unit is
116 free to carry algorithm-specific metadata (a single parameter and gradient,
117 a Muon HSDP assignment, collective ordering, and so on) as long as it
118 exposes ``slots``.
119 """
121 slots: List[SwapSlot]
124@dataclass
125class UpdateUnit:
126 """Generic per-parameter update unit: parameter, gradient and swap slots.
128 Kept as a concrete unit for optimizers whose atomic update is one parameter,
129 and as the runtime's default unit type. ``adapter_index`` identifies the
130 owning parameter group. Optimizers with a richer atomic unit (Muon HSDP
131 assignments, for example) define their own unit exposing ``slots``.
132 """
134 adapter_index: int
135 param: Any
136 grad: Any
137 slots: List[SwapSlot]
140class PipelineSwapRuntime:
141 """Torch tensor storage/copy runtime and pipeline orchestration."""
143 def __init__(self, config: Any) -> None:
144 """Initialize instance-owned streams and packed storage metadata."""
145 self.config = config
146 self._copy_stream: Optional[Any] = None
147 self._packed_enabled = bool(getattr(config, "packed_swap", True))
148 self._host_buffers: Dict[Any, Any] = {}
149 self._host_layout_signature = ()
150 self._packed_batch_plans: List[_PackedBatchPlan] = []
151 self._staging_arenas: List[Optional[_StagingArena]] = [None, None]
152 self._packed_ready_events: Dict[int, Any] = {}
153 self._packed_offload_events: Dict[int, Any] = {}
154 self._packed_tail_event: Optional[Any] = None
155 self._packed_device_views: Dict[Any, Any] = {}
156 self._packed_slot_bind_hook: Optional[Callable[[SwapSlot], None]] = None
158 @property
159 def packed_enabled(self) -> bool:
160 """Return whether this runtime may build packed state candidates."""
161 return self._packed_enabled
163 @property
164 def copy_stream(self) -> Optional[Any]:
165 """Return the transfer stream after its lazy initialization."""
166 return self._copy_stream
168 def set_packed_slot_bind_hook(self, hook: Callable[[SwapSlot], None]) -> None:
169 """Install the adapter hook that republishes a packed slot's live view."""
170 self._packed_slot_bind_hook = hook
172 def _bind_packed_slot(self, slot: SwapSlot, tensor: Any) -> None:
173 """Bind a packed view and notify the owning state adapter."""
174 slot.bind_tensor(tensor)
175 if self._packed_slot_bind_hook is not None:
176 self._packed_slot_bind_hook(slot)
178 def partition(
179 self,
180 units: Sequence[SwapUpdateUnit],
181 unit_costs: Optional[Sequence[int]] = None,
182 ) -> List[List[SwapUpdateUnit]]:
183 """Partition consecutive units using explicit costs or local state bytes.
185 Args:
186 units: Atomic updates in execution order, including empty local units.
187 unit_costs: Optional non-negative integer costs, one per unit. A
188 distributed adapter supplies rank-independent planning weights.
190 Returns:
191 Consecutive batches retaining every unit in its original order.
193 Raises:
194 ValueError: If explicit costs have the wrong length or invalid values.
195 """
196 ordered_units = list(units)
197 if unit_costs is None:
198 unit_costs = [self._unit_cost(unit) for unit in ordered_units]
199 else:
200 unit_costs = list(unit_costs)
201 if len(unit_costs) != len(ordered_units):
202 raise ValueError("unit_costs must contain exactly one cost per update unit.")
203 if any(isinstance(cost, bool) or not isinstance(cost, int) or cost < 0 for cost in unit_costs):
204 raise ValueError("unit_costs must contain non-negative integers.")
205 if not ordered_units:
206 return []
207 swap_times = max(1, min(int(self.config.swap_times), len(ordered_units)))
208 unit_costs = [max(1, cost) for cost in unit_costs]
209 remaining_cost = sum(unit_costs)
210 batches: List[List[SwapUpdateUnit]] = []
211 start = 0
212 for batch_index in range(swap_times - 1):
213 remaining_batches = swap_times - batch_index
215 # Compare against the remaining average without introducing floating-point rounding.
216 # Bind the loop state as defaults so the closure captures this iteration's values.
217 def should_stop_growing(
218 end: int,
219 current_cost: int,
220 next_cost: int,
221 remaining_batches: int = remaining_batches,
222 remaining_cost: int = remaining_cost,
223 start: int = start,
224 ) -> bool:
225 """Return whether adding the next unit would not improve balance."""
226 current_distance = abs(remaining_cost - current_cost * remaining_batches)
227 next_distance = abs(remaining_cost - (current_cost + next_cost) * remaining_batches)
228 return end > start and current_distance <= next_distance
230 max_end = len(ordered_units) - remaining_batches + 1
231 end, current_cost = self._grow_balanced_segment(unit_costs, start, max_end, should_stop_growing)
232 batches.append(ordered_units[start:end])
233 remaining_cost -= current_cost
234 start = end
235 batches.append(ordered_units[start:])
236 return batches
238 def run_pipeline(
239 self,
240 batches: Sequence[Sequence[UpdateUnit]],
241 step_context: Any,
242 step_batch: Callable[[List[UpdateUnit], Any], Any],
243 ) -> List[Any]:
244 """Run one-batch-ahead prefetch while releasing completed offloads before widening the window."""
245 results = []
246 batch_lists = [list(batch) for batch in batches]
247 if not batch_lists:
248 return results
250 if self.supports_packed_pipeline(batch_lists):
251 return self._run_packed_pipeline(batch_lists, step_context, step_batch)
253 try:
254 self.prefetch(batch_lists[0]) # prefetch 0
255 for index, batch_list in enumerate(batch_lists):
256 self.wait_prefetch(batch_list) # wait_prefetch n
258 previous_index = index - 1
259 if previous_index >= 0:
260 self.wait_offload(batch_lists[previous_index]) # wait_offload n-1
262 next_index = index + 1
263 if next_index < len(batch_lists):
264 self.prefetch(batch_lists[next_index]) # prefetch n+1
266 results.append(step_batch(batch_list, step_context)) # update n
267 self.refresh_swappable_slots(batch_list)
268 self.offload(batch_list) # offload n
270 self.wait_offload(batch_lists[-1])
271 finally:
272 self._drain_pending_transfers(batch_lists)
273 return results
275 def _drain_pending_transfers(self, batch_lists: Sequence[Sequence[UpdateUnit]]) -> None:
276 """Settle slots left in an intermediate ``h2d``/``d2h`` state.
278 Called from ``run_pipeline``'s ``finally`` so that an exception raised
279 mid-pipeline (for example in ``step_batch``) cannot leave prefetched or
280 offloaded slots holding device storage and a pending stream event. Slots
281 that already reached a terminal state (``device`` or ``host``) are left
282 untouched.
283 """
284 for batch_list in batch_lists:
285 self.wait_prefetch(batch_list)
286 self.wait_offload(batch_list)
288 def _run_packed_pipeline(
289 self,
290 batches: Sequence[List[UpdateUnit]],
291 step_context: Any,
292 step_batch: Callable[[List[UpdateUnit], Any], Any],
293 ) -> List[Any]:
294 """Run updates with two reusable staging buffers.
296 Buffer parity is fixed by batch index. D2H batch ``n`` and H2D batch
297 ``n + 2`` share one copy-stream chain while the other arena updates.
298 """
299 results = []
300 try:
301 self.begin_packed_step(batches)
302 self.enqueue_packed_prefetch(0, 0)
303 if len(batches) > 1:
304 self.enqueue_packed_prefetch(1, 1)
305 for batch_index, batch in enumerate(batches):
306 staging_index = batch_index % 2
307 self.wait_packed_prefetch(batch_index, staging_index)
308 completed_index = batch_index - 2
309 if completed_index >= 0:
310 self.wait_packed_offload(completed_index)
311 self.finish_packed_offload(completed_index)
312 self.activate_packed_batch(batch_index, staging_index)
313 results.append(step_batch(batch, step_context))
314 self.refresh_swappable_slots(batch)
315 next_index = batch_index + 2
316 self.enqueue_packed_offload_prefetch(
317 batch_index,
318 next_index if next_index < len(batches) else None,
319 staging_index,
320 )
321 drain_start = max(0, len(batches) - 2)
322 for batch_index in range(drain_start, len(batches)):
323 self.wait_packed_offload(batch_index)
324 self.finish_packed_offload(batch_index)
325 finally:
326 self.release_packed_step_results(results)
327 self.end_packed_step()
328 return results
330 # Both hooks below are dispatched through ``self`` from the packed pipeline
331 # above and are overridden per backend, so they stay instance
332 # methods even though the base bodies only discard their arguments.
333 def release_packed_step_results(self, results: List[Any]) -> None:
334 """Release backend-specific update outputs before staging teardown."""
335 del results
337 def refresh_swappable_slots(self, batch: Sequence[UpdateUnit]) -> None:
338 """Refresh slots that become swappable after the optimizer update."""
339 del batch
341 def synchronize_cpu_mirrors(self, slots: Iterable[SwapSlot]) -> None:
342 """Ensure CPU mirrors contain latest data for checkpointing."""
343 slot_list = [slot for slot in _iter_unique_slot_objects(slots) if slot.swappable]
344 if not slot_list:
345 return
347 compute_stream = self.current_stream()
348 with self.stream_context(compute_stream):
349 for event in _iter_unique_events(slot_list):
350 self.wait_event(event, compute_stream)
352 for slot in slot_list:
353 if slot.state == "host":
354 if slot.cpu_tensor is None:
355 raise RuntimeError(f"Swap slot {slot.name!r} is host-resident but has no CPU mirror.")
356 continue
357 if slot.state != "d2h":
358 self.copy_to_cpu(slot)
359 self.wait_offload_slot(slot)
361 checkpoint_event = self.record_event(compute_stream) if compute_stream is not None else None
363 # Storage release only needs stream ordering above. Host completion is
364 # required separately because checkpoint_state_dict reads CPU mirrors.
365 self.wait_event(checkpoint_event, None)
366 for slot in slot_list:
367 slot.event = None
368 self.restore_packed_host_views(slot_list)
370 def materialize_slots(self, slots: Iterable[SwapSlot]) -> None:
371 """Materialize host-resident slots on their original device.
373 HSDP checkpoint broadcasts operate on device tensors. Swap slots are
374 normally host-resident between optimizer steps, so checkpoint code uses
375 this helper to reuse the normal H2D stream/event ordering before issuing
376 a collective.
377 """
378 slot_list = [
379 slot for slot in _iter_unique_slot_objects(slots)
380 if slot.swappable and slot.state == "host"
381 ]
382 if not slot_list:
383 return
384 for slot in slot_list:
385 if slot.packed:
386 device_tensor = torch.empty(slot.shape, dtype=slot.dtype, device=slot.device)
387 self._bind_packed_slot(slot, device_tensor)
388 self.prefetch([UpdateUnit(0, None, None, slot_list)])
389 self.wait_prefetch([UpdateUnit(0, None, None, slot_list)])
391 def restore_packed_host_views(self, slots: Iterable[SwapSlot]) -> None:
392 """Restore persistent packed host views after temporary materialization."""
393 for slot in _iter_unique_slot_objects(slots):
394 if (slot.packed and slot.state == "host" and slot.cpu_tensor is not None
395 and slot.tensor is not slot.cpu_tensor):
396 self._bind_packed_slot(slot, slot.cpu_tensor)
398 def prefetch(self, batch: Sequence[UpdateUnit]) -> None:
399 """Prefetch batch slots from CPU to device."""
400 slots = [slot for slot in _iter_unique_slots(batch) if slot.swappable and slot.state == "host"]
401 if not slots:
402 return
404 for slot in slots:
405 if slot.cpu_tensor is None:
406 raise RuntimeError(f"Swap slot {slot.name!r} is host-resident but has no CPU mirror.")
408 restored_slots = []
409 try:
410 for slot in slots:
411 self.restore_device_storage(slot)
412 restored_slots.append(slot)
413 except Exception:
414 for restored_slot in reversed(restored_slots):
415 self.release_device_storage(restored_slot)
416 raise
418 copy_stream = self._get_copy_stream()
419 compute_event = self._record_current_stream_event() if copy_stream is not None else None
420 with self.stream_context(copy_stream):
421 self.wait_event(compute_event, copy_stream)
422 for slot in slots:
423 self.copy_to_device(slot)
424 slot.state = "h2d"
425 copy_event = self.record_event(copy_stream) if copy_stream is not None else None
426 for slot in slots:
427 slot.event = copy_event
429 def wait_prefetch(self, batch: Sequence[UpdateUnit]) -> None:
430 """Wait for batch prefetch copies."""
431 slots = [slot for slot in _iter_unique_slots(batch) if slot.swappable and slot.state == "h2d"]
432 if not slots:
433 return
435 compute_stream = self.current_stream()
436 with self.stream_context(compute_stream):
437 for event in _iter_unique_events(slots):
438 self.wait_event(event, compute_stream)
439 for slot in slots:
440 self.wait_prefetch_slot(slot)
441 slot.event = None
443 def offload(self, batch: Sequence[UpdateUnit]) -> None:
444 """Offload batch slots from device to CPU."""
445 slots = [slot for slot in _iter_unique_slots(batch) if slot.swappable and slot.state == "device"]
446 if not slots:
447 return
449 self._enqueue_offload_slots(slots)
451 def offload_initial_slots(self, slots: Iterable[SwapSlot]) -> None:
452 """Offload existing device-resident slots before the first optimizer update."""
453 slot_list = [
454 slot for slot in _iter_unique_slot_objects(slots)
455 if slot.swappable and slot.state == "device"
456 ]
457 if not slot_list:
458 return
460 self._enqueue_offload_slots(slot_list)
461 # Waiting here also releases device storage. Deferring this until the first
462 # prefetch would leave cold optimizer states resident during forward/backward.
463 self._wait_offload_slots(slot_list)
465 def _enqueue_offload_slots(self, slots: Sequence[SwapSlot]) -> None:
466 """Enqueue D2H copies for device-resident slots."""
467 copy_stream = self._get_copy_stream()
468 compute_event = self._record_current_stream_event() if copy_stream is not None else None
469 with self.stream_context(copy_stream):
470 self.wait_event(compute_event, copy_stream)
471 for slot in slots:
472 self.copy_to_cpu(slot)
473 slot.state = "d2h"
474 copy_event = self.record_event(copy_stream) if copy_stream is not None else None
475 for slot in slots:
476 slot.event = copy_event
478 def wait_offload(self, batch: Sequence[UpdateUnit]) -> None:
479 """Wait for batch offload copies."""
480 slots = [slot for slot in _iter_unique_slots(batch) if slot.swappable and slot.state == "d2h"]
481 if not slots:
482 return
484 self._wait_offload_slots(slots)
486 def _wait_offload_slots(self, slots: Sequence[SwapSlot]) -> None:
487 """Wait for D2H copies and release device storage for copied slots."""
488 compute_stream = self.current_stream()
489 with self.stream_context(compute_stream):
490 for event in _iter_unique_events(slots):
491 self.wait_event(event, compute_stream)
492 for slot in slots:
493 self.wait_offload_slot(slot)
494 slot.event = None
496 @staticmethod
497 def _unit_cost(unit: UpdateUnit) -> int:
498 return sum(slot.storage_nbytes for slot in unit.slots if slot.swappable)
500 @staticmethod
501 def _grow_balanced_segment(
502 unit_costs: Sequence[int],
503 start: int,
504 max_end: int,
505 should_stop: Callable[[int, int, int], bool],
506 ) -> Tuple[int, int]:
507 """Extend a segment from ``start`` while ``should_stop`` stays false.
509 Returns the exclusive end index and the accumulated cost of ``[start, end)``.
510 """
511 end = start
512 current_cost = 0
513 while end < max_end:
514 next_cost = unit_costs[end]
515 if should_stop(end, current_cost, next_cost):
516 break
517 current_cost += next_cost
518 end += 1
519 return end, current_cost
521 def _get_copy_stream(self) -> Any:
522 if self._copy_stream is None:
523 self._copy_stream = self.new_stream()
524 return self._copy_stream
526 def _record_current_stream_event(self) -> Any:
527 current_stream = self.current_stream()
528 return self.record_event(current_stream)
530 def populate_slot_metadata(self, slot: SwapSlot, template: Any) -> None:
531 """Populate stable logical tensor metadata without allocating device state."""
532 storage_tensor = self._storage_tensor(template)
533 slot.shape = tuple(storage_tensor.shape)
534 slot.dtype = storage_tensor.dtype
535 slot.device = storage_tensor.device
536 slot.numel = int(storage_tensor.numel())
537 slot.storage_nbytes = slot.numel * int(storage_tensor.element_size())
539 def is_packable_template(self, tensor: Any, min_numel: int) -> bool:
540 """Return whether a state shaped like ``tensor`` can use packed staging."""
541 if not self._packed_enabled:
542 return False
543 return self.is_swappable_tensor(tensor, min_numel)
545 @staticmethod
546 def is_distributed_tensor(tensor: Any) -> bool:
547 """Return whether ``tensor`` exposes a DTensor local shard."""
548 return tensor is not None and callable(getattr(tensor, "to_local", None))
550 def prepare_packed_host(self, slots: Sequence[SwapSlot]) -> None:
551 """Pack logical optimizer states into persistent pinned buffers by dtype."""
552 if not self._packed_enabled:
553 return
555 packed_slots = [slot for slot in slots if slot.swappable and slot.packed]
556 signature = self._validate_packed_host_layout(packed_slots)
557 if signature == self._host_layout_signature:
558 return
560 self._prepare_packed_host_sources(packed_slots)
561 slots_by_dtype: Dict[Any, List[SwapSlot]] = {}
562 for slot in packed_slots:
563 slots_by_dtype.setdefault(slot.dtype, []).append(slot)
565 self._host_buffers = {
566 dtype: self._pack_dtype_slots(dtype, dtype_slots)
567 for dtype, dtype_slots in slots_by_dtype.items()
568 }
569 self._host_layout_signature = signature
571 def _validate_packed_host_layout(
572 self, packed_slots: Sequence[SwapSlot],
573 ) -> Tuple[Tuple[Any, ...], ...]:
574 """Validate candidates and return their persistent layout signature."""
575 if self._packed_batch_plans or self._packed_tail_event is not None:
576 raise RuntimeError("Cannot rebuild packed optimizer layout during an active transfer.")
577 slot_ids = [id(slot) for slot in packed_slots]
578 if len(slot_ids) != len(set(slot_ids)):
579 raise RuntimeError("Packed optimizer layout contains a duplicate state slot.")
580 return tuple(
581 (id(slot), slot.dtype, slot.shape, slot.numel, slot.device)
582 for slot in packed_slots
583 )
585 def _prepare_packed_host_sources(self, packed_slots: Sequence[SwapSlot]) -> None:
586 """Validate source layouts and settle copies before rebuilding host buffers."""
587 for slot in packed_slots:
588 source = slot.checkpoint_tensor
589 if source is not None:
590 local = self._storage_tensor(source)
591 if local.layout != torch.strided or not local.is_contiguous():
592 raise ValueError("Packed optimizer states must be dense and contiguous before host packing.")
593 # A layout rebuild reads CPU mirrors on the host. Stream ordering alone
594 # is insufficient, including after an initial-state offload.
595 if self._copy_stream is not None:
596 self._copy_stream.synchronize()
597 for event in _iter_unique_events(packed_slots):
598 self.wait_event(event, None)
599 if any(slot.state in ("h2d", "d2h") for slot in packed_slots):
600 self.synchronize_cpu_mirrors(packed_slots)
602 def invalidate_packed_layout(self) -> None:
603 """Discard an idle packed layout after state slots have been rebound.
605 Persistent slot mirrors retain their values until the next host packing.
606 The instance's copy stream remains available for subsequent steps.
607 """
608 if self._packed_batch_plans or self._packed_tail_event is not None:
609 raise RuntimeError("Cannot reset optimizer parameters during a packed transfer.")
610 for arena in self._staging_arenas:
611 if arena is not None:
612 arena.dtype_views.clear()
613 arena.raw_buffer.untyped_storage().resize_(0)
614 self._staging_arenas = [None, None]
615 self._packed_device_views.clear()
616 self._packed_ready_events.clear()
617 self._packed_offload_events.clear()
618 self._host_buffers = {}
619 self._host_layout_signature = ()
621 def _pack_dtype_slots(self, dtype: Any, dtype_slots: Sequence[SwapSlot]) -> Any:
622 """Pack every slot of one dtype into a single pinned host buffer.
624 Each slot is re-pointed at its own view inside the new buffer and marked
625 host-resident, so the packed buffer becomes the slot's CPU mirror from
626 here on. Any device copy that fed the mirror is released immediately
627 afterwards, because the host view now owns the only valid values.
628 """
629 total_numel = sum(slot.numel for slot in dtype_slots)
630 host_buffer = torch.empty(total_numel, dtype=dtype, device="cpu", pin_memory=True)
631 host_offset = 0
632 for slot in dtype_slots:
633 flat_view = host_buffer.narrow(0, host_offset, slot.numel)
634 host_view = flat_view.view(slot.shape)
635 source = slot.tensor if slot.state == "device" else slot.checkpoint_tensor
636 if source is None:
637 host_view.zero_()
638 else:
639 source_tensor = self._storage_tensor(source)
640 host_view.copy_(source_tensor.detach().reshape(-1).view(slot.shape), non_blocking=False)
641 if source_tensor.device.type != "cpu" and source is slot.tensor:
642 self.release_device_storage(slot)
643 slot.host_offset = host_offset
644 slot.cpu_tensor = host_view
645 slot.bind_tensor(host_view)
646 slot.state = "host"
647 slot.event = None
648 host_offset += slot.numel
649 return host_buffer
651 def is_swappable_tensor(self, tensor: Any, min_numel: int) -> bool:
652 """Return whether ``tensor`` can participate in swap."""
653 storage_tensor = self._storage_tensor(tensor)
654 if not isinstance(storage_tensor, torch.Tensor):
655 return False
656 if not storage_tensor.is_floating_point():
657 return False
658 if int(storage_tensor.numel()) < int(min_numel):
659 return False
660 if storage_tensor.is_sparse:
661 return False
662 if not storage_tensor.is_contiguous():
663 return False
664 if storage_tensor.device.type == "cpu":
665 return False
666 try:
667 storage_size = int(storage_tensor.untyped_storage().size())
668 expected_size = int(storage_tensor.numel()) * int(storage_tensor.element_size())
669 if storage_size != expected_size:
670 return False
671 except RuntimeError:
672 return False
673 return True
675 def storage_nbytes(self, tensor: Any) -> int:
676 """Return storage bytes for a Torch tensor."""
677 storage_tensor = self._storage_tensor(tensor)
678 if not isinstance(storage_tensor, torch.Tensor):
679 return 0
680 try:
681 return int(storage_tensor.untyped_storage().size())
682 except RuntimeError:
683 return int(storage_tensor.numel()) * int(storage_tensor.element_size())
685 def make_cpu_tensor(self, tensor: Any) -> Any:
686 """Create a CPU mirror tensor."""
687 storage_tensor = self._storage_tensor(tensor)
688 if isinstance(storage_tensor, torch.Tensor):
689 source = storage_tensor.detach()
690 try:
691 cpu_tensor = torch.empty_like(source, device="cpu", pin_memory=True)
692 except RuntimeError:
693 cpu_tensor = torch.empty_like(source, device="cpu")
694 cpu_tensor.copy_(source, non_blocking=True)
695 return cpu_tensor
696 raise ValueError(f"Expected torch.Tensor for CPU mirror, got {type(tensor)!r}.")
698 def make_zero_cpu_tensor_like(self, tensor: Any) -> Any:
699 """Create a zero-valued CPU mirror without materializing device state."""
700 storage_tensor = self._storage_tensor(tensor)
701 if not isinstance(storage_tensor, torch.Tensor):
702 raise ValueError(f"Expected torch.Tensor for CPU mirror, got {type(tensor)!r}.")
703 try:
704 cpu_tensor = torch.empty_like(storage_tensor, device="cpu", pin_memory=True)
705 except RuntimeError:
706 cpu_tensor = torch.empty_like(storage_tensor, device="cpu")
707 cpu_tensor.zero_()
708 return cpu_tensor
710 @staticmethod
711 def make_device_tensor_like(param: Any, saved_tensor: Any) -> Any:
712 """Create a live state tensor on the parameter device."""
713 if not isinstance(saved_tensor, torch.Tensor):
714 raise ValueError(f"Expected torch.Tensor in optimizer state, got {type(saved_tensor)!r}.")
715 return saved_tensor.detach().to(device=param.device, dtype=saved_tensor.dtype).clone()
717 def make_empty_device_tensor_like(self, param: Any, saved_tensor: Any) -> Any:
718 """Create an uninitialized live state tensor shell on the parameter device."""
719 if not isinstance(saved_tensor, torch.Tensor):
720 raise ValueError(f"Expected torch.Tensor in optimizer state, got {type(saved_tensor)!r}.")
721 if self.is_distributed_tensor(param):
722 return self.make_empty_state_like(param)
723 return torch.empty_like(saved_tensor, device=param.device, dtype=saved_tensor.dtype)
725 def make_empty_state_like(self, param: Any) -> Any:
726 """Allocate a state shell with the parameter's local and logical metadata."""
727 local = torch.empty_like(self._storage_tensor(param), memory_format=torch.preserve_format)
728 # Explicit from_local preserves layout even inside SkipDTensorDispatch,
729 # where empty_like(param) would otherwise return a plain local tensor.
730 return _restore_dtensor_metadata(param, local)
732 def copy_to_device(self, slot: SwapSlot) -> None:
733 """Copy one CPU mirror to device tensor."""
734 if slot.state != "host":
735 return
736 if slot.cpu_tensor is None:
737 return
738 self._storage_tensor(slot.tensor).copy_(slot.cpu_tensor, non_blocking=True)
739 slot.state = "h2d"
741 def wait_prefetch_slot(self, slot: SwapSlot) -> None:
742 """Torch fallback copies are synchronous on CPU and stream-ordered on device."""
743 slot.state = "device"
745 def copy_to_cpu(self, slot: SwapSlot) -> None:
746 """Copy one device tensor to CPU mirror."""
747 source = self._storage_tensor(slot.tensor)
748 if slot.cpu_tensor is None:
749 slot.cpu_tensor = self.make_cpu_tensor(source)
750 else:
751 slot.cpu_tensor.copy_(source.detach(), non_blocking=True)
752 slot.state = "d2h"
754 def wait_offload_slot(self, slot: SwapSlot) -> None:
755 """Release device storage after D2H copy completes."""
756 if self._storage_tensor(slot.tensor).device.type != "cpu":
757 self.release_device_storage(slot)
758 slot.state = "host"
760 def restore_device_storage(self, slot: SwapSlot) -> None:
761 """Restore device tensor storage before H2D."""
762 storage_tensor = self._storage_tensor(slot.tensor)
763 if storage_tensor.device.type == "cpu":
764 return
765 storage = storage_tensor.untyped_storage()
766 if storage.size() != slot.storage_nbytes:
767 storage.resize_(slot.storage_nbytes)
769 def release_device_storage(self, slot: SwapSlot) -> None:
770 """Release device storage for a swappable tensor."""
771 storage_tensor = self._storage_tensor(slot.tensor)
772 if storage_tensor.device.type == "cpu":
773 return
774 storage = storage_tensor.untyped_storage()
775 if storage.size() != 0:
776 storage.resize_(0)
778 @staticmethod
779 def device_handle() -> Any:
780 """Return the Torch device module (``torch.npu``) that owns streams and events."""
781 try:
782 return getattr(torch, _DEVICE_TYPE)
783 except AttributeError as error:
784 raise RuntimeError(
785 f"Swap optimizer expects device handle: 'torch.{_DEVICE_TYPE}' failed."
786 ) from error
788 def current_stream(self) -> Any:
789 """Return the current compute stream."""
790 return self.device_handle().current_stream()
792 def new_stream(self) -> Any:
793 """Create the copy stream."""
794 return self.device_handle().Stream()
796 def stream_context(self, stream: Any) -> Any:
797 """Return a stream context for ``stream``."""
798 if stream is None:
799 return contextlib.nullcontext()
800 return self.device_handle().stream(stream)
802 def record_event(self, stream: Any = None) -> Any:
803 """Record an event on ``stream``."""
804 event = self.device_handle().Event()
805 if stream is None:
806 event.record()
807 else:
808 event.record(stream)
809 return event
811 @staticmethod
812 def wait_event(event: Any, stream: Any = None) -> None:
813 """Make ``stream`` wait for ``event``."""
814 if event is None:
815 return
816 if stream is None:
817 event.synchronize()
818 return
819 event.wait(stream)
821 def supports_packed_pipeline(self, batches: Sequence[Sequence[UpdateUnit]]) -> bool:
822 """Make the final, step-specific decision to use the packed pipeline."""
823 if not self._packed_enabled or not batches:
824 return False
825 units = itertools.chain.from_iterable(batches)
826 slots = itertools.chain.from_iterable(unit.slots for unit in units)
827 swappable_slots = [slot for slot in slots if slot.swappable]
828 if not swappable_slots:
829 return False
830 if len({id(slot) for slot in swappable_slots}) != len(swappable_slots):
831 raise RuntimeError("Packed optimizer batches contain a duplicate state slot.")
832 devices = {slot.device for slot in swappable_slots}
833 eligible = len(devices) == 1 and None not in devices and all(
834 slot.packed
835 and slot.cpu_tensor is not None
836 and slot.dtype in self._host_buffers
837 for slot in swappable_slots
838 )
839 # A packed slot always carries a host mirror and a host-buffer dtype once
840 # ``prepare_packed_host`` has run, so the only remaining way for
841 # ``eligible`` to be false is a device mismatch. Per-tensor swap cannot
842 # recover from that -- the slot tensor is already a CPU host view -- so
843 # every packed slot must fail loudly rather than fall through silently.
844 if not eligible and any(slot.packed for slot in swappable_slots):
845 raise RuntimeError(
846 "Packed optimizer states cannot fall back to per-tensor swap after "
847 f"host packing; packed pipeline eligibility failed for local devices {sorted(map(str, devices))}."
848 )
849 return eligible
851 @staticmethod
852 def _storage_tensor(tensor: Any) -> Any:
853 """Return the local tensor whose storage is managed by this runtime."""
854 to_local = getattr(tensor, "to_local", None)
855 return to_local() if callable(to_local) else tensor
857 @staticmethod
858 def _first_swappable_slot(batches: Sequence[Sequence[UpdateUnit]]) -> SwapSlot:
859 """Return the first swappable slot across ``batches``, in iteration order."""
860 units = itertools.chain.from_iterable(batches)
861 slots = itertools.chain.from_iterable(unit.slots for unit in units)
862 swappable = (slot for slot in slots if slot.swappable)
863 return next(swappable)
865 def begin_packed_step(self, batches: Sequence[Sequence[UpdateUnit]]) -> None:
866 """Build batch transfer plans and materialize two raw staging buffers."""
867 first_slot = self._first_swappable_slot(batches)
868 # FSDP may leave gradient reductions/reshards queued on auxiliary
869 # streams. A current-stream event cannot cover those streams, so the
870 # device-wide synchronization is required before optimizer reads state.
871 getattr(torch, first_slot.device.type).synchronize(first_slot.device)
872 self._packed_tail_event = None
873 self._packed_batch_plans = [self._build_packed_batch_plan(batch) for batch in batches]
874 dtype_layouts, total_bytes, device = self._packed_layout(self._packed_batch_plans)
876 for staging_index in range(2):
877 arena = self._materialize_staging_arena(staging_index, total_bytes, device)
878 layout_signature = tuple(
879 (dtype, offset, num_bytes) for dtype, (offset, num_bytes) in dtype_layouts.items()
880 )
881 if arena.layout_signature != layout_signature:
882 arena.dtype_views = {
883 dtype: arena.raw_buffer.narrow(0, offset, num_bytes).view(dtype)
884 for dtype, (offset, num_bytes) in dtype_layouts.items()
885 }
886 arena.layout_signature = layout_signature
887 self._drop_packed_views(staging_index)
888 self._packed_ready_events = {}
889 self._packed_offload_events = {}
891 def enqueue_packed_prefetch(self, batch_index: int, staging_index: int) -> None:
892 """Enqueue one packed H2D and record its ready event."""
893 copy_stream = self._get_copy_stream()
894 with self.stream_context(copy_stream):
895 self._copy_packed_to_device(batch_index, staging_index)
896 ready_event = self.record_event(copy_stream)
897 self._packed_ready_events[batch_index] = ready_event
898 self._packed_tail_event = ready_event
900 def wait_packed_prefetch(self, batch_index: int, staging_index: int) -> None:
901 """Order the compute stream after the batch's packed transfer chain."""
902 del staging_index
903 ready_event = self._packed_ready_events.get(batch_index)
904 if ready_event is None:
905 raise RuntimeError(f"Packed optimizer batch {batch_index} has no ready event.")
906 self.wait_event(ready_event, self.current_stream())
908 def activate_packed_batch(self, batch_index: int, staging_index: int) -> None:
909 """Bind each swap slot to its slice of one staging arena."""
910 batch_plan = self._packed_batch_plans[batch_index]
911 arena = self._require_staging_arena(staging_index)
912 for dtype, region in batch_plan.regions.items():
913 dtype_view = arena.dtype_views[dtype]
914 for slot in region.slots:
915 relative_offset = slot.host_offset - region.host_offset
916 cache_key = (
917 staging_index,
918 id(arena.raw_buffer),
919 id(dtype_view),
920 id(slot),
921 dtype,
922 relative_offset,
923 slot.numel,
924 slot.shape,
925 )
926 device_view = self._packed_device_views.get(cache_key)
927 if device_view is None:
928 device_view = dtype_view.narrow(0, relative_offset, slot.numel).view(slot.shape)
929 self._packed_device_views[cache_key] = device_view
930 self._bind_packed_slot(slot, device_view)
931 slot.state = "device"
932 slot.event = None
934 def enqueue_packed_offload_prefetch(
935 self,
936 batch_index: int,
937 next_index: Optional[int],
938 staging_index: int,
939 ) -> None:
940 """Serialize current D2H and next same-parity H2D on the copy stream."""
941 copy_stream = self._get_copy_stream()
942 compute_event = self._record_current_stream_event()
943 with self.stream_context(copy_stream):
944 self.wait_event(compute_event, copy_stream)
945 self._copy_packed_to_host(batch_index, staging_index)
946 # This arena may be overwritten by the following H2D. Exception
947 # cleanup must use the host result once its D2H has been submitted.
948 for region in self._packed_batch_plans[batch_index].regions.values():
949 for slot in region.slots:
950 slot.state = "d2h"
951 if next_index is not None:
952 self._copy_packed_to_device(next_index, staging_index)
953 chain_event = self.record_event(copy_stream)
954 self._packed_offload_events[batch_index] = chain_event
955 self._packed_tail_event = chain_event
956 if next_index is not None:
957 self._packed_ready_events[next_index] = chain_event
959 def wait_packed_offload(self, batch_index: int) -> None:
960 """Order the compute stream after a packed transfer chain."""
961 offload_event = self._packed_offload_events.get(batch_index)
962 if offload_event is None:
963 raise RuntimeError(f"Packed optimizer batch {batch_index} has no offload event.")
964 self.wait_event(offload_event, self.current_stream())
966 def finish_packed_offload(self, batch_index: int) -> None:
967 """Make persistent pinned views authoritative after D2H completion."""
968 batch_plan = self._packed_batch_plans[batch_index]
969 for region in batch_plan.regions.values():
970 for slot in region.slots:
971 self._bind_packed_slot(slot, slot.cpu_tensor)
972 slot.state = "host"
973 slot.event = None
975 def end_packed_step(self) -> None:
976 """Drain the packed copy chain and release step-local staging storage."""
977 if self._copy_stream is not None:
978 # A submission can fail after enqueueing a copy but before recording
979 # its tail event, even when no batch has activated a device view yet.
980 self._copy_stream.synchronize()
981 elif self._packed_tail_event is not None:
982 self.wait_event(self._packed_tail_event, None)
983 active_slots: Dict[int, Any] = {}
984 for batch_plan in self._packed_batch_plans:
985 for region in batch_plan.regions.values():
986 for slot in region.slots:
987 if slot.state in ("device", "d2h"):
988 active_slots[id(slot)] = slot
989 if active_slots:
990 compute_stream = self.current_stream()
991 synchronize = getattr(compute_stream, "synchronize", None)
992 if synchronize is not None:
993 synchronize()
994 for slot in active_slots.values():
995 if slot.state == "device":
996 slot.cpu_tensor.copy_(self._storage_tensor(slot.tensor).detach(), non_blocking=False)
997 self._bind_packed_slot(slot, slot.cpu_tensor)
998 slot.state = "host"
999 slot.event = None
1000 for arena in self._staging_arenas:
1001 if arena is None:
1002 continue
1003 # Drop views before shrinking the raw storage; otherwise a view
1004 # can keep the device allocation alive after the step ends.
1005 arena.dtype_views = {}
1006 arena.layout_signature = None
1007 storage = arena.raw_buffer.untyped_storage()
1008 if storage.size() != 0:
1009 storage.resize_(0)
1010 self._packed_device_views = {}
1011 self._packed_batch_plans = []
1012 self._packed_ready_events = {}
1013 self._packed_offload_events = {}
1014 self._packed_tail_event = None
1016 @staticmethod
1017 def _build_packed_batch_plan(batch: Sequence[UpdateUnit]) -> _PackedBatchPlan:
1018 """Group a batch's packed slots into contiguous host regions by dtype."""
1019 slots_by_dtype: Dict[Any, List[SwapSlot]] = {}
1020 seen_slots = set()
1021 for unit in batch:
1022 for slot in unit.slots:
1023 if not slot.swappable or not slot.packed:
1024 continue
1025 if id(slot) in seen_slots:
1026 raise RuntimeError("Packed optimizer batch contains a duplicate state slot.")
1027 slots_by_dtype.setdefault(slot.dtype, []).append(slot)
1028 seen_slots.add(id(slot))
1030 regions = {}
1031 for dtype, slots in slots_by_dtype.items():
1032 host_offset = slots[0].host_offset
1033 expected_offset = host_offset
1034 for slot in slots:
1035 if slot.host_offset != expected_offset:
1036 raise RuntimeError(
1037 f"Packed optimizer batch has a non-contiguous {dtype} host range at slot {slot.name!r}."
1038 )
1039 expected_offset += slot.numel
1040 regions[dtype] = _PackedBatchRegion(dtype, host_offset, expected_offset - host_offset, slots)
1041 return _PackedBatchPlan(regions)
1043 def _materialize_staging_arena(self, staging_index: int, total_bytes: int, device: Any) -> _StagingArena:
1044 """Allocate or resize one packed device staging arena for the requested layout."""
1045 arena = self._staging_arenas[staging_index]
1046 if arena is None or arena.raw_buffer.device != device:
1047 raw_buffer = torch.empty(total_bytes, dtype=torch.uint8, device=device)
1048 arena = _StagingArena(raw_buffer)
1049 self._staging_arenas[staging_index] = arena
1050 self._drop_packed_views(staging_index)
1051 return arena
1053 raw_buffer = arena.raw_buffer
1054 storage = raw_buffer.untyped_storage()
1055 if storage.size() < total_bytes:
1056 storage.resize_(total_bytes)
1057 arena.dtype_views = {}
1058 arena.layout_signature = None
1059 self._drop_packed_views(staging_index)
1060 raw_buffer.set_(storage, 0, (total_bytes,), (1,))
1061 return arena
1063 def _drop_packed_views(self, staging_index: int) -> None:
1064 """Drop cached views for one arena after its storage/layout changes."""
1065 self._packed_device_views = {
1066 key: value for key, value in self._packed_device_views.items() if key[0] != staging_index
1067 }
1069 def _copy_packed_to_device(self, batch_index: int, staging_index: int) -> None:
1070 batch_plan = self._packed_batch_plans[batch_index]
1071 arena = self._require_staging_arena(staging_index)
1072 for dtype, region in batch_plan.regions.items():
1073 host_view = self._host_buffers[dtype].narrow(0, region.host_offset, region.numel)
1074 arena.dtype_views[dtype].narrow(0, 0, region.numel).copy_(host_view, non_blocking=True)
1076 def _copy_packed_to_host(self, batch_index: int, staging_index: int) -> None:
1077 batch_plan = self._packed_batch_plans[batch_index]
1078 arena = self._require_staging_arena(staging_index)
1079 for dtype, region in batch_plan.regions.items():
1080 host_view = self._host_buffers[dtype].narrow(0, region.host_offset, region.numel)
1081 host_view.copy_(arena.dtype_views[dtype].narrow(0, 0, region.numel), non_blocking=True)
1083 def _require_staging_arena(self, staging_index: int) -> _StagingArena:
1084 arena = self._staging_arenas[staging_index]
1085 if arena is None:
1086 raise RuntimeError(f"Packed optimizer staging arena {staging_index} is not materialized.")
1087 return arena
1089 @staticmethod
1090 def _align_bytes(num_bytes: int) -> int:
1091 return ((num_bytes + _PACKED_ALIGNMENT_BYTES - 1) // _PACKED_ALIGNMENT_BYTES) * _PACKED_ALIGNMENT_BYTES
1093 def _packed_layout(self, batch_plans: Sequence[_PackedBatchPlan]) -> Tuple[Dict[Any, Tuple[int, int]], int, Any]:
1094 """Return the dtype -> (offset, bytes) staging layout shared by every batch."""
1095 max_numel_by_dtype: Dict[Any, int] = {}
1096 device = None
1097 for batch_plan in batch_plans:
1098 for dtype, region in batch_plan.regions.items():
1099 max_numel_by_dtype[dtype] = max(max_numel_by_dtype.get(dtype, 0), region.numel)
1100 if device is None and region.slots:
1101 device = region.slots[0].device
1102 if device is None:
1103 raise RuntimeError("Packed optimizer pipeline has no device-resident state metadata.")
1105 dtype_layouts = {}
1106 byte_offset = 0
1107 for dtype in sorted(max_numel_by_dtype, key=str):
1108 byte_offset = self._align_bytes(byte_offset)
1109 element_size = int(self._host_buffers[dtype].element_size())
1110 num_bytes = max_numel_by_dtype[dtype] * element_size
1111 dtype_layouts[dtype] = (byte_offset, num_bytes)
1112 byte_offset += num_bytes
1113 return dtype_layouts, self._align_bytes(byte_offset), device
1116@dataclass
1117class _PackedBatchRegion:
1118 """One dtype-contiguous host range transferred for a pipeline batch."""
1120 dtype: Any
1121 host_offset: int
1122 numel: int
1123 slots: List[SwapSlot]
1126@dataclass
1127class _PackedBatchPlan:
1128 """Packed transfer regions for one optimizer pipeline batch."""
1130 regions: Dict[Any, _PackedBatchRegion] = field(default_factory=dict)
1133@dataclass
1134class _StagingArena:
1135 """One raw device allocation and its dtype-specific views."""
1137 raw_buffer: Any
1138 dtype_views: Dict[Any, Any] = field(default_factory=dict)
1139 layout_signature: Any = None
1142def _iter_unique_slots(units: Sequence[UpdateUnit]) -> Iterable[SwapSlot]:
1143 """Yield slots once by object identity."""
1144 slots = itertools.chain.from_iterable(unit.slots for unit in units)
1145 return _iter_unique_slot_objects(slots)
1148def _iter_unique_slot_objects(slots: Iterable[SwapSlot]) -> Iterable[SwapSlot]:
1149 """Yield slots once by tensor object identity."""
1150 unique_slots: Dict[int, SwapSlot] = {}
1151 for slot in slots:
1152 unique_slots.setdefault(id(slot.tensor), slot)
1153 return unique_slots.values()
1156def _iter_unique_events(slots: Iterable[SwapSlot]) -> Iterable[Any]:
1157 """Yield non-empty events once by object identity."""
1158 seen = set()
1159 for slot in slots:
1160 event = slot.event
1161 if event is None:
1162 continue
1163 key = id(event)
1164 if key in seen:
1165 continue
1166 seen.add(key)
1167 yield event
1170def _slot_tensor(unit: SwapUpdateUnit, key: str, fallback: Any) -> Any:
1171 """Return an active swap slot tensor for ``key``, or ``fallback``.
1173 A slot is authoritative only while it is swappable, still resident on the
1174 local device and bound to a live tensor; otherwise the optimizer state entry
1175 is the source of truth for this step.
1176 """
1177 for slot in unit.slots:
1178 if slot.name == key and slot.swappable and slot.state == "device" and slot.tensor is not None:
1179 return slot.tensor
1180 return fallback
1183def validate_state_keys(state_keys: Optional[Sequence[str]]) -> Optional[tuple[str, ...]]:
1184 """Validate user-provided optimizer state keys.
1186 Only lexical validity is checked here: a key must appear in
1187 :data:`KNOWN_STATE_KEYS`. Whether the wrapped optimizer family actually
1188 owns that state is decided later by the family adapter, which fails during
1189 ``validate()``/wrap time.
1190 """
1191 if state_keys is None:
1192 return None
1193 normalized = tuple(state_keys)
1194 invalid = sorted(set(normalized) - set(KNOWN_STATE_KEYS))
1195 if invalid:
1196 raise ValueError(
1197 "SwapOptimizerConfig.state_keys only supports optimizer state slots "
1198 f"{KNOWN_STATE_KEYS}, but got {invalid}."
1199 )
1200 return normalized
1203#: Builds the family adapter for a wrapped optimizer.
1204AdapterFactory = Callable[[Any, Any, PipelineSwapRuntime], "StateSwapAdapter"]
1207class StateSwapAdapter:
1208 """Algorithm-independent optimizer-state swap adapter base.
1210 Owns the slot map, the configured/default state-key selection framework, the
1211 checkpoint strip/export/restore path and packed-state publication. The
1212 numerical update and the per-parameter unit construction belong to the
1213 optimizer family subclasses, which implement :meth:`default_state_keys`,
1214 :meth:`prepare_step`, :meth:`iter_update_units` and :meth:`step_batch`.
1215 """
1217 @staticmethod
1218 def default_state_keys() -> Tuple[str, ...]:
1219 """Return the family's default swap state keys.
1221 Subclasses must override this; the base class deliberately assumes
1222 nothing about an optimizer's state layout, so there is no default.
1224 Raises:
1225 NotImplementedError: Always, unless a subclass overrides it.
1226 """
1227 raise NotImplementedError
1229 def __init__(self, optimizer: Any, config: Any, runtime: Any) -> None:
1230 """Bind a leaf optimizer to its state slots and transfer runtime."""
1231 self.optimizer = optimizer
1232 self.config = config
1233 self.runtime = runtime
1234 self._slots: Dict[Tuple[int, str], SwapSlot] = {}
1235 self._packed_state_bindings: Dict[int, Tuple[Any, str]] = {}
1236 self._checkpoint_active = False
1237 self._checkpoint_broadcasted = False
1238 self._checkpoint_state_bindings: List[Tuple[Any, str, Any]] = []
1239 self.runtime.set_packed_slot_bind_hook(self.on_packed_slot_bound)
1241 def iter_update_units(self, step_context: Dict[str, Any]) -> List[UpdateUnit]:
1242 """Return units collected in ``prepare_step``."""
1243 return step_context["units"]
1245 def partition_costs(
1246 self,
1247 units: Sequence[SwapUpdateUnit],
1248 step_context: Any,
1249 ) -> Optional[Sequence[int]]:
1250 """Return optional planning costs; default to the runtime's local bytes."""
1251 del units, step_context
1253 def prepare_pipeline(
1254 self,
1255 batches: Sequence[Sequence[SwapUpdateUnit]],
1256 step_context: Any,
1257 ) -> None:
1258 """Pack local slots in execution order after logical partitioning."""
1259 del step_context
1260 if not self.runtime.packed_enabled:
1261 return
1262 self._refresh_packed_state_bindings()
1263 active_slots = [slot for batch in batches for unit in batch for slot in unit.slots]
1264 active_ids = {id(slot) for slot in active_slots}
1265 # Inactive state stays on host at the layout tail. It must not leave
1266 # holes inside active dtype regions or keep old packed buffers alive.
1267 inactive_slots = [slot for slot in self._ordered_slots() if id(slot) not in active_ids]
1268 self.runtime.prepare_packed_host(active_slots + inactive_slots)
1269 self.publish_packed_state()
1271 def _refresh_packed_state_bindings(self) -> None:
1272 """Index packed slots by their live optimizer-state owner."""
1273 bindings = {}
1274 for group in self.optimizer.param_groups:
1275 for param in group["params"]:
1276 for key in self._configured_state_keys():
1277 slot = self._slots.get((id(param), key))
1278 if slot is not None and slot.packed:
1279 bindings[id(slot)] = (param, key)
1280 self._packed_state_bindings = bindings
1282 def on_packed_slot_bound(self, slot: SwapSlot) -> None:
1283 """Publish a packed host or staging view into optimizer state."""
1284 binding = self._packed_state_bindings.get(id(slot))
1285 if binding is None:
1286 raise RuntimeError(f"Packed optimizer slot {slot.name!r} has no live state binding.")
1287 param, key = binding
1288 state = self.optimizer.state.get(param)
1289 if state is None:
1290 raise RuntimeError(f"Packed optimizer slot {slot.name!r} has no optimizer state entry.")
1291 state[key] = slot.tensor
1293 def initial_slots(self) -> Iterable[SwapSlot]:
1294 """Discover optimizer states materialized before the swap wrapper was created."""
1295 slots = []
1296 for group in self.optimizer.param_groups:
1297 for param in group["params"]:
1298 state = self.optimizer.state.get(param)
1299 if state:
1300 slots.extend(self._build_slots(param, state))
1301 return tuple(slots)
1303 def finish_step(self, step_context: Any) -> Any:
1304 """Finish one outer optimizer step."""
1305 del step_context
1307 def _checkpoint_broadcast_state_keys(self) -> Tuple[str, ...]:
1308 """Return the wrapped optimizer state fields consumed by its broadcast."""
1309 get_state_keys = getattr(self.optimizer, "_state_keys_for_checkpoint_broadcast", None)
1310 if not callable(get_state_keys):
1311 return ()
1312 return tuple(get_state_keys())
1314 def _make_checkpoint_cpu_tensor(self, param: Any) -> torch.Tensor:
1315 """Allocate one host receiver matching the optimizer broadcast contract."""
1316 local_param = self.runtime._storage_tensor(param)
1317 try:
1318 return torch.empty_like(local_param, dtype=torch.float32, device="cpu", pin_memory=True)
1319 except RuntimeError:
1320 return torch.empty_like(local_param, dtype=torch.float32, device="cpu")
1322 def _publish_checkpoint_cpu_state(self) -> None:
1323 """Expose CPU mirrors and host receivers to the fused state broadcast.
1325 The wrapped broadcast packs these tensors into one bounded device buffer
1326 per batch. Existing swap slots therefore stay host-resident, while ranks
1327 without owner state receive directly into temporary CPU tensors.
1328 """
1329 broadcast_info = getattr(self.optimizer, "_param_to_broadcast_info", None)
1330 state_keys = self._checkpoint_broadcast_state_keys()
1331 if not broadcast_info or not state_keys:
1332 return
1334 for param in broadcast_info:
1335 param_state = self.optimizer.state.get(param)
1336 if param_state is None:
1337 param_state = {}
1338 for key in state_keys:
1339 slot = self._slots.get((id(param), key))
1340 replacement = None
1341 if slot is not None and slot.swappable:
1342 if slot.state != "host" or slot.cpu_tensor is None:
1343 raise RuntimeError(
1344 f"Checkpoint state slot {key!r} must have an authoritative CPU mirror."
1345 )
1346 replacement = slot.cpu_tensor
1347 elif key not in param_state:
1348 replacement = self._make_checkpoint_cpu_tensor(param)
1350 if replacement is None:
1351 continue
1352 if param not in self.optimizer.state:
1353 self.optimizer.state[param] = param_state
1354 previous = param_state.get(key, _MISSING_CHECKPOINT_STATE)
1355 self._checkpoint_state_bindings.append((param, key, previous))
1356 param_state[key] = replacement
1358 def _restore_checkpoint_state_bindings(self) -> None:
1359 """Restore live state entries after exporting the CPU checkpoint values."""
1360 for param, key, previous in reversed(self._checkpoint_state_bindings):
1361 param_state = self.optimizer.state.get(param)
1362 if param_state is None:
1363 continue
1364 if previous is _MISSING_CHECKPOINT_STATE:
1365 param_state.pop(key, None)
1366 if not param_state:
1367 self.optimizer.state.pop(param, None)
1368 else:
1369 param_state[key] = previous
1370 self._checkpoint_state_bindings.clear()
1372 def prepare_checkpoint(self) -> None:
1373 """Publish host-resident swap state for an optional HSDP broadcast."""
1374 if self._checkpoint_active:
1375 return
1376 self._checkpoint_active = True
1377 self._checkpoint_broadcasted = False
1378 try:
1379 self._refresh_packed_state_bindings()
1380 # A previous optimizer step may still have a completed D2H event or
1381 # a live device state. Make the host mirror authoritative first;
1382 # the broadcast streams it through its bounded device buffer.
1383 self.runtime.synchronize_cpu_mirrors(self.all_slots())
1384 self._publish_checkpoint_cpu_state()
1385 except Exception:
1386 try:
1387 self.runtime.synchronize_cpu_mirrors(self.all_slots())
1388 except RuntimeError: # pragma: no cover - best effort after a transfer failure
1389 logger.exception("Failed to restore swap slots after checkpoint preparation failed")
1390 finally:
1391 self._restore_checkpoint_state_bindings()
1392 self._checkpoint_active = False
1393 self._checkpoint_broadcasted = False
1394 raise
1396 def finish_checkpoint(self) -> None:
1397 """Return checkpoint-broadcast state to the normal host-resident form."""
1398 if not self._checkpoint_active:
1399 return
1400 try:
1401 self.runtime.synchronize_cpu_mirrors(self.all_slots())
1402 finally:
1403 self._restore_checkpoint_state_bindings()
1404 self._checkpoint_active = False
1405 self._checkpoint_broadcasted = False
1407 def all_slots(self) -> Iterable[SwapSlot]:
1408 """Iterate known swap slots."""
1409 return tuple(self._ordered_slots())
1411 def tracked_slots(self) -> Iterable[SwapSlot]:
1412 """Return all registered slots, including parameters removed from groups."""
1413 return tuple({id(slot): slot for slot in self._slots.values()}.values())
1415 def rebind_parameter_slots(self, new_to_old_param: Dict[Any, Any]) -> None:
1416 """Rebuild state bindings using an explicit parameter identity mapping.
1418 Args:
1419 new_to_old_param: New-to-old identities; unrelated leaf entries are
1420 ignored. CPU mirrors must already be synchronized.
1422 Raises:
1423 ValueError: If a state cannot be mapped to the new local shape.
1424 """
1425 params = [param for group in self.optimizer.param_groups for param in group["params"]]
1426 local_ids = {id(param) for param in self._params_with_local_state(params)}
1427 rebound_slots = {}
1428 rebound_state = {}
1429 for param in params:
1430 if id(param) not in local_ids:
1431 continue
1432 old_param = new_to_old_param.get(param, param)
1433 state = self.optimizer.state.get(param, self.optimizer.state.get(old_param))
1434 if state is None:
1435 continue
1436 state = dict(state)
1437 for key in self._configured_state_keys():
1438 if key not in state:
1439 continue
1440 slot = self._slots.get((id(param), key))
1441 if slot is None:
1442 slot = self._slots.get((id(old_param), key))
1443 slot = self._rebind_state_slot(param, key, state[key], slot, param is not old_param)
1444 state[key] = slot.tensor
1445 rebound_slots[(id(param), key)] = slot
1446 rebound_state[param] = state
1448 # Publish only after all sources have been read. A failed conversion must
1449 # leave the old slot map (and its authoritative CPU mirrors) recoverable.
1450 for param in list(self.optimizer.state):
1451 if isinstance(param, torch.Tensor) and param not in rebound_state:
1452 self.optimizer.state.pop(param)
1453 self.optimizer.state.update(rebound_state)
1454 self._slots = rebound_slots
1455 self._refresh_packed_state_bindings()
1457 def _rebind_state_slot(
1458 self, param: Any, key: str, tensor: Any, slot: Optional[SwapSlot], replaced: bool,
1459 ) -> SwapSlot:
1460 """Preserve a compatible binding or recreate it from authoritative values."""
1461 local_param = self.runtime._storage_tensor(param)
1462 if slot is not None and (tensor is slot.tensor or tensor is slot.cpu_tensor):
1463 if (not replaced and slot.shape == tuple(local_param.shape)
1464 and slot.dtype == local_param.dtype and slot.device == local_param.device
1465 and not self.runtime.is_distributed_tensor(param)):
1466 return slot
1467 source = slot.checkpoint_tensor
1468 else:
1469 # Identity verifies a missing map; it never infers a replacement.
1470 if any(tensor is tracked.tensor or tensor is tracked.cpu_tensor for tracked in self._slots.values()):
1471 raise ValueError("Replacing a swap parameter requires an explicit new_to_old_param mapping.")
1472 source = tensor
1473 source = self.runtime._storage_tensor(source)
1474 if source.numel() != local_param.numel():
1475 raise ValueError(
1476 f"Cannot rebind {key!r} with local shape {tuple(source.shape)} "
1477 f"to {tuple(local_param.shape)}; restore state for the new layout first."
1478 )
1479 cpu_tensor = self._cast_swappable_tensor_to_cpu(param, source).reshape(local_param.shape)
1480 device_tensor = self.runtime.make_empty_state_like(param)
1481 rebound = self._make_slot(key, device_tensor)
1482 if rebound.swappable:
1483 rebound.cpu_tensor = self.runtime.make_cpu_tensor(cpu_tensor)
1484 rebound.state = "host"
1485 self.runtime.release_device_storage(rebound)
1486 if rebound.packed:
1487 rebound.bind_tensor(rebound.cpu_tensor)
1488 else:
1489 self.runtime._storage_tensor(device_tensor).copy_(cpu_tensor)
1490 return rebound
1492 def checkpoint_state_dict(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
1493 """Return optimizer checkpoint state using CPU mirrors for swapped slots."""
1494 del args, kwargs
1495 self._refresh_packed_state_bindings()
1496 self.runtime.synchronize_cpu_mirrors(self.all_slots())
1497 return self.export_swappable_state(self.optimizer.state_dict())
1499 @contextlib.contextmanager
1500 def checkpoint_load_context(self, state_dict: Dict[str, Any], flatten: bool) -> Iterable[None]:
1501 """Provide state keys for FQN loading without an initialization step.
1503 Args:
1504 state_dict: Chain checkpoint keyed by model FQN.
1505 flatten: Whether the checkpoint uses flattened FQN entries.
1506 """
1507 previous_state = self.optimizer.state
1508 schema = copy.copy(previous_state)
1509 schema.clear()
1510 for group in self.optimizer.param_groups:
1511 for param in group["params"]:
1512 name = param.model_name
1513 if flatten:
1514 prefix = f"state.{name}."
1515 schema[param] = {
1516 key[len(prefix):].split(".", 1)[0]: None
1517 for key in state_dict if key.startswith(prefix)
1518 }
1519 else:
1520 schema[param] = dict(state_dict.get("state", {}).get(name, {}))
1521 # DCP uses live state keys to unflatten and calls step() if state is
1522 # empty. Empty owner ranks must never initiate that extra collective.
1523 self.optimizer.state = schema
1524 try:
1525 yield
1526 finally:
1527 if self.optimizer.state is schema:
1528 self.optimizer.state = previous_state
1530 def load_checkpoint_state_dict(
1531 self,
1532 state_dict: Dict[str, Any],
1533 *args: Any,
1534 **kwargs: Any,
1535 ) -> None:
1536 """Load optimizer checkpoint state while restoring swap-managed slots."""
1537 del args, kwargs
1538 stripped, removed = self.strip_swappable_state(state_dict)
1539 self.optimizer.load_state_dict(stripped)
1540 self.load_swappable_state(state_dict, removed)
1542 def _params_with_local_state(self, params: Sequence[Any]) -> Iterable[Any]:
1543 """Return parameters whose swapped state belongs to this rank.
1545 The default restores every checkpoint entry. Distributed optimizer
1546 adapters may filter the sequence to preserve rank-local state ownership.
1547 """
1548 return params
1550 def publish_packed_state(self) -> None:
1551 """Publish persistent packed CPU mirrors to the wrapped optimizer state."""
1552 if not self.runtime.packed_enabled:
1553 return
1554 self._refresh_packed_state_bindings()
1555 for group in self.optimizer.param_groups:
1556 for param in group["params"]:
1557 state = self.optimizer.state.get(param)
1558 for key in self._configured_state_keys():
1559 slot = self._slots.get((id(param), key))
1560 if slot is None or not slot.packed or slot.cpu_tensor is None:
1561 continue
1562 if state is None:
1563 state = self.optimizer.state[param]
1564 state[key] = slot.cpu_tensor
1566 def export_swappable_state(self, state_dict: Dict[str, Any]) -> Dict[str, Any]:
1567 """Build a checkpoint-safe Torch optimizer state dict.
1569 Torch optimizer state dicts are keyed by saved parameter ids, while the
1570 adapter tracks live swap slots by the current parameter objects. This
1571 method walks both orders together and exports each parameter's optimizer
1572 state with the data source that currently owns the valid tensor values.
1574 If an Adam state tensor such as ``exp_avg`` or ``exp_avg_sq`` has been
1575 offloaded, the live device tensor may only be a placeholder with its
1576 storage released. In that case, write a cloned CPU mirror into the
1577 exported state dict so checkpoints contain the real optimizer values.
1578 Non-swappable state, metadata, and tensors that are still resident on
1579 device are deep-copied from the original Torch state dict unchanged.
1580 """
1581 exported = {
1582 key: self._copy_checkpoint_value(value)
1583 for key, value in state_dict.items()
1584 if key not in ("state", "param_groups")
1585 }
1586 exported["param_groups"] = self._copy_checkpoint_value(
1587 state_dict.get("param_groups", [])
1588 )
1589 exported["state"] = {}
1590 saved_groups = exported.get("param_groups", [])
1591 params_in_order = []
1592 for group in self.optimizer.param_groups:
1593 params_in_order.extend(group["params"])
1594 ids_in_order = []
1595 for group in saved_groups:
1596 ids_in_order.extend(group["params"])
1597 for param, param_id in zip(params_in_order, ids_in_order):
1598 saved_state = state_dict.get("state", {}).get(param_id)
1599 if not saved_state:
1600 continue
1601 exported_state = {}
1602 for key in self._state_keys_for_param(param):
1603 slot = self._slots.get((id(param), key))
1604 if self.has_host_mirror(slot, key, saved_state):
1605 local_copy = slot.cpu_tensor.detach().clone()
1606 exported_state[key] = _restore_dtensor_metadata(param, local_copy)
1607 for key, value in saved_state.items():
1608 if key not in exported_state:
1609 exported_state[key] = self._copy_checkpoint_value(value)
1610 exported["state"][param_id] = exported_state
1611 return exported
1613 @staticmethod
1614 def _copy_checkpoint_value(value: Any) -> Any:
1615 """Clone checkpoint values without dispatching ``deepcopy`` on DTensor.
1617 HSDP state broadcasts may leave a distributed tensor in the wrapped
1618 optimizer state. Hyper DTensor intentionally rejects unregistered
1619 ``__deepcopy__`` dispatch, so clone its local shard and restore the
1620 global layout explicitly. Nested containers are copied recursively;
1621 scalar metadata still uses regular ``copy.deepcopy`` semantics.
1622 """
1623 if isinstance(value, torch.Tensor):
1624 to_local = getattr(value, "to_local", None)
1625 local = to_local() if callable(to_local) else value
1626 if local is not value:
1627 local_copy = local.detach().cpu().clone()
1628 return _restore_dtensor_metadata(value, local_copy)
1629 return copy.deepcopy(value)
1630 if isinstance(value, dict):
1631 return {
1632 key: StateSwapAdapter._copy_checkpoint_value(item)
1633 for key, item in value.items()
1634 }
1635 if isinstance(value, list):
1636 return [StateSwapAdapter._copy_checkpoint_value(item) for item in value]
1637 if isinstance(value, tuple):
1638 return tuple(StateSwapAdapter._copy_checkpoint_value(item) for item in value)
1639 return copy.deepcopy(value)
1641 @staticmethod
1642 def has_host_mirror(slot: Optional[SwapSlot], key: str, saved_state: Dict[str, Any]) -> bool:
1643 """Return whether ``slot`` currently owns the authoritative CPU mirror for ``key``.
1645 A slot qualifies only when it has been offloaded to host *and* still
1646 holds a CPU mirror: after :meth:`export_swappable_state`'s device-side
1647 fast path, the live tensor may be a storage-released placeholder, so the
1648 CPU mirror is the only source of the real optimizer values.
1649 """
1650 return (
1651 slot is not None
1652 and slot.state == "host"
1653 and slot.cpu_tensor is not None
1654 and key in saved_state
1655 )
1657 def strip_swappable_state(self, state_dict: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[int, Dict[str, Any]]]:
1658 """Split Adam state tensors out before delegating to Torch loading.
1660 PyTorch's ``load_state_dict`` eagerly restores tensors into the
1661 optimizer state. For swap-managed Adam buffers, that would bypass the
1662 adapter/runtime bookkeeping and can place large tensors directly on the
1663 device. This method therefore copies the checkpoint state dict's
1664 *containers*, removes the Adam buffers that may be swap-managed, and
1665 returns them in a side table keyed by the checkpoint parameter id.
1667 Only the container skeleton is rebuilt: leaf tensors are shared with the
1668 caller's checkpoint instead of being duplicated. The removed buffers are
1669 consumed by :meth:`load_swappable_state`, which either takes a reference
1670 to the checkpoint storage (CPU checkpoints, where ``to(device="cpu")`` is
1671 a no-op) or copies it once into a host mirror. Deep-copying the whole
1672 state dict would hold one extra full copy of every optimizer tensor on
1673 host memory for the duration of the load, which for large models is a
1674 needless doubling of the checkpoint-loading peak.
1676 The stripped state dict is safe to pass to the wrapped optimizer's
1677 ``load_state_dict`` for ordinary fields such as parameter groups and
1678 step counters: Torch does not mutate the tensors it reads from it. The
1679 removed tensors must be handed to ``load_swappable_state`` afterwards so
1680 they can be restored with the correct CPU mirror/device placeholder
1681 layout.
1682 """
1683 stripped = self._copy_state_dict_containers(state_dict)
1684 removed: Dict[int, Dict[str, Any]] = {}
1685 swappable_keys = self._configured_state_keys()
1686 for param_id, saved_state in list(stripped.get("state", {}).items()):
1687 if not isinstance(saved_state, dict):
1688 continue
1689 for key in swappable_keys:
1690 if key in saved_state:
1691 removed.setdefault(param_id, {})[key] = saved_state.pop(key)
1692 return stripped, removed
1694 @staticmethod
1695 def _copy_state_dict_containers(state_dict: Dict[str, Any]) -> Dict[str, Any]:
1696 """Shallow-copy a checkpoint state dict, sharing every leaf tensor.
1698 Rebuilds only the nesting Torch's ``load_state_dict`` walks: the top
1699 level, the ``state`` per-parameter mappings, and the ``param_groups``
1700 entries. Do not add tensor copies here: this helper exists to keep the
1701 checkpoint-loading host peak free of a duplicate optimizer state.
1702 """
1703 state = state_dict.get("state")
1704 copied_state = (
1705 {param_id: dict(saved_state) if isinstance(saved_state, dict) else saved_state
1706 for param_id, saved_state in state.items()}
1707 if isinstance(state, dict)
1708 else state
1709 )
1710 param_groups = state_dict.get("param_groups")
1711 copied_groups = (
1712 [dict(group) if isinstance(group, dict) else group for group in param_groups]
1713 if isinstance(param_groups, list)
1714 else param_groups
1715 )
1716 copied = {key: value for key, value in state_dict.items() if key not in ("state", "param_groups")}
1717 if state is not None:
1718 copied["state"] = copied_state
1719 if param_groups is not None:
1720 copied["param_groups"] = copied_groups
1721 return copied
1723 def load_swappable_state(self, original_state_dict: Dict[str, Any], removed: Dict[int, Dict[str, Any]]) -> None:
1724 """Restore removed Adam buffers under swap runtime control.
1726 This is the second half of checkpoint loading. After Torch has loaded
1727 the stripped state dict, this method maps checkpoint parameter ids back
1728 to the current parameter objects by walking saved and current parameter
1729 groups in order. Each removed Adam buffer is then recreated as an
1730 optimizer state entry and registered as a ``SwapSlot``.
1732 Packed runtimes place checkpoint values directly in persistent pinned
1733 host views. Legacy runtimes retain an empty device placeholder whose
1734 storage is restored only during prefetch. Buffers that do not meet the
1735 runtime's swappability criteria are materialized directly on the
1736 parameter's device and tracked as normal device-resident slots.
1737 """
1738 saved_groups = original_state_dict.get("param_groups", [])
1739 self._slots = {}
1740 # Walk saved and current parameter groups in order to map checkpoint ids
1741 # back to the live parameter objects.
1742 saved_ids = [
1743 saved_id
1744 for saved_group, current_group in zip(saved_groups, self.optimizer.param_groups)
1745 for saved_id in saved_group["params"]
1746 ]
1747 groups = self.optimizer.param_groups
1748 current_params = [param for group in groups for param in group["params"]]
1749 local_param_ids = {
1750 id(param) for param in self._params_with_local_state(current_params)
1751 }
1752 for saved_id, param in zip(saved_ids, current_params):
1753 if id(param) not in local_param_ids:
1754 self.optimizer.state.pop(param, None)
1755 continue
1756 for key, saved_tensor in removed.get(saved_id, {}).items():
1757 self._restore_swappable_entry(param, saved_id, saved_groups, key, saved_tensor)
1758 if self.runtime.packed_enabled:
1759 self.runtime.prepare_packed_host(self._ordered_slots())
1760 self.publish_packed_state()
1762 def _restore_swappable_entry(
1763 self,
1764 param: Any,
1765 saved_id: int,
1766 saved_groups: List[Dict[str, Any]],
1767 key: str,
1768 saved_tensor: Any,
1769 ) -> None:
1770 """Restore one removed checkpoint buffer as a host or device-resident slot."""
1771 state = self.optimizer.state[param]
1772 cpu_tensor = self._cast_swappable_tensor_to_cpu(param, saved_tensor)
1773 if self.runtime.packed_enabled and self.runtime.is_packable_template(param, self.config.min_numel):
1774 if self.runtime.is_distributed_tensor(param):
1775 logical_tensor = self.runtime.make_empty_state_like(param)
1776 slot = self._make_slot(key, logical_tensor)
1777 state[key] = logical_tensor
1778 self.runtime.release_device_storage(slot)
1779 else:
1780 slot = self._make_slot(key, None, template=param)
1781 state[key] = cpu_tensor
1782 slot.tensor = cpu_tensor
1783 slot.cpu_tensor = cpu_tensor
1784 slot.state = "host"
1785 self._slots[(id(param), key)] = slot
1786 return
1787 device_tensor = self.runtime.make_empty_device_tensor_like(param, cpu_tensor)
1788 slot = self._make_slot(key, device_tensor)
1789 if slot.swappable:
1790 state[key] = device_tensor
1791 slot.cpu_tensor = self.runtime.make_cpu_tensor(cpu_tensor)
1792 slot.state = "host"
1793 self._slots[(id(param), key)] = slot
1794 self.runtime.release_device_storage(slot)
1795 else:
1796 device_tensor = self._cast_state_tensor_like_torch(
1797 param,
1798 saved_tensor,
1799 saved_id,
1800 saved_groups,
1801 key,
1802 )
1803 state[key] = device_tensor
1804 self._slots[(id(param), key)] = self._make_slot(key, device_tensor)
1806 @staticmethod
1807 def _cast_state_tensor_like_torch(
1808 param: Any,
1809 saved_tensor: Any,
1810 saved_id: int,
1811 saved_groups: List[Dict[str, Any]],
1812 key: str,
1813 ) -> Any:
1814 """Cast a loaded state tensor using PyTorch optimizer load semantics."""
1815 if not isinstance(saved_tensor, torch.Tensor):
1816 raise ValueError(f"Expected torch.Tensor in optimizer state, got {type(saved_tensor)!r}.")
1817 process = getattr(torch.optim.Optimizer, "_process_value_according_to_param_policy", None)
1818 if process is not None:
1819 return process(param, saved_tensor, saved_id, saved_groups, key).detach().clone()
1820 if key == "step":
1821 return saved_tensor.detach().clone()
1822 if param.is_floating_point():
1823 return saved_tensor.detach().to(dtype=param.dtype, device=param.device).clone()
1824 return saved_tensor.detach().to(device=param.device).clone()
1826 @staticmethod
1827 def _cast_swappable_tensor_to_cpu(param: Any, saved_tensor: Any) -> Any:
1828 """Cast swappable state dtype like PyTorch while keeping values on CPU."""
1829 if not isinstance(saved_tensor, torch.Tensor):
1830 raise ValueError(f"Expected torch.Tensor in optimizer state, got {type(saved_tensor)!r}.")
1831 to_local = getattr(saved_tensor, "to_local", None)
1832 source = to_local() if callable(to_local) else saved_tensor
1833 if param.is_floating_point():
1834 return source.detach().to(dtype=param.dtype, device="cpu")
1835 return source.detach().to(device="cpu")
1837 def _register_present_slots(self, param: Any, state: Dict[str, Any]) -> None:
1838 """Register configured state tensors that already exist in an optimizer state mapping."""
1839 for key in self._configured_state_keys():
1840 tensor = state.get(key)
1841 if tensor is None or (id(param), key) in self._slots:
1842 continue
1843 self._slots[(id(param), key)] = self._make_slot(key, tensor)
1845 def _build_slots(self, param: Any, state: Dict[str, Any]) -> List[SwapSlot]:
1846 """Return swap slots associated with the current parameter state."""
1847 slots = []
1848 for key in self._state_keys_for_param(param):
1849 tensor = state.get(key)
1850 if tensor is None:
1851 continue
1852 slot = self._slots.get((id(param), key))
1853 if slot is None:
1854 slot = self._make_slot(key, tensor)
1855 self._slots[(id(param), key)] = slot
1856 elif slot.tensor is not tensor and slot.cpu_tensor is not tensor:
1857 slot = self._make_slot(key, tensor)
1858 self._slots[(id(param), key)] = slot
1859 slots.append(slot)
1860 return slots
1862 def _make_slot(self, key: str, tensor: Any, template: Optional[Any] = None) -> SwapSlot:
1863 """Create a swap slot for a state tensor or a packed-state template."""
1864 metadata_tensor = tensor if tensor is not None else template
1865 if metadata_tensor is None:
1866 raise ValueError(f"Cannot build swap slot {key!r} without a tensor or template.")
1867 if tensor is None:
1868 swappable = self.runtime.is_packable_template(metadata_tensor, self.config.min_numel)
1869 else:
1870 swappable = self.runtime.is_swappable_tensor(tensor, self.config.min_numel)
1871 packed = bool(self.runtime.packed_enabled and swappable)
1872 slot = SwapSlot(
1873 name=key,
1874 tensor=tensor,
1875 cpu_tensor=None,
1876 swappable=swappable,
1877 state="device" if tensor is not None else "pending",
1878 packed=packed,
1879 logical_tensor=tensor if packed and self.runtime.is_distributed_tensor(tensor) else None,
1880 )
1881 self.runtime.populate_slot_metadata(slot, metadata_tensor)
1882 return slot
1884 def _ordered_slots(self) -> List[SwapSlot]:
1885 """Return slots in optimizer parameter and configured state-key order."""
1886 slots = []
1887 seen_slots = set()
1888 for group in self.optimizer.param_groups:
1889 for param in group["params"]:
1890 for key in self._configured_state_keys():
1891 slot = self._slots.get((id(param), key))
1892 if slot is not None and id(slot) not in seen_slots:
1893 slots.append(slot)
1894 seen_slots.add(id(slot))
1895 return slots
1897 def _state_keys_for_param(self, param: Any) -> Tuple[str, ...]:
1898 state = self.optimizer.state[param]
1899 keys = self._configured_state_keys()
1900 result = []
1901 for key in keys:
1902 if key in state:
1903 result.append(key)
1904 elif self.config.state_keys is not None:
1905 raise ValueError(f"Requested state key '{key}' is not present for parameter.")
1906 return tuple(result)
1908 def _configured_state_keys(self) -> Tuple[str, ...]:
1909 """Return optimizer state keys selected for swap by the current config."""
1910 return tuple(self.config.state_keys or self.default_state_keys())
1913def _check_step_hooks(optimizer: Any) -> None:
1914 """Reject step hooks that the swap pipeline bypasses."""
1915 if (getattr(optimizer, "_optimizer_step_pre_hooks", None)
1916 or getattr(optimizer, "_optimizer_step_post_hooks", None)
1917 or torch_optimizer._global_optimizer_pre_hooks
1918 or torch_optimizer._global_optimizer_post_hooks):
1919 raise ValueError(
1920 "Optimizer-state swap does not support registered step hooks; "
1921 "remove the hooks or disable swap"
1922 )
1925class SwapOptimizer(torch.optim.Optimizer):
1926 """Torch optimizer wrapper around a family-provided swap adapter.
1928 The wrapper is algorithm agnostic: it owns the pipeline runtime, exposes the
1929 wrapped optimizer's ``param_groups``/``state``/``defaults`` and drives one
1930 step as ``prepare_step`` -> ``partition`` -> ``run_pipeline`` ->
1931 ``finish_step``. The concrete adapter is built by the family factory passed
1932 as ``adapter_factory`` (see ``swap_adam.py``/``swap_muon.py``).
1933 """
1935 _is_swap_optimizer = True
1937 def __init__(self, optimizer: Any, config: Any, adapter_factory: AdapterFactory) -> None:
1938 """Wrap a leaf with its family adapter and independent transfer runtime."""
1939 _check_step_hooks(optimizer)
1940 # Do not call ``torch.optim.Optimizer.__init__``: the wrapped base
1941 # optimizer already owns param_groups/state/defaults. Inheriting keeps
1942 # PyTorch LR schedulers and isinstance checks happy while this wrapper
1943 # delegates all optimizer state to ``self.optimizer``.
1944 self.optimizer = optimizer
1945 self.config = config
1946 self._step_active = False
1947 self._reset_active = False
1948 self._reset_required = False
1949 self._export_active = False
1950 self.runtime = PipelineSwapRuntime(config)
1951 self.adapter = adapter_factory(optimizer, config, self.runtime)
1952 self.adapter.validate()
1953 # Torch Adam states are normally lazy, but callers may materialize them
1954 # before wrapping to avoid first-step initialization in the measured loop.
1955 initial_slots = tuple(self.adapter.initial_slots())
1956 self.runtime.offload_initial_slots(initial_slots)
1957 self.runtime.prepare_packed_host(initial_slots)
1958 self.adapter.publish_packed_state()
1960 def __getattr__(self, name: str) -> Any:
1961 """Delegate unknown attributes to the base optimizer."""
1962 return getattr(self.optimizer, name)
1964 @property
1965 def param_groups(self) -> List[Dict[str, Any]]:
1966 """Proxy parameter groups."""
1967 return self.optimizer.param_groups
1969 @param_groups.setter
1970 def param_groups(self, value: List[Dict[str, Any]]) -> None:
1971 """Replace the underlying optimizer's parameter groups."""
1972 self.optimizer.param_groups = value
1974 @property
1975 def state(self) -> Dict[Any, Any]:
1976 """Proxy optimizer state."""
1977 return self.optimizer.state
1979 @property
1980 def defaults(self) -> Dict[str, Any]:
1981 """Proxy optimizer defaults."""
1982 return self.optimizer.defaults
1984 def add_param_group(self, param_group: Dict[str, Any]) -> None:
1985 """Proxy param group addition."""
1986 self.optimizer.add_param_group(param_group)
1988 def zero_grad(self, set_to_none: bool = True) -> None:
1989 """Proxy gradient clearing."""
1990 self.optimizer.zero_grad(set_to_none=set_to_none)
1992 def reset_optimizer_parameters(
1993 self, new_to_old_param: Optional[Dict[torch.nn.Parameter, torch.nn.Parameter]] = None,
1994 ) -> None:
1995 """Rebuild optimizer metadata and rebind state at an idle step boundary.
1997 Args:
1998 new_to_old_param: Explicit new-to-old identities after param_groups
1999 replacement. Omit when parameter identities are unchanged.
2001 Raises:
2002 RuntimeError: If a step, checkpoint or transfer is still active.
2003 """
2004 slots = tuple(self.adapter.tracked_slots())
2005 if (self._step_active or self._reset_active or self._export_active or self.adapter._checkpoint_active
2006 or self.runtime._packed_batch_plans or self.runtime._packed_tail_event is not None
2007 or any(slot.state in ("h2d", "d2h") for slot in slots)):
2008 raise RuntimeError("Cannot reset optimizer parameters during a step, checkpoint or transfer.")
2009 self._reset_active = True
2010 self._reset_required = True
2011 try:
2012 with self._no_grad_context():
2013 self.runtime.synchronize_cpu_mirrors(slots)
2014 reset = getattr(self.optimizer, "reset_optimizer_parameters", None)
2015 if callable(reset):
2016 reset()
2017 self.adapter.rebind_parameter_slots(new_to_old_param or {})
2018 self.runtime.invalidate_packed_layout()
2019 self._reset_required = False
2020 finally:
2021 self._reset_active = False
2023 def step(self, closure: Optional[Any] = None) -> Any:
2024 """Run one optimizer step with pipeline state swap."""
2025 if closure is not None:
2026 raise ValueError("Swap optimizer does not support closure.")
2027 _check_step_hooks(self.optimizer)
2028 self._check_reset_complete()
2029 self._step_active = True
2030 try:
2031 with self._no_grad_context():
2032 step_context = self.adapter.prepare_step()
2033 units = self.adapter.iter_update_units(step_context)
2034 unit_costs = self.adapter.partition_costs(units, step_context)
2035 batches = self.runtime.partition(units, unit_costs=unit_costs)
2036 self.adapter.prepare_pipeline(batches, step_context)
2037 self.runtime.run_pipeline(batches, step_context, self.adapter.step_batch)
2038 return self.adapter.finish_step(step_context)
2039 finally:
2040 self._step_active = False
2042 def _check_reset_complete(self) -> None:
2043 """Prevent using a partially rebuilt optimizer after reset failed."""
2044 if self._reset_required:
2045 raise RuntimeError("Optimizer parameter reset has not completed; retry reset before step or checkpoint.")
2047 def state_dict(self) -> Dict[str, Any]:
2048 """Return optimizer state dict using CPU mirrors for swappable tensors."""
2049 self._check_reset_complete()
2050 self._export_active = True
2051 try:
2052 return self._export_state_dict()
2053 finally:
2054 self._export_active = False
2056 def _export_state_dict(self) -> Dict[str, Any]:
2057 """Reuse an outer chain checkpoint window or open a leaf-owned window."""
2058 broadcast = getattr(self.optimizer, "_broadcast_state_fused_for_ckpt", None)
2059 if not callable(broadcast):
2060 return self.adapter.checkpoint_state_dict()
2062 owns_checkpoint_window = not self.adapter._checkpoint_active
2063 try:
2064 if owns_checkpoint_window:
2065 self._prepare_swap_checkpoint()
2066 if not self.adapter._checkpoint_broadcasted:
2067 self._broadcast_state_fused_for_ckpt()
2068 return self.adapter.checkpoint_state_dict()
2069 finally:
2070 if owns_checkpoint_window:
2071 self.cleanup_synced_state()
2073 def _prepare_swap_checkpoint(self) -> None:
2074 """Publish CPU swap state needed by an HSDP checkpoint broadcast."""
2075 self._check_reset_complete()
2076 broadcast = getattr(self.optimizer, "_broadcast_state_fused_for_ckpt", None)
2077 if callable(broadcast):
2078 self.adapter.prepare_checkpoint()
2080 @contextlib.contextmanager
2081 def _swap_checkpoint_load_context(self, state_dict: Dict[str, Any], flatten: bool) -> Iterable[None]:
2082 """Let chain loading discover state keys without updating any parameters."""
2083 previous_active = self._export_active
2084 self._export_active = True
2085 try:
2086 with self.adapter.checkpoint_load_context(state_dict, flatten):
2087 yield
2088 finally:
2089 self._export_active = previous_active
2091 def _finish_swap_checkpoint(self) -> None:
2092 """Return checkpoint-broadcast state to its normal swapped form."""
2093 self.adapter.finish_checkpoint()
2095 def _broadcast_state_fused_for_ckpt(self) -> None:
2096 """Broadcast wrapped optimizer state through bounded device buffers."""
2097 broadcast = getattr(self.optimizer, "_broadcast_state_fused_for_ckpt", None)
2098 if not callable(broadcast):
2099 return
2100 prepared_here = not self.adapter._checkpoint_active
2101 try:
2102 if prepared_here:
2103 self.adapter.prepare_checkpoint()
2104 broadcast()
2105 self.adapter._checkpoint_broadcasted = True
2106 except Exception:
2107 if prepared_here:
2108 self.cleanup_synced_state()
2109 raise
2111 def cleanup_synced_state(self) -> None:
2112 """Release HSDP checkpoint copies and restore host-resident slots."""
2113 try:
2114 self.adapter.finish_checkpoint()
2115 finally:
2116 cleanup = getattr(self.optimizer, "cleanup_synced_state", None)
2117 if callable(cleanup):
2118 cleanup()
2120 def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
2121 """Load optimizer state dict while keeping swappable tensors on CPU mirrors."""
2122 previous_active = self._export_active
2123 self._export_active = True
2124 try:
2125 self.adapter.load_checkpoint_state_dict(state_dict)
2126 finally:
2127 self._export_active = previous_active
2129 @contextlib.contextmanager
2130 def _no_grad_context(self) -> Iterable[None]:
2131 with torch.no_grad():
2132 yield