Non-reentrant activation checkpoint CPU offload - #8181
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e09c9d5e62
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| self._pending_cpu_buffers: list[tuple[_BufferKey, torch.Tensor, torch.Event]] = [] | ||
|
|
||
| # cpu-like accelerators have no streams; _pack_tensor never offloads there | ||
| self.s0 = get_accelerator().current_stream() |
There was a problem hiding this comment.
When the manager is constructed outside the stream that later runs the training step, or is reused under a different stream, self.s0 remains the construction-time stream. Consequently, _pack_tensor makes the offload stream wait on the wrong producer stream, while _restore_tensor makes the wrong consumer stream wait for the H2D copy; the actual training stream can therefore read data before a transfer completes and silently produce incorrect results. Resolve the current compute stream when entering the context or whenever a tensor is packed or unpacked.
Useful? React with 👍 / 👎.
| if (gpu_tensor.storage_offset() != storage_offset or gpu_tensor.stride() != stride): | ||
| gpu_tensor = torch.as_strided( | ||
| gpu_tensor, | ||
| size=shape, | ||
| stride=stride, | ||
| storage_offset=storage_offset, |
There was a problem hiding this comment.
For a marked hidden state that is a view with a nonzero storage_offset and has fallen out of the forward stash, _empty_cpu_like allocates backing storage starting at offset zero, but this restoration reapplies the original positive offset. For common slices such as base[1:], the resulting as_strided view extends beyond the newly allocated storage and raises a runtime error during backward. Allocate backing storage that includes the prefix or normalize the restored tensor to offset zero.
Useful? React with 👍 / 👎.
| @@ -0,0 +1,333 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
There was a problem hiding this comment.
The reviewed commit message has no Signed-off-by: trailer, so it does not satisfy this repository's mandatory commit-signoff requirement. Recreate the commit with git commit --signoff using the configured author identity.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
saved_tensors_hooks context manager that offloads checkpoint hidden-state inputs to a pinned CPU buffer pool on a side stream; works with use_reentrant=False. Adapted from axolotl PR 3776. Signed-off-by: Wing Lian <wing@axolotl.ai>
Address review: the compute stream can differ from the construction-time stream, so resolve it via a property at each sync point; restoring an offloaded view reapplied its storage_offset onto offset-0 pool storage. Signed-off-by: Wing Lian <wing@axolotl.ai>
e09c9d5 to
f32c5d1
Compare
|
Thank you very much, @winglian, for porting it to Deepspeed @tohtana, @tjruwase - for a better context - when Wing and I talked I said my implementation of activation offload was a very inefficient blocking prototype implementation. When Wing said he created an efficient 0-cost (full overlap) I have asking him if he could contribute his code to Deepspeed. |
|
@winglian, I think the last transformer block's activation should not be offloaded as it's almost immediately used on the first backward call and it could be exposed and blocking operation if the head's bwd compute is very fast - do you have that already and I missed it in the code? Do you think there should be an arg controlling how many last activations to keep, probably with 1 being the default - is there a use-case where more than 1 would be beneficial? |
2 participants
upstreaming feature to support non-reentrant checkpointing from axolotl: axolotl-ai-cloud/axolotl#3776
@sfc-gh-sbekman
saved_tensors_hooks context manager that offloads checkpoint hidden-state inputs to a pinned CPU buffer pool on a side stream; works with use_reentrant=False.