Skip to content
Soarr01Public

About

Peft of linearization: convert pretrained HF transformers to gated linear attention (GLA)

Resources

Stars

0 stars

Watchers

0 watching

Forks

Latest commit

 

History

1 Commit

Folders and files

Repository files navigation

relinear — convert pretrained transformers to linear/hybrid-attention models

The peft of linearization. One pipeline: HF transformer in → linear or hybrid model out. The 2024–26 conversion recipes exist only as fragmented, mostly-abandoned single-method paper-code; the primitives layer (fla-org) is alive and excellent but ships layers, not conversion. relinear is the missing layer between them.

Born from a data-driven gap hunt over a 96k-repo GitHub corpus (2026-07): the conversion cell is all single-paper code — MambaInLlama (242★, stale 266d), phi-mamba (dead 662d), Liger, LoLCATs, Taylor-Calibrate (2★) — while hybrid models ship in production (Nemotron-H, Qwen-hybrid, Falcon-H1) and no maintained toolkit unifies the recipes.

The pipeline (what this scaffold already runs, pure torch, CPU)

surgery        swap softmax attention -> GLA, REUSING the teacher's q/k/v/o   (LoLCATs move)
transfer       stage 1: train ONLY feature maps + gates to mimic each softmax
               layer's output on teacher hidden states                        (layerwise distill)
select         HALO-style: layers with the largest post-transfer residual
               resist linearization -> keep them softmax (hybrid), no grid search
distill        stage 2: brief end-to-end KL to teacher logits
eval           associative recall — THE capability linearization loses (Zoology:
               ~82% of the SSM-vs-attention gap is recall). No recall eval = lying.
python -m relinear.cli demo --hybrid-k 1     # full ladder end-to-end, ~5 min CPU (synthetic teacher)
python -m tests.run_tests                    # the 5 discriminating tests, ~9 min (synthetic)

python -m relinear.cli convert-hf --hybrid-k 1     # same pipeline on a REAL pretrained model (arch auto-detected)
python -m tests.run_hf_tests                       # 3 tests, GPT-Neo (no RoPE), needs GPU/network
python -m tests.run_hf_llama_tests                 # 6 tests, Llama (RoPE + GQA), needs GPU/network

