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