Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/core/activation

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

7# http://www.apache.org/licenses/LICENSE-2.0

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"""Eager non-reentrant activation checkpointing for the Torch backend.

17The saved-tensor hook algorithm is adapted from ``torch.utils.checkpoint`` in

18PyTorch release/2.9. Hyper owns the scheduling and session extensions here so

19the implementation can run consistently on PyTorch 2.6, 2.7, and 2.9 without

20patching the installed framework.

21"""

22import contextlib

23import contextvars

24import threading

25import uuid

26import warnings

27import weakref

28from collections import defaultdict

29from dataclasses import dataclass, field

30from functools import partial

31from typing import Any, Callable, DefaultDict, Dict, Generator, Iterator, List, Optional, Tuple

33import torch

34from torch.utils._pytree import tree_map

35from torch.utils.checkpoint import DefaultDeviceType

36from torch.utils.checkpoint import checkpoint as torch_checkpoint

37from torch.utils.checkpoint import set_checkpoint_early_stop

40_DEFAULT_DETERMINISM_MODE = "default"

41_RECOMPUTE_COLLECTOR = contextvars.ContextVar("hyper_recompute_collector", default=None)

42_RECOMPUTE_SESSION = contextvars.ContextVar("hyper_recompute_session", default=None)

43_SESSION_FRAMES: DefaultDict[Any, weakref.WeakSet] = defaultdict(weakref.WeakSet)

44_SESSION_FRAMES_LOCK = threading.RLock()

47class CheckpointError(RuntimeError):

48 """Raised when checkpoint forward and recomputation are inconsistent."""

51class _Handle:

52 """Identity key for one recomputed saved tensor."""

55class _Holder:

56 """Saved-tensor placeholder containing handles keyed by recompute session."""

58 def __init__(self) -> None:

59 """Initialize an empty per-session handle mapping."""

60 self.handles: Dict[Any, Optional[_Handle]] = {}

63class _StopRecomputationError(Exception):

64 """Internal control-flow exception used by early-stop recomputation."""

67class _SessionActivation:

68 """Control-plane state shared by checkpoint frames in one session scope."""

70 def __init__(self, session_id: Any, retain_on_unpack: bool) -> None:

71 """Initialize one scoped session activation."""

72 self.session_id = session_id

73 self.retain_on_unpack = retain_on_unpack

74 self.frames: weakref.WeakSet = weakref.WeakSet()

77class _NoopSaveInputs(torch.autograd.Function):

78 """Save checkpoint inputs without adding a meaningful forward operation."""

80 @staticmethod

81 def forward(*args: Any) -> Any:

82 """Return a dummy output whose grad function retains checkpoint inputs."""

83 del args

84 return torch.empty((0,))

86 @staticmethod

87 def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:

88 """Save tensor inputs while retaining non-tensor input structure."""

89 del output

90 tensor_pairs = [(index, value) for index, value in enumerate(inputs) if isinstance(value, torch.Tensor)]

91 tensor_indices, tensors = zip(*tensor_pairs)

92 index_to_saved_index = {input_index: saved_index for saved_index, input_index in enumerate(tensor_indices)}

93 stored_args = [None if isinstance(value, torch.Tensor) else value for value in inputs]

95 def get_args(saved_tensors: Tuple[Any, ...]) -> List[Any]:

96 """Reconstruct the original checkpoint arguments."""

97 restored_args = [

98 saved_tensors[index_to_saved_index[index]] if index in tensor_indices else value

99 for index, value in enumerate(stored_args)

100 ]

101 return restored_args[1:]

102

103 ctx.get_args = get_args

104 ctx.save_for_backward(*tensors)

105

106 @staticmethod

107 def backward(ctx: Any, *grad_outputs: Any) -> None:

108 """Reject direct backward through the internal input saver."""

109 del ctx, grad_outputs

110 raise CheckpointError("The internal checkpoint input saver must not be backwarded directly.")

111

112

113class _CheckpointFrame:

114 """State shared by one checkpoint forward and its recomputations."""

115

116 def __init__(self, recompute_fn: Callable, early_stop: bool, metadata_fn: Optional[Callable]) -> None:

117 """Initialize frame state captured by the saved-tensor hooks."""

118 self.recompute_fn = recompute_fn

119 self.input_saver = None

120 self.weak_holders: List[weakref.ReferenceType] = []

121 self.recomputed: DefaultDict[Any, weakref.WeakKeyDictionary] = defaultdict(weakref.WeakKeyDictionary)

122 self.recomp_counter: DefaultDict[Any, int] = defaultdict(int)

123 self.is_recomputed: DefaultDict[Any, bool] = defaultdict(bool)

124 self.early_stop = early_stop

125 self.metadata_fn = metadata_fn

126 self.x_metadatas: List[Any] = []

127 self.forward_completed = False

128 self.ignore_saved_mismatch = False

129 self.active_session: Optional[_SessionActivation] = None

130

131 def check_recomputed_tensors_match(self, session_id: Any) -> None:

132 """Validate saved-tensor count and metadata after recomputation."""

133 if self.ignore_saved_mismatch:

134 return

135 if len(self.weak_holders) != self.recomp_counter[session_id]:

136 raise CheckpointError(

137 "Hyper checkpoint saved a different number of tensors during forward and recomputation. "

138 f"Forward saved {len(self.weak_holders)} tensors, but recomputation saved "

139 f"{self.recomp_counter[session_id]} tensors."

140 )

141

142 mismatches = []

143 for index, weak_holder in enumerate(self.weak_holders):

144 holder = weak_holder()

145 if holder is None:

146 continue

147 handle = holder.handles.get(session_id)

148 _internal_assert(handle is not None, "Missing recomputed tensor handle during metadata validation.")

149 _internal_assert(

150 handle in self.recomputed[session_id],

151 "Missing recomputed tensor during metadata validation.",

152 )

153 recomputed_tensor = self.recomputed[session_id][handle]

154 recomputed_metadata = self.metadata_fn(recomputed_tensor)

155 if self.x_metadatas[index] != recomputed_metadata:

156 mismatches.append((index, self.x_metadatas[index], recomputed_metadata))

157

158 if mismatches:

159 details = "\n".join(

160 f"tensor {index}: forward={forward_metadata}, recompute={recomputed_metadata}"

161 for index, forward_metadata, recomputed_metadata in mismatches

162 )

163 raise CheckpointError(

164 "Hyper checkpoint detected different tensor metadata during recomputation:\n" + details

165 )

166

167 def clear_session(self, session_id: Any) -> None:

168 """Release all tensors and handles associated with one session."""

169 for weak_holder in self.weak_holders:

170 holder = weak_holder()

171 if holder is not None:

172 holder.handles.pop(session_id, None)

173 self.recomputed.pop(session_id, None)

174 self.recomp_counter.pop(session_id, None)

175 self.is_recomputed.pop(session_id, None)

176

177

178def _bind_session_activation(frame: _CheckpointFrame, activation: _SessionActivation) -> None:

179 """Bind one activation to a frame outside the unpack hot path."""

180 if frame.active_session is activation:

181 return

182 if frame.active_session is not None:

183 raise CheckpointError("Concurrent recompute sessions on the same checkpoint frame are not supported.")

184 frame.active_session = activation

185 activation.frames.add(frame)

186

187

188def _register_session_frame(

189 frame: _CheckpointFrame,

190 session_id: Any,

191 activation: Optional[_SessionActivation] = None,

192) -> None:

193 """Register a frame for cleanup and bind its current activation when present."""

194 with _SESSION_FRAMES_LOCK:

195 _SESSION_FRAMES[session_id].add(frame)

196 if activation is not None:

197 _internal_assert(activation.session_id == session_id, "Session activation key does not match its frame.")

198 _bind_session_activation(frame, activation)

199

200

201def _activate_registered_frames(activation: _SessionActivation) -> None:

202 """Install an activation on every frame already registered for its session."""

203 with _SESSION_FRAMES_LOCK:

204 for frame in list(_SESSION_FRAMES.get(activation.session_id, ())):

205 _bind_session_activation(frame, activation)

206

207

208def _deactivate_session(activation: _SessionActivation) -> None:

209 """Remove one activation from every frame bound at context entry."""

210 with _SESSION_FRAMES_LOCK:

211 for frame in list(activation.frames):

212 if frame.active_session is activation:

213 frame.active_session = None

214 activation.frames.clear()

215

216

217def _internal_assert(condition: bool, message: str) -> None:

218 if not condition:

219 raise CheckpointError(message)

220

221

222def _noop_context_fn() -> Tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]:

223 return contextlib.nullcontext(), contextlib.nullcontext()

224

225

226def _default_metadata_fn(tensor: Any) -> Dict[str, Any]:

227 return {"shape": tensor.shape, "dtype": tensor.dtype, "device": tensor.device}

228

229

230def _infer_device_type(*args: Any) -> str:

231 """Return the preferred non-CPU device type found in checkpoint inputs."""

232 device_types = []

233

234 def add_device_type(value: Any) -> None:

235 """Record one non-CPU tensor device type."""

236 if isinstance(value, torch.Tensor) and value.device.type != "cpu":

237 device_types.append(value.device.type)

238

239 tree_map(add_device_type, args)

240 device_types_set = set(device_types)

241 if len(device_types_set) > 1:

242 warnings.warn(

243 "Hyper checkpoint received tensors on multiple non-CPU device types. RNG state is preserved only for "

244 "one device type; CUDA is preferred when present.",

245 stacklevel=3,

246 )

247 if not device_types:

248 return DefaultDeviceType.get_device_type()

249 if "cuda" in device_types_set:

250 return "cuda"

251 return device_types[0]

252

253

254def _get_device_module(device_type: str) -> Any:

255 if device_type == "meta":

256 return torch.device("meta")

257 return getattr(torch, device_type)

258

259

260def _get_device_states(device_type: str, *args: Any) -> Tuple[List[int], List[Any]]:

261 """Capture RNG states for non-CPU input devices of the requested type."""

262 device_ids = []

263

264 def add_device_id(value: Any) -> None:

265 """Record one non-CPU tensor device index."""

266 if isinstance(value, torch.Tensor) and value.device.type not in {"cpu", "meta"}:

267 device_ids.append(value.get_device())

268

269 tree_map(add_device_id, args)

270 device_module = _get_device_module(device_type)

271 states = []

272 for device_id in device_ids:

273 with device_module.device(device_id):

274 states.append(device_module.get_rng_state())

275 return device_ids, states

276

277

278def _set_device_states(device_type: str, devices: List[int], states: List[Any]) -> None:

279 if device_type == "meta":

280 return

281 device_module = _get_device_module(device_type)

282 for device, state in zip(devices, states):

283 with device_module.device(device):

284 device_module.set_rng_state(state)

285

286

287def _get_autocast_kwargs(device_type: str) -> Tuple[Optional[Dict[str, Any]], Dict[str, Any]]:

288 """Return active autocast settings for the selected device and CPU."""

289 device_kwargs = None

290 if torch.amp.is_autocast_available(device_type):

291 device_kwargs = {

292 "enabled": torch.is_autocast_enabled(device_type),

293 "dtype": torch.get_autocast_dtype(device_type),

294 "cache_enabled": torch.is_autocast_cache_enabled(),

295 }

296 cpu_kwargs = {

297 "enabled": torch.is_autocast_enabled("cpu"),

298 "dtype": torch.get_autocast_dtype("cpu"),

299 "cache_enabled": torch.is_autocast_cache_enabled(),

300 }

301 return device_kwargs, cpu_kwargs

302

303

304def _autocast_context(device_type: str, kwargs: Optional[Dict[str, Any]]) -> Any:

305 """Create an autocast context only when another thread enabled it."""

306 if kwargs is None or not kwargs.get("enabled", False):

307 return contextlib.nullcontext()

308 return torch.amp.autocast(device_type=device_type, **kwargs)

309

310

311def _create_recomputation_hooks(frame: _CheckpointFrame, session_id: Any) -> Any:

312 """Create saved-tensor hooks that retain tensors from one recomputation."""

313 frame_ref = weakref.ref(frame)

314

315 def pack_hook(tensor: Any) -> Any:

316 """Store recomputed tensors in their forward holders."""

317 tensor = tensor.detach() if tensor.requires_grad else tensor

318 target_frame = frame_ref()

319 _internal_assert(target_frame is not None, "Checkpoint frame was released during recomputation.")

320 recompute_index = target_frame.recomp_counter[session_id]

321 target_frame.recomp_counter[session_id] += 1

322

323 if recompute_index >= len(target_frame.weak_holders):

324 if not target_frame.early_stop and not target_frame.forward_completed:

325 target_frame.ignore_saved_mismatch = True

326 return tensor

327 raise CheckpointError(

328 "Hyper checkpoint tried to save more tensors during recomputation than during forward."

329 )

330

331 holder = target_frame.weak_holders[recompute_index]()

332 if holder is not None:

333 _internal_assert(

334 holder.handles.get(session_id) is None,

335 "A recomputed tensor handle already exists for this session.",

336 )

337 handle = _Handle()

338 holder.handles[session_id] = handle

339 target_frame.recomputed[session_id][handle] = tensor

340

341 if target_frame.early_stop and target_frame.recomp_counter[session_id] == len(target_frame.weak_holders):

342 raise _StopRecomputationError

343 return tensor

344

345 def unpack_hook(tensor: Any) -> Any:

346 """Return tensors saved by operations inside the recomputation."""

347 return tensor

348

349 return torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook)

350

351

352# PyTorch exposes this tracing guard only as a private decorator.

353@torch._disable_dynamo # pylint: disable=protected-access

354def _run_fn_with_dynamo_disabled(function: Callable, *args: Any, **kwargs: Any) -> Any:

355 """Run recomputation without tracing the saved-tensor unpack hook with Dynamo."""

356 return function(*args, **kwargs)

357

358

359def _run_recomputation(frame: _CheckpointFrame, session_id: Any) -> None:

360 """Run and validate a frame recomputation for the given session."""

361 if frame.is_recomputed[session_id]:

362 return

363

364 activation = frame.active_session

365 if activation is not None:

366 _internal_assert(activation.session_id == session_id, "Active session key does not match recomputation key.")

367 previous_activation = _RECOMPUTE_SESSION.get()

368 token = None

369 if activation is not None and previous_activation is not activation:

370 token = _RECOMPUTE_SESSION.set(activation)

371 try:

372 input_context = frame.input_saver.grad_fn

373 args = input_context.get_args(input_context.saved_tensors)

374 try:

375 with _create_recomputation_hooks(frame, session_id), torch.autograd.enable_grad():

376 _run_fn_with_dynamo_disabled(frame.recompute_fn, *args)

377 except _StopRecomputationError:

378 pass

379 finally:

380 if token is not None:

381 _RECOMPUTE_SESSION.reset(token)

382 frame.is_recomputed[session_id] = True

383 frame.check_recomputed_tensors_match(session_id)

384

385

386def _create_checkpoint_hooks(frame: _CheckpointFrame) -> Any:

387 """Create hooks that lazily recompute tensors saved during forward."""

388 def pack_hook(tensor: Any) -> _Holder:

389 """Replace a forward saved tensor with an opaque holder."""

390 holder = _Holder()

391 frame.weak_holders.append(weakref.ref(holder))

392 if frame.metadata_fn is not None:

393 with torch.no_grad():

394 frame.x_metadatas.append(frame.metadata_fn(tensor))

395 return holder

396

397 def unpack_hook(holder: _Holder) -> Any:

398 """Return the corresponding tensor from lazy or prefired recomputation."""

399 activation = frame.active_session

400 if activation is not None:

401 session_id = activation.session_id

402 retain_on_unpack = activation.retain_on_unpack

403 else:

404 session_id = torch._C._current_graph_task_id() # pylint: disable=W0212

405 if session_id == -1:

406 session_id = int(uuid.uuid4())

407 retain_on_unpack = False

408

409 _run_recomputation(frame, session_id)

410 _internal_assert(session_id in holder.handles, "No recomputed tensor was saved for this checkpoint value.")

411 handle = holder.handles[session_id]

412 if handle is None:

413 raise CheckpointError("A checkpoint tensor was unpacked more than once in the same recompute session.")

414 _internal_assert(handle in frame.recomputed[session_id], "The recomputed tensor has already been released.")

415 tensor = frame.recomputed[session_id][handle]

416 if not retain_on_unpack:

417 holder.handles[session_id] = None

418 return tensor

419

420 return torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook)

421

422

423def _is_compiling() -> bool:

424 compiler = getattr(torch, "compiler", None)

425 return bool(compiler is not None and compiler.is_compiling())

426

427

428def _native_checkpoint(

429 function: Callable,

430 *args: Any,

431 context_fn: Callable,

432 preserve_rng_state: bool,

433 determinism_check: str,

434 debug: bool,

435 early_stop: bool,

436 **kwargs: Any,

437) -> Any:

438 """Use the public native API for compile, adapting 2.6/2.7 early-stop."""

439 with set_checkpoint_early_stop(early_stop):

440 return torch_checkpoint(

441 function,

442 *args,

443 use_reentrant=False,

444 context_fn=context_fn,

445 preserve_rng_state=preserve_rng_state,

446 determinism_check=determinism_check,

447 debug=debug,

448 **kwargs,

449 )

450

451

452@dataclass

453class _ExecutionState:

454 """Device, RNG, and autocast state captured before a checkpointed forward.

455

456 Bundles the values that both the forward pass and every later recomputation

457 must agree on, so the generator that owns the checkpoint lifetime can hand

458 them around as one object instead of a dozen loose locals.

459 """

460

461 device_type: str

462 device_module: Any

463 forward_context: Any

464 recompute_context: Any

465 device_autocast_kwargs: Optional[Dict[str, Any]]

466 cpu_autocast_kwargs: Dict[str, Any]

467 device_was_initialized: bool = False

468 device_ids: List[int] = field(default_factory=list)

469 device_states: List[Any] = field(default_factory=list)

470 cpu_state: Any = None

471

472 @classmethod

473 def capture(cls, context_fn: Callable, preserve_rng_state: bool, args: Tuple[Any, ...]) -> "_ExecutionState":

474 """Resolve contexts and capture RNG state before the forward pass runs."""

475 device_type = _infer_device_type(*args)

476 device_module = _get_device_module(device_type)

477 forward_context, recompute_context = _resolve_contexts(context_fn)

478 device_autocast_kwargs, cpu_autocast_kwargs = _get_autocast_kwargs(device_type)

479

480 # Device RNG is captured only when the device is already initialized;

481 # otherwise the device indexes its state lazily at first use.

482 device_initialized = preserve_rng_state and getattr(device_module, "_initialized", False)

483 device_ids, device_states = _get_device_states(device_type, *args) if device_initialized else ([], [])

484

485 return cls(

486 device_type=device_type,

487 device_module=device_module,

488 forward_context=forward_context,

489 recompute_context=recompute_context,

490 device_autocast_kwargs=device_autocast_kwargs,

491 cpu_autocast_kwargs=cpu_autocast_kwargs,

492 device_was_initialized=device_initialized,

493 device_ids=device_ids,

494 device_states=device_states,

495 cpu_state=torch.get_rng_state() if preserve_rng_state else None,

496 )

497

498 def check_device_state_preserved(self, preserve_rng_state: bool) -> None:

499 """Fail when the device was first initialized inside the checkpointed forward."""

500 device_initialized = getattr(self.device_module, "_initialized", False)

501 if preserve_rng_state and not self.device_was_initialized and device_initialized:

502 raise RuntimeError(

503 "The device state was initialized inside a Hyper checkpoint forward, so its initial RNG state "

504 "could not be preserved. Initialize the device before entering checkpoint."

505 )

506

507 def recompute_fn(self, function: Callable, *inputs: Any) -> None:

508 """Restore execution state and rerun the checkpointed ``function``."""

509 function_kwargs, *function_args = inputs

510 with torch.random.fork_rng(

511 devices=self.device_ids,

512 enabled=self.cpu_state is not None,

513 device_type=self.device_type,

514 ):

515 if self.cpu_state is not None:

516 torch.set_rng_state(self.cpu_state)

517 if self.device_was_initialized:

518 _set_device_states(self.device_type, self.device_ids, self.device_states)

519

520 device_autocast_context = _autocast_context(self.device_type, self.device_autocast_kwargs)

521 cpu_autocast_context = _autocast_context("cpu", self.cpu_autocast_kwargs)

522 with device_autocast_context, cpu_autocast_context, self.recompute_context:

523 function(*function_args, **function_kwargs)

524

525

526def _resolve_contexts(context_fn: Callable) -> Tuple[Any, Any]:

527 """Validate and unpack the (forward, recompute) context tuple from ``context_fn``."""

528 contexts = context_fn()

529 if not isinstance(contexts, tuple) or len(contexts) != 2:

530 raise ValueError("context_fn must return a (forward_context, recompute_context) tuple.")

531 return contexts[0], contexts[1]

532

533

534def _checkpoint_without_reentrant_generator(

535 function: Callable,

536 preserve_rng_state: bool,

537 context_fn: Callable,

538 determinism_check: str,

539 early_stop: bool,

540 *args: Any,

541 **kwargs: Any,

542) -> Generator[None, None, None]:

543 """Set up eager checkpoint state around the caller's forward execution."""

544 metadata_fn = _resolve_metadata_fn(determinism_check)

545 state = _ExecutionState.capture(context_fn, preserve_rng_state, args)

546

547 frame = _CheckpointFrame(partial(state.recompute_fn, function), early_stop, metadata_fn)

548 dummy = torch.empty((0,), requires_grad=True)

549 frame.input_saver = _NoopSaveInputs.apply(dummy, kwargs, *args)

550

551 if frame.input_saver.grad_fn is None:

552 yield

553 return

554

555 if _RECOMPUTE_SESSION.get() is not None:

556 raise CheckpointError("Nested checkpoint is not supported during scheduled recomputation.")

557

558 collector = _RECOMPUTE_COLLECTOR.get()

559 if collector is not None:

560 collector.append(frame)

561 try:

562 with _create_checkpoint_hooks(frame), state.forward_context:

563 yield

564 frame.forward_completed = True

565 state.check_device_state_preserved(preserve_rng_state)

566 except BaseException:

567 if collector is not None and frame in collector:

568 collector.remove(frame)

569 raise

570

571

572def _resolve_metadata_fn(determinism_check: str) -> Callable:

573 """Return the metadata extractor implementing ``determinism_check``."""

574 metadata_functions = {_DEFAULT_DETERMINISM_MODE: _default_metadata_fn, "none": lambda tensor: None}

575 if determinism_check not in metadata_functions:

576 raise ValueError(

577 f"determinism_check must be one of {list(metadata_functions)}, but got {determinism_check!r}."

578 )

579 return metadata_functions[determinism_check]

580

581

582def checkpoint(

583 function: Callable,

584 *args: Any,

585 use_reentrant: bool = False,

586 context_fn: Callable = _noop_context_fn,

587 preserve_rng_state: bool = True,

588 determinism_check: str = _DEFAULT_DETERMINISM_MODE,

589 debug: bool = False,

590 early_stop: bool = True,

591 **kwargs: Any,

592) -> Any:

593 """Run Hyper's non-reentrant checkpoint implementation.

594

595 Eager execution always uses this implementation. Compile execution falls

596 back to PyTorch's public non-reentrant checkpoint API.

597 """

598 if use_reentrant is not False:

599 raise ValueError("Hyper checkpoint only supports use_reentrant=False.")

600 if not isinstance(early_stop, bool):

601 raise ValueError(f"early_stop must be bool, but got {type(early_stop).__name__}.")

602 if not isinstance(preserve_rng_state, bool):

603 raise ValueError(

604 f"preserve_rng_state must be bool, but got {type(preserve_rng_state).__name__}."

605 )

606 if not callable(context_fn):

607 raise ValueError("context_fn must be callable.")

608 if _is_compiling():

609 return _native_checkpoint(

610 function,

611 *args,

612 context_fn=context_fn,

613 preserve_rng_state=preserve_rng_state,

614 determinism_check=determinism_check,

615 debug=debug,

616 early_stop=early_stop,

617 **kwargs,

618 )

619 if debug:

620 raise ValueError("debug=True is not supported by Hyper eager checkpoint yet.")

621

622 generator = _checkpoint_without_reentrant_generator(

623 function,

624 preserve_rng_state,

625 context_fn,

626 determinism_check,

627 early_stop,

628 *args,

629 **kwargs,

630 )

631 next(generator)

632 try:

633 result = function(*args, **kwargs)

634 except BaseException:

635 generator.close()

636 raise

637 try:

638 next(generator)

639 except StopIteration:

640 return result

641 generator.close()

642 raise CheckpointError("The internal checkpoint generator yielded more than once.")

643

644

645@contextlib.contextmanager

646def recompute_handle_collector_ctx() -> Iterator[List[Any]]:

647 """Collect opaque checkpoint handles created in this context."""

648 handles = []

649 token = _RECOMPUTE_COLLECTOR.set(handles)

650 try:

651 yield handles

652 finally:

653 _RECOMPUTE_COLLECTOR.reset(token)

654

655

656def recompute_handle(handle: Any, session_id: Any) -> None:

657 """Run one collected checkpoint recomputation ahead of backward."""

658 if not isinstance(handle, _CheckpointFrame):

659 raise ValueError("handle must be produced by recompute_handle_collector_ctx().")

660 _validate_session_id(session_id)

661 activation = _RECOMPUTE_SESSION.get()

662 if (

663 activation is not None

664 and activation.session_id == session_id

665 and activation.retain_on_unpack

666 ):

667 _register_session_frame(handle, session_id, activation)

668 _run_recomputation(handle, session_id)

669 return

670 if activation is not None:

671 raise CheckpointError("recompute_handle cannot enter another active recompute session.")

672

673 _register_session_frame(handle, session_id)

674 with recompute_session_ctx(session_id=session_id, retain_on_unpack=True):

675 _run_recomputation(handle, session_id)

676

677

678def _validate_session_id(session_id: Any) -> None:

679 if session_id is None:

680 raise ValueError("session_id must not be None.")

681 try:

682 hash(session_id)

683 except TypeError as error:

684 raise ValueError("session_id must be hashable.") from error

685

686

687@contextlib.contextmanager

688def recompute_session_ctx(session_id: Any, retain_on_unpack: bool = False) -> Iterator[Any]:

689 """Select the key and retention policy used by checkpoint unpack hooks."""

690 _validate_session_id(session_id)

691 if not isinstance(retain_on_unpack, bool):

692 raise ValueError(f"retain_on_unpack must be bool, but got {type(retain_on_unpack).__name__}.")

693 if _RECOMPUTE_SESSION.get() is not None:

694 raise CheckpointError("Nested recompute session contexts are not supported.")

695 activation = _SessionActivation(session_id, retain_on_unpack)

696 token = _RECOMPUTE_SESSION.set(activation)

697 try:

698 _activate_registered_frames(activation)

699 yield session_id

700 finally:

701 try:

702 _deactivate_session(activation)

703 finally:

704 _RECOMPUTE_SESSION.reset(token)

705

706

707def clear_recompute_session(session_id: Any) -> None:

708 """Release retained recomputation data for a session; repeated calls are safe."""

709 _validate_session_id(session_id)

710 with _SESSION_FRAMES_LOCK:

711 registered_frames = _SESSION_FRAMES.pop(session_id, None)

712 frames = list(registered_frames) if registered_frames is not None else []

713 for frame in frames:

714 frame.clear_session(session_id)

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / checkpoint.py: 85%

400 statements