Wire async CPU offload into native cpu_checkpointing
Extract the copy machinery from CheckpointHiddenStatesOffload into a
reusable _ActivationOffloadEngine (pinned buffer pool, side stream,
forward stash, keep-last, tracker, offload_input/restore_input). The HF
saved_tensors_hooks class now composes the engine, so its behavior and
tests are unchanged.
DeepSpeed native cpu_checkpointing (non-partition path) uses the shared
engine to offload checkpoint inputs on a side stream instead of blocking
.to('cpu')/.to(cuda). Inputs are offloaded after forward and their GPU
storage emptied; backward restores in place onto the saved tensor to
preserve autograd identity and requires_grad. The engine is process-wide
and lazily built so keep-last/stash span the whole forward and the stream
binds to the selected device; reset() drains it for eval.
Also fix two pre-existing bugs that made non_reentrant_checkpoint +
cpu_checkpointing incorrect: inputs were emptied before the forward (layer
ran on scalars) and saved_data was never restored during recompute. The
save now happens after forward and saved_data is rebuilt before recompute.
Add native cpu_checkpointing correctness tests for both reentrant and
non-reentrant paths and update the cpu_checkpointing config docs.
Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
