build(worker): port bidirectionalization seam to transformers 5.14.1 - #17
Merged
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
Scope and process constraints the user set:
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
transformers==4.51.3to5.14.1and rewrote Decision 2's bidirectionalization onto v5's mask interface. v5 removed all three eager seams the old code patched (GPT-2's per-layerself.biasbuffer, the RoPE family's per-model_update_causal_mask, Gemma 2/3's per-decoder-layer sliding-window re-mask);install_mask_annealnow registers a per-AnnealStatekey (a2d_annealed_eager_<id(state)>) inAttentionMaskInterface/AttentionInterfaceand points only that model'sconfig._attn_implementationat it, so an un-patched sibling in the same process keeps HF's own causal mask. The annealed reveal derives frommasking_utils.eager_mask's own output (masked cells atfinfo.min, only revealed when the key is a real token), and_a2d_eager_attentiondelegates tosys.modules[type(module).__module__].eager_attention_forwardso each family's eager maths is unchanged.load_modelmoves fromtorch_dtype=todtype=.attn.full/attn.gqa/attn.swastay 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_seamreadsconfig.layer_types, the newis_rope_familyreadsnum_key_value_heads, GPT-2 is still keyed onGPT2Attention.gqa_attention.pyandswa_attention.pydrop their hand-rolled mask wrapper and signature-agnostic layer-forward wrap entirely and delegate to the shared install.install_mask_annealraises ifset_attn_implementationleaves the config off the a2d key (a silently-ignored key would look causal at every alpha while still passing the D13 identity gate), andinstall_anneal_patchrejectsreorder_and_upcast_attn=Trueby name rather than as a baremax_abs_diffmismatch.tests/test_bidir.pyasserts the key is live inALL_MASK_ATTENTION_FUNCTIONS, that a sibling model stays on"eager", and that only the patched model shifts earlier-position logits atalpha=1; thetiny_gpt2fixture 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 excludinguv.lock): it knowingly leaves 4.x prose in the touched modules,docs/, andAGENTS.md, which the stacked docs-only PR refreshes.Risk Assessment
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) andcargo test --workspace(47 passed) as the baseline, then went past unit level: built the Rust CLI and drove the actual end-user patha2d detect->a2d convert->a2d sampleover hermetic tiny fixtures for all three attention seams, where every conversion reportedidentity 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=0Evidence: 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 OKEvidence: 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'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 passedEvidence: Evidence index + reproduction scripts
Pipeline
Updates from git push no-mistakes
✅ **intent** - passed
✅ No issues found.
✅ **Rebase** - passed
✅ No issues found.
packages/a2d-worker-hf/src/a2d_core/transform/attention.py:137-set_attn_implementationis not guaranteed to apply the key. Inmodeling_utils.py(v5) it doesif not self._can_set_attn_implementation(): logger.warning(...)and then falls through without assigning — no exception._can_set_attn_implementationis a source-inspection heuristic: it returns False whensys.modules.get(cls.__module__)is None, wheninspect.getsourceraises OSError/TypeError (source-less deployment: frozen binary, zipimport, .pyc-only wheel), or when the regex finds an*Attention*(nn.Module)class but"ALL_ATTENTION_FUNCTIONS.get_interface("is absent from the module text — a string HF can change in any patch release. When that happens the model stays on"eager", 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_seamcatches it for GPT-2 in this version only, and cannot catch the source-less case because tests always run from source. Make it structural: aftermodel.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"eager"changes which branch GPT-2 takes.modeling_gpt2.pycomputesusing_eager = self.config._attn_implementation == "eager"and thenif using_eager and self.reorder_and_upcast_attn: ... _upcast_and_reordered_attn(...) else: attention_interface(...). For a GPT-2 checkpoint withreorder_and_upcast_attn=True(a publicGPT2Configflag, used by Megatron-derived checkpoints), the un-patched reference copy loaded in the same process still takes_upcast_and_reordered_attnwhile 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 numericmax_abs_diffand no explanation of why. Detect classifies such a config asattn.full, which is non-blocking, so it reaches the worker. Pre-existing in a different form (in 4.51.3 the_upcast_and_reordered_attnpath never called the patchedeager_attention_forwardeither), but this change makes the divergence structural rather than incidental. Suggest either rejectingreorder_and_upcast_attn=Trueininstall_anneal_patchwith 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 seetransform/gqa_attention.pyandtransform/attention.py"for the v5 seams", but after the portgqa_attention.pyholds no seam code at all — it is a six-line capability gate (is_rope_familyplus a delegatinginstall_gqa_anneal_patch). The entire v5 seam isannealed_eager_mask/_a2d_eager_attention/install_mask_annealintransform/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()readsmodel.config.layer_types, the same fieldhas_sliding_window_seamdispatches 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 fromconfig.layer_types) that would restore the module-side check the oldlayer.is_slidinggave. Low impact: the behavioural assertions in the same tests (_shift(..., k=0, alpha=0.0) == 0.0and> 1e-6at 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)raisesAttributeError: 'NoneType' object has no attribute 'startswith'whenconfig._attn_implementationisNone, instead of the intended ValueError naming Decision 2.PreTrainedConfig.__init__sets_attn_implementation = kwargs.pop("attn_implementation", 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_modeland every conftest fixture passattn_implementation="eager"- so this is error-message quality only, butapply_transforms/generate_canvasaccept a caller-supplied model. Fix:(implementation or "").startswith(...).packages/a2d-worker-hf/src/a2d_core/transform/identity.py:36-check_identity's docstring says "patchedis the patched (possibly grown) model whose_a2d_annealisstate", but this PR deletes the_a2d_annealmodule tag entirely - the state is now reached through the per-state key inALL_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- Eachinstall_mask_annealadds two permanent entries to HF's process-globalAttentionMaskInterface/AttentionInterface_global_mappingdicts, and the mask closure keeps itsAnnealStatealive forever. That strong ref is load-bearing (it is what stopsid(state)from being recycled onto a stale key), so it is a deliberate tradeoff, not a defect. Worth noting becausesample/denoiser.py:48re-installs with a freshAnnealState(alpha=1.0)on everygenerate_canvascall, 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 --frozenthenuv run python -c "import transformers, torch"(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-clithen manual E2E:./target/debug/a2d detect <fixture>,./target/debug/a2d convert <fixture> --out <run> --data corpus.jsonl --seq-len 8 --max-steps 3 --anneal-steps 3 --per-device-batch-size 2 --device cpu --dtype float32and./target/debug/a2d sample <run> -p 'w1 w2 w3' --canvas-len 8 --num-steps 4 --device cpufor gpt2 (attn.full), gemma (attn.gqa) and gemma3 (attn.swa)Manual per-family verification script drivingresolve_capabilities+apply_transformsover 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 intransformers.masking_utils.ALL_MASK_ATTENTION_FUNCTIONS, and an un-patched sibling model's shiftManual reload of each converted checkpoint withAutoModelForCausalLM.from_pretrained, plus a grep of the savedconfig.jsonfora2d_annealed_eager_/_attn_implementationdocs/PLAN-PHASE2.md:191- docs/PLAN-PHASE2.md still documents Decision 2 as the 4.48.3 GPT-2self.biasmonkeypatch (lines 58, 191-197, 220, 267, 341-342), includingPin transformers==4.48.3andtorch_dtype=float32. Left alone on purpose: the file is stampedStatus: approved plan, pre-implementationand 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_upcastappears nowhere in crates/, docs/, or README.md. A GPT-2 config with that flag passesa2d detectand 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-2self.bias, per-layeris_slidingand 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 needsgit 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 forreorder_and_upcast_attnis 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.