Skip to content

DistABLPLoader: vectorized label remap + optional AnchorLabels edge-list output - #681

Draft
kmontemayor2-sc wants to merge 20 commits into
mainfrom
kmonte/ablp-vectorized-labels-and-list-output
Draft

DistABLPLoader: vectorized label remap + optional AnchorLabels edge-list output#681
kmontemayor2-sc wants to merge 20 commits into
mainfrom
kmonte/ablp-vectorized-labels-and-list-output

Conversation

@kmontemayor2-sc

@kmontemayor2-sc kmontemayor2-sc commented Jun 25, 2026

Copy link
Copy Markdown
Collaborator

Summary

Reworks how DistABLPLoader remaps Anchor-Based Link Prediction (ABLP) labels
from global node ids to subgraph-local indices, and adds an optional edge-list
output format for those labels.

At prod scale, we see 60% speedup for ABLP next time for colcated and 30% for graph store (e.g. 40 -> 15 and 60 -> 40).

Internally, we see small but

  • Vectorized label remap. Label remapping now runs as a single
    sorted-membership join (searchsorted over the node map) instead of a
    per-anchor Python loop. The result is loss-equivalent to the previous output
    (same (anchor, label) pairs; within-anchor order is unspecified, and the
    ABLP contrastive loss is order-invariant). We observe meaningful sampling/
    collation performance improvements at production scale.

  • AnchorLabels edge-list container + use_list_output flag.
    DistABLPLoader(..., use_list_output=True) returns labels as an
    AnchorLabels edge list (two co-indexed [E] tensors,
    anchor_index / label_index) that the loss can index directly with no
    padding or per-anchor Python loop. use_list_output=False (default) preserves
    the existing ragged dict[int, torch.Tensor] output, so this is backward
    compatible. AnchorLabels.to_dict() recovers the dict form.

  • Public API. AnchorLabels is exported from gigl.distributed.

  • Examples. The link-prediction training examples (homogeneous +
    heterogeneous, colocated + graph-store) are updated to consume the
    AnchorLabels edge-list output.

Tests

  • Equivalence of the vectorized remap output against constructed expected values,
    covering empty / fully-padded / duplicate-label / multi-edge-type cases.
  • use_list_output=True vs =False produce equivalent labels.
  • CUDA device-placement regression test for the remap kernel.
  • Guard: a non-unique node→global map raises ValueError.

Notes

  • Labels are loss-equivalent to the prior output, not order-identical.
  • Documentation-only follow-ups: worked examples for both output formats in the
    DistABLPLoader docstring and AnchorLabels shape docs.

kmontemayor and others added 12 commits June 25, 2026 14:57
…racle

Factor the inline per-anchor label-remap loop in DistABLPLoader._set_labels
into a module-level function _loop_set_labels. The new function is a
behavior-preserving extraction: _set_labels delegates to it, producing
identical output. _loop_set_labels will serve as the equivalence oracle
for the vectorized kernel added in the next task.

Also imports PADDING_NODE from gigl.utils.data_splitters (used by the
vectorized kernel in the next task) and adds the contract test file
tests/unit/distributed/vectorized_set_labels_test.py.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…oop oracle

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ls kernel

Add frozen dataclass AnchorLabels (anchor_index, label_index, num_anchors)
with to_dict() bridge; thin wrapper _remap_one_label_tensor_edge_list over
the shared _membership_remap; and edge_list_set_labels driver. Proves
to_dict() reproduces _loop_set_labels bit-for-bit via parametrized equivalence
tests (11 new tests, all classes pass).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…eous training

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…eneous training

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…test

Line-wrapping only (ruff format) for the edge-list label reads in the four
link-prediction training examples and the loader equivalence test. No logic
change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@kmontemayor2-sc
kmontemayor2-sc force-pushed the kmonte/ablp-vectorized-labels-and-list-output branch from 292f4d4 to bd97d03 Compare June 25, 2026 14:57
kmontemayor and others added 2 commits June 25, 2026 16:18
…ctors

**Behavioral change:** Drop the composite-key stable argsort from
`_membership_remap` (the `composite_key = anchor_kept * (num_nodes + 1) +
local_idx` / `torch.argsort(composite_key, stable=True)` block). The pair
stream is now emitted in column-visit (row-major masked flatten) order rather
than ascending-local-index (torch.nonzero) order.

**Why this is safe:** `RetrievalLoss` (`gigl/nn/loss.py`) is
`CrossEntropyLoss(reduction="sum")` over a diagonal-targeted score matrix that
masks collisions by id VALUE, not position. It is invariant to any joint
permutation of `(anchor_index, label_index)` pairs. The only constraint is
co-indexing (pair k stays intact) and per-anchor grouping, both of which are
preserved. Within-anchor label order was an implementation artifact of the loop
oracle -- never a loss requirement.

