Skip to content

build(worker): port bidirectionalization seam to transformers 5.14.1 - #17

Merged
undeemed merged 4 commits into
mainfrom
fm/a2d-dep-transformers-t1
Jul 29, 2026
Merged

build(worker): port bidirectionalization seam to transformers 5.14.1#17
undeemed merged 4 commits into
mainfrom
fm/a2d-dep-transformers-t1

Conversation

@undeemed

@undeemed undeemed commented Jul 29, 2026

Copy link
Copy Markdown
Owner

PR 1 of 2. Merge this one first; the stacked docs PR is #18 (#18), based on this branch.

Intent

Port the a2d worker from transformers 4.51.3 to 5.14.1. This is a real major-version port, not a version-string edit: v5 removed all three eager attention seams a2d's bidirectionalization (Decision 2) patched - GPT-2's per-layer self.bias buffer, the RoPE family's per-model _update_causal_mask method, and Gemma 2/3's per-decoder-layer sliding-window re-mask - and unified them into one documented extension point, ALL_MASK_ATTENTION_FUNCTIONS keyed by config._attn_implementation. On the 4.51.3 code the CI python job was 62 failed / 29 passed; it is now 92 passed.

Deliberate decisions a reviewer reading only the diff would not know:

  • Register into HF's mask/attention registries rather than monkeypatch library internals. The key is per-AnnealState (f'a2d_annealed_eager_{id(state)}') and only that model's config is pointed at it, so isolation is per model, not per process. That is load-bearing: the D13 identity gate holds an un-patched sibling copy of the same model in the same process and it must keep HF's own causal mask. Earlier attempts at a WeakKeyDictionary keyed by config failed (PretrainedConfig defines eq without hash) and at tagging config._a2d_anneal failed (save_pretrained then raises 'Object of type AnnealState is not JSON serializable'); the registry key was chosen over both. The registered closure holds a strong ref to the state so id() cannot be recycled, and _attn_implementation is verified not to be written into config.json.
  • The reveal derives FROM HF's own eager_mask output rather than rebuilding the mask, which is what makes alpha=0 bit-identical to base by construction, independent of dtype/cache offset/padding. Genuine 2D padding is preserved: a masked cell is only revealed when its key is a real token.
  • _a2d_eager_attention dispatches via sys.modules[type(module).module].eager_attention_forward instead of registering one global eager function, because each family's eager maths differs (GPT-2 attention scaling, Gemma 2 logit softcapping).
  • attn.full / attn.gqa / attn.swa stay three DISTINCT capabilities even though v5 collapsed them onto one seam, because they are a contract shared with the Rust detect crate and the handler registry is keyed on them. They are now three structural gates on one install; resolve_capabilities reads config.layer_types then num_key_value_heads then GPT2Attention, from the model itself, never from job tags. Mistral was verified to correctly stay attn.gqa (it folds its window into one model-level mask and declares no layer_types).
  • Verified empirically across gpt2, llama, qwen2, gemma, gemma2, gemma3, mistral: alpha=0 max abs diff exactly 0.0, alpha=1 non-zero. Mistral's alpha=1 shift equals Llama's exactly, independently confirming the sliding window fully opened.
  • Acceptance criterion from the task: a silently-not-applied patch is the dangerous failure mode now that the seam is a registry key rather than a monkeypatch, so test_bidir.py::test_install_routes_the_model_and_only_it_through_the_annealed_seam asserts the key is live in ALL_MASK_ATTENTION_FUNCTIONS, that a sibling model stays on 'eager', and that only the patched model opens at alpha=1.

Scope and process constraints the user set:

  • The captain's standing rule is every PR under 500 changed lines excluding lockfiles. The honest port was 759. The user was asked and ruled explicitly for option B: split into two stacked green PRs rather than ship one oversized PR. THIS IS PR 1 OF 2 (483 lines excl uv.lock): the bump, the code port, and the pyproject pin comment. PR 2 (branch fm/a2d-dep-transformers-t1-docs, already committed, opens immediately after this one) is docs-only and refreshes the prose.
  • Therefore, BY DESIGN AND BY THE USER'S EXPLICIT RULING, this PR knowingly leaves stale prose describing 4.x internals (_update_causal_mask, self.bias, per-layer is_sliding, the old signature-agnostic layer-forward wrap) in transform/attention.py, transform/gqa_attention.py, transform/swa_attention.py, transform/apply.py, transform/handlers/{gqa,swa}attention.py, worker.py, tests/conftest.py, tests/test{gqa,swa}_attention.py, the root pyproject.toml mypy comment and AGENTS.md. That is not an oversight; it is the agreed split seam, it is called out at the end of the commit message, and PR 2 fixes every one of those files. Do not flag stale docstrings in those files as a defect of this PR.
  • A code-only split (port before bump) was ruled out because 4.51.3 GPT-2 does not route through masking_utils, so any intermediate commit would be red, and a split leaving an intermediate broken is not a split.
  • CPU only: this container is Linux aarch64 with no CUDA/MPS/Metal. Tests are hermetic tiny random-weight fixtures; no weights are ever downloaded (Gemma is gated, CI is CPU/no-network).
  • Do not push to dependabot PR build(deps): bump transformers from 4.51.3 to 5.14.1 #12's branch; firstmate closes build(deps): bump transformers from 4.51.3 to 5.14.1 #12 as superseded once this lands.

Also in scope and done: the task asked to re-check whether the root pyproject.toml comment claiming transformers 4.51.3 ships an empty py.typed is still true for 5.14.1. It is - the marker is still empty - and follow_imports = 'skip' is still required (removing it reintroduces exactly 68 mypy errors, because v5's inline annotations disagree with the runtime on **kwargs config fields and Trainer.compute_loss). The override is kept; the corrected comment explaining why lands in PR 2 with the rest of the prose.

Full local gate is green on this branch: ruff check, ruff format --check, mypy strict, 92 pytest passed, cargo fmt --check, cargo clippy --workspace --all-targets -D warnings, cargo test --workspace.

What Changed

  • Bumped the worker pin from transformers==4.51.3 to 5.14.1 and rewrote Decision 2's bidirectionalization onto v5's mask interface. v5 removed all three eager seams the old code patched (GPT-2's per-layer self.bias buffer, the RoPE family's per-model _update_causal_mask, Gemma 2/3's per-decoder-layer sliding-window re-mask); install_mask_anneal now registers a per-AnnealState key (a2d_annealed_eager_<id(state)>) in AttentionMaskInterface/AttentionInterface and points only that model's config._attn_implementation at it, so an un-patched sibling in the same process keeps HF's own causal mask. The annealed reveal derives from masking_utils.eager_mask's own output (masked cells at finfo.min, only revealed when the key is a real token), and _a2d_eager_attention delegates to sys.modules[type(module).__module__].eager_attention_forward so each family's eager maths is unchanged. load_model moves from torch_dtype= to dtype=.
  • attn.full / attn.gqa / attn.swa stay three distinct capabilities (they are the contract shared with the Rust detect crate and the handler registry key) but become three structural gates over that one install: has_sliding_window_seam reads config.layer_types, the new is_rope_family reads num_key_value_heads, GPT-2 is still keyed on GPT2Attention. gqa_attention.py and swa_attention.py drop their hand-rolled mask wrapper and signature-agnostic layer-forward wrap entirely and delegate to the shared install.
  • Added two failure guards and their tests: install_mask_anneal raises if set_attn_implementation leaves the config off the a2d key (a silently-ignored key would look causal at every alpha while still passing the D13 identity gate), and install_anneal_patch rejects reorder_and_upcast_attn=True by name rather than as a bare max_abs_diff mismatch. tests/test_bidir.py asserts the key is live in ALL_MASK_ATTENTION_FUNCTIONS, that a sibling model stays on "eager", and that only the patched model shifts earlier-position logits at alpha=1; the tiny_gpt2 fixture takes config overrides. Worker and Rust suites are green (93 pytest, cargo test --workspace). This is PR 1 of 2 by the 500-line rule (483 lines excluding uv.lock): it knowingly leaves 4.x prose in the touched modules, docs/, and AGENTS.md, which the stacked docs-only PR refreshes.

Risk Assessment

⚠️ Medium: A real major-version dependency port that rewrites the correctness-critical bidirectionalization seam and reshapes the transitive dependency graph (requests dropped for httpx, typer/rich added), but every load-bearing claim checks out against the installed 5.14.1 source, the dangerous silent-no-op failure mode is now a hard raise, and only three low-severity informational items remain.

Testing

Synced the pinned env (transformers 5.14.1, torch 2.13.0+cpu, CPU aarch64), ran the full worker pytest suite twice (93 passed both times, confirming the new per-AnnealState registry key is not order-sensitive) and cargo test --workspace (47 passed) as the baseline, then went past unit level: built the Rust CLI and drove the actual end-user path a2d detect -> a2d convert -> a2d sample over hermetic tiny fixtures for all three attention seams, where every conversion reported identity gate: PASS (max_abs_diff=0.00e0), annealed 0.000 -> 0.667 across training steps, wrote a checkpoint, exited 0, and then decoded text from the converted checkpoint. I also ran a per-family matrix (gpt2, llama, qwen2, gemma, gemma2, gemma3, mistral) proving alpha=0 is exactly bit-identical to base, alpha=1 genuinely opens the mask, a2d's key is live in HF's mask registry, and an un-patched sibling in the same process stays causal, and verified the saved checkpoints reload with plain from_pretrained with no anneal registry key persisted into config.json. No UI surface exists in this change (Rust CLI plus a Python worker emitting JSONL), so the reviewer-visible evidence is CLI transcripts and rendered tables rather than screenshots. Everything passed; no code or test changes were required and the worktree is clean.

Evidence: a2d convert -> a2d sample CLI transcript, all three seams (transformers 5.14.1)

$ a2d detect models/gemma3 SUPPORTED: gemma3_text capabilities: paradigm.ar-transformer, attn.gqa, attn.swa, pos.learned, ffn.dense, norm.rms $ a2d convert models/gemma3 --out runs/gemma3 --data corpus.jsonl --seq-len 8 --max-steps 3 --anneal-steps 3 --per-device-batch-size 2 --device cpu job started (worker: a2d-worker-hf 0.1.0) progress: patch 3/7 progress: identity 4/7 identity gate: PASS (max_abs_diff=0.00e0, tolerance=1e-6) step 1: loss=11.0271 anneal=0.000 lr=1.00e-4 tokens=16 step 2: loss=5.4356 anneal=0.333 lr=6.67e-5 tokens=32 step 3: loss=5.2019 anneal=0.667 lr=3.33e-5 tokens=48 checkpoint @ step 3: .../runs2/gemma3/checkpoints/checkpoint-3 job completed exit=0 $ a2d sample runs/gemma3 -p 'w1 w2 w3' --canvas-len 8 --num-steps 4 --device cpu w1 w2 w3 w1 w1 w1 w24 w1 exit=0

# a2d convert -> a2d sample, real CLI, transformers 5.14.1 on CPU aarch64
# hermetic tiny fixtures saved from tests/conftest.py (no weights downloaded)

$ a2d detect models/gpt2
SUPPORTED: gpt2
  layers: 2  d_model: 16  heads: 2/2 (kv)  vocab: 64
  capabilities: paradigm.ar-transformer, attn.full, pos.learned, ffn.dense
  plan: objective=mdlm
  estimated memory: 0.00 GB (approx, weights-only)

$ a2d convert models/gpt2 --out runs/gpt2 --data corpus.jsonl --seq-len 8 --max-steps 3 --anneal-steps 3 --per-device-batch-size 2 --device cpu
job started (worker: a2d-worker-hf 0.1.0)
progress: ingest 0/7
progress: materialize 1/7
progress: grow 2/7
progress: patch 3/7
progress: identity 4/7
identity gate: PASS (max_abs_diff=0.00e0, tolerance=1e-6)
progress: train 5/7
step 1: loss=10.9538 anneal=0.000 lr=1.00e-4 tokens=16
step 2: loss=5.4769 anneal=0.333 lr=6.67e-5 tokens=32
step 3: loss=5.2855 anneal=0.667 lr=3.33e-5 tokens=48
checkpoint @ step 3: /tmp/no-mistakes-evidence/01KYPMN1Q9YKTNQVTJBMY28BW1/work/runs2/gpt2/checkpoints/checkpoint-3
progress: save 6/7
job completed
run 3c731ad2-feea-4c74-aab0-b8e00dae10f5 completed -> /tmp/no-mistakes-evidence/01KYPMN1Q9YKTNQVTJBMY28BW1/work/runs2/gpt2
exit=0

$ a2d sample runs/gpt2 -p 'w1 w2 w3' --canvas-len 8 --num-steps 4 --device cpu
w1 w2 w3 w53 w36 w52 w61 w61
exit=0

$ a2d detect models/gemma
SUPPORTED: gemma
  layers: 2  d_model: 16  heads: 4/1 (kv)  vocab: 64
  capabilities: paradigm.ar-transformer, attn.gqa, pos.learned, ffn.dense, norm.rms
  plan: objective=mdlm
  estimated memory: 0.00 GB (approx, weights-only)

$ a2d convert models/gemma --out runs/gemma --data corpus.jsonl --seq-len 8 --max-steps 3 --anneal-steps 3 --per-device-batch-size 2 --device cpu
job started (worker: a2d-worker-hf 0.1.0)
progress: ingest 0/7
progress: materialize 1/7
progress: grow 2/7
progress: patch 3/7
progress: identity 4/7
identity gate: PASS (max_abs_diff=0.00e0, tolerance=1e-6)
progress: train 5/7
step 1: loss=10.9807 anneal=0.000 lr=1.00e-4 tokens=16
step 2: loss=5.4707 anneal=0.333 lr=6.67e-5 tokens=32
step 3: loss=5.2158 anneal=0.667 lr=3.33e-5 tokens=48
checkpoint @ step 3: /tmp/no-mistakes-evidence/01KYPMN1Q9YKTNQVTJBMY28BW1/work/runs2/gemma/checkpoints/checkpoint-3
progress: save 6/7
job completed
run e89f8cf7-87a3-4a7d-829b-71fbff390a32 completed -> /tmp/no-mistakes-evidence/01KYPMN1Q9YKTNQVTJBMY28BW1/work/runs2/gemma
exit=0

$ a2d sample runs/gemma -p 'w1 w2 w3' --canvas-len 8 --num-steps 4 --device cpu
w1 w2 w3 w40 w40 w40 w40 w40
exit=0

$ a2d detect models/gemma3
SUPPORTED: gemma3_text
  layers: 4  d_model: 16  heads: 4/1 (kv)  vocab: 64
  capabilities: paradigm.ar-transformer, attn.gqa, attn.swa, pos.learned, ffn.dense, norm.rms
  plan: objective=mdlm
  estimated memory: 0.00 GB (approx, weights-only)

$ a2d convert models/gemma3 --out runs/gemma3 --data corpus.jsonl --seq-len 8 --max-steps 3 --anneal-steps 3 --per-device-batch-size 2 --device cpu
job started (worker: a2d-worker-hf 0.1.0)
progress: ingest 0/7
progress: materialize 1/7
progress: grow 2/7
progress: patch 3/7
progress: identity 4/7
identity gate: PASS (max_abs_diff=0.00e0, tolerance=1e-6)
progress: train 5/7
step 1: loss=11.0271 anneal=0.000 lr=1.00e-4 tokens=16
step 2: loss=5.4356 anneal=0.333 lr=6.67e-5 tokens=32
step 3: loss=5.2019 anneal=0.667 lr=3.33e-5 tokens=48
checkpoint @ step 3: /tmp/no-mistakes-evidence/01KYPMN1Q9YKTNQVTJBMY28BW1/work/runs2/gemma3/checkpoints/checkpoint-3
progress: save 6/7
job completed
run b76e71cd-5541-465f-ae35-51d8290d17fd completed -> /tmp/no-mistakes-evidence/01KYPMN1Q9YKTNQVTJBMY28BW1/work/runs2/gemma3
exit=0

$ a2d sample runs/gemma3 -p 'w1 w2 w3' --canvas-len 8 --num-steps 4 --device cpu
w1 w2 w3 w1 w1 w1 w24 w1
exit=0
Evidence: Per-family anneal seam matrix (D13 identity + bidirectionality + per-model isolation)

transformers 5.14.1 torch 2.13.0+cpu family capability seam live a0 max|diff| a1 shift sibling stays causal ------------------------------------------------------------------------------------------------------------ gpt2 attn.full True 0.0 0.0048 yes (shift 0.0) OK llama attn.gqa True 0.0 0.0197 yes (shift 0.0) OK qwen2 attn.gqa True 0.0 0.0221 yes (shift 0.0) OK gemma attn.gqa True 0.0 0.0059 yes (shift 0.0) OK gemma2 attn.swa True 0.0 0.0534 yes (shift 0.0) OK gemma3 attn.swa True 0.0 0.1535 yes (shift 0.0) OK mistral (sliding, folded into one mask) attn.gqa True 0.0 0.0197 yes (shift 0.0) OK RESULT: ALL FAMILIES OK

transformers 5.14.1  torch 2.13.0+cpu

family                                   capability  seam live a0 max|diff|   a1 shift  sibling stays causal
------------------------------------------------------------------------------------------------------------
gpt2                                     attn.full   True               0.0     0.0048  yes (shift 0.0)      OK
llama                                    attn.gqa    True               0.0     0.0197  yes (shift 0.0)      OK
qwen2                                    attn.gqa    True               0.0     0.0221  yes (shift 0.0)      OK
gemma                                    attn.gqa    True               0.0     0.0059  yes (shift 0.0)      OK
gemma2                                   attn.swa    True               0.0     0.0534  yes (shift 0.0)      OK
gemma3                                   attn.swa    True               0.0     0.1535  yes (shift 0.0)      OK
mistral (sliding, folded into one mask)  attn.gqa    True               0.0     0.0197  yes (shift 0.0)      OK

D13: alpha=0 is bit-identical to base (0.0 exactly) for every family; alpha=1 moves an
earlier position when a later token changes (bidirectional); an un-patched sibling of the
same model in the same process stays causal (0.0 shift) - per-model isolation holds.

RESULT: ALL FAMILIES OK
Evidence: Converted checkpoints reload cleanly; no anneal registry key persisted

gemma3/model/config.json model_type : gemma3_text vocab_size : 65 (base 64 + grown mask row) a2d : {"objective": "mdlm", "mask_token_id": 64, "final_alpha": 0.6666666666666666, "sampler": {"canvas_len": 8, "num_steps": 8, "temperature": 1.0}} attn-impl keys present: none 'a2d_annealed_eager_' in config.json: False reload from_pretrained : Gemma3ForCausalLM | attn_implementation = 'sdpa'

# Converted checkpoints produced by 'a2d convert' above, loaded back with plain HF from_pretrained.
# Checks the load-bearing claim: the per-model anneal registry key must NOT be persisted into
# config.json, or the shipped checkpoint would be unloadable by anyone without a2d in-process.

gpt2/model/config.json
  model_type            : gpt2
  vocab_size            : 65 (base 64 + grown mask row)
  a2d                   : {"objective": "mdlm", "mask_token_id": 64, "final_alpha": 0.6666666666666666, "sampler": {"canvas_len": 8, "num_steps": 8, "temperature": 1.0}}
  attn-impl keys present: none
  'a2d_annealed_eager_' in config.json: False
  reload from_pretrained : GPT2LMHeadModel | attn_implementation = 'sdpa'
  files                  : ['config.json', 'generation_config.json', 'model.safetensors', 'tokenizer.json', 'tokenizer_config.json']

gemma/model/config.json
  model_type            : gemma
  vocab_size            : 65 (base 64 + grown mask row)
  a2d                   : {"objective": "mdlm", "mask_token_id": 64, "final_alpha": 0.6666666666666666, "sampler": {"canvas_len": 8, "num_steps": 8, "temperature": 1.0}}
  attn-impl keys present: none
  'a2d_annealed_eager_' in config.json: False
  reload from_pretrained : GemmaForCausalLM | attn_implementation = 'sdpa'
  files                  : ['config.json', 'generation_config.json', 'model.safetensors', 'tokenizer.json', 'tokenizer_config.json']

gemma3/model/config.json
  model_type            : gemma3_text
  vocab_size            : 65 (base 64 + grown mask row)
  a2d                   : {"objective": "mdlm", "mask_token_id": 64, "final_alpha": 0.6666666666666666, "sampler": {"canvas_len": 8, "num_steps": 8, "temperature": 1.0}}
  attn-impl keys present: none
  'a2d_annealed_eager_' in config.json: False
  reload from_pretrained : Gemma3ForCausalLM | attn_implementation = 'sdpa'
  files                  : ['config.json', 'generation_config.json', 'model.safetensors', 'tokenizer.json', 'tokenizer_config.json']
Evidence: Seam hook-guard tests passing

test_bidir.py::test_future_token_reaches_earlier_positions_only_when_bidirectional PASSED test_bidir.py::test_install_routes_the_model_and_only_it_through_the_annealed_seam PASSED test_bidir.py::test_install_rejects_reorder_and_upcast_attn PASSED 3 passed

rootdir: /home/fleet/.no-mistakes/worktrees/5ae58bdd975c/01KYPMN1Q9YKTNQVTJBMY28BW1/packages/a2d-worker-hf
configfile: pyproject.toml
plugins: typeguard-4.5.2, anyio-4.14.2
collecting ... collected 3 items

packages/a2d-worker-hf/tests/test_bidir.py::test_future_token_reaches_earlier_positions_only_when_bidirectional PASSED [ 33%]
packages/a2d-worker-hf/tests/test_bidir.py::test_install_routes_the_model_and_only_it_through_the_annealed_seam PASSED [ 66%]
packages/a2d-worker-hf/tests/test_bidir.py::test_install_rejects_reorder_and_upcast_attn PASSED [100%]

============================== 3 passed in 1.43s ===============================
Evidence: Evidence index + reproduction scripts
# Evidence: a2d worker ported to transformers 5.14.1 (PR 1 of 2)

Environment: Linux aarch64, CPU only, `transformers 5.14.1`, `torch 2.13.0+cpu`.
All models are hermetic tiny random-weight fixtures built in-process from the repo's own
`packages/a2d-worker-hf/tests/conftest.py` savers. No weights were downloaded.

| file                          | what it shows                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               |
| ----------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `cli-convert-sample.txt`      | Real `a2d detect` -> `a2d convert` -> `a2d sample` transcript for all three attention seams (GPT-2 `attn.full`, Gemma 1 `attn.gqa`, Gemma 3 `attn.swa`). Every run: identity gate PASS with `max_abs_diff=0.00e0`, anneal climbing 0.000 -> 0.667 over the training steps, checkpoint written, `job completed`, exit 0, then a diffusion decode off the converted checkpoint.                                                                                                                                               |
| `family-matrix.txt`           | Per-family matrix through `resolve_capabilities` + `apply_transforms` (the same dispatch `worker.py` runs) for gpt2, llama, qwen2, gemma, gemma2, gemma3, mistral: capability resolved, a2d's key live in `ALL_MASK_ATTENTION_FUNCTIONS`, `alpha=0` max abs diff exactly `0.0`, `alpha=1` shift non-zero, and an un-patched sibling of the same model in the same process still causal (0.0 shift). Mistral stays `attn.gqa` and its `alpha=1` shift (0.0197) equals Llama's exactly, i.e. its sliding window opened fully. |
| `checkpoint-config-check.txt` | The converted checkpoints reloaded with plain `AutoModelForCausalLM.from_pretrained`: no `a2d_annealed_eager_*` key and no `_attn_implementation` persisted into `config.json`, `vocab_size` grown 64 -> 65, `a2d` block present.                                                                                                                                                                                                                                                                                           |
| `seam-guard-tests.txt`        | The new hook-guard tests (`test_install_routes_the_model_and_only_it_through_the_annealed_seam`, `test_install_rejects_reorder_and_upcast_attn`) passing.                                                                                                                                                                                                                                                                                                                                                                   |
| `work/`                       | Reproduction scripts (`make_models.py`, `family_matrix.py`) plus the fixture model dirs and the run dirs the CLI produced.                                                                                                                                                                                                                                                                                                                                                                                                  |

Pipeline

Updates from git push no-mistakes

✅ **intent** - passed

✅ No issues found.

✅ **Rebase** - passed

✅ No issues found.

⚠️ **Review** - 3 infos
  • ⚠️ packages/a2d-worker-hf/src/a2d_core/transform/attention.py:137 - set_attn_implementation is not guaranteed to apply the key. In modeling_utils.py (v5) it does if not self._can_set_attn_implementation(): logger.warning(...) and then falls through without assigning — no exception. _can_set_attn_implementation is a source-inspection heuristic: it returns False when sys.modules.get(cls.__module__) is None, when inspect.getsource raises OSError/TypeError (source-less deployment: frozen binary, zipimport, .pyc-only wheel), or when the regex finds an *Attention*(nn.Module) class but &#34;ALL_ATTENTION_FUNCTIONS.get_interface(&#34; is absent from the module text — a string HF can change in any patch release. When that happens the model stays on &#34;eager&#34;, keeps HF's own causal mask at every alpha, and the D13 gate still passes (identity.py says so explicitly), so conversion ships a causal model labelled bidirectional. test_install_routes_the_model_and_only_it_through_the_annealed_seam catches it for GPT-2 in this version only, and cannot catch the source-less case because tests always run from source. Make it structural: after model.set_attn_implementation(key), if model.config._attn_implementation != key: raise ValueError(...).
  • ⚠️ packages/a2d-worker-hf/src/a2d_core/transform/attention.py:137 - Moving the config off the literal string &#34;eager&#34; changes which branch GPT-2 takes. modeling_gpt2.py computes using_eager = self.config._attn_implementation == &#34;eager&#34; and then if using_eager and self.reorder_and_upcast_attn: ... _upcast_and_reordered_attn(...) else: attention_interface(...). For a GPT-2 checkpoint with reorder_and_upcast_attn=True (a public GPT2Config flag, used by Megatron-derived checkpoints), the un-patched reference copy loaded in the same process still takes _upcast_and_reordered_attn while the patched copy now takes _a2d_eager_attention, so base and patched@alpha=0 compute attention by different code paths. The identity gate then fails with a bare numeric max_abs_diff and no explanation of why. Detect classifies such a config as attn.full, which is non-blocking, so it reaches the worker. Pre-existing in a different form (in 4.51.3 the _upcast_and_reordered_attn path never called the patched eager_attention_forward either), but this change makes the divergence structural rather than incidental. Suggest either rejecting reorder_and_upcast_attn=True in install_anneal_patch with an explicit message, or folding it into the seam.
  • ℹ️ packages/a2d-worker-hf/pyproject.toml:12 - The pin comment added by this commit says to see transform/gqa_attention.py and transform/attention.py "for the v5 seams", but after the port gqa_attention.py holds no seam code at all — it is a six-line capability gate (is_rope_family plus a delegating install_gqa_anneal_patch). The entire v5 seam is annealed_eager_mask / _a2d_eager_attention / install_mask_anneal in transform/attention.py. This comment is newly authored in this PR, not part of the known-stale prose deferred to PR 2, so it is worth pointing at the one file that actually holds the seam.
  • ℹ️ packages/a2d-worker-hf/tests/test_swa_attention.py:30 - _sliding() reads model.config.layer_types, the same field has_sliding_window_seam dispatches on, so the four "sanity: the stack really does mix local and global layers" assertions no longer confirm anything independent of the production gate. v5 does expose a per-module equivalent (Gemma2DecoderLayer.layer_type / Gemma3DecoderLayer.layer_type, set at init from config.layer_types) that would restore the module-side check the old layer.is_sliding gave. Low impact: the behavioural assertions in the same tests (_shift(..., k=0, alpha=0.0) == 0.0 and &gt; 1e-6 at alpha=1) are the real proof and are unchanged, so this is a note rather than a gap.

🔧 Fix: verify anneal seam install and reject reorder_and_upcast_attn
3 infos still open:

  • ℹ️ packages/a2d-worker-hf/src/a2d_core/transform/attention.py:128 - implementation.startswith(_KEY_PREFIX) raises AttributeError: &#39;NoneType&#39; object has no attribute &#39;startswith&#39; when config._attn_implementation is None, instead of the intended ValueError naming Decision 2. PreTrainedConfig.__init__ sets _attn_implementation = kwargs.pop(&#34;attn_implementation&#34;, None), so a config never dispatched through a model carries None (verified: GPT2Config(...)._attn_implementation is None). The old code compared with != and degraded gracefully. Unreachable via a2d's own paths - load_model and every conftest fixture pass attn_implementation=&#34;eager&#34; - so this is error-message quality only, but apply_transforms/generate_canvas accept a caller-supplied model. Fix: (implementation or &#34;&#34;).startswith(...).
  • ℹ️ packages/a2d-worker-hf/src/a2d_core/transform/identity.py:36 - check_identity's docstring says "patched is the patched (possibly grown) model whose _a2d_anneal is state", but this PR deletes the _a2d_anneal module tag entirely - the state is now reached through the per-state key in ALL_MASK_ATTENTION_FUNCTIONS. This is the only remaining live reference to a removed symbol outside the files the author explicitly deferred to the docs-only follow-up (attention.py, gqa_attention.py, swa_attention.py, apply.py, handlers/*, worker.py, conftest.py, test_{gqa,swa}_attention.py, root pyproject, AGENTS.md), so it will be missed unless identity.py is added to that list.
  • ℹ️ packages/a2d-worker-hf/src/a2d_core/transform/attention.py:134 - Each install_mask_anneal adds two permanent entries to HF's process-global AttentionMaskInterface/AttentionInterface _global_mapping dicts, and the mask closure keeps its AnnealState alive forever. That strong ref is load-bearing (it is what stops id(state) from being recycled onto a stale key), so it is a deliberate tradeoff, not a defect. Worth noting because sample/denoiser.py:48 re-installs with a fresh AnnealState(alpha=1.0) on every generate_canvas call, so a long sampling or eval loop in one process grows both registries unboundedly (~2 entries + one dataclass per sample). Entries are tiny and the worker is one-job-per-process, so no action needed; it only matters if sampling ever moves into a long-lived service.
✅ **Test** - passed

✅ No issues found.

  • uv sync --frozen then uv run python -c &#34;import transformers, torch&#34; (confirmed transformers 5.14.1 / torch 2.13.0+cpu on aarch64 CPU)
  • uv run pytest -q (full worker suite, 93 passed; run twice to check the per-AnnealState registry key does not make results order-dependent)
  • uv run pytest packages/a2d-worker-hf/tests/test_bidir.py -v (the new seam hook-guard + reorder_and_upcast_attn rejection tests)
  • cargo test --workspace (Rust detect/contracts/run/cli suites, 47 passed)
  • cargo build -p a2d-cli then manual E2E: ./target/debug/a2d detect &lt;fixture&gt;, ./target/debug/a2d convert &lt;fixture&gt; --out &lt;run&gt; --data corpus.jsonl --seq-len 8 --max-steps 3 --anneal-steps 3 --per-device-batch-size 2 --device cpu --dtype float32 and ./target/debug/a2d sample &lt;run&gt; -p &#39;w1 w2 w3&#39; --canvas-len 8 --num-steps 4 --device cpu for gpt2 (attn.full), gemma (attn.gqa) and gemma3 (attn.swa)
  • Manual per-family verification script driving resolve_capabilities + apply_transforms over the repo's tiny fixtures for gpt2, llama, qwen2, gemma, gemma2, gemma3 and mistral: alpha=0 max abs diff vs base, alpha=1 shift at an earlier position, a2d key membership in transformers.masking_utils.ALL_MASK_ATTENTION_FUNCTIONS, and an un-patched sibling model's shift
  • Manual reload of each converted checkpoint with AutoModelForCausalLM.from_pretrained, plus a grep of the saved config.json for a2d_annealed_eager_ / _attn_implementation
⚠️ **Document** - 3 infos
  • ℹ️ docs/PLAN-PHASE2.md:191 - docs/PLAN-PHASE2.md still documents Decision 2 as the 4.48.3 GPT-2 self.bias monkeypatch (lines 58, 191-197, 220, 267, 341-342), including Pin transformers==4.48.3 and torch_dtype=float32. Left alone on purpose: the file is stamped Status: approved plan, pre-implementation and was already stale against the 4.51.3 pin before this change, so the repo convention appears to be that plan docs are frozen historical records. Needs a human call: freeze as history, add a superseded-by note, or rewrite to the v5 seam.
  • ℹ️ packages/a2d-worker-hf/src/a2d_core/transform/attention.py:156 - The port adds a new convert-time hard rejection (reorder_and_upcast_attn=True) that no user-facing doc mentions and that the Rust detect gate does not screen for - reorder_and_upcast appears nowhere in crates/, docs/, or README.md. A GPT-2 config with that flag passes a2d detect and only fails at convert, which cuts against the SPEC's "gate before GPU" goal. I documented it in the code doc comments, but whether it belongs in the README limits table (capability-keyed today) or should instead be gated in detect is a design call, not a doc edit.

🔧 Fix: revert docs resync; defer prose to stacked docs PR
3 infos still open:

  • ℹ️ packages/a2d-worker-hf/src/a2d_core/transform/attention.py:1 - Per the user's ruling I reverted commit 9c434ac wholesale instead of syncing docs, so PR 1 knowingly ships stale transformers-4.x prose: _update_causal_mask, GPT-2 self.bias, per-layer is_sliding and the old layer-forward wrap in transform/{attention,gqa_attention,swa_attention,apply}.py, transform/handlers/{full,gqa,swa}_attention.py, transform/identity.py, worker.py step-4 comment, tests/{conftest,test_gqa_attention,test_swa_attention}.py, the root pyproject.toml mypy py.typed comment, AGENTS.md, docs/CONCEPTS.md and docs/SPEC-HANDOFF.md. Not a defect of this PR - it is the agreed split seam. The resynced text is preserved verbatim in the reverted commit 9c434ac; note that branch fm/a2d-dep-transformers-t1-docs does not exist in this gate repo, so PR 2 needs git cherry-pick 9c434ac (or the equivalent) or that prose is lost.
  • ℹ️ docs/PLAN-PHASE2.md:191 - Frozen as history per the user's ruling; no edit made in PR 1. The one-line superseded-by note lands in the docs-only PR 2.
  • ℹ️ packages/a2d-worker-hf/src/a2d_core/transform/attention.py:152 - No change in PR 1 per the user's ruling: detect-side screening for reorder_and_upcast_attn is a Rust contract change tracked as its own task, not a doc line. The convert-time rejection stays undocumented in README/docs by that decision.
✅ **Lint** - passed

✅ No issues found.

✅ **Push** - passed

✅ No issues found.

undeemed added 4 commits July 29, 2026 09:52
1 of 2. The docs-only follow-up (PR 2) refreshes the prose this PR leaves
stale; see the note at the end.

v5 unified the three eager causal seams a2d patched. GPT-2 no longer bakes
causality into a per-layer `self.bias` buffer, the RoPE family no longer builds
its mask in a per-model `_update_causal_mask` method, and Gemma 2/3 no longer
apply their sliding window inside each local decoder layer's forward. Every
family now routes causality through `transformers.masking_utils`, which builds
the 4D additive mask via `ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation]`.

Re-target all three onto that one seam. `install_mask_anneal` registers a mask
function (and the matching attention function) under a key private to its
`AnnealState` and points only that model's config at it, so isolation is per
model rather than per process - the D13 identity gate's un-patched reference
copy keeps HF's own causal mask. The registered function calls HF's `eager_mask`
and re-reveals every cell that mask masked for a real, non-padded key, so
alpha=0 is bit-identical to base by construction and alpha=1 opens the future,
Mistral's single-mask window and Gemma 2/3's separate sliding mask alike.

`attn.full` / `attn.gqa` / `attn.swa` stay distinct capabilities (detect's
contract, the handler registry) but are now three structural gates on one
install: `resolve_capabilities` reads `config.layer_types` for
`sliding_attention`, then `num_key_value_heads`, then `GPT2Attention`.

A registry key that HF silently ignores would leave the model fully causal at
every alpha, so add `test_install_routes_the_model_and_only_it_through_the_annealed_seam`
asserting the key is live in `ALL_MASK_ATTENTION_FUNCTIONS`, that a sibling
model stays on `eager`, and that only the patched one opens at alpha=1.

Also: `from_pretrained(torch_dtype=)` -> `dtype=`.

KNOWN-STALE PROSE, fixed in PR 2: module docstrings and comments in
transform/attention.py, transform/gqa_attention.py, transform/swa_attention.py,
transform/apply.py, transform/handlers/{gqa,swa}_attention.py, worker.py,
tests/conftest.py, tests/test_{gqa,swa}_attention.py, the root pyproject.toml
mypy override comment and AGENTS.md still describe 4.x internals
(`_update_causal_mask`, `self.bias`, per-layer `is_sliding`). The code is the
truth in this PR; those docstrings are not.
@undeemed
undeemed merged commit 6d66642 into main Jul 29, 2026
3 checks passed
@undeemed
undeemed deleted the fm/a2d-dep-transformers-t1 branch July 29, 2026 14:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant