Repository navigation
Kmonte/memexp e6 - #794
Draft
kmontemayor2-sc wants to merge 6 commits into
Draft
Kmonte/memexp e6#794kmontemayor2-sc wants to merge 6 commits into
kmontemayor2-sc wants to merge 6 commits into
Conversation
Port of the loader change from #783. The ids, features, quantized features and labels are now copied batch by batch into one preallocated tensor, and each source batch is released after it is copied. Before, `tf.concat` held every batch plus a second full-size output at the same time. The helper takes an `axis`, so the edge ids, which load as (2, num_edges) batches, use the same path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: dsaini2-sc <dsaini2-sc@snapchat.com> Co-authored-by: swang7-sc <swang7-sc@snapchat.com> Co-authored-by: mkolodner-sc <mkolodner-sc@snapchat.com>
Add `_concatenate_partitioned_chunks`. It builds each output in one preallocated tensor, copies the chunks into it one by one, and frees each chunk after the copy. Before, `torch.cat` (and `torch.stack` over two cats) kept every chunk, the cat temporaries and the stacked copy alive at the same time. Zero-row chunks are skipped, the same as `torch.cat` ignores legacy 1-D empty tensors. Sites changed: - `DistRangePartitioner._partition_edge_index_and_edge_features`: the edge index (stack of two cats → one (2, N) tensor), edge features, quantized edge features and edge weights, in one pass. - `DistPartitioner._partition_node_features_and_labels`: the node ids, node features, quantized node features and labels, in one pass. Also drop `input_parts` / `input_data`, which kept the unpartitioned tensors alive after `del node_features`. - `DistPartitioner._partition_label_edge_index`: the stack of two cats → one (2, N) tensor. Left alone: - The range node-feature reorder (`feats[argsort(ids)]`) still briefly holds about 2x the features. It is a reorder, not a concat. - The non-range `DistPartitioner._partition_edge_index_and_edge_features`, which has the same pattern but is not on the range path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: dsaini2-sc <dsaini2-sc@snapchat.com> Co-authored-by: swang7-sc <swang7-sc@snapchat.com> Co-authored-by: mkolodner-sc <mkolodner-sc@snapchat.com>
`_concatenate_tf_tensors_to_torch` now allocates a CPU output with `allocate_preshared`. A large output goes straight into shared memory, so the `share_memory_()` that later hands it to the partitioner does not copy it. Before, the full feature tensor briefly existed twice: the private copy and its new shared copy. Small outputs stay plain tensors, and non-CPU outputs still use `torch.empty`. Tests check that: - the output is shared - `share_memory()` keeps its storage - the values match `torch.cat` Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: dsaini2-sc <dsaini2-sc@snapchat.com> Co-authored-by: swang7-sc <swang7-sc@snapchat.com> Co-authored-by: mkolodner-sc <mkolodner-sc@snapchat.com>
- Remove the spawn-child hand-off test. It passed without the shared-memory change, because sending a tensor to another process moves it into shared memory anyway, so it could not catch a regression. - Add a test for all-zero-row chunks. A rank that receives nothing for a type shapes its outputs from the first chunk, and nothing tested that path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: dsaini2-sc <dsaini2-sc@snapchat.com> Co-authored-by: swang7-sc <swang7-sc@snapchat.com> Co-authored-by: mkolodner-sc <mkolodner-sc@snapchat.com>
Cut comments that restated the code, keep the reasons, and list the test chunk fields by index. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: dsaini2-sc <dsaini2-sc@snapchat.com> Co-authored-by: swang7-sc <swang7-sc@snapchat.com> Co-authored-by: mkolodner-sc <mkolodner-sc@snapchat.com>
…oners Move the two near-identical helpers into `gigl.utils.concat.concatenate_chunks`. It consumes a list of chunks, preallocates each output once, copies the chunks in, and frees each chunk after its copy. It also: - skips zero-row chunks and shapes the outputs from the first chunk with rows, like `torch.cat`; - checks each field's dtype and shape off the concat axis, because `copy_` would silently cast or broadcast; - stacks a tuple of fields into one tensor, e.g. `(0, 1)` builds a `(2, N)` edge index; - with `preshare=True`, allocates CPU outputs with `allocate_preshared`. The loader keeps presharing its outputs, and the partitioners keep plain `torch.empty` outputs, so behavior does not change. The partitioner call sites shrink to one call each. The helper tests move to `tests/unit/utils/concat_test.py`. A new loader-level test checks that `load_as_torch_tensors` returns shared outputs that `share_memory` keeps in place. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: dsaini2-sc <dsaini2-sc@snapchat.com> Co-authored-by: swang7-sc <swang7-sc@snapchat.com> Co-authored-by: mkolodner-sc <mkolodner-sc@snapchat.com>
This branch has not been deployed
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.
Scope of work done
Where is the documentation for this feature?: N/A
Did you add automated tests or write a test plan?
Updated Changelog.md? NO
Ready for code review?: NO