HF integration (roadmap #1, in progress) — real model, real text

relinear/hfmodels.py + hfconvert.py + hfeval.py run the same surgery → transfer → select → distill ladder on a real pretrained causal LM instead of the synthetic teacher, using held-out perplexity on real prose (relinear/data/sample_corpus.txt, public-domain Conan Doyle) as the capability signal instead of the recall task (there's no recall vocabulary on natural text).

Verified target: EleutherAI/gpt-neo-125m. GPTNeoSelfAttention exposes exactly the q_proj/k_proj/v_proj/out_proj surface surgery already assumes — no GQA, no RoPE (learned position embeddings, like TinyTransformer) — so it needed an adapter, not a redesign. GPT-Neo alternates global (full causal) / local (sliding-window) layers; only global layers are linearization candidates here — local layers are already linear-time and are left untouched, so the hybrid selection is scoped to the 6 global layers.

Real numbers (measured, --transfer-steps 40 --distill-steps 30, RTX 2080 Ti): held-out perplexity teacher 19.24 → raw swap (no training) 8174.24 → converted, hybrid_k=1 (kept layer 8, the highest-residual global layer) 25.33. Same shape as the synthetic ladder — raw swap destroys the model, the pipeline recovers almost all of it — now on a real 125M model and real text, not a toy recall task.

Second target: JackFram/llama-160m (Llama architecture, no GQA — MHA + RoPE). Adding RoPE needed exactly one new hook in the shared base class: GLAAttention._qk_transform(q, k) (models.py), called post-split, pre-feature-map, identity by default (so GPT-Neo's path is byte-for-byte unchanged). LlamaGLAAttention overrides it to apply the SAME rotation the teacher uses — reusing the teacher's own LlamaRotaryEmbedding + apply_rotary_pos_emb directly rather than reimplementing the math, the same "reuse, don't rebuild" move as q/k/v/o. The recurrence itself doesn't change: RoPE is a per-token transform the recurrence never has to know about.

This also forced a real generalization, not just a second copy-paste: hfconvert.py's attention_transfer_hf/logit_distill_hf/convert_hf now take an arch module (hfmodels.py or hfllama.py) exposing surgery, candidate_layer_indices, get_attn — adding an architecture means adding one such module, not touching the training loop. relinear.cli convert-hf auto-detects the arch from config.model_type.

Real numbers, Llama path (--transfer-steps 60 --distill-steps 40, RTX 2080 Ti): held-out perplexity teacher 63.0 (llama-160m is a small wikitext-only-pretrained model, evaluated out-of-domain on Conan Doyle — higher than GPT-Neo's is expected, not a bug) → raw swap 774.5 → converted, hybrid_k=1 (kept layer 2) 94.0. Same ladder shape, now with RoPE in the loop.

Third target: HuggingFaceTB/SmolLM2-135M (GQA — 9 query heads / 3 kv heads, n_rep=3). GLAAttention now takes an n_kv_heads param: k_proj/v_proj/Fk are sized to the physical kv-head count, and a new _repeat_kv expands the (feature-mapped) k and (raw) v to the query-head count before the recurrence loop — same contiguous grouping as HF's own repeat_kv (query head qh reads kv head qh // n_rep; verified bit-identical against transformers.models.llama.modeling_llama.repeat_kv in test_repeat_kv_matches_hf), so the recurrence itself still never has to know about GQA. n_kv_heads=None (the default) reproduces the old MHA shapes exactly — GPT-Neo and llama-160m are unaffected byte-for-byte.

Real numbers, GQA path (--transfer-steps 20 --distill-steps 15, RTX 2080 Ti): raw swap 1420.9 → converted, hybrid_k=1 (kept layer 7) 56.9. Same ladder shape, now with a real grouped-query-attention model.

Fourth target: NousResearch/Llama-3.2-1B (16 layers, 32 query / 8 kv heads, hidden 2048 — an ungated mirror of Meta's gated original, same architecture) — the roadmap's literal 1B milestone. Getting here needed two more fixes, both real bugs, not 1B-specific hacks:

  1. --dtype {fp32,fp16,bf16} (new CLI flag): two full-precision copies of a 1.24B-param model don't fit fp32 on an 11GB GPU (~10GB just for weights, before activations). But fp16 silently produced NaN residuals in the GLA recurrence — the accumulate-then-divide state update (num / den, den can get very small) needs more dynamic range than fp16's narrow exponent gives it; this is the same well-known fp16-vs-bf16 instability LLM practitioners hit elsewhere, now inside a custom recurrence instead of a transformer forward pass. bf16 (same exponent range as fp32) has no such issue. Use --dtype bf16 for anything ≥1B, not fp16.
  2. from_hf only moved device, never dtype — .to(sa.q_proj.weight.device) left the newly-constructed GLA module's Linears/Fq/Fk/gate in default fp32 even when the teacher loaded in bf16/fp16, crashing on the first matmul. Fixed to .to(device=..., dtype=...) in both hfmodels.py and hfllama.py. Same bug, same fix, second time — see the gotcha in memory about checking dtype whenever device is checked.
  3. init_state's zeros defaulted to fp32 regardless of the surrounding model's dtype — same class of bug as #2, one layer down (the recurrent state tensors, not the module params). Now takes a dtype param, threaded from _scan's input tensor.

Parameter-efficient distillation (also needed to fit the 1B optimizer, and just correct regardless): stage 2 used to run Adam(student.parameters()) — fine at 125M, but Adam's fp32 moment buffers over a full 1B-param model don't fit an 11GB GPU alongside the model itself. hfconvert._freeze_non_gla now freezes everything except the converted GLA modules' own parameters before distillation, so only the (tiny) converted-attention params get an optimizer — actually the "peft" in "peft of linearization," not just an efficiency hack. Verified to not regress the smaller models (GPT-Neo/llama-160m/SmolLM2 numbers unchanged within noise after this change).

Real numbers, 1B path (--dtype bf16 --transfer-steps 30 --distill-steps 20 --batch-size 4 --seq-len 64, RTX 2080 Ti, ~8GB peak): held-out perplexity (eval seq_len=256) teacher 14.80 → raw swap 12898.73 → converted, hybrid_k=1 (kept layer 10) 483.22 — proving the mechanism runs correctly at 1B scale. Evaluated at the SAME seq_len used for training (64) instead of the default 256, the same converted model scores 288.31 — notably better, and the more methodologically honest number (see below).

Investigated "does more training help?" — result: no, not with this recipe, and that's worth documenting rather than hiding. Pushing transfer/distill steps well past 30+20 (tried 200+150, and separately 50+40 with larger seq_len for more per-step data diversity) made held-out perplexity consistently WORSE (812–1114), not better, across every variant tried — different layer selections, forced vs auto-selected keep-layer, and more diverse per-step batches all showed the same direction. Isolating the two candidate explanations:

  • Train/eval seq_len mismatch is real but only partially responsible. Training windows are 64–96 tokens; hfeval.perplexity's default eval window is 256. Re-evaluating the 30+20 run at seq_len=64 (matching training) improved its score from 483→288 — the recurrent state hadn't been asked to generalize past what it was trained on, so the default 256-token eval was quietly understating this config's real quality.
  • But re-evaluating the 200+150 run at ITS matching training length (64) still scored ~1104 — essentially unchanged from its 256-eval score of ~1078, and far worse than the 30-step run's 288 either way. So the seq_len mismatch does NOT explain why more steps is worse — that part is a genuine optimization/overfitting effect of this recipe (plausibly stage-2's LR being too high to run for hundreds of steps without drifting, or the small trainable param set overfitting the KL loss on a narrow set of repeatedly-sampled corpus windows). Not yet root-caused further — logged as an open question, not fixed.

Practical takeaway (documented, not (yet) automated): for this model/recipe, ~30 transfer / ~20 distill steps is closer to a local optimum than "more is better," and perplexity should be reported at the seq_len actually used for training, not a fixed default. Early-stopping against held-out perplexity during training would be the principled fix (not yet implemented — see scope limits). Not claiming LoLCATs-quality reproduction yet either way — what's proven here is that the pipeline, memory footprint, and dtype handling all work correctly end-to-end on a real ~1B model on consumer hardware smaller than the 24GB roadmap target; closing the remaining quality gap is future work.

v1 scope limits (fast-follow, not yet done): fused-qkv attention (GPT-2's c_attn) is not supported. Streaming step() doesn't work for RoPE models — it would need to rotate the recurrent state by the per-step incremental angle (RetNet/xpos-style), not just rotate q/k, which isn't implemented; LlamaGLAAttention.step() raises rather than silently returning wrong numbers. No padding-mask / KV-cache support (fixed-length windows only — HF's use_cache must be False at every call site, see hfconvert.py/hfeval.py). No --recipe/--select flags yet (roadmap #2). No checkpoint export yet (roadmap #4) — the converted model exists only in memory. 1B-scale conversion quality (not just mechanism) still needs work to reach LoLCATs parity — "just add more steps" was tried and made it worse (see above); the likely next lever is held-out-monitored early stopping and/or a lower stage-2 learning rate, not raw step count. No eval-time seq_len auto-matching yet — hfeval.perplexity still defaults to 256 regardless of training seq_len, so callers must pass a matching seq_len explicitly for an honest number.

Recipe registry (roadmap #2) — real algorithmic differences, not renamed flags

--recipe lolcats|liger|mohawk and --select halo|kl|fixed|recall on convert-hf. Each recipe changes real architecture or loss, not just a label:

  • lolcats (default): the pipeline described above — output-MSE attention-transfer only.
  • liger: the gate is reused from the ALREADY-COMPUTED key projection (gate_mode="key" in GLAAttention, models.py) instead of a dedicated Linear(d, n_heads) — a per-head learned vector [h_kv, dk] reusing kraw rather than a fresh d×h weight matrix. Verified: 126× fewer gate parameters (520 vs 65,568 on a 32-head/2048-hidden config) with comparable quality on both the synthetic recall task (0.670 vs 0.672 post-distill) and a real-model smoke test.
  • mohawk: adds a matrix-alignment auxiliary loss on top of lolcats' output-MSE — align GLAAttention.implicit_attn_weights(x) (the causal, row-normalized T×T attention weights the recurrence implicitly computes — verified bit-identical, via matrix-multiply, to what _scan's O(T) loop actually produces) against the teacher's REAL softmax attention weights, captured by forcing output_attentions=True on just that one teacher module via a forward hook. MOHAWK-inspired (a real matrix-orientation-stage signal), not the original paper's full 3-stage curriculum — documented as such, not oversold.
  • halo (default select): the existing residual proxy — near-free, reuses attention-transfer's own per-layer residuals.
  • kl: ablation-based — for each candidate layer, temporarily restore JUST that layer to softmax (via arch.set_attn, new symmetric counterpart to arch.get_attn) and measure KL divergence to the teacher's logits; keep the k layers whose restoration helps most. Costs one extra full-model forward per candidate layer per eval batch (real cost, no shortcut).
  • fixed: no data-driven selection — --fixed-layers 2,5,9 names the kept layers directly (select_fixed, convert.py).
  • recall — the fix for the project's former "central open problem": a NEW greedy-forward ablation search that directly targets MQAR + needle recall, instead of an aggregate output-matching proxy (residual or KL) that turned out to be blind to which layers are recall-critical. See "Closing the recall gap" below for the full story and numbers — this is not an incremental option, it's the resolution of what was previously documented as an open architectural/objective-level risk.

Smoke-tested on EleutherAI/gpt-neo-125m (small step counts — mechanism validation, not a quality benchmark): all three recipes and all four select strategies ran end-to-end and landed well below raw-swap perplexity (8174.24). Numbers updated 2026-07-09 after fixing the _freeze_non_gla bug (see HANDOVER.md — stage 2 now genuinely trains ONLY Fq/Fk/gate, not q/k/v/o too): lolcats 509.23, liger 209.40, mohawk 270.64 (all much higher than the 61.25/36.67/52.29 previously reported — expected, not a regression: true peft-style training has much less capacity to fit with at this tiny 15+10-step smoke-test budget than the accidentally-also-fine-tuning-q/k/v/o runs did; still "finite and well below raw-swap" is all this smoke test asserts, and it still passes); select=kl, select=fixed, and select=recall also finite and sane. tests/run_recipe_tests.py, 4/4 PASS (extended to cover select=recall, which needed a tokenizer param threaded through convert_hf — unused by every other select, fully backward compatible). Not yet re-verified at 1B scale.

Eval: MQAR + needle-in-haystack (roadmap #3) — an important honest finding

Perplexity is an average over a whole sequence; a model can have decent average perplexity while having completely lost a specific, hard-to-reconstruct-by-guessing capability like multi-key recall. relinear/evalmqar.py (MQAR — multi-query associative recall, generalizing evalharness.py's toy single-query task to real tokenizer vocab AND multiple queries per sequence, the actual "MQ") and relinear/evalneedle.py (needle-in-haystack: insert a fact in real background text at a controllable depth, ask about it, manually greedy-decode a few tokens with NO KV-cache — this pipeline's GLA doesn't support streaming for RoPE models yet — and check whether the answer comes back) are now reported by default after every convert-hf conversion, alongside perplexity (--skip-mqar-needle to opt out for a fast iteration loop).

Both harnesses discriminate correctly (tests/run_eval_tests.py, verified): teacher MQAR 0.562, needle pass-rate 1.00 at every depth (0/25/50/75/100%) → raw swap (untrained) 0.000 / 0.00 — exactly the expected collapse.

But converting with a step count that gives a "good-looking" perplexity number (EleutherAI/gpt-neo-125m, hybrid_k=1, 40 transfer + 30 distill steps → perplexity 30.62, respectably close to teacher's 19.24) leaves MQAR at 0.031–0.047 and needle pass-rate at 0.00 across every depth — essentially UNRECOVERED, despite the perplexity number looking fine. This is not a bug in the eval; it's the eval working exactly as intended, and it's the most important finding from this roadmap item: every perplexity number reported earlier in this README for the HF path should be read as "average next-token statistics recovered," not "recall capability recovered" — those are measurably different things for this pipeline at these step counts. This is precisely Zoology's "~82% of the gap is recall" thesis showing up empirically in relinear's own numbers, not just cited from a paper. Closing this gap (not just the perplexity gap) is now the real quality bar for roadmap #1's 1B milestone and is NOT yet met at any scale tested.

Checkpoint export (roadmap #4) — real weight round-trip, incl. a live HF Hub push/pull

--save-dir PATH on convert-hf saves a converted model; relinear.cli load-checkpoint PATH reloads it. This is deliberately NOT a plain HF checkpoint — GLAAttention reuses the original q_proj/k_proj/v_proj/o_proj shapes exactly (only NEW parameters — Fq/Fk/gate — are added), so a naive AutoModelForCausalLM.from_pretrained(path) on the saved directory would not error, it would silently reconstruct the ORIGINAL, UNCONVERTED model — HF warns about (doesn't fail on) the extra Fq/Fk/gate tensors it doesn't recognize, so the wrong model loads with no crash to catch it. relinear/export.py saves a manifest (base model_id, arch adapter, recipe, gate_mode, kept-softmax layers) plus the converted state_dict, and load_converted rebuilds the exact module structure via arch.surgery(...) with the recorded settings before loading weights into it — always use this, never bare AutoModelForCausalLM.

Verified: a real, bit-identical local round-trip (tests/run_export_tests.py, 2/2 PASS) — save a converted GPT-Neo-125m student, reload it fresh (re-downloading the base model and re-running surgery), and the reloaded model's logits match the original exactly (max diff 0.0, not just "close"), for both lolcats and liger (different gate_mode, different param shapes — both round-trip correctly).

HF Hub upload/download — real, live-verified, not just unit-tested. relinear.export.push_to_hub/load_converted_from_hub (via huggingface_hub) push a saved checkpoint to a real HF Hub model repo (writes a model card that also warns against bare AutoModelForCausalLM) and pull it back. cli.py exposes this as --push-to-hub REPO_ID [--hub-public] on convert-hf, a standalone push-to-hub <path> <repo_id> subcommand, and load-checkpoint auto-detecting a repo id vs a local directory. Live- tested end to end against a real (private) repo on the developer's own HF account: pushed a converted GPT-2 checkpoint, downloaded it back fresh via load_converted_from_hub, and its logits matched the original local checkpoint's exactly (max diff 0.0) — the same bit-identical bar as the local round-trip test, now proven across a real network upload/download, not just a local file copy. The load-checkpoint <repo_id> CLI path was verified too (round-trip perplexity reported correctly after an auto-detected Hub download). Default private=True — the test repo was private, matching this default.

vLLM serving (roadmap #4) — real custom backend, GPT-2 AND Llama (incl. true GQA) verified end-to-end

relinear/vllm_backend.py is a genuine vLLM model backend for relinear-converted checkpoints — not a config file, not a stub. vLLM already ships infrastructure for exactly this class of model: MambaBase (constant-size recurrent state instead of a growing KV cache) + LinearAttentionBackend, used today by real linear-attention models (MiniMax-Text-01 was the reference pattern this follows). Hybrid models — mixing full softmax attention layers with linear-attention layers, exactly relinear's hybrid_k/ keep_softmax design — are natively supported too: vLLM's model runner auto-discovers every attention-family layer registered per-index into its compilation context and asks each one for its own cache spec, so a hybrid model just needs to instantiate the RIGHT layer type per index (vLLM's stock Attention vs a new RelinearGLALayer), mirroring relinear.hfgpt2.surgery()/relinear.hfllama.surgery() exactly.

State is one tensor per layer, [h, dk, dk+1] — reusing the "augmented value" trick already validated in models.py's _scan_fla (the last column IS the z-normalizer's own recurrent state) rather than inventing new math for serving. Kernels are fla-org's own chunk_gla (prefill) and fused_recurrent_gla (decode) — not vLLM's vendored fla ops, which only implement Gated DeltaNet's different (delta-rule) recurrence. Weight loading pulls from two sources per layer: the base HF checkpoint (kept-softmax layers — Conv1D-transposed the same way vLLM's own GPT2 loader already does for GPT-2, or merged via the same stacked_params_mapping vLLM's own LlamaForCausalLM.load_weights uses for Llama's fused qkv_proj) and relinear's own saved converted state_dict (GLA layers, already in the right convention — no further transpose needed).

RelinearGLALayer (the shared GLA layer both architectures use) supports GQA (n_kv_heads, _repeat_kv ported verbatim from relinear.models.GLAAttention) and RoPE (rotary_emb — reuses the TEACHER's own LlamaRotaryEmbedding module directly, applied to q/k right after the head split and before the feature maps, matching models.py's _qk_transform hook's exact position — same "don't reimplement the teacher's module" philosophy as relinear/hfllama.py, just using vLLM's real per-token positions tensor instead of hfllama.py's training-only arange(T) convention, a genuine correctness improvement for real serving where a request isn't always a fresh window starting at 0).

Verified end-to-end on three checkpoints, two architectures:

  • GPT-2 (openai-community/gpt2, hybrid_k=1, layer 5 kept softmax, 11 other layers GLA): generated coherent text ("The quick brown fox" → " is a very good friend.\n\nThe best part of the game is"), checked qualitatively.
  • Llama, MHA (JackFram/llama-160m, hybrid_k=1, layer 2 kept softmax, 11 other layers GLA): generated "The quick brown fox" → ".\nThe first time I saw the sign.\nThe sign was a" — cross-checked against relinear.export.load_converted's own greedy decode on the identical prompt and it matched character-for-character, real numerical evidence the new RoPE+GQA serving-time math is correct, not just plausible-looking.
  • Llama, real GQA (NousResearch/Llama-3.2-1B, 32 query / 8 kv heads, n_rep=4, hybrid_k=1, layer 10 kept softmax): the 160m case above is actually MHA (n_kv_heads==n_heads), so it exercised the _repeat_kv code path without ever testing a real kv_heads < heads shape — this checkpoint closes that gap. Same cross-check methodology, same result: vLLM's fla-org-kernel decode and load_converted's pure-torch decode matched character-for-character on "The quick brown fox" (both produced the same low-quality-but-identical continuation — this was a low-step-budget conversion for shape verification, not a quality claim).

Log inspection on all three confirmed fla-org's real Triton kernels (chunk_gla_fwd_kernel_o/fused_recurrent_fwd_kernel, plus the RoPE path also exercising chunk_local_cumsum_vector_kernel/chunk_fwd_kernel_h/ chunk_gla_fwd_A_kernel_intra_sub_inter/_intra) actually executed during inference — not just that the process didn't crash.

torch.compile + CUDA graphs — done, and it took two real bugs to get right, not just "remove enforce_eager=True." RelinearGPT2Model/RelinearLlamaModel are now @support_torch_compile-decorated, and it's the DEFAULT for both loader functions (no caller-side flag needed). Two bugs found via systematic-debugging, not guessed: (1) @support_torch_compile's dynamic-dims auto-detection silently found zero dynamic dims, because this file's from __future__ import annotations makes annotations strings at runtime, breaking the detector's type-object comparison — fixed by passing dynamic_arg_dims explicitly; (2) once compiling, output SILENTLY DEGENERATED (no crash, no warning — eager mode: coherent text, compiled mode: repeated-token garbage) because vllm::relinear_gla (our custom op) wasn't in vLLM's compilation_config.splitting_ops, so the CUDA-graph capture incorrectly folded our op's genuinely data-dependent control flow (a variable-length per-sequence prefill loop) into a single static region — fixed by auto-injecting our op into splitting_ops (see HANDOVER.md item 13 for the full root-cause story). Verified after the fix, default settings, no manual flags: GPT-2 and Llama (RoPE+GQA) compiled-mode output matched eager mode's character-for-character on the same checkpoints/prompts used above. Measured speedup: 100-token decode throughput, same GPU/checkpoint, warmup excluded — eager 41.3 tok/s vs compiled 273.4 tok/s, 6.6x at this small (batch=1) scale.

Batched cu_seqlens prefill — done. All concurrent prefill sequences in a step now go through ONE chunk_gla(..., cu_seqlens=...) call (fla-org's own FlashAttention-style packed-variable-length API — the same convention vLLM's query_start_loc already uses), not a per-sequence Python loop. Verified at the kernel level (tests/ run_fla_kernel_tests.py, 4 variable-length sequences with a mix of zero/nonzero initial states, batched vs looped: output/state diff 0.0/1.5e-08) and end-to-end (4 concurrent prompts in one llm.generate() call, no crash; single-request output confirmed bit-identical before/after the refactor — see HANDOVER.md item 14).

v1 scope, honestly: GPT-Neo not supported (vLLM ships no native GPT-Neo model file to build against — GPT-NeoX is a different architecture); no tensor parallelism (single-GPU); LoRA- adapted and liger (gate_mode="key") checkpoints explicitly rejected rather than silently mishandled. Building this required an isolated conda env (vLLM's dependencies conflicted with the shared base environment other projects on the development machine use) and surfaced several vLLM-version-specific API gotchas — see HANDOVER.md item 10 and its Gotchas section for the full list (useful if the installed vLLM version differs from 0.24.0, since real API drift was hit even against the version this was researched against).

fla-org kernel integration (roadmap #1) — verified 50x+ speedup, forward AND gradient

GLAAttention(use_fla_kernel=True) (--fla-kernel on convert-hf) swaps the parallel- forward path from a T-step pure-torch Python loop to fla-org's chunk_gla Triton kernel — the real production GLA implementation this project always intended to build on (see "Honest positioning" below), rather than shipping only the reference loop indefinitely.

relinear's recurrence and fla's kernel contract differ in three ways, each bridged without changing relinear's actual math:

  • Gate shape: relinear's forget gate is a single scalar per (batch, head, timestep), multiplying the ENTIRE [dk,dk] state uniformly. fla's chunk_gla expects a gate VECTOR over the key dimension (the general GLA formulation — relinear's scalar gate is a degenerate special case of it, not an approximation of it). Broadcasting the scalar to a length-dk vector reproduces relinear's exact semantics. fla also expects the gate in log-space (exponentiates internally), so log(gate) is passed, not the raw gate.
  • Normalizer: relinear divides by q·z (a linear-attention-style running normalizer, computed via its own parallel recurrence) — fla's "vanilla" GLA kernel has no built-in equivalent. Recovered via the standard "augmented value" trick: append a column of ones to v before the kernel call. That extra output channel's state update is exactly z's own recurrence, so its output value at each step IS q·z — one kernel call yields both the numerator and the denominator, no second call needed.
  • Scale: fla's kernel defaults to a softmax-attention-style 1/sqrt(dk) scale on q; relinear's math has no such scale, so scale=1.0 is passed explicitly to disable it.

Verified (tests/run_fla_kernel_tests.py): forward-value match for both MHA and GQA (max abs diff ~3e-4, floating-point-level) — but a forward-only check is not enough for a training-time kernel swap, since a wrong backward pass would silently corrupt every downstream training run without showing up in any forward comparison. So gradients (w.r.t. input AND every trainable parameter) were checked too: max diff ~3e-6, even tighter than the forward match. Measured speedup: 50-55x at production-relevant scale (B=8,T=512,d=2048 — roughly 1B-model-sized: 218.7ms → 4.0ms per forward call; B=4,T=256, d=1024: 107.5ms → 2.2ms). A full end-to-end smoke test (convert-hf --model EleutherAI/gpt-neo-125m --fla-kernel) also completed successfully.

Scope: only the bulk parallel-forward path is accelerated (what attention-transfer/ distillation/eval actually call — always with no carried state). step()'s O(1) single- token streaming (only exercised by the synthetic non-HF test suite; RoPE streaming is separately unsupported per the scope limits above) stays on the pure-torch path unchanged — chunked kernels aren't naturally suited to single-token stepping and it isn't the performance bottleneck this item targets. Default remains use_fla_kernel=False (the pure-torch path stays the CPU-compatible, zero-extra-dependency default and correctness reference); passing use_fla_kernel=True without flash-linear-attention installed and a CUDA GPU available raises a clear error rather than silently falling back to the slow path and misreporting what actually ran.

Fused-qkv support (roadmap #1) — GPT-2's c_attn, the fifth architecture

relinear/hfgpt2.py extends the arch-adapter abstraction to GPT-2-family models, whose GPT2Attention has no separate q_proj/k_proj/v_proj — a single c_attn module projects to 3*embed_dim in one matmul, split into q/k/v along the output dimension at forward time. c_attn (and c_proj, the output projection) are also Conv1D, not nn.Linear — a GPT/GPT-2-era module with a TRANSPOSED weight convention ([in,out] instead of nn.Linear's [out,in]) and forward(x) = x @ weight + bias (no transpose in the matmul). from_hf reconstructs q/k/v/o as 4 separate nn.Linears by slicing c_attn's weight/bias into 3 equal column blocks — matching .split(self.split_size, dim=2)'s order and boundaries exactly — and transposing each slice to nn.Linear's convention; c_proj is handled the same way for o_proj.

GPT-2 is also the first tested architecture with BIASED q/k/v (every other adapter so far — GPT-Neo, Llama-family — has bias-free q/k/v) — dropping the bias would silently produce a WRONG raw-swap comparison, not just a slower one, so GLAAttention gained a qkv_bias param (default False, preserving every existing adapter's behavior exactly; GPT-2's adapter is the only caller that passes True).

Verified (tests/run_gpt2_tests.py, 4/4 PASS, openai-community/gpt2): the q/k/v/o reconstruction matches the teacher's real Conv1D forward to floating-point precision (~1e-5 — torch.addmm vs F.linear use different matmul kernels for the identical computation, not a translation bug; weights themselves are copied bit-exact). Full ladder: teacher perplexity 44.0 → raw-swap 14962.9 (raw swap badly hurts it, as expected) → converted 81.0 (hybrid_k=1, kept softmax layer 5) — same shape as every other architecture's ladder. GPT-2 has no RoPE (learned absolute position embeddings, same category as GPT-Neo) and no GQA (predates it) and no global/local layer split (unlike GPT-Neo — every layer is a candidate).

Known limitation: the mohawk recipe needs the teacher's real softmax attention weights (output_attentions=True); GPT-2's default _attn_implementation is "sdpa", which does not materialize attention weights even when requested (unlike GPT-Neo/Llama's "eager" default) — pass attn_implementation="eager" explicitly when loading a GPT-2 teacher if using --recipe mohawk. The default lolcats/liger recipes don't need output_attentions and are unaffected.

--init taylor-calibrate (roadmap #1) — inconclusive, kept opt-in

relinear/hfconvert.py:calibrate_taylor_init replaces Fq/Fk's default identity init with a short (default 30-step) Adam pre-calibration, run once before attention_transfer_hf. The math: softmax attention's pre-normalization score exp(q·k/sqrt(dk)) has a numerically-safe first-order Taylor expansion 1 + q·k/sqrt(dk); the GLA kernel phi(Fq@q)·phi(Fk@k) is approximately affine near Fx=0 too, so calibrating Fq/Fk to match phi(Fq@q)·phi(Fk@k) ≈ 1 + q·k/sqrt(dk) on the teacher's real hidden states (same _layer_io source, same RoPE/GQA pipeline order attention_transfer_hf itself uses) targets a better-than-identity starting point using real data, not just an analytic guess. Explicitly disclosed as Taylor-Calibrate-INSPIRED, not a reproduction of the cited paper — same honesty pattern as the mohawk recipe entry in the table below. Selected via convert_hf(..., init="identity"|"taylor-calibrate") / cli.py's --init flag.

Result, measured on GPT-Neo-125m (hybrid_k=1, kept layer 8), matched step budgets: short budget (15 transfer / 10 distill steps) taylor-calibrate is WORSE (perplexity 187.22 vs identity's 169.73 — calibration steps are pure overhead too small a subsequent budget can't recoup); larger budget (60/40 steps) it's roughly at parity (23.51 vs 24.07 — within single-run noise). At both budgets per-layer transfer residuals are consistently lower with taylor-calibrate (the calibration measurably improves the attention-transfer FIT), but that doesn't reliably translate into better final distilled perplexity. Not a proven win — identity remains the default; taylor-calibrate is available for further experimentation (multi-seed runs, larger-budget-only use) rather than recommended by default. Full regression suite re-run and confirmed unaffected on the default identity path.

Closing the recall gap, part 1: data augmentation (negative — see part 2 for what worked)

Hypothesis: natural prose (the training corpus) rarely exercises exact multi-key lookup, so attention-transfer never gets a signal that teaches real recall no matter how long you train on it alone. --augment none|mqar-short|mqar-long|mixed|curriculum (relinear/augment.py) mixes synthetic MQAR-structured batches into training to test this directly.

Confirmed, with a real tradeoff: on GPT-Neo-125m (hybrid_k=1, lolcats), mixing in SHORT synthetic MQAR sequences recovers short-range MQAR sharply (0.031 → 0.531, near teacher's 0.562) at a small perplexity cost (30.62 → 35.26). A naive 3-way mix (short + long + corpus) in one fixed step budget diluted all targets below what either single-augment run achieved on its own — a genuine capacity/step-budget tradeoff, not something that "just works" by throwing more data types at it. Bigger step budgets, curricula (sequential instead of interleaved phases), weighted/blocked curricula, and combining the mohawk recipe with augmentation were all tried subsequently — full numbers in relinear/augment.py's docstring and HANDOVER.md.

★ Important correction: an earlier draft of this section reported needle-in-haystack recovering partially (0/5 → 3/5) via long-synthetic-MQAR augmentation. That result does NOT reproduce — it turned out to be an artifact of a real bug in hfconvert._freeze_non_gla (stage 2 was accidentally also fine-tuning q/k/v/o, not just Fq/Fk/gate as documented; now fixed). With the bug fixed, that same config scores needle 0/5. As of this fix, TRUE peft-style stage 2 (only Fq/Fk/gate trainable, as designed) has not recovered needle-in-haystack in any configuration tried — recipe (lolcats, mohawk) × augmentation (none/short/long/mixed/curriculum in several variants), all 0/5. See HANDOVER.md for the full corrected history.

Every training-side lever in this part (augmentation data-mix/curriculum, mohawk's matrix-alignment loss, LoRA, fully unfreezing q/k/v/o) was exhausted and found negative or inconclusive. That's what motivated the architectural pivot in part 2 below.

Closing the recall gap, part 2: ★★★ RESOLVED — it was layer selection, not training

Per systematic-debugging discipline (3+ failed fixes ⇒ question the architecture, not keep patching), the investigation pivoted from "more training" to "which layers get kept softmax." This found the actual root cause within hours, no training changes needed.

A hybrid_k sweep on GPT-Neo-125m using the DEFAULT --select halo showed MQAR creeping up slowly (0.062→0.055→0.133 across k=1,2,3) with needle stuck at 0/5 — then BOTH suddenly jumped to teacher-level at k=4 (MQAR 0.453, needle 5/5). A sharp phase transition, not a smooth trend — the signature of a threshold/combinatorial effect, not "a little more training capacity helps a little at a time."

Combinatorial probing (--select fixed) then found real synergy, not additivity: layer 8 alone drives most of MQAR (0.062); layer 6 ALONE contributes literally nothing (0.023, identical to the broken baseline); but layers {6,8} TOGETHER unlock partial needle recall, and adding layer 4 extends it to full depth (5/5). No single-layer "importance" ranking could ever surface this — layer 6's value is invisible until layer 8 is already present.

Critically, the ALREADY-SHIPPED --select kl (ablation + KL-to-teacher — more principled than HALO's residual proxy) was tested against this same combination and also failed: it picked [2,4,6], completely missing layer 8 (MQAR 0.023, needle 0/5 — no better than HALO). This is the key finding: it's not that HALO's specific proxy is flawed, it's that BOTH existing selection metrics are aggregate output-matching signals (residual MSE, KL divergence) dominated by a model's typical, high-frequency behavior — associative recall depends on a handful of specific tokens in a specific context, a signal that barely registers in either metric averaged over ordinary text. Layer 8's own KL- restoration score was literally WORSE than the layers KL-selection preferred, even though it's the single most recall-critical layer.

The fix: a new --select recall (relinear/hfconvert.py:select_recall) — a GREEDY FORWARD ablation search (not single-layer-at-once like select_kl): each round, every remaining candidate is trial-restored to softmax IN ADDITION to layers already picked in previous rounds, scored by actual MQAR accuracy + a needle probe, and the best addition is kept. This structure is what makes synergy detectable — testing layer 6 "with 8 already kept" gives a real signal where testing it alone gives none.

Verified working, fully automatically, no hand-picking: --select recall --hybrid-k 3 on GPT-Neo-125m independently discovered [6,8,10] — an even better combination than the hand-found {4,6,8} — via a live, logged, round-by-round greedy trace. Final pipeline result: converted perplexity 20.91, MQAR 0.430 (77% of teacher's 0.559), needle pass-rate 5/5 — FULL recovery, identical to the teacher's own depth curve — at the SAME hybrid_k=3 budget where HALO's default gets needle 0/5.

What this means: the recall gap was never evidence of an architectural or objective- level limitation of GLA itself. It was that neither shipped selection method ever targeted recall directly. --select recall costs more (needle probes involve real generation) but is now the recommended choice whenever recall matters, not just perplexity — see HANDOVER.md item 15 for the full investigation and numbers.

Generalization check (Llama-160m) — the mechanism generalizes, the "full recovery" number doesn't (yet). At hybrid_k=6 (50% of Llama-160m's 12 layers — matching GPT- Neo's successful proportion), halo still fully fails (MQAR 0.047, needle 0/5) while select_recall clearly beats it (MQAR 0.133, needle 3/5) with a real, steadily-climbing signal in the live search trace — confirming the core mechanism (greedy-forward search finding recall-critical layer synergies invisible to residual/KL proxies) is not GPT-Neo-specific. But the recovery is gradual and partial here (60% needle, 36%-of- teacher MQAR), not GPT-Neo's sharp full snap to 5/5 — likely due to Llama-160m's RoPE architecture, its weaker teacher checkpoint (MQAR only 0.371 to begin with), or a more diffusely-spread recall circuit needing a larger budget than tried so far. See HANDOVER.md item 15 for the full data and honestly-open follow-ups (a larger-budget Llama sweep, 1B-scale testing, the single-depth needle probe's simplification, and whether the greedy search is near-optimal in general).

What is verified today (5/5 PASS, synthetic pipeline)

  1. test_surgery_swaps_and_reuses_weights — every attention module swapped; q/k/v/o bit-exact from the teacher; keep_softmax honored.
  2. test_attention_transfer_reduces_residual — per-layer softmax-matching residual: mean 0.48 → 0.16 (best layer 0.03).
  3. test_recall_recovery_ladder — teacher 1.000 | raw swap 0.25 → +transfer 0.47 → +distill 0.67. Every pipeline stage earns its keep, measured on the capability that actually breaks.
  4. test_selection_finds_recall_critical_layer — the residual proxy picks the layer whose softmax must survive: keeping it → recall 1.000; keeping the anti-selected layer → 0.66 (no better than pure linear). HALO's thesis, demonstrated without any conversion-training search.
  5. test_streaming_o1_state — converted model streams in O(1): step==forward parity (2e-7), state size constant in T while a KV cache grows linearly.

Honest positioning (competitors are ingredients)

project what it is what it is not
fla-org/flash-linear-attention (5.3k★, alive) the kernels/layers this should build on in production a conversion pipeline
HazyResearch/lolcats attention-transfer + LoRA recipe, Llama-family multi-recipe, multi-target, maintained product
MambaInLlama / phi-mamba (MOHAWK) one distillation recipe each, stale/dead a toolkit
OpenSparseLLMs/Linearization (Liger) key-matrix-reuse gating recipe general
HALO / RAD / KL-guided / Taylor-Calibrate layer-selection & init papers (2026) shipped anywhere usable — relinear's own --init taylor-calibrate (see above) is INSPIRED by the latter, tried, and found inconclusive

relinear's bet is the peft bet: none of these die because the method is wrong — they die because single-method paper-code has no reason to be maintained. A recipe-agnostic pipeline does.

Roadmap to product (this scaffold = the mechanism proof)

  1. HF integration — in progress, recall gap now resolved on GPT-Neo-125m: convert-hf proves the ladder on five real models end to end, up to and including 1B scale (GPT-Neo-125m; Llama-160m with RoPE; SmolLM2-135M with RoPE+GQA; Llama-3.2-1B with RoPE+GQA+bf16; GPT-2 with fused q/k/v). The real quality bar, per roadmap #3's finding, is MQAR/needle recovery, not perplexity — every training-side lever tried (augmentation, LoRA, full unfreeze, mohawk) was negative, but the actual fix turned out to be layer SELECTION, not training: the new --select recall (see "Closing the recall gap, part 2" above) achieves needle 5/5 and MQAR 0.430 (77% of teacher) on GPT-Neo-125m at hybrid_k=3 — the SAME budget the default halo selector gets needle 0/5 at. fla-org kernels as the production GLA target — done (--fla-kernel, 50-55x measured speedup, verified forward AND gradient equivalent — see HANDOVER.md item 6). LoRA in stage 2 — done (--lora-rank, stable but no capability gain at rank 8 — see HANDOVER.md item 8). Fused-qkv (GPT-2 c_attn) — done (relinear/hfgpt2.py; the first architecture with biased q/k/v, needed a new qkv_bias param on GLAAttention — see HANDOVER.md). Remaining: verify select_recall generalizes beyond GPT-Neo-125m (Llama, 1B scale) — the real blocker for calling this roadmap item fully done. Milestone: one 1B conversion reproducing LoLCATs-quality numbers (perplexity AND recall) on a single 24GB GPU.
  2. Recipe registry — done (see above): --recipe lolcats|liger|mohawk, --select halo|kl|fixed|recall. --init identity|taylor-calibrate — done, inconclusive (see above; kept opt-in, identity remains the default). Remaining: re-verify recipes' MQAR/needle recovery using --select recall now that it actually surfaces recall-critical layers (previously blocked on "roadmap #1's recall gap," which is now understood to be a selection problem, not a recipe problem).
  3. Eval — done (see above): MQAR + needle + perplexity all reported by default after every conversion. Revealed that perplexity recovery does NOT imply recall recovery for this pipeline at the step counts tried so far — the most important finding from this item, now the real target for roadmap #1.
  4. Export — done (see above): --save-dir/load-checkpoint, a real bit-identical round-trip via relinear/export.py — but NOT a plain HF checkpoint (needs load_converted, not AutoModelForCausalLM). HF Hub upload/download — done, live-verified (push_to_hub/load_converted_from_hub, --push-to-hub/ push-to-hub CLI, real private-repo push+pull tested against the developer's own HF account, bit-identical result — see above). vLLM serving — done for GPT-2 AND Llama-family, v1 (relinear/vllm_backend.py; verified end-to-end on both — GPT-2 checked qualitatively (coherent text), Llama checked more rigorously: vLLM's output matched relinear.export.load_converted's own greedy decode CHARACTER-FOR-CHARACTER on the same prompt, real evidence the new RoPE+GQA math is correct, not just plausible-looking. fla-org's real kernels confirmed running for both prefill and decode on both architectures — see HANDOVER.md item 10 for full detail, scope limits, and version-specific vLLM gotchas). GPT-Neo vLLM support NOT done (vLLM ships no native GPT-Neo model file to build against); no tensor parallelism (single-GPU only). Llama-3.2-1B's real GQA shape (32 query / 8 kv heads) — now verified end-to-end (see above): character-for-character match against the plain-transformers reference, same rigor bar as llama-160m's MHA case. torch.compile/ CUDA-graph support — now done, default for both loaders, 6.6x measured decode speedup, two real bugs found and fixed (see above and HANDOVER.md item 13).
  5. Model zoo: publish converted checkpoints (the distribution/stars engine). Was blocked on roadmap #1's recall gap; the mechanism is now demonstrated to work (--select recall, GPT-Neo-125m), but still blocked pending broader verification (Llama, 1B scale) — no point publishing checkpoints on the strength of a single 125M-model result.

Honest risks (from the hunt, kept in view)

  • Star ceiling is arch-tooling scale (fla-org 5.3k★), not agent-tooling scale — this is a credibility+utility play, not a hype play.
  • Scoop risk: fla-org itself is best-positioned to ship conversion. Mitigation: build AS a fla companion (their modules as targets, PR-friendly), move fast.
  • Quality bar: converted checkpoints must actually hold up on recall/long-context evals or this joins the graveyard it was mined from — hence the eval harness is not optional.
  • ★ The recall gap — formerly this project's central open problem — is RESOLVED as of 2026-07-11, but not yet generalized. Every training-side lever (augmentation data-mix/ curriculum, the mohawk matrix-alignment recipe, fully unfreezing q/k/v/o, LoRA) was genuinely negative — but that's because the gap was never a training or architectural capacity problem at all. It was a LAYER-SELECTION problem: both shipped selectors (halo's residual proxy, kl's ablation-KL-to-teacher) are aggregate output-matching metrics blind to which layers are recall-critical, because recall lives in a handful of specific, SYNERGISTIC layer combinations invisible to any metric averaged over ordinary text. A new --select recall (greedy-forward MQAR+needle ablation search) found a combination on GPT-Neo-125m achieving needle 5/5 (full recovery) and MQAR 0.430 (77% of teacher's 0.559) at hybrid_k=3 — the same layer budget where HALO gets needle 0/5. See HANDOVER.md item 15 for the full investigation. Generalization check on Llama-160m — the MECHANISM generalizes (beats HALO at matched budget on a second architecture), but the "clean full recovery" number is GPT-Neo-125m-specific, not yet reproduced elsewhere. At hybrid_k=6 (50% of Llama-160m's 12 layers — the same proportion where GPT-Neo's phase transition happened), halo still completely fails (MQAR 0.047, needle 0/5, confirming again it's blind to recall regardless of budget), while select_recall clearly beats it (MQAR 0.133 vs 0.047, needle 3/5 vs 0/5) with a real, steadily-climbing signal in the live search trace (not noise, unlike an earlier hybrid_k=2 attempt). But Llama-160m's improvement was GRADUAL, not GPT-Neo's sharp 0/5→5/5 snap, and tops out at needle 3/5 (60%) / MQAR 0.133 (36% of teacher) at this budget — not full recovery. Plausible reasons: Llama-160m is RoPE-based (a real architectural difference), is a comparatively weak community checkpoint (teacher MQAR only 0.371, vs GPT-Neo's 0.559 — less recall signal to transfer in the first place), or its recall circuit may simply need more than 6 of 12 layers to fully snap into place. Read "the recall gap is closeable via better selection" as the generalized claim; "full recovery at hybrid_k=4" as a GPT-Neo-125m-specific instance of it, not a universal constant. See HANDOVER.md item 15 for the full data (both hybrid_k=2 and hybrid_k=6 on Llama) and remaining honest gaps: no larger-budget Llama sweep tried yet, no 1B-scale test, the single mid-depth needle probe is a simplification, and whether the greedy-forward search is near-optimal in general is unestablished. This is real, substantial progress, not a "fully closed, ship it" result — the model-zoo blocker (roadmap #5) should still wait for broader verification.

Layout

file role
relinear/models.py TinyTransformer (teacher stand-in), SoftmaxAttention, GLAAttention (pure-torch reference target; O(1) step(); _qk_transform hook for RoPE; n_kv_heads/_repeat_kv for GQA; gate_mode/gate_parameters() for the liger recipe; implicit_attn_weights for the mohawk recipe; use_fla_kernel/_scan_fla for the fla-org-backed fast path)
relinear/convert.py surgery / attention_transfer / select_hybrid / select_fixed / logit_distill / convert
relinear/evalharness.py associative-recall data + accuracy + teacher training
relinear/cli.py demo (synthetic ladder) + convert-hf (real-model ladder, arch auto-detected, --recipe/--select)
tests/run_tests.py the 5 discriminating tests (synthetic)
relinear/hfmodels.py GPT-Neo arch adapter: HFGLAAttention, surgery/get_attn/set_attn/candidate_layer_indices, ATTN_WEIGHTS_INDEX
relinear/hfllama.py Llama arch adapter: LlamaGLAAttention (RoPE + GQA), surgery/get_attn/set_attn/candidate_layer_indices, ATTN_WEIGHTS_INDEX
relinear/hfgpt2.py GPT-2 arch adapter: GPT2GLAAttention (fused Conv1D c_attn split into q/k/v/o, biased), surgery/get_attn/set_attn/candidate_layer_indices, ATTN_WEIGHTS_INDEX
relinear/hfconvert.py arch-agnostic attention_transfer_hf / logit_distill_hf / convert_hf / select_kl, via forward hooks; _freeze_non_gla for peft-style stage-2; _unfreeze_qkvo/_init_lora-driven LoRA for extra stage-2 capacity; RECIPES/SELECTS registries
relinear/hfeval.py corpus loading, window batches, held-out perplexity
relinear/evalmqar.py multi-query associative recall on real tokenizer vocab: make_mqar_batch, mqar_accuracy
relinear/evalneedle.py needle-in-haystack on real text: needle_probe, needle_sweep, manual uncached greedy decode
relinear/data/sample_corpus.txt public-domain text (Conan Doyle) for real-text train/eval
tests/run_hf_tests.py 3 discriminating tests, GPT-Neo (GPU/network)
tests/run_hf_llama_tests.py 6 discriminating tests: llama-160m (RoPE) + SmolLM2-135M (RoPE+GQA), incl. a repeat_kv-vs-HF unit test (GPU/network)
tests/run_gpt2_tests.py 4 discriminating tests: surgery, Conv1D-split reconstruction accuracy, raw-swap perplexity, pipeline recovery (GPU/network)
tests/run_recipe_tests.py 4 discriminating tests: liger param count, implicit-matrix correctness, all recipes/selects run and beat raw-swap (GPU/network)
tests/run_eval_tests.py 2 discriminating tests (harness sanity) + 1 documented-not-asserted finding: perplexity recovery ≠ recall recovery (GPU/network)
relinear/export.py save_converted / load_converted — manifest + state_dict round-trip, NOT a plain HF checkpoint
tests/run_export_tests.py 2 discriminating tests: bit-identical round-trip, manifest correctness incl. liger's different param shapes (GPU/network)
tests/run_fla_kernel_tests.py 3 discriminating tests + 1 informational: forward equivalence (MHA + GQA), gradient equivalence (input + all params), speed benchmark (GPU + flash-linear-attention)
relinear/augment.py make_batch_fn — mixes synthetic MQAR batches (and, for curriculum, sequential/blocked phases) into attention-transfer/distillation training (--augment); docstring has the full recall-gap investigation + numbers
relinear/vllm_backend.py vLLM serving backend, GPT-2 + Llama v1: RelinearGLALayer (MambaBase, fla-org kernels, GQA + RoPE support), RelinearGPT2ForCausalLM/RelinearLlamaForCausalLM (hybrid softmax+GLA), load_relinear_{gpt2,llama}_for_vllm entry points
HANDOVER.md continuation guide for a fresh session/machine: status, key finding, exact next-step commands in priority order, gotchas

About

Peft of linearization: convert pretrained HF transformers to gated linear attention (GLA)

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages