Add LGATr-slim and LLoCa - #31
Open
hqucms wants to merge 13 commits into
Open
Conversation
Port the L-GATr-slim network (lgatr v2.0.0 nets/slim.py + layers/slim_layers.py) and the tagging-guide LGATrSlimWrapper dense path into a self-contained weaver/nn/model/LGATrSlim.py: - Slim layers (SlimLinear/SlimRMSNorm/SlimGLU/SlimSelfAttention/SlimMLP/SlimBlock) and the LGATrSlim core net, with fp32 precision islands under autocast (minimum_autocast_precision / naive_amp), gradient checkpointing, optional torch.compile, and dead-tail parameter freezing. Bit-exact parity with the reference implementation verified via state_dict transfer. - LGATrSlimTagger: weaver-facing wrapper with four-momentum embedding (E, px, py, pz), symmetry-breaking spurions, global class token, native SDPA key masking, optional masked-mean aggregation, and SequenceTrimmer reuse. Defaults follow the tagging-guide tag_slim config (2.03M params). - movedim replaced by a normalized-permute helper: negative dims produce invalid Transpose nodes in the TorchScript ONNX exporter. - Example network config (test/networks/example_LGATrSlim.py) and unit tests covering forward shapes, padding/permutation invariance, for_inference softmax, backward, and the weaver.train.onnx export path with dynamic axes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015cGUWDpH3eU9dXgb3DrgpK
Port of the LLoCa (Lorentz Local Canonicalization) transformer from the lloca package (dev branch, commit 8a5bb43) and the tagging-guide wrapper: transformer_v2 backbone with frame-to-frame LLoCa attention (including preserve_variance), LearnedPDFrames frames-net, and the TensorReps transform machinery, wrapped as a weaver-facing LLoCaTransformerTagger. Everything runs on weaver's dense zero-padded layout; the equivectors edge convolution is a masked dense reimplementation of the upstream torch_geometric MessagePassing module, so there is no torch_geometric dependency. Attention is fixed to native scaled_dot_product_attention, and the frames-net edge standardization initializes lazily from the first batch into checkpointed buffers. Tests cover shapes, padding/permutation invariance, Lorentz invariance (without spurions), backward, ONNX export (float32 paths), and cross-checks against the upstream lloca implementation. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015cGUWDpH3eU9dXgb3DrgpK
Port the sparse-padding flash attention support from tagging-guide (the LGATrWrapper/TransformerWrapper _forward_sparse paths) and the lgatr/lloca attention_backends packages. Both taggers gain an `attention_backend` option: - "native" (default, unchanged): dense zero-padded SDPA; required for ONNX. - "varlen": torch's native flash-attention varlen kernel (torch.nn.attention.varlen.varlen_attn, torch >= 2.10; no extra deps). - "flash": the flash-attn package, preferring the FlashAttention-3 interface and falling back to FlashAttention-2. For the varlen/flash backends the taggers switch to the packed sparse layout: padding is dropped (dense_to_sparse), per-event cu_seqlens drive block-diagonal varlen attention, global class tokens are interleaved at each event's segment start, and readout is per segment (class-token gather or segment mean). The backend is selected from the attention kwargs, as upstream. On CPU, where the varlen kernels are unavailable, the packed layout falls back to a materialized block-diagonal SDPA mask. The LLoCa frames-net stays dense (mathematically equivalent either way). The kernel wrapper zero-pads the head dim to a multiple of 8 and passes the softmax scale explicitly so it keeps the un-padded 1/sqrt(channels) (LGATr-slim head dims, e.g. 28, are not multiples of 8); fp32 inputs are cast to the autocast half dtype and back. LLoCaAttention additionally accepts token-resolved p_ref for the packed layout. Tested on CPU (packed matches dense to ~1e-7) and on an RTX 5090 with both kernels (matches dense fp32 to <=5e-4 in fp16; fwd+bwd speedup at batch 512 / maxlen 128 jagged: ~2.5x for LGATr-slim, ~1.15x for LLoCa). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QZ3Uh4TGsZvcjKMR6Huxyx
Port the xformers memory-efficient attention support from tagging-guide: attention_backend="xformers" runs block-diagonal attention over the packed tokens via xformers.ops.memory_efficient_attention with a BlockDiagonalMask built from the packed sequence lengths (host-side, so this backend incurs a device-to-host sync that the varlen/flash backends avoid). The wrapper casts fp32 inputs to the autocast dtype and zero-pads the head dim to a multiple of 8 with the softmax scale pinned to the un-padded value, as _run_varlen_kernel already does: on recent GPU architectures (e.g. sm_120) xformers has no fp32 kernel and only the fp16/bf16 flash kernels remain. Validated on an RTX 5090 (torch 2.11.0+cu128, xformers 0.0.35): packed xformers vs dense native outputs agree to 4e-6..2.4e-4 for both models, fwd+bwd, with and without bf16 amp; bit-identical to the varlen backend. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ExsDt5azhwa5gJSMNrCEb8
Port of the tagging-guide "Transformer" baseline (model=tr): the LLoCa transformer_v2 backbone run with identity frames, reusing the existing LLoCaTransformer port. Computes the seven standardized kinematic tagging features internally, supports class-token or masked-mean readout, the packed varlen/flash/xformers attention backends, and ONNX export (with momentum_float64=False). The upstream spurions survive the identity-frames path as featureless valid tokens, reproduced as num_register_tokens. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01X8kqxhVmL1tef9t3pbeFPM
Add an auxiliary_scalars option to LGATrSlimTagger and make the existing ones actually skip the computation instead of computing all seven features and slicing them away: - PlainTransformerTagger: auxiliary_scalars=None now bypasses the momentum path entirely, so pf_vectors is unused and v may be None. - LLoCaTransformerTagger: auxiliary_scalars=None skips the global-frame features fed to the frames-net; the new local_auxiliary_scalars selects the local-frame features fed to the transformer (previously hardcoded to "all"), and None also skips the local-frame four-momentum transforms. - LGATrSlimTagger: gains auxiliary_scalars/momentum_float64 to compute the features internally from pf_vectors instead of taking them from the data config. PlainTransformerTagger and LGATrSlimTagger default to None, i.e. the features come from pf_features via the weaver data config as for ParticleTransformer; note this changes PlainTransformerTagger's default away from the upstream behaviour, which is auxiliary_scalars="all". LLoCaTransformerTagger keeps "all" on both switches. The shared helpers (get_auxiliary_scalars and the four-momentum primitives) move from LLoCaTransformer.py to a new weaver/nn/model/kinematics.py, since LLoCaTransformer imports from LGATrSlim and could not be imported back. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018SaqkBKBrGSDjHD8QPdhWi
`ParticleTransformer`'s default `Embed` applies `BatchNorm1d` + `LayerNorm` + `Linear` + activation. `use_plain_embed=True` reduces it to a bare stack of `Linear` layers, reproducing the single `linear_in` of `PlainTransformerTagger` so the two models can be compared without the input-normalization confound. Exposed on `ParticleTransformer` and both tagger wrappers, so it can be set from the command line with `-o use_plain_embed True`. It overrides `normalize_input`, is rejected together with `use_conv_embed`, and warns when `embed_dims` has more than one entry (a plain stack of linear layers collapses to a single one). Also fixes `Embed.forward`: the (N, C, P) -> (N, P, C) transpose used to sit inside the `input_bn is not None` branch, so `normalize_input=False` applied the embedding across the sequence axis instead of the channel axis. That path was never exercised before, but `use_plain_embed` needs it. `conv_embed` is `Identity` unless `use_conv_embed`, which asserts `normalize_input`, so the default path is unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018SaqkBKBrGSDjHD8QPdhWi
…kends
The varlen/flash/xformers kernels have no fp32 path, so `_run_varlen_kernel`
and `_xformers_attention` cast fp32 inputs to half. `native` does not: its SDPA
is pinned to fp32 by `minimum_autocast_precision`, even under autocast. That
makes the packed backends measurably less accurate, which was not documented
anywhere and is surprising when a run only changes `attention_backend`.
Measured on JetClass (tag-slim config, batch 512, 1000 steps/epoch, 6 epochs,
no AMP), val accuracy, run-to-run noise floor ~0.1-0.3 pp from repeats:
dense + fp32 0.8179 / 0.8188 <- native
dense + fp16 0.8077
dense + bf16 0.8056
packed + fp16 0.8046
packed + bf16 0.7961 / 0.7929
The ~1 pp cost is the precision, not the packed layout: forcing the *dense*
path to fp16 reproduces it, while packed-vs-dense at matched fp16 costs
<=0.3 pp, at the noise floor. An fp32-forced packed run matches dense logits
to 3e-7 and gradients to 5e-6, so the layout itself is exact.
Switching the cast to bf16 was tried and is worse by ~1 pp. fp16 does underflow
the backward pass -- attention output gradients sit near 1e-7, below fp16's
smallest subnormal, so ~96% of gradient entries flush to zero and the gradient
cosine against fp32 drops to 0.996 -- but those entries carry only ~1% of the
gradient norm, while bf16's three fewer mantissa bits degrade the forward
values that matter (fwd relnorm 2.1e-4 vs 3.3e-5). Gradient cosine is the wrong
proxy for the training outcome here.
So keep fp16, and instead:
- extract the choice into `_half_cast_dtype()`, documenting the measurements,
and warn once that the packed backends cost ~1% absolute accuracy;
- give `_autocast_active()` an optional device_type, so a CPU autocast region
is no longer misread as a CUDA one when picking the cast dtype.
No behaviour change for either code path that occurs in practice.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GQN5bNk6ptZahQCy4JYYXg
The varlen/flash/xformers backends drop the padding but have no fp32 kernel, so
they run attention in half precision and lose ~1 pp of tagging accuracy against
`native` (see the previous commit). flex_attention does have an fp32 path, and
its BlockMask makes the block-diagonal attention O(tokens x block) rather than
the O(tokens^2) of a materialized SDPA mask -- so it can drop the padding *and*
keep native's precision. A dense fp32 block-diagonal mask is not an option: at
batch 512 the packed sequence is ~35k tokens, i.e. a ~25 GiB mask.
`_flex_attention` is pinned to fp32 by `minimum_autocast_precision` exactly as
`_sdpa_attention` is, so `flex` and `native` share one precision policy. Both
flex_attention and create_block_mask are compiled with dynamic=True, since the
packed token count changes from batch to batch.
Numerically it reproduces `native`: logits relnorm 2.7e-7, gradient relnorm
2.8e-6, gradient cosine 1.00000000 (`varlen`, for contrast: 3.3e-5 / 7.5e-2 /
0.9972). The new CUDA test asserts flex matches the dense result to rtol=1e-4,
where the flash-derived backends only manage 2e-2.
Confirmed in training -- JetClass, tag-slim config, batch 512, 1000 steps/epoch,
6 epochs, no AMP, val accuracy:
native 0.72057 0.77595 0.78761 0.80172 0.81171 0.81794
native 0.73627 0.77612 0.78705 0.80039 0.81516 0.81877 (repeat)
flex 0.72303 0.75555 0.77970 0.79357 0.80537 0.81738
varlen 0.73275 0.76306 0.76876 0.78706 0.79567 0.80461 (fp16)
flex lands within 0.06-0.14 pp of native, i.e. inside the run-to-run noise
measured from the native repeats, and 1.28 pp above varlen.
The BlockMask tile is set to 64 rather than flex's default 128: a 128-token tile
is wider than a whole jet (JetClass averages ~70 constituents), so most of each
diagonal tile is masked-out work. That cut a compiled step from 90 ms to 70 ms.
32 is rejected by the kernel.
One fwd+bwd step at batch 512 on an RTX 5090, compiled / eager:
native 81 ms / 12.0 GiB 142 ms / 14.1 GiB
flex 70 ms / 5.8 GiB 88 ms / 6.9 GiB
varlen 39 ms / 5.2 GiB 66 ms / 6.3 GiB
So flex dominates native outright -- same numbers, 1.16x faster, 2.1x less
memory -- while varlen stays 1.8x faster than flex for its ~1 pp.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GQN5bNk6ptZahQCy4JYYXg
`--compile` compiled the model twice: train.py sets network_options["compile_model"] = True, which the LGATr-family taggers forwarded into their backbone as `compile=`, and train.py then wrapped the whole model in torch.compile as well. The outer wrap failed for every backend, so --compile was unusable for these models. `compile_model` is not only a "compile yourself" switch -- ParticleNet and ParticleTransformer use it to *adapt* to being compiled (SequenceTrimmer round_to_32, sparse_eval). So the flag keeps flowing; what changes is that the three taggers no longer self-compile, leaving train.py's end-to-end wrap as the single compilation. `compile_kwargs` is consequently unused and now warns rather than being silently dropped; pass --compiler-option instead. Two bugs had to be fixed for the outer wrap to work: - The packed path built the global-token flag by mutating a view in place, `new_s[:, -1].index_fill_(0, global_idxs, 1.0)`. Functionalization turns that into select -> index_put -> copy_(empty, ...) -> select_scatter, and Inductor asserts on a copy_ whose target is an in-graph empty. Rebuilt functionally (index_put + index_fill + cat). The same line was present in all three models. - flex's block mask is built from derived symbols like (tokens + k) // 64, which Inductor's statically_known_multiple_of cannot factor, so they reached the scheduler as CantSplit. `_build_flex_block_mask` is now decorated with @torch.compiler.disable(); create_block_mask is compiled on its own in _get_flex_ops, so nothing is lost by keeping it out of the enclosing graph. Also adopt ParticleTransformer's trimmer settings in the three taggers (round_to_32=compile_model, num_extra_tokens = the prepended spurion/register tokens plus the global token), so a compiled run sees sequence lengths that are clean multiples of 32 once the prepended tokens are added. On torch 2.11 `native` + `--compile` still fails, on an Inductor miscodegen of a split-reduction workspace slice under dynamic shapes; this is unrelated to the change here (it failed before it too). Use `--compiler-option dynamic False` there, or torch 2.13, where all backends work. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GQN5bNk6ptZahQCy4JYYXg
…m count An `--compile` run with `attention_backend='flex'` died on the first training step with `InductorError: AssertionError: flex_attention is not an OpOverload`, i.e. Inductor had no lowering registered for the flex HOP. Root cause: `flops_counter.patch_tensor_ops` monkeypatches `torch.mm` & co. with `torch_function_wrapper` objects that have no `__name__`. The flops pass ran inside `model_setup`, before anything else had imported Inductor, and the flex backend calls `torch.compile` from inside the model's own forward -- so the process's first `import torch._inductor` happened while those names were patched. It died in `kernel/mm.py` (`ExternKernelChoice(torch.mm, ...)` reads `kernel.__name__`) and Python rolled `compile_fx`, `select_algorithm`, `lowering` and `kernel` back out of `sys.modules`. `kernel/flex/*` had already finished and stayed cached, so the retry re-ran `lowering.py` against a fresh `lowerings` dict but skipped the flex submodule -- the flex_attention lowering only ever existed in the discarded dict. The flops exception is caught and logged, so the only visible symptom was the later, unrelated-looking failure. - train.py: move the flops/params report out of `model_setup` into a `report_flops()` closure in `_main`, called after `torch.compile(model)` (so Inductor is already imported) and on `orig_model` rather than the compiled / DataParallel wrapper. Predict-only runs report it in the test block. - LGATrSlim.py: stop pre-compiling `flex_attention` in `_get_flex_ops` -- under an enclosing `torch.compile` the HOP is lowered by that graph, so the wrapper only nested one compile inside another. `_flex_attention` now picks the raw callable when `torch.compiler.is_compiling()` and a separately compiled one (`_get_compiled_flex_attention`) otherwise, keeping the eager / predict / ONNX path fused. `create_block_mask` stays compiled: uncompiled it materializes a dense (tokens, tokens) mask. - flops_counter.py: also log the exact (unrounded) parameter count, before the forward pass so it is still reported when the flops estimation itself fails. Verified on 2x RTX 5090 / torch 2.13.0+cu130: the full JetOrigin command trains and validates (exit 0); compiled and eager flex paths give identical losses. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YQeHLHwCD7iTGWAZvY6Gxy
…r block Every encoder block passed both the pair bias and the key-padding mask into `Attention.forward`, which added them together itself. Both operands are the same for all blocks, so an 8-layer model built eight identical (batch, num_heads, seq_len, seq_len) tensors and -- since each is an input to its block's softmax -- kept all eight alive for the backward pass. `_forward_encoder` now does that sum once and hands every block the result. Outputs are bit-identical; on an RTX 5090 (torch 2.12.1+cu130, 8 layers, batch 256, 128 constituents) a fwd+bwd step goes 62.8 -> 61.6 ms / 7.32 -> 6.44 GiB in fp32 and 48.4 -> 45.1 ms / 4.36 -> 3.92 GiB under bf16 autocast. Under `--compile` it changes nothing, as Inductor already CSEs the repeated add; the win is for eager runs. Two kinds of block keep the per-block path. One excluded by `block_ids_with_attn_mask` takes no bias at all, so it never built such a tensor and still gets the cheap (batch, seq_len) mask. One with `scale_attn_mask` scales the bias by its own `c_mask`, which must not reach the padding `-inf` -- a negative `c_mask` would flip it to `+inf` and point the attention straight at the padded tokens. Those two exclusions differ in an easy-to-miss way, and both matter for peak memory, so they are derived once in `__init__` rather than re-tested in the loop: - The un-merged bias can be dropped after the merge whenever no block scales it. Requiring `all(block_ids_with_attn_mask)` as well keeps it alive for a block that was never given it: 6.82 -> 6.45 GiB instead of 6.82 GiB on a 4-of-8 config. - If `scale_attn_mask` is on for every block, nothing can read the shared mask and building it pins a dense tensor for nothing (+0.10 GiB, +1.4% peak). Both are covered by the new tests, which count the dense masks a forward actually materializes (4-D `aten.add.Tensor` calls under a dispatch mode): the count must not grow with `num_layers` on the shared path, and must be exactly `num_layers` -- no shared mask on top -- when every block scales the bias. The rest of the file pins the equivalences the shared mask rests on and the padding invariant itself, including that an event padded out to the batch length gives the same answer as that event alone. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BG6VLqVb7NST9KnBT3M63B
… to 64
`get_sparse_attention_kwargs` built the flex `BlockMask` with a hardcoded
64-token tile. Inductor rejects a tile that is not a multiple of the BLOCK_M /
BLOCK_N of the Triton config it picks, and without `max_autotune` there is
exactly one candidate config, so the mismatch is fatal rather than one lost
choice:
LoweringException: ValueError: Q and KV block size must be divisible by
BLOCK_M and BLOCK_N. We got Q_BLOCK_SIZE=64 and KV_BLOCK_SIZE=64.
That config comes from a table keyed on (capability, dtype, head_dim) with
entries only at head_dim 64/128/192/256; every other head_dim falls through to a
(64, 64) default, which is why the pin had worked so far. So the round head dims
are the broken ones -- on sm_120, fp32 head_dim 64 asks for BLOCK_M=128. In
L-GATr-slim it takes `-o hidden_v_channels 128 -o hidden_s_channels 512
-o num_heads 16` (head_dim = 4*8 + 32 = 64) to land there, but PlainTransformer
hits it at the entirely ordinary embed_dim 512 / 8 heads.
`_flex_block_size` now reads the config Inductor will actually use and mirrors
its tiling, which divides by construction. Both simpler options are worse: a
tile of 128 is valid for every table entry but costs 33% of a compiled fwd+bwd
step on the default config (45 -> 60 ms, batch 256, RTX 5090), and a square tile
throws away the asymmetry -- at head_dim 64, (128, 32) beats (128, 128) by 19%
(319 vs 393 ms). All three taggers pass their head dimension down for this;
`_flex_head_dim` is read off the constructed attention module rather than
re-derived from the constructor arguments, and a test pins it against the real
q/k/v, since a silent drift would resurface as the same crash.
This also fixes a pre-existing `AssertionError: failed to set ranges` in the
same path, which had made the eager flex path unusable and left
`test_flex_backend_cuda` failing. `create_block_mask` was compiled with
`dynamic=True`, which also sets `specialize_int=False` and unspecializes every
int Dynamo sees -- closed-over ones included -- so the block size went symbolic
and `ceil(tokens / block)` became a ratio of two unknowns Inductor cannot
factor. It is now one compiled builder cached per block size, with the block
size closed over and the token count taken from `batch.shape[0]`, a tensor dim
the caller marks dynamic. Marking it explicitly rather than leaning on automatic
dynamic matters: automatic dynamic escalates after enough distinct token counts
and unspecializes the block size anyway, which would bring the failure back a
few steps into training rather than on the first one.
Verified on an RTX 5090 (torch 2.12.1+cu130): packed matches dense to <= 2.4e-7
for head dims 16/24/28/64/128 across all three taggers, 25 varying token counts
run without a recompile storm both eager and under `--compile`, and the suite
passes (80 tests on CPU, 74 on GPU).
One pre-existing failure is left alone, having been confirmed identical before
these changes: the eager path raises `NoValidChoicesError` on a single static
shape whose token count exceeds one Q tile when head_dim is not a power of two.
It clears once shapes go dynamic, so training is unaffected -- which is why the
new dynamic-shape test uses training-shaped batches rather than one fixed shape.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QGwcN9N1nUGrPPbT6b5V54
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.
No description provided.