Repository navigation
Compute a missing hop_distance on the padded Graph Transformer k-hop path - #850
Draft
kmontemayor2-sc wants to merge 99 commits into
Draft
kmontemayor2-sc wants to merge 99 commits into
kmontemayor2-sc wants to merge 99 commits into
Conversation
…eroData GraphTransformerEncoder.forward built a new HeteroData from the projected features and called heterodata_to_graph_transformer_input, which converts it with to_homogeneous. Dynamo cannot trace either step, so torch.compile of the whole encoder failed. Add tensors_to_graph_transformer_input, which takes x/edge_index/edge_attr dicts and the graph-level relative-encoding matrices and concatenates them in the same node and edge type order to_homogeneous uses. heterodata_to_graph_transformer_input is now a thin wrapper over it, and the encoder calls it directly. Node features are gathered with a masked gather instead of bool-mask indexing, which drops a host sync. Outputs are unchanged: torch.equal against the previous implementation over 432 transform configurations (k-hop/PPR, in/out, hop order, relation edges, explicit anchors, multiple node and edge types). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Brings in the per-seed-role PPR relation filter (#799). The encoder no longer builds a projected HeteroData, so the expansion flags and typed-PPR channel masks now reach the transform as explicit tensors_to_graph_transformer_input arguments (ppr_is_expanded_dict, ppr_relation_channel_masks); heterodata_to_graph_transformer_input reads them from the batch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A missing-attribute lookup on HeteroData runs a regex, and torch.compile guards on the regex cache, which broke the compiled-forward test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Matches the HeteroData wrapper, which already skips stores that carry only edge_attr instead of raising AttributeError. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- On CUDA, khop / anchor_only configs with no attention bias build their sequences with a fixed-shape layout (gigl/transforms/static_khop.py) and run the layers on one [tokens, dim] buffer with varlen attention (gigl/nn/packed_attention.py), so no FLOPs go to padding. - The encoder picks the path per forward: packed when the config supports it, on CUDA, and not training with attention dropout; otherwise the existing padded path runs unchanged. Both give the same embeddings and gradients up to float error. - GraphTransformerEncoderLayer.forward_packed is public; the packed path costs two host reads per batch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Require a head dim divisible by 8 and hop_distance <= 2 for the packed path - Accept a string device in GraphTransformerEncoder.forward - Make varlen_attention, padded_packed_attention, build_neighbor_table and khop_sequence_layout private; make pack_sequences generic over the grid - Trim the encoder notes and CHANGELOG entry to the full condition set - Tests: fold the CUDA predicate check into the CPU one, run the compile test on CUDA, parity over hop_distance 1 and 2, drop covered static_khop cases Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
HeteroData.to moves the batch in place, so the CPU encoder was fed CUDA data. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…parse layout Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
… blockers, layout budget - Clamp the varlen kernel's length bounds to the token counts so a training step on a batch with fewer tokens than max_seq_len no longer fails in backward. - List why a config cannot run packed (_packed_backend_blockers) and log it at construction; attention_dropout_rate > 0 is a blocker so training and eval take the same path. - Fall back to the padded path for a batch whose k-hop candidate set exceeds a budget. - The k-hop layout accepts int32 edge indices and anchors; the homogeneous edge index helper places the result on the device and keeps the dtype. The shared helpers are public, and PackedSequences/pack_sequences live next to the layout. - Add the hop embedding directly on the packed path and read the batch's relative-encoding attributes only on the padded path. - once_differentiable backward, grad-mode aware needs_grad, and validation of max_seq_len and hop_distance. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
The memory-efficient attention kernels were called as aten ops from an autograd.Function. Inductor applies layout rules to those ops that need concrete strides, and it fails on the [1, tokens, ...] views when the token count is only known at run time (whole-encoder torch.compile). Wrap forward and backward as torch.library custom ops, gigl::varlen_attention_forward and gigl::varlen_attention_backward, with fake implementations, so Inductor treats them as opaque. Eager numerics are unchanged.
…put projection _encode_packed contains graph breaks (data-dependent sizes), and a graph break inside a call leaves the bound method on the interpreter stack. torch.compile guards on that fresh bound-method object, so nesting the call in _project_output(...) recompiles a whole-encoder compile on every batch. The result is assigned to a local first; add a comment saying so. Add tests that a compiled encoder reuses its graphs across batches (CPU, eager backend) and that the whole encoder compiles with Inductor on CUDA and matches eager.
The packed backend runs without attention dropout, so the forward op returns only (output, logsumexp) and the backward kernel builds the dropout-free Philox scalars itself. The Python layer wraps the forward result in VarlenAttentionOutput. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Annotations only, no runtime checking. The registered custom-op kernels and fakes keep plain Tensor types because custom_op infers the op schema from the annotations and rejects jaxtyping types; their docstrings carry the shapes. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
…ecting Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
A batch whose k-hop candidate set is too large for the packed layout runs on the padded path. The count and the warning live in a torch.compiler.disable helper so the counter does not become a Dynamo guard. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
…_layout and the logger Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
…axtyping shapes; add at-scale results to the docstring Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
…e hop attribute on both paths, clamp pack_sequences lengths - Reject a hop_distance embedding with <= hop_distance entries on configs that can run packed, and document that the transform's h_max and direction must match the encoder. - Raise the padded path's missing-attribute error from the packed path too. - pack_sequences clamps lengths to max_seq_len instead of returning inconsistent offsets. - State that the varlen dropout is untested rather than claiming how the kernel behaves. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
…tead of custom ops Aten already has an autograd formula and meta registration for _efficient_attention_forward, so the two gigl:: custom ops, their fakes, the autograd.Function and the registration helper are removed. The logsumexp is requested only when a gradient is needed. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
khop_layout takes prioritize_hop_order and re-sorts the kept neighbours by (hop, node id) before truncation. The candidate budget already bounds memory for deeper hops, and oversized batches fall back to the padded path. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
…ment the asymmetry Requiring the attribute on the packed path broke pipelines that skip the all-pairs hop-distance transform on purpose, which is the point of computing hops from the layout. Document that the padded path raises and the packed path does not read it. Add CUDA gradient tests for hop_distance 3 and prioritize_hop_order. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
…op, with a TODO for torch >= 2.10
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Every contracted function shares the code object of Jaxtyping's wrapper, so Dynamo counted each one that started a frame as a recompile of the same code. The launcher now runs that frame eagerly with skip_code, and Jaxtyping's config reads jaxtyping_disable as true while Dynamo traces, so inlined wrappers call the undecorated function. Eager calls are still checked. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The error and the encoder docstring now show how to compile each layer's forward and forward_packed in place, besides whole-encoder compile. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The Shape Contract wrapper shares one code object across every contracted function, and its typeguard checks ran inside compiled frames. Dynamo now skips that code object, and the checks are off while Dynamo traces. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The nested-equality test helper raises AssertionError, so python -O keeps the checks. The encoder docstring lists the AddHeteroHopDistanceEncoding settings that a supplied attribute must use. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…d_attention packed_attention and _padded_packed_attention take the PackedLayout as layout, the name the rest of the packed backend uses. The module docstring now says that the attention does not depend on how the sequences are built; k-hop is one example. The two CUDA parity tests call packed_attention instead of _varlen_attention, so they also cover the first-token offsets and length clamps in the wrapper. The meta-tensor test still calls the private op. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
The check ran for every config that could use the packed backend, which moved more configs into it as the backend grew and wrongly covered PPR configs, whose hops come from the batch. Gate it on khop instead. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Conflicts resolved: gigl/nn/graph_transformer.py (kept the hop branch's hop_distance docstring, with the embedding-size sentence now saying khop).
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
#815 landed on main as a squash commit. Resolved gigl/nn/packed_attention.py to this branch's version, which is main's file plus the non-empty sequence precondition in the packed_attention docstring. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Resolved with the branch side: gigl/nn/graph_transformer.py, tests/unit/nn/packed_graph_transformer_test.py Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.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
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.
Summary
The packed Graph Transformer path computes
hop_distancefrom the sequence layout, so it runs withoutAddHeteroHopDistanceEncoding. A batch whose k-hop layout is too large falls back to the padded path. That path read thehop_distanceattribute and raisedValueErrorwhen it was missing. So the transform had to stay in the pipeline even when nearly every batch ran packed.With this change, the padded path computes
hop_distancefrom the anchors when the batch has none. A fallback batch runs, andAddHeteroHopDistanceEncodingis optional for k-hop sequences.Approach
_get_anchor_hop_distance_matrixingigl/transforms/graph_transformer.py. It does a sparse BFS from the unique anchors over the direction-oriented edge index. It records the first-reach hop (1 tohop_distance) for each (anchor, node) pair, in the CSR format ofAddHeteroHopDistanceEncoding, for the anchor rows only.tensors_to_graph_transformer_inputuses the helper only for k-hop sequences, and only whenhop_distanceis an anchor-based input or attention bias and the batch does not have it. It works on a copy, so the caller's attribute dict does not change.GraphTransformerEncoderno longer raises on a fallback batch without the attribute.User-facing changes
AddHeteroHopDistanceEncodingis optional for k-hop Graph Transformer encoders that usehop_distanceas an anchor-based input or bias.hop_distanceattribute unchanged.hop_distance(pairwise_attention_bias_attr_names) still raises.Verification
make unit_test_py(includesty): 1458 tests OK, 16 skipped.AddHeteroHopDistanceEncoding(is_undirected=False, sampling_direction=..., h_max=hop_distance). They cover both directions, hops 1 to 3, hop-ordered and unordered sequences, and the input and bias roles.packed_graph_transformer_test.py(52),packed_attention_test.py(4),nn/graph_transformer_test.py(76),transforms/graph_transformer_test.py(69), andstatic_khop_test.py(14).Performance
One L4 GPU, padded path forced, hub-heavy batch. "With attribute" times the transform plus the forward. "Without attribute" times the forward only, which computes the hop itself. The outputs match exactly.
Known limitation, not caused by this change: at 256 anchors and 60000 nodes, both variants run out of CUDA memory (29.3 GiB) in the existing
_lookup_csr_values_and_found.🤖 Generated with Claude Code