Add pin_empty helper for empty scratch destinations - #8332
Add pin_empty helper for empty scratch destinations#8332sfc-gh-truwase wants to merge 3 commits into
Conversation
The repeating pin_memory(torch.empty(...), make_copy=False) idiom hid a full-size pageable template; pin_empty takes shape/dtype only. Signed-off-by: tunji-ruwase_snow <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
![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: bd38fec2d4
ℹ️ 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".
| pins = get_active_native_pinned_memory() | ||
| if pins is not None: | ||
| return pins.pin_empty(template, shape) | ||
| return self._torch_empty_pinned(template, shape) |
There was a problem hiding this comment.
When MPS is active without the native backend, both new helpers reach this inherited _torch_empty_pinned path, which calls tensor.new_empty(..., pin_memory=True). MPS_Accelerator intentionally overrides _torch_pin_memory because PyTorch pinning raises on Apple Silicon (accelerator/mps_accelerator.py:233-236), but it does not override _torch_empty_pinned; consequently pin_empty and the newly routed ZeRO-1/2 state-offload paths now raise instead of returning an ordinary unified-memory CPU buffer. Add an MPS-specific no-op allocation override analogous to the CPU implementation.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
This MPS issue might better be addressed in separate PR. @sfc-gh-truwase if you believe this could be an issue, maybe open an issue for better tracking.
@PKUWZP for MPS related.
There was a problem hiding this comment.
@delock I agree to defer to a separate PR, especially given the ongoing MPS work stream.
pin() sizes the allocation from the input tensor, so a 0-element dtype template cannot go through pin(); both paths now use _new_locked, and pin_empty is a thin pin_memory(alloc_shape=) wrapper. Signed-off-by: tunji-ruwase_snow <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
alloc_shape silently dropped make_copy; pin_empty now uses a private _pin_fresh path, requires a host device and dtype, and skips CPU torch accounting. Signed-off-by: tunji-ruwase_snow <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
2 participants
Summary
get_accelerator().pin_empty()/pin_empty_like()so the repeatingpin_memory(torch.empty(...), make_copy=False)idiom from Route ZeRO/SuperOffload pin sites through accelerator pin_memory #8256 lives in one place (review note).offload_stateshp/lp pins,offload_optimizer_states, and the SuperOffload worker grad buffer through the helper.Test plan
pre-commit run --fileson touched pathspytest tests/unit/v1/accelerator/test_accelerator.py tests/unit/v1/pin_memory/test_offload_route.py tests/unit/v1/pin_memory/test_pin_memory.py— 37 passed, 5 skipped (CPU host; GPU-only offload-route cases skipped)Made with Cursor