**Why the dict path is unaffected:** `anchor_of_entry` is built as
`arange(N).repeat_interleave(M)` (row-major), so the masked flatten is already
non-decreasing in anchor_index. `bincount`/`split` requires only contiguous
grouping by anchor, which row-major flatten guarantees without any argsort.

**Readability refactors in `_membership_remap`:**
- Rename terse tensors: `valid` -> `is_present`, `found` -> `is_exact_match`,
  `positions` -> `sorted_positions`, `local_idx` -> `local_index`,
  `anchor_kept` -> `anchor_of_matched`
- Add numbered step comments explaining the searchsorted membership lookup

**Readability refactors in outer kernels:**
- Add `_remap_group` helper to collapse the ~40-line duplicated positive/
  negative per-edge-type loops in both `vectorized_set_labels` and
  `edge_list_set_labels` (uses TypeVar for generic return type)
- `_sorted_for` closure pattern preserved, inlined per-kernel (memoizes
  torch.sort across pos/neg edge types of the same node type)

**Other improvements:**
- Add doctest to `AnchorLabels` showing a 2-anchor case + `to_dict()` round-trip
- Update `AnchorLabels` docstring: document column-visit order and
  loss-permutation-invariance rationale
- Update all docstrings to drop "bit-for-bit"/"torch.nonzero order" language

**Test contract relaxation (tests remain non-vacuous):**
- `vectorized_set_labels_test.py` / `edge_list_set_labels_test.py`:
  `_assert_label_dicts_equal` -> `_assert_label_dicts_set_equal` (uses
  `sorted()` per anchor; still catches membership errors + multiplicity)
- `test_unsorted_node_map_exact_order` -> `test_unsorted_node_map_correct_membership`:
  asserts SET {0, 2} instead of sequence [0, 2]; docstring explains column
  vs ascending-local order difference
- `RemapOneEdgeListTest.test_unsorted_node_map_nontrivial_sort_perm` ->
  `test_unsorted_node_map_correct_membership`: pins column order [2, 0] and
  explains it is SET-equal to [0, 2]; verifies sort_perm mapping is correct
- `dist_ablp_neighborloader_test.py`: `_ordered_global_pairs` ->
  `_global_pair_set` (sorts within anchor); `_collect_homogeneous_labels`
  docstring updated; the in-process exact-tensor assertion
  (`label_index == cat(dict values)`) is preserved -- both paths draw from the
  same `_membership_remap` pair stream so they remain identical
- CUDA device test docstring: remove "stable-argsort tie-break" reference;
  keep the duplicate [15,15] row (still tests duplicate handling + device placement)

Device fix preserved: `anchor_of_entry` is still built on `label_tensor.device`.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…er-identical

After dropping the order-reproduction argsort, per-anchor label order is
column-visit order, not the loop's nonzero order. The (query, label) pairs are
unchanged and the contrastive loss is order-invariant, so the read is equivalent
for training; the comments now say so instead of implying byte-identical order.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@kmontemayor2-sc
kmontemayor2-sc force-pushed the kmonte/ablp-vectorized-labels-and-list-output branch from 486e670 to d42d0df Compare June 26, 2026 19:24
GiGL main resolved ABLP labels with a per-anchor Python loop. This replaces it
with one vectorized kernel and a single output type, plus a thin dict view for
backward compatibility.

- Resolve labels with a sorted-membership join (_membership_remap: sort the node
  map once, searchsorted the label ids, keep exact matches) instead of the
  O(N_anchors * M * N_nodes) per-anchor loop. Delete _loop_set_labels (no
  production callers) and the redundant dict-producing kernel.
- Collapse the remap to two functions, _membership_remap and edge_list_set_labels;
  the dict path is just AnchorLabels.to_dict(), selected by use_list_output on
  DistABLPLoader (default False keeps the ragged dict). Drop the _remap_group
  callback indirection, the _LabelT TypeVar, and the vestigial
  supervision_edge_types parity param.
- AnchorLabels stores labels as two parallel (anchor_index, label_index) tensors
  so the loss can index them directly; within-anchor order is unspecified because
  the ABLP contrastive loss is order-invariant over the pairs.
- Make the __debug__ unique-node-map check a cheap adjacent-difference test on the
  already-sorted map (it ran torch.unique every batch, and GiGL is not launched
  with -O).
