Size CSR/CSC indptr to the local partition node range in distributed builds - #746
Draft
kmontemayor2-sc wants to merge 3 commits into
Draft
Size CSR/CSC indptr to the local partition node range in distributed builds#746kmontemayor2-sc wants to merge 3 commits into
kmontemayor2-sc wants to merge 3 commits into
Conversation
…builds Distributed partitioned graph builds size the CSR/CSC index pointer from the global node-id ceiling: GLT's coo_to_csr/coo_to_csc infer the compressed-dimension size as max(id) + 1 when node_sizes is not supplied, and nothing supplies it. Under range partitioning each rank owns a contiguous id slice [lo, hi), so every rank's indptr carries an all-zero prefix of lo entries — the cost grows linearly with partition index and adding partitions cannot shrink it. On a ~1B-node graph over 80 range partitions, the top rank's indptr per edge type is ~960M int64 entries (~7.15 GiB) where ~12M (~92 MiB) suffice. Design convention: the compressed dimension is rebased onto the local range while everything else stays global. - OffsetTopology (gigl/distributed/utils/topology.py): a GLT Topology subclass built from COO with an explicit offset and local node count. indptr has exactly num_nodes + 1 entries; neighbor ids, edge ids, and edge weights keep global values, realigned by one stable sort on the rebased compressed dimension (this also drops stock GLT's second SparseTensor build on the weighted path). to_coo() restores global ids; to_csr()/to_csc() raise rather than leak the local-based pointer. Empty partitions build a zero indptr of the local size. Plain attributes keep ForkingPickler round trips working. - DistDataset._initialize_graph receives the node partition books explicitly (build() initializes the graph before assigning them) and builds one topology per edge type, releasing each edge type's COO tensors as soon as its topology owns the data. Message-passing edge types are validated against the rank's range bounds all-or-nothing (violation raises) because id_select is node-type-scoped; label edge types are decided per type, falling back to a global topology when misaligned. - id_select is installed unconditionally in DistDataset.__init__ so every sampling worker rebuilt from the ipc handle gets it back. Activation is detected at call time from the topology types; active translation subtracts each selected id's own partition lower bound (searchsorted on the RangePartitionBook bounds, partition 0 -> 0), which keeps both the local and the remote RPC sampling branches correct. - Degree aggregation sizes its all-reduce as MAX(offset + len) across ranks and scatters each rank's slice at its own offset (also for the heterogeneous per-anchor merge), so degree tensors stay globally indexed. - _get_padded_labels indexes label indptr by (anchor - offset), raising on anchors below the offset (torch negative indices wrap silently) and padding out-of-range anchors on the allow_non_existant path. - The sampling producer rejects SamplingType.SUBGRAPH for rebased graphs: GLT's subgraph path sends global ids straight to the native samplers, bypassing id_select. Intra-row neighbor order is no longer column-sorted (per-row multisets and edge-id/weight pairing are preserved); GLT APIs relying on that ordering, such as strict negative sampling, are documented unsupported for this topology. Tests cover the local sizing against stock Topology's global sizing, per-row multiset/pairing equivalence, per-element id_select translation, ForkingPickler round trips of Graph(OffsetTopology) and the dataset ipc handle, degree scatter across unequal ranges, offset label lookups, and a two-rank end-to-end sampling test where rank 1 owns a nonzero range and both local and remote one-hop sampling return correct global neighbor ids, edge ids, and weighted-edge choices. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ght storage Rebasing the compressed dimension onto the local partition range leaves a few consumers that still assume globally-sized topologies, and the build still holds edge weights twice. This closes them. The graph-store server reported the rebased topology's native row count, which covers only the local partition range, while get_edge_index returns global ids; get_edge_size now adds the topology offset so both speak the same id space. The shared graph-store sampling backend accepted SamplingType.SUBGRAPH, whose GLT path sends global node ids straight to the native samplers, bypassing the id_select translation and silently returning wrong subgraphs on rebased topologies; the backend now refuses it at construction, matching the colocated producer's rejection. The dataset kept the partition-order edge weights alongside the topology-ordered copy every conversion produces, doubling weight memory for the lifetime of the dataset. It now references the topology-owned storage, so weights are stored once per edge type; dataset.edge_weights is therefore ordered to match each topology's edge order, with the per-edge id-to-weight pairing preserved. The build also initialized each native graph eagerly, though the building process never samples from it and sampling workers rebuild the graph from its ipc handle before initializing their own copy on demand; the eager init only spent an edge-scale transient (the native CPU init runs unique over the indices) per edge type. Native initialization is now deferred to first use. Tests pin the global row bound reported for a rebased partition, the SUBGRAPH rejection, the storage identity and id-to-weight pairing of dataset.edge_weights on inputs whose topology order differs from input order, and the deferred native initialization. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
kmontemayor2-sc
changed the title
Collaborator
Author
|
/all_test |
Contributor
GiGL Automation@ 21:15:33UTC : 🔄 @ 23:19:48UTC : ❌ Workflow failed. |
Contributor
GiGL Automation@ 21:15:35UTC : 🔄 @ 22:35:51UTC : ✅ Workflow completed successfully. |
Contributor
GiGL Automation@ 21:15:36UTC : 🔄 @ 21:17:28UTC : ✅ Workflow completed successfully. |
Contributor
GiGL Automation@ 21:15:36UTC : 🔄 @ 22:54:40UTC : ✅ Workflow completed successfully. |
Contributor
GiGL Automation@ 21:15:37UTC : 🔄 @ 21:24:10UTC : ✅ Workflow completed successfully. |
Contributor
GiGL Automation@ 21:15:40UTC : 🔄 @ 21:24:59UTC : ✅ Workflow completed successfully. |
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.
Problem
When building a distributed graph from a range-partitioned COO edge index,
coo_to_csr/coo_to_cscinfernode_sizesfromint(row.max()) + 1wheneverthe caller does not pass it — and no caller in the graph-build path does. Under
range partitioning, partition r's rows occupy
[lo_r, hi_r), so the resultingindex-pointer array is correctly truncated at the top of the local range but
still carries an all-zero head prefix spanning
[0, lo_r).The waste therefore grows linearly with partition index, and adding partitions
does not reduce it, because the prefix is a function of the global node-id
space rather than of the local edge count.
Measured on an 80-partition build of a ~960M-node heterogeneous graph, per-partition
peak build memory fits:
i.e. roughly 36 GB more on the highest-numbered partition than the lowest, purely
in zero-filled pointer entries. Per edge type on the top partition that is
~7.15 GiB (959,816,917 x 8 B) of array for ~12M meaningful entries. Two nearly
empty edge types in that graph each paid the full ~7 GiB. In a container with a
fixed memory limit this is the difference between a build that fits and one that
is OOM-killed.
What this changes
OffsetTopology, agraphlearn_torch.data.Topologysubclass that sizes theindex-pointer array to the partition's own node range and records the offset it
was rebased by. Nothing is vendored: GLT already exposes
DistDataset(id_select=...)and
Graphaccepts any object satisfying the topology protocol, so the nativeCSR path picks up the smaller array unchanged and its footprint shrinks 1:1.
Sampling ids arrive in the global space, so an
id_selecttranslation isinstalled unconditionally in
__init__rather than passed as a constructorargument — the IPC handles drop constructor state, and every sampling worker
rebuilds the dataset through
ForkingPickler, so a constructor-only installwould silently revert workers to untranslated lookups. Activation is derived at
call time from the topologies themselves so no extra state has to survive IPC.
Rebasing is all-or-nothing for message-passing edge types: a compressed node type
whose range does not align with the partition raises rather than silently falling
back, because
id_selectis scoped per node type. Edge types that can never entera sampling fanout (label edges) fall back to global sizing per type. The subclass
is restricted to
input_layout="COO".The optimisation applies only under range partitioning. Under hash partitioning
owned ids span the global range, rebasing buys nothing, and the code skips it.
The second commit closes the consumers that assumed global-space topologies and
removes two unrelated memory costs found during review:
SUBGRAPHsampling is rejected on a rebased graph in the graph-store backend(the colocated producer was already guarded).
get_edge_sizereturned local-range values whileget_edge_indexreturnedglobal ids; both now report the same space.
copy, permanently doubling weight memory. The topology-owned storage is now the
single copy, and the
edge_weightsproperty documents the ordering contract.lazy_init()is removed. The building process neversamples, workers rebuild the native graph from the IPC handle, and every native
consumer initialises on demand — so it only bought an edge-scale transient.
Correctness
The native samplers index
row_ptr[v]with the incoming id and return zeroneighbors for out-of-range values without a lower-bound check, so a wrong
translation would produce wrong neighborhoods rather than an error. Tests target
that directly rather than asserting wiring:
topology_test.py(19) — local sizing asserted asindptr.numel() == num_nodes + 1for rows in
[lo, hi)withlo >> 0, and demonstrates in the same module thatstock
Topologyallocatesrows.max()+-sized pointers, so the assertion failsagainst pre-change behaviour. Per-row neighbor multisets and edge-id/weight
pairing match stock across CSR/CSC x weighted/unweighted; empty partitions and
global
to_coo()covered.range_partition_rebase_test.py(16) — two-rankmp.spawnend-to-end test whererank 1 owns a nonzero range: local and remote RPC sampling return correct global
neighbor and edge ids, weighted fanout-1 selects the deterministic edge (any
misalignment would pick a zero-weight one), degrees come back globally indexed.
Plus
ForkingPicklerround trips, the all-or-nothing raise, and per-type labelfallback.
degree_test.py(17),data_splitters_test.py(53),dist_server_test.py(44).make format_py,make check_format_py,make type_checkpass.distributed_weighted_sampling_test.pywas updated for the new weight ordering andrelies on CI; the loader and dataset integration suites also run there.
Known limitations
Within-row column ordering is not preserved. This is documented on the subclass;
the strict negative sampler that would depend on it is not reachable from this
codebase, where sampling configs are constructed with
with_neg=False.Deferred, reasoned in review rather than skipped: sizing by max-active compressed
id (trims within-partition slack, but would decouple pointer length from the
partition book that degree aggregation relies on);
_node_idsmaterialization(pre-existing, orthogonal); dropping
edge_idswhen no edge features exist.Separately, the native CSR build eagerly computes a column count via
at::_uniqueover the whole indices tensor, which the CPU path never reads.Avoiding it needs an optional argument on the C++ binding and a native rebuild,
so it is left for a follow-up.