Repository navigation
Assert offload-state removal only when offload leaves the accelerator - #8738
Merged
Merged
Conversation
validate_device requires the contiguous grad buffer to be gone once it is reported on cpu, since a cpu device means offload_states() moved it to the host tier. On the cpu accelerator the offload target is the accelerator itself, so the buffer is never moved and the assertion fails for the parametrizations that include it. Gate the check on the offload actually crossing memory tiers. The new flag sits next to allocator_backed_memory_stats, which gates the memory-delta assertions for the same reason (deepspeedai#8409). On CUDA it is always true, so the assertion behaves exactly as before; on cpu the state falls through to the device comparison, which holds. Measured on cpu/gloo with 2 ranks: 112 of the 128 parametrizations pass without this change and the 16 failures are all this one assertion; with the guard the whole file passes. Reported in deepspeedai#8381, where the cpu run exercises these tests. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
sfc-gh-truwase
approved these changes
Oct 6, 2026
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Oct 6, 2026
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Oct 9, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
validate_deviceintests/unit/v1/zero/test_offload_states.pyrequires the contiguous grad buffer to be gone once it is reported on cpu, because a cpu device meansoffload_states()moved it to the host tier:On the cpu accelerator the offload target is the accelerator, so
state_device[state] == torch.device("cpu")holds for everything, yet the buffer was never moved anywhere — the assertion fails for every parametrization that includes it.This is the device-state counterpart of #8409, which gated the memory-delta assertions on
allocator_backed_memory_statsfor the same reason. Add a siblingoffload_leaves_acceleratorflag (next to that one) and require it for the removal check.Behavior
offload_leaves_acceleratorTrueFalseValidation (executed on real hardware)
20-core x86_64 CPU without AVX512-FP16, gloo backend, 2 ranks (
LOCAL_SIZE=2),-n 4:16 failed, 112 passed— and all 16 failures are this single assertion (State OffloadStateTypeEnum.contiguous_grad_buffer must be removed after offload_states()).128 passed, 0 failed.So the cpu accelerator keeps the coverage of the other 112 parametrizations instead of losing them to a blanket skip.
The enclosing cpu CI run is #8381, which exercises this file at multi-rank and reports the same result (
128 PASSED).