- Tests assert against constructed expected values and per-anchor label SETS, not
  within-anchor order. Examples consume the edge-list directly.

Behavior-preserving: per-anchor label sets are unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@kmontemayor2-sc
kmontemayor2-sc force-pushed the kmonte/ablp-vectorized-labels-and-list-output branch from d42d0df to 641b24e Compare June 26, 2026 21:13
kmontemayor and others added 3 commits June 26, 2026 21:47
…ions

Comments and docstrings only -- no behavior change (verified: the AST with
docstrings stripped is identical to the prior commit for every touched file).

- Docstrings lead with why over what; the membership-remap algorithm is shown as
  a small worked example (node map + [N_anchors, M] padded label tensor -> pairs)
  with labeled steps instead of dense prose.
- Inline comments reference those docstring steps rather than floating free.
- Tensor dimensions annotated throughout (N_anchors / M / N_nodes / K / E),
  including the K (non-padding candidates) vs E (matched pairs) distinction.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…uniqueness guard

- Add a use_list_output=True worked example to DistABLPLoader.__init__ that
  mirrors the use_list_output=False example (same graph, AnchorLabels values).
- Document AnchorLabels.to_dict()'s non-decreasing anchor_index precondition.
- Promote the _membership_remap duplicate-node-map guard from a __debug__
  assert to an always-on ValueError so it still fires under `python -O`; update
  the test to expect ValueError.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Comments should describe the code as it is, not as a delta from a prior or
never-committed version. Reword/remove comments that only parse if the reader
knows the code's development history:

- Drop the "Always-on (not __debug__) ... python -O" framing on the node-map
  uniqueness guard; keep the precondition + empty-slice safety rationale.
- Fix a stale _membership_remap docstring line still describing a __debug__
  assertion (it is now an always-on ValueError).
- Remove "with no argsort" / "no argsort needed" (presupposed the dropped
  order-reproduction argsort); state the grouped-by-anchor property directly.
- Reword "a sorted-membership join rather than a per-anchor scan ... trades a
  broadcast-compare for a search" to describe the algorithm and its complexity
  as-is.
- Drop "rather than maintaining a second kernel" in _set_labels.
- Same scrub in edge_list_set_labels_test.py (module docstring + guard test).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@kmontemayor2-sc kmontemayor2-sc changed the title Kmonte/ablp vectorized labels and list output DistABLPLoader: vectorized label remap + optional AnchorLabels edge-list output Jun 29, 2026

@mkolodner-sc mkolodner-sc left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks a lot for looking into this Kyle! This is a super exciting change, and looking forward to realizing the throughput and general complexity improvements from this work.

Did a first pass here through the production logic and left some comments.

Comment on lines +113 to +114
num_anchors (int): Total number of anchors ``N_anchors`` (rows of the
source padded label tensor), including anchors with no labels.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need the num_anchors data_class? This should already be available in Data/HeteroData object right

Comment on lines +111 to +112
anchor_index (torch.Tensor): ``[E]`` long tensor of local anchor rows.
label_index (torch.Tensor): ``[E]`` long tensor of local label node ids.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm wondering if we can represent this instead as a [2, E] edge tensor, since that is essentially what it seems to be used for here. This would make it so we'd only need one .to(device) call. If we can make this just one tensor instead of two and we can remove the num_anchors field as well (which already exists in the HeteroData), I'm wondering if we can remove the need for a dataclass here all together.

The benefit is that the resulting labeled object would just be an edge tensor and would feel more PyG-esque and natural to use, rather than a GiGL-specific dataclass for labeled objects. A [2, E] (labeled) edge index tensor also feels similar to what we use for other objects such as the message passing edge index tensor or PPR edges, which we don't have custom dataclasses for.

logger = Logger()


@dataclass(frozen=True)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we move the newly introduced remapping logic to a separate utility file, maybe creating a new ablp_loader utility file? Might help make our code a bit easier to read/follow-through s.t. if users are clicking through and trying to understand what/how to use ABLP loader at a high level, they don't have to be exposed to this internal label mapping logic as the first thing they see.

Comment on lines +124 to +125
Every anchor ``0..num_anchors-1`` receives a key; anchors with no labels
map to an empty ``long`` tensor on the same device as ``label_index``.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we add a note here that this we may eventually want to remove this path in the future?

# Set self._shutdowned right away, that way if we throw here, and __del__ is called,
# then we can properly clean up and don't get extraneous error messages.
self._shutdowned = True
self._use_list_output = use_list_output

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we add a warning here that this will be deprecated in a future release? Unless you see some benefit to us keeping this around.


