## Summary
- Continues [@winglian](https://github.com/winglian)'s work from #8181:
a `saved_tensors_hooks` context manager that offloads non-reentrant
checkpoint hidden-state inputs to a pinned CPU buffer pool on a side
stream (`use_reentrant=False`).
- The first two commits are **authored and signed off by Wing Lian**
(`wing@axolotl.ai`); they are the original #8181 patches, rebased onto
current `master`. This follow-up commit addresses review without
rewriting those commits.
- Review follow-up: restore `GradientCheckpointingLayer.__call__` when
no manager is active (HybridEngine train/rollout), skip offloading the
last checkpoint input (`keep_last_count=1`), rename `*_size` knobs to
`*_bytes` / `*_count`, and add tests for the HF signature contract and
keep-last behavior.
- Wires the async offload into DeepSpeed native `cpu_checkpointing`: the
copy machinery is factored into a reusable `_ActivationOffloadEngine`,
which both the HF hooks class and native `CheckpointFunction` /
`non_reentrant_checkpoint` share. Also fixes two pre-existing
`non_reentrant_checkpoint` + `cpu_checkpointing` bugs (inputs emptied
before forward; `saved_data` never restored during recompute).
Original upstreaming context: axolotl-ai-cloud/axolotl#3776, requested
in #8181.
## Test plan
- [x] `pytest
tests/unit/runtime/activation_checkpointing/test_offload_activations.py`
on H200 — HF `saved_tensors_hooks` path (26 passed)
- [x] `pytest
tests/unit/runtime/activation_checkpointing/test_activation_checkpointing.py`
on H200 — native reentrant + new `cpu_checkpointing` offload tests (30
passed)
- [x] `pytest
tests/unit/runtime/activation_checkpointing/test_activation_checkpointing_non_reentrant.py`
on H200 — native non-reentrant + new `cpu_checkpointing` offload test
(49 passed)
- [x] H200 microbenchmark (activation-dominant): async CPU offload
matches blocking's 58% peak-memory reduction at ~5.7% step-time overhead
(vs ~6.8x for blocking), i.e. 6.4x faster than blocking offload
- [ ] CI unit tests for activation checkpointing
Made with [Cursor](https://cursor.com)
---------
Signed-off-by: Wing Lian <wing@axolotl.ai>
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Wing Lian <wing@axolotl.ai>
Co-authored-by: Cursor <cursoragent@cursor.com>
Olatunji Ruwase
858e91eea0