Chunk the degree-tensor all-reduce to bound peak device memory - #766
Draft
mkolodner-sc wants to merge 1 commit into
Draft
Chunk the degree-tensor all-reduce to bound peak device memory#766mkolodner-sc wants to merge 1 commit into
mkolodner-sc wants to merge 1 commit into
Conversation
The single-shot all-reduce moved the whole padded int64 degree tensor to the process-group device. For large graphs that tensor is several GiB, so together with the collective buffer it can OOM small-memory GPUs during dataloader setup, even though the reduced result is only ever used on CPU. Reduce one fixed-size slice at a time: keep the padded tensor on CPU and, per slice, copy it to the device, all-reduce, and copy the result back. A chunked element-wise SUM is bit-identical to a whole-tensor SUM (every rank pads to the same size and uses identical slice bounds); the over-counting correction is unchanged. This runs once at loader setup, so there is no per-batch cost.
mkolodner-sc
force-pushed
the
mkolodner-sc/chunk-degree-allreduce
branch
from
September 9, 2026 21:27
fcff48d to
7c24c08
Compare
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
1 participant
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.
What
_all_reduce_single_degree_tensorpadded the degree tensor to the global maxsize, moved the whole
int64tensor to the process-group device, andall-reduced it in a single call.
Why
For large graphs that tensor is several GiB. Together with the collective's
buffer it can OOM small-memory GPUs during dataloader setup, even though the
reduced result is only ever used on CPU.
How
Reduce one fixed-size slice (
_DEGREE_ALLREDUCE_CHUNK) at a time: keep thepadded tensor on CPU and, per slice, copy to device -> all-reduce -> copy back.
A chunked element-wise SUM is bit-identical to the whole-tensor SUM (every rank
pads to the same size and uses identical slice bounds); the over-counting
correction is unchanged. Runs once at loader setup, so there is no per-batch
throughput cost.