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? |
| manager.mark(hidden_states) | ||
| return _ORIG_GRADIENT_CHECKPOINTING_LAYER_CALL(self, *args, **kwargs) | ||
|
|
||
| GradientCheckpointingLayer.__call__ = _checkpoint_offload_call |
There was a problem hiding this comment.
__call__ will stay changed after exit from the context manager. What will happen when the model is forwarded again without this context manager? I'm thinking about the scenario that same model is used for both training and rollout with HybridEngineRollout.
| manager = _current_manager() | ||
| if (manager is not None and self.training and torch.is_grad_enabled() | ||
| and getattr(self, "gradient_checkpointing", False)): | ||
| hidden_states = args[0] if args else kwargs.get("hidden_states") |
There was a problem hiding this comment.
Here args[0] depends on signature of GradientCheckpointingLayer. Is it possible to add a comment about this assumption, and also add a test to check whether signature changed?
def test_gradient_checkpointing_layer_signature_contract():
"""Pin the HF signature _checkpoint_offload_call depends on.
If HF moves/renames hidden_states, fail loudly so the marker patch gets updated."""
transformers = pytest.importorskip("transformers")
if not hasattr(transformers, "GradientCheckpointingLayer"):
pytest.skip("...")
import inspect
from transformers import GradientCheckpointingLayer
sig = inspect.signature(GradientCheckpointingLayer.forward)
params = [p for p in sig.parameters.values() if p.name != "self"]
assert params, "GradientCheckpointingLayer.forward has no params"
assert params[0].name == "hidden_states", (
f"_checkpoint_offload_call assumes args[0] is hidden_states, but "
f"GradientCheckpointingLayer.forward now starts with '{params[0].name}'. "
f"Update the marker patch in patch_gradient_checkpointing_layer_marker."
)
| loss.backward() | ||
| ``` | ||
|
|
||
| Re-use one manager and wrap each training step (forward+backward) in it. Requires `transformers`: the marker identifying checkpoint inputs is installed on `GradientCheckpointingLayer`, and all other saved tensors pass through untouched. Tune with `use_pin_memory`, `use_streams`, `min_offload_size`, `max_fwd_stash_size` and `max_cpu_buffer_pool_size`. |
There was a problem hiding this comment.
The _size here does not have same meaning, suggest to use _bytes and _count instead for less ambiguity.
| self._tracker: dict[int, tuple[torch.Tensor, torch.device, torch.Size, tuple[int, ...], _BufferKey]] = {} | ||
| self._fwd_stash: dict[int, tuple[torch.Tensor, torch.Event]] = {} | ||
| self._cpu_buffer_pool: dict[_BufferKey, list[torch.Tensor]] = {} | ||
| self._cpu_buffer_pool_size = 0 |
There was a problem hiding this comment.
same here, is it bytes or count? Better clarify by name.
3 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.