Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / swap_optimizer_base.py: 85%
316 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +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"""Common abstractions for optimizer state swap."""
17from __future__ import annotations
19import contextlib
20from dataclasses import dataclass
21from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence
24STATE_KEYS = ("exp_avg", "exp_avg_sq", "max_exp_avg_sq")
25MASTER_PARAM_KEY = "master_param"
26SUPPORTED_STATE_KEYS = STATE_KEYS + (MASTER_PARAM_KEY,)
29@dataclass
30class SwapSlot:
31 """One logical optimizer state tensor that may be swapped."""
33 name: str
34 tensor: Any
35 cpu_tensor: Optional[Any] = None
36 storage_nbytes: int = 0
37 swappable: bool = True
38 state: str = "device"
39 event: Optional[Any] = None
40 shape: tuple[int, ...] = ()
41 dtype: Optional[Any] = None
42 device: Optional[Any] = None
43 numel: int = 0
44 host_offset: int = 0
45 packed: bool = False
46 logical_tensor: Optional[Any] = None
48 def bind_tensor(self, tensor: Any) -> None:
49 """Bind the slot to a host or staging tensor view."""
50 device = getattr(tensor, "device", None)
51 device_type = getattr(device, "type", None)
52 if self.logical_tensor is not None and device_type != "cpu":
53 # DTensor exposes no public setter for replacing its local shard.
54 setattr(self.logical_tensor, "_local_tensor", tensor)
55 self.logical_tensor.data = tensor
56 self.tensor = self.logical_tensor
57 else:
58 self.tensor = tensor
61 @property
62 def checkpoint_tensor(self) -> Any:
63 """Return the CPU tensor when present, otherwise the live tensor."""
64 return self.cpu_tensor if self.cpu_tensor is not None else self.tensor
67@dataclass
68class UpdateUnit:
69 """Per-parameter optimizer update unit used by the pipeline runtime.
71 ``adapter_index`` is interpreted by the backend adapter: it identifies a
72 Torch parameter group or a MindSpore parameter/state entry.
73 """
75 adapter_index: int
76 param: Any
77 grad: Any
78 slots: List[SwapSlot]
81class OptimizerSwapAdapter:
82 """Backend optimizer adapter interface."""
84 def __init__(self, optimizer: Any, config: Any, runtime: Any) -> None:
85 self.optimizer = optimizer
86 self.config = config
87 self.runtime = runtime
89 @classmethod
90 def matches(cls, optimizer: Any) -> bool:
91 """Return whether this adapter supports ``optimizer``."""
92 raise NotImplementedError
94 def validate(self) -> None:
95 """Validate optimizer flags and unsupported modes."""
96 raise NotImplementedError
98 def prepare_step(self, *args: Any, **kwargs: Any) -> Any:
99 """Prepare one outer optimizer step."""
100 raise NotImplementedError
102 def iter_update_units(self, step_context: Any) -> List[UpdateUnit]:
103 """Return update units for the current outer step."""
104 raise NotImplementedError
106 def step_batch(self, batch: List[UpdateUnit], step_context: Any) -> Any:
107 """Update one pipeline batch."""
108 raise NotImplementedError
110 def finish_step(self, step_context: Any) -> Any:
111 """Finish one outer optimizer step."""
112 del step_context
114 def all_slots(self) -> Iterable[SwapSlot]:
115 """Iterate all known swap slots."""
116 return ()
118 def initial_slots(self) -> Iterable[SwapSlot]:
119 """Build or return slots that can be offloaded before the first update."""
120 return self.all_slots()
122 def checkpoint_state_dict(self, *args: Any, **kwargs: Any) -> Any:
123 """Return a checkpoint view when the backend supports it."""
124 raise NotImplementedError
126 def load_checkpoint_state_dict(self, state_dict: Any, *args: Any, **kwargs: Any) -> None:
127 """Load a checkpoint view when the backend supports it."""
128 raise NotImplementedError
131class PipelineSwapRuntime:
132 """Backend-neutral pipeline orchestration.
134 Concrete backend runtimes provide tensor copy/storage methods. This class
135 owns batching and the common state machine order.
136 """
138 def __init__(self, config: Any) -> None:
139 self.config = config
140 self._copy_stream: Optional[Any] = None
142 def partition(self, units: Sequence[UpdateUnit]) -> List[List[UpdateUnit]]:
143 """Partition update units into balanced batches by swappable state bytes."""
144 non_empty_units = list(units)
145 if not non_empty_units:
146 return []
147 swap_times = max(1, min(int(self.config.swap_times), len(non_empty_units)))
148 unit_costs = [max(1, self._unit_cost(unit)) for unit in non_empty_units]
149 remaining_cost = sum(unit_costs)
150 batches: List[List[UpdateUnit]] = []
151 start = 0
152 for batch_index in range(swap_times - 1):
153 remaining_batches = swap_times - batch_index
154 max_end = len(non_empty_units) - remaining_batches + 1
155 end = start
156 current_cost = 0
157 while end < max_end:
158 next_cost = unit_costs[end]
159 # Compare against the remaining average without introducing floating-point rounding.
160 current_distance = abs(remaining_cost - current_cost * remaining_batches)
161 next_distance = abs(remaining_cost - (current_cost + next_cost) * remaining_batches)
162 if end > start and current_distance <= next_distance:
163 break
164 current_cost += next_cost
165 end += 1
166 batches.append(non_empty_units[start:end])
167 remaining_cost -= current_cost
168 start = end
169 batches.append(non_empty_units[start:])
170 return batches
172 def run_pipeline(
173 self,
174 batches: Sequence[Sequence[UpdateUnit]],
175 step_context: Any,
176 step_batch: Callable[[List[UpdateUnit], Any], Any],
177 ) -> List[Any]:
178 """Run one-batch-ahead prefetch while releasing completed offloads before widening the window."""
179 results = []
180 batch_lists = [list(batch) for batch in batches]
181 if not batch_lists:
182 return results
184 if self.supports_packed_pipeline(batch_lists):
185 return self._run_packed_pipeline(batch_lists, step_context, step_batch)
187 self.prefetch(batch_lists[0]) # prefetch 0
188 for index, batch_list in enumerate(batch_lists):
189 self.wait_prefetch(batch_list) # wait_prefetch n
191 previous_index = index - 1
192 if previous_index >= 0:
193 self.wait_offload(batch_lists[previous_index]) # wait_offload n-1
195 next_index = index + 1
196 if next_index < len(batch_lists):
197 self.prefetch(batch_lists[next_index]) # prefetch n+1
199 results.append(step_batch(batch_list, step_context)) # update n
200 self.refresh_swappable_slots(batch_list)
201 self.offload(batch_list) # offload n
203 self.wait_offload(batch_lists[-1])
204 return results
206 def _run_packed_pipeline(
207 self,
208 batches: Sequence[List[UpdateUnit]],
209 step_context: Any,
210 step_batch: Callable[[List[UpdateUnit], Any], Any],
211 ) -> List[Any]:
212 """Run updates with two reusable staging buffers.
214 Buffer parity is fixed by batch index. D2H batch ``n`` and H2D batch
215 ``n + 2`` share one copy-stream chain while the other arena updates.
216 """
217 results = []
218 try:
219 self.begin_packed_step(batches)
220 self.enqueue_packed_prefetch(0, 0)
221 if len(batches) > 1:
222 self.enqueue_packed_prefetch(1, 1)
223 for batch_index, batch in enumerate(batches):
224 staging_index = batch_index % 2
225 self.wait_packed_prefetch(batch_index, staging_index)
226 completed_index = batch_index - 2
227 if completed_index >= 0:
228 self.wait_packed_offload(completed_index)
229 self.finish_packed_offload(completed_index)
230 self.activate_packed_batch(batch_index, staging_index)
231 results.append(step_batch(batch, step_context))
232 self.refresh_swappable_slots(batch)
233 next_index = batch_index + 2
234 self.enqueue_packed_offload_prefetch(
235 batch_index,
236 next_index if next_index < len(batches) else None,
237 staging_index,
238 )
239 drain_start = max(0, len(batches) - 2)
240 for batch_index in range(drain_start, len(batches)):
241 self.wait_packed_offload(batch_index)
242 self.finish_packed_offload(batch_index)
243 finally:
244 self.release_packed_step_results(results)
245 self.end_packed_step()
246 return results
248 def release_packed_step_results(self, results: List[Any]) -> None:
249 """Release backend-specific update outputs before staging teardown."""
250 del results
252 def supports_packed_pipeline(self, batches: Sequence[Sequence[UpdateUnit]]) -> bool:
253 """Return whether this runtime can use two packed staging buffers."""
254 del batches
255 return False
257 def begin_packed_step(self, batches: Sequence[Sequence[UpdateUnit]]) -> None:
258 """Allocate and prepare packed staging storage for one optimizer step."""
259 del batches
260 raise NotImplementedError
262 def enqueue_packed_prefetch(self, batch_index: int, staging_index: int) -> None:
263 """Enqueue a standalone H2D into one staging buffer."""
264 del batch_index, staging_index
265 raise NotImplementedError
267 def wait_packed_prefetch(self, batch_index: int, staging_index: int) -> None:
268 """Make compute wait until a packed batch is ready for update."""
269 del batch_index, staging_index
270 raise NotImplementedError
272 def activate_packed_batch(self, batch_index: int, staging_index: int) -> None:
273 """Bind batch optimizer states to views in a staging buffer."""
274 del batch_index, staging_index
275 raise NotImplementedError
277 def enqueue_packed_offload_prefetch(
278 self,
279 batch_index: int,
280 next_index: Optional[int],
281 staging_index: int,
282 ) -> None:
283 """Enqueue D2H and the next H2D serially through one staging buffer."""
284 del batch_index, next_index, staging_index
285 raise NotImplementedError
287 def wait_packed_offload(self, batch_index: int) -> None:
288 """Synchronize a trailing D2H before staging storage is released."""
289 del batch_index
290 raise NotImplementedError
292 def finish_packed_offload(self, batch_index: int) -> None:
293 """Rebind an offloaded batch to its persistent CPU state views."""
294 del batch_index
295 raise NotImplementedError
297 def end_packed_step(self) -> None:
298 """Finish backend staging cleanup after all transfers complete."""
299 raise NotImplementedError
301 def refresh_swappable_slots(self, batch: Sequence[UpdateUnit]) -> None:
302 """Refresh slots that become swappable after the optimizer update."""
303 del batch
305 def synchronize_cpu_mirrors(self, slots: Iterable[SwapSlot]) -> None:
306 """Ensure CPU mirrors contain latest data for checkpointing."""
307 slot_list = [slot for slot in _iter_unique_slot_objects(slots) if slot.swappable]
308 if not slot_list:
309 return
311 compute_stream = self.current_stream()
312 with self.stream_context(compute_stream):
313 for event in _iter_unique_events(slot_list):
314 self.wait_event(event, compute_stream)
316 for slot in slot_list:
317 if slot.state == "host":
318 if slot.cpu_tensor is None:
319 raise RuntimeError(f"Swap slot {slot.name!r} is host-resident but has no CPU mirror.")
320 continue
321 if slot.state != "d2h":
322 self.copy_to_cpu(slot)
323 self.wait_offload_slot(slot)
325 checkpoint_event = self.record_event(compute_stream) if compute_stream is not None else None
327 # Storage release only needs stream ordering above. Host completion is
328 # required separately because checkpoint_state_dict reads CPU mirrors.
329 self.wait_event(checkpoint_event, None)
330 for slot in slot_list:
331 slot.event = None
333 def prefetch(self, batch: Sequence[UpdateUnit]) -> None:
334 """Prefetch batch slots from CPU to device."""
335 slots = [slot for slot in _iter_unique_slots(batch) if slot.swappable and slot.state == "host"]
336 if not slots:
337 return
339 for slot in slots:
340 if slot.cpu_tensor is None:
341 raise RuntimeError(f"Swap slot {slot.name!r} is host-resident but has no CPU mirror.")
343 for slot in slots:
344 self.restore_device_storage(slot)
346 copy_stream = self._get_copy_stream()
347 compute_event = self._record_current_stream_event() if copy_stream is not None else None
348 with self.stream_context(copy_stream):
349 self.wait_event(compute_event, copy_stream)
350 for slot in slots:
351 self.copy_to_device(slot)
352 slot.state = "h2d"
353 copy_event = self.record_event(copy_stream) if copy_stream is not None else None
354 for slot in slots:
355 slot.event = copy_event
357 def wait_prefetch(self, batch: Sequence[UpdateUnit]) -> None:
358 """Wait for batch prefetch copies."""
359 slots = [slot for slot in _iter_unique_slots(batch) if slot.swappable and slot.state == "h2d"]
360 if not slots:
361 return
363 compute_stream = self.current_stream()
364 with self.stream_context(compute_stream):
365 for event in _iter_unique_events(slots):
366 self.wait_event(event, compute_stream)
367 for slot in slots:
368 self.wait_prefetch_slot(slot)
369 slot.event = None
371 def offload(self, batch: Sequence[UpdateUnit]) -> None:
372 """Offload batch slots from device to CPU."""
373 slots = [slot for slot in _iter_unique_slots(batch) if slot.swappable and slot.state == "device"]
374 if not slots:
375 return
377 self._enqueue_offload_slots(slots)
379 def offload_initial_slots(self, slots: Iterable[SwapSlot]) -> None:
380 """Offload existing device-resident slots before the first optimizer update."""
381 slot_list = [
382 slot for slot in _iter_unique_slot_objects(slots)
383 if slot.swappable and slot.state == "device"
384 ]
385 if not slot_list:
386 return
388 self._enqueue_offload_slots(slot_list)
389 # Waiting here also releases device storage. Deferring this until the first
390 # prefetch would leave cold optimizer states resident during forward/backward.
391 self._wait_offload_slots(slot_list)
393 def _enqueue_offload_slots(self, slots: Sequence[SwapSlot]) -> None:
394 """Enqueue D2H copies for device-resident slots."""
395 copy_stream = self._get_copy_stream()
396 compute_event = self._record_current_stream_event() if copy_stream is not None else None
397 with self.stream_context(copy_stream):
398 self.wait_event(compute_event, copy_stream)
399 for slot in slots:
400 self.copy_to_cpu(slot)
401 slot.state = "d2h"
402 copy_event = self.record_event(copy_stream) if copy_stream is not None else None
403 for slot in slots:
404 slot.event = copy_event
406 def wait_offload(self, batch: Sequence[UpdateUnit]) -> None:
407 """Wait for batch offload copies."""
408 slots = [slot for slot in _iter_unique_slots(batch) if slot.swappable and slot.state == "d2h"]
409 if not slots:
410 return
412 self._wait_offload_slots(slots)
414 def _wait_offload_slots(self, slots: Sequence[SwapSlot]) -> None:
415 """Wait for D2H copies and release device storage for copied slots."""
416 compute_stream = self.current_stream()
417 with self.stream_context(compute_stream):
418 for event in _iter_unique_events(slots):
419 self.wait_event(event, compute_stream)
420 for slot in slots:
421 self.wait_offload_slot(slot)
423 def _unit_cost(self, unit: UpdateUnit) -> int:
424 return sum(slot.storage_nbytes for slot in unit.slots if slot.swappable)
426 def _get_copy_stream(self) -> Any:
427 if self._copy_stream is None:
428 self._copy_stream = self.new_stream()
429 return self._copy_stream
431 def _record_current_stream_event(self) -> Any:
432 current_stream = self.current_stream()
433 return self.record_event(current_stream)
435 def current_stream(self) -> Any:
436 """Return the current compute stream for the active backend."""
437 raise NotImplementedError
439 def new_stream(self) -> Any:
440 """Create a copy stream for the active backend."""
441 raise NotImplementedError
443 def stream_context(self, stream: Any):
444 """Return a context manager that makes ``stream`` current."""
445 del stream
446 return contextlib.nullcontext()
448 def restore_device_storage(self, slot: SwapSlot) -> None:
449 """Restore device storage before enqueueing H2D copy."""
450 del slot
452 def record_event(self, stream: Any = None) -> Any:
453 """Record an event on ``stream`` when the backend supports events."""
454 del stream
456 def wait_event(self, event: Any, stream: Any = None) -> None:
457 """Make ``stream`` wait for ``event`` when the backend supports events."""
458 del event, stream
460 def make_cpu_tensor(self, tensor: Any) -> Any:
461 """Create a CPU mirror tensor."""
462 raise NotImplementedError
464 def copy_to_device(self, slot: SwapSlot) -> None:
465 """Copy one slot CPU mirror to its device tensor."""
466 raise NotImplementedError
468 def wait_prefetch_slot(self, slot: SwapSlot) -> None:
469 """Wait for one slot prefetch."""
470 raise NotImplementedError
472 def copy_to_cpu(self, slot: SwapSlot) -> None:
473 """Copy one slot device tensor to its CPU mirror."""
474 raise NotImplementedError
476 def wait_offload_slot(self, slot: SwapSlot) -> None:
477 """Wait for one slot offload."""
478 raise NotImplementedError
481def _iter_unique_slots(units: Sequence[UpdateUnit]) -> Iterable[SwapSlot]:
482 """Yield slots once by object identity."""
483 return _iter_unique_slot_objects(slot for unit in units for slot in unit.slots)
486def _iter_unique_slot_objects(slots: Iterable[SwapSlot]) -> Iterable[SwapSlot]:
487 """Yield slots once by tensor object identity."""
488 unique_slots: Dict[int, SwapSlot] = {}
489 for slot in slots:
490 unique_slots.setdefault(id(slot.tensor), slot)
491 return unique_slots.values()
494def _iter_unique_events(slots: Iterable[SwapSlot]) -> Iterable[Any]:
495 """Yield non-empty events once by object identity."""
496 seen = set()
497 for slot in slots:
498 event = slot.event
499 if event is None:
500 continue
501 key = id(event)
502 if key in seen:
503 continue
504 seen.add(key)
505 yield event
508def validate_state_keys(state_keys: Optional[Sequence[str]]) -> Optional[tuple[str, ...]]:
509 """Validate user-provided logical state keys."""
510 if state_keys is None:
511 return None
512 normalized = tuple(state_keys)
513 invalid = sorted(set(normalized) - set(SUPPORTED_STATE_KEYS))
514 if invalid:
515 raise ValueError(
516 "SwapOptimizerConfig.state_keys only supports Adam/AdamW logical slots "
517 f"{SUPPORTED_STATE_KEYS}, but got {invalid}."
518 )
519 return normalized