num_labels = int(label_tensor.size(1))
flat = label_tensor.reshape(-1) # [N_anchors * M] before the pad mask
# step 0 (see docstring example): tag each flattened entry with its anchor row.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similarly, in-line comments are useful, but these feel a bit too long and hard to follow for me (we have step 0, step 0 continued, precondition for step 1, step 1 ...)

"""
sorted_cache: dict[NodeType, tuple[torch.Tensor, torch.Tensor]] = {}

def _sorted_for(node_type: NodeType) -> tuple[torch.Tensor, torch.Tensor]:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sorted_for is a bit ambiguous of a name, can we make this clearer? Ditto for the other functions here (i.e. "remap" is also vague.

Comment on lines +247 to +249
# searchsorted returns N_nodes for an id larger than every entry, which would
# gather out of bounds at step 2; clamp it back into range.
sorted_positions = sorted_positions.clamp_(max=num_nodes - 1)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In what scenarios do we expect this to happen?

Comment on lines +264 to +265
anchor_of_matched.to(to_device).to(torch.long),
local_index.to(to_device).to(torch.long),

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need .to(to_device).to(torch.long)? Asking since we also call to_device in the training example loops

)


def edge_list_set_labels(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: I think some of the new function and var names introduced in this PR are a bit non-intuitive for me. A few examples

  • edge_list_set_labels might be easier to make sense of if we put the action first i.e. set_edge_list_labels
  • node_local_to_global_by_type could be better and more aligned with other vars as local_id_to_global_id_by_node_type
  • sorted_node and sorted_perm could have easier names to follow as well

@kmontemayor2-sc

Copy link
Copy Markdown
Collaborator Author

/all_test

@github-actions

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 16:48:34UTC : 🔄 Integration Test started.

@ 18:05:25UTC : ✅ Workflow completed successfully.

@github-actions

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 16:48:35UTC : 🔄 Lint Test started.

@ 16:56:39UTC : ✅ Workflow completed successfully.

@github-actions

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 16:48:38UTC : 🔄 E2E Test started.

@ 18:17:36UTC : ❌ Workflow failed.
Please check the logs for more details.

@github-actions

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 16:48:39UTC : 🔄 Python Unit Test started.

@ 17:56:56UTC : ❌ Workflow failed.
Please check the logs for more details.

@github-actions

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 16:48:41UTC : 🔄 C++ Unit Test started.

@ 16:50:36UTC : ✅ Workflow completed successfully.

@github-actions

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 16:48:42UTC : 🔄 Scala Unit Test started.

@ 16:57:26UTC : ✅ Workflow completed successfully.

…dge types

_set_labels built the local->global node map from
self._supervision_edge_types[2]. Those edge types are stored reversed when
edge_dir="in" (dist_ablp_neighborloader.py:486-490), so that index is the anchor
node type there, while the label edge types reaching collation stay
outward-facing. Every heterogeneous edge_dir="in" batch therefore raised
KeyError on its supervision node type.

The pre-refactor code inserted both endpoints of each supervision edge type,
which masked the direction ambiguity. Deriving the map from the label edge type
keys that _remap_labels_by_edge_type actually looks up drops the edge_dir
coupling entirely.

Add heterogeneous coverage for both output formats across both edge dirs. The
existing equivalence test is homogeneous and edge_dir="out", so it could not
catch this. Anchors and supervision nodes use distinct node types, so a label
remapped against the wrong node map yields a wrong global id instead of passing
silently.

Fixes the 5 unit errors and the het-dblp-sup{,-gs} e2e trainer failures in
https://github.com/Snapchat/GiGL/actions/runs/30379971583

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@kmontemayor2-sc

Copy link
Copy Markdown
Collaborator Author

/all_test

@github-actions

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 23:42:16UTC : 🔄 Python Unit Test started.

@ 01:01:08UTC : ✅ Workflow completed successfully.

@github-actions

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 23:42:18UTC : 🔄 Scala Unit Test started.

@ 23:51:08UTC : ✅ Workflow completed successfully.

@github-actions

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 23:42:19UTC : 🔄 Lint Test started.

@ 23:51:09UTC : ✅ Workflow completed successfully.

@github-actions

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 23:42:21UTC : 🔄 C++ Unit Test started.

@ 23:44:57UTC : ✅ Workflow completed successfully.

@github-actions

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 23:42:22UTC : 🔄 E2E Test started.

@ 01:07:18UTC : ✅ Workflow completed successfully.

@github-actions

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

GiGL Automation

@ 23:42:26UTC : 🔄 Integration Test started.

@ 01:05:49UTC : ✅ Workflow completed successfully.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants