diff --git a/AGENTS.md b/AGENTS.md index ea83902..9cf79a8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -4,8 +4,8 @@ This file is the project's committed home for project-intrinsic agent knowledge: - **What a2d is + the extension playbook:** `docs/SPEC-HANDOFF.md` (esp. §3 open-closed model, §5 repo layout) and `docs/ARCHITECTURE.md` (the ML recipe: mask annealing, MDLM, D13 identity gate). Adding a model/capability is additive by policy — one file per quirk, never edit an existing adapter/handler. - **Detect (Rust, config-only, no GPU):** `crates/a2d-detect/`. New `model_type` quirk = one file in `src/adapters/` delegating to `generic::detect` + `inferred=false`, one `mod` line in `adapters/mod.rs`, and one `fixtures/configs//{config.json,expected.json}` (auto-picked up by `tests/corpus.rs`). Capabilities are the unit of support, not model names; the gate blocks only conversion-blocking caps. -- **Convert handlers (Python worker):** `packages/a2d-worker-hf/src/a2d_core/transform/`. Three eager causal seams — GPT-2 bakes causality into per-layer `self.bias` (`attn.full`, `attention.py`); the RoPE family (Llama/Qwen2/Gemma 1) routes it through the 4D mask `_update_causal_mask` builds (`attn.gqa`, `gqa_attention.py` - the reveal opens EVERY base-masked real-key cell, so single-mask windowed families like Mistral v0.1 or Qwen2 with active `use_sliding_window` also unwindow through this seam); Gemma 2/3 add a per-layer sliding window on their *local* decoder layers on top of that same full mask (`attn.swa`, `swa_attention.py` - anneals BOTH the full-causal future-reveal AND each sliding layer's far-past window via a signature-agnostic layer-forward wrap; `resolve_capabilities` checks SWA before GQA because Gemma 2/3 own both seams). `worker.py` picks the handler by structural introspection (`resolve_capabilities` in `apply.py`), not from job tags — the `ConversionJob` carries no capability set. D13 rule: patched@`alpha=0` must equal base logits bit-for-bit. `attn.swa` is a SUPPORTED, non-blocking capability (`crates/a2d-contracts/src/lib.rs` `blocking()`); `attn.sink`/`attn.mla`/SSM still reject. -- **HF transformers is pinned to `4.51.3`** (`packages/a2d-worker-hf/pyproject.toml`), the first pin shipping Gemma 3; the GPT-2 eager seam and Gemma 1 `_update_causal_mask` are unchanged from the old `4.48.3`. Two gotchas: 4.51.3 ships an (empty) `py.typed`, so the root `pyproject.toml` mypy override keeps `transformers`/`tokenizers` untyped via `follow_imports = "skip"` (not just `ignore_missing_imports`); and in 4.51.3 Gemma 3's sliding window is applied per-decoder-layer in `Gemma3DecoderLayer.forward`, NOT via a model-level mask-mapping dict (that is a later-transformers refactor) - verify the seam before assuming. +- **Convert handlers (Python worker):** `packages/a2d-worker-hf/src/a2d_core/transform/`. `transformers>=5` unified what used to be three distinct eager seams (GPT-2's `self.bias`, the RoPE family's `_update_causal_mask`, Gemma 2/3's per-layer window) into ONE: every family's 4D additive mask is built by `transformers.masking_utils` through `ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation]`. `attention.py` `install_mask_anneal` registers a per-`AnnealState` key in that registry (plus `ALL_ATTENTION_FUNCTIONS`, same key) and points only that model's config at it, so isolation is per model — the identity gate's un-patched reference copy is unaffected. The reveal derives FROM `eager_mask`'s output (every base-masked cell with a real, non-padded key), which is why `alpha=0` is bit-identical by construction; it opens strictly-future cells, Mistral-style single-mask windows, and Gemma 2/3's separate `sliding_attention` mask alike. `attn.full`/`attn.gqa`/`attn.swa` stay distinct capabilities (detect's contract, registry keys) but are now three gates on one install: `resolve_capabilities` (`apply.py`) checks `config.layer_types` for `sliding_attention` (swa) before `num_key_value_heads` (gqa) before `GPT2Attention` (full), reading the model's own structure, not job tags — the `ConversionJob` carries no capability set. D13 rule: patched@`alpha=0` must equal base logits bit-for-bit. `attn.swa` is a SUPPORTED, non-blocking capability (`crates/a2d-contracts/src/lib.rs` `blocking()`); `attn.sink`/`attn.mla`/SSM still reject. +- **HF transformers is pinned to `5.14.1`** (`packages/a2d-worker-hf/pyproject.toml`). Two gotchas: it still ships an (empty) `py.typed`, and its inline annotations disagree with the runtime (`**kwargs` config fields, `Trainer.compute_loss`), so the root `pyproject.toml` mypy override keeps `transformers`/`tokenizers` untyped via `follow_imports = "skip"` (not just `ignore_missing_imports`) — dropping it reintroduces 68 errors. And a silently-not-applied mask patch is the dangerous failure mode now that the seam is a registry key rather than a monkeypatch: `test_bidir.py::test_install_routes_the_model_and_only_it_through_the_annealed_seam` is the guard. - **Hermetic tests only:** never download weights (Gemma is gated; CI is CPU/no-network). Build tiny random-weight configs in-process (see `tests/conftest.py` `tiny_gpt2`/`tiny_gqa`/`tiny_gemma2`/`tiny_gemma3`/`tiny_mistral`). - **Full CI gate before shipping:** `cargo fmt --all --check`, `cargo clippy --workspace --all-targets -- -D warnings`, `cargo test --workspace`; `uv run ruff check .`, `uv run ruff format --check .`, `uv run mypy` (strict), `uv run pytest`; contracts are generated — after touching `crates/a2d-contracts/` run `bash scripts/codegen.sh` (CI fails on `schema/`+`packages/a2d-contracts/` drift). diff --git a/crates/a2d-contracts/src/lib.rs b/crates/a2d-contracts/src/lib.rs index c7630b0..2293fff 100644 --- a/crates/a2d-contracts/src/lib.rs +++ b/crates/a2d-contracts/src/lib.rs @@ -180,12 +180,11 @@ impl Capability { /// True ONLY for the five caps Phase 1 cannot handle. Every implemented /// conversion cap - now including `AttnSwa`, whose sliding-window anneal the - /// worker performs for both flavors (per-layer Gemma 2/3 via the `attn.swa` - /// handler; single-mask Mistral/Qwen2 via the shared `attn.gqa` mask reveal) - - /// and every fidelity cap - /// returns false, so fidelity tags cannot block by construction. Flipping - /// `AttnSink` here later is the remaining "enable GPT-OSS" change - /// (ARCHITECTURE 5's "flip the gate"). + /// worker performs for both flavors (a per-layer sliding mask, Gemma 2/3, via + /// the `attn.swa` handler; single-mask Mistral via the shared `attn.gqa` mask + /// reveal) - and every fidelity cap returns false, so fidelity tags cannot + /// block by construction. Flipping `AttnSink` here later is the remaining + /// "enable GPT-OSS" change (ARCHITECTURE 5's "flip the gate"). pub fn blocking(self) -> bool { matches!( self, diff --git a/docs/CONCEPTS.md b/docs/CONCEPTS.md index 771e57c..675e314 100644 --- a/docs/CONCEPTS.md +++ b/docs/CONCEPTS.md @@ -13,7 +13,7 @@ It predicts the next token left-to-right, and a **causal mask** enforces the rul The block of positions it works on is the **canvas**; turning masks back into tokens over N steps is **denoising**. 3. **a2d's whole job** is the conversion between them: take AR weights, **anneal** the causal mask off (the `attn.full` / `attn.gqa` / `attn.swa` transforms), and briefly retrain with a masking objective (**MDLM**). -One transform per attention seam plus one objective. That is the core. +One transform per attention capability - all three sharing one mask seam - plus one objective. That is the core. Everything else in the vocabulary exists to decide *which models a2d will touch* - the complexity lives in honest **detection**, not in the small **conversion**. @@ -24,7 +24,7 @@ Everything else in the vocabulary exists to decide *which models a2d will touch* - **Bidirectional** - a token sees left *and* right. What diffusion needs. - **MDLM** - the masked-diffusion objective a2d uses: mask a random fraction, predict them, repeat. - **Canvas** - the span of positions a diffusion model denoises (its workspace). -- **Anneal (`attn.full`, `attn.gqa`, `attn.swa`)** - a2d's trick: slowly turn the causal mask off so AR weights adapt to bidirectional. One transform per attention seam: GPT-2 bakes causality per layer (`attn.full`); the RoPE family (Gemma 1 / Qwen2 / Llama / Mistral) routes it through one shared mask (`attn.gqa`); Gemma 2/3 add a per-layer sliding window on top of that mask (`attn.swa`). +- **Anneal (`attn.full`, `attn.gqa`, `attn.swa`)** - a2d's trick: slowly turn the causal mask off so AR weights adapt to bidirectional. `transformers` v5 builds every family's mask through one shared interface, so there is ONE anneal seam with three capability gates: GPT-2's dense attention (`attn.full`); the RoPE family (Gemma 1 / Qwen2 / Llama / Mistral), which routes causality - and Mistral's window - through that single mask (`attn.gqa`); models whose `config.layer_types` names per-layer sliding layers - Gemma 2/3 - which additionally take a sliding-window mask from the same interface (`attn.swa`). - **Identity gate** - hard correctness check: at `anneal=0` the patched model must match the base model's logits, or convert aborts. - **GQA** - grouped-query attention: fewer key/value heads than query heads (a memory trick). The mechanism rides along in HF's forward; the `attn.gqa` tag also names the RoPE-family anneal transform (see Anneal). - **RoPE** - rotary position encoding (modern position scheme). Passthrough. diff --git a/docs/PLAN-PHASE2.md b/docs/PLAN-PHASE2.md index 3bb0fb2..89c6c2c 100644 --- a/docs/PLAN-PHASE2.md +++ b/docs/PLAN-PHASE2.md @@ -1,6 +1,7 @@ # a2d Phase 2 Implementation Plan - Conversion core: dense happy path **Status:** approved plan, pre-implementation. +**Superseded in part:** Decision 2's GPT-2 `self.bias` monkeypatch and the `transformers==4.48.3` pin below are historical - the shipped anneal installs through the transformers v5 mask interface in [`transform/attention.py`](../packages/a2d-worker-hf/src/a2d_core/transform/attention.py) - and this plan stays frozen as approved rather than being rewritten. **Source spec:** [`SPEC-HANDOFF.md`](SPEC-HANDOFF.md) §4.2 (conversion pipeline: mask annealing, shift removal, identity gate, MDLM), §4.3 (run-dir target), §6 (Phase 2 scope and exit criteria), §9.4 (dllm wrap-vs-implement open decision); [`ARCHITECTURE.md`](ARCHITECTURE.md) M0 recipe (the three things conversion touches) and D13 (the identity test). ## Context @@ -385,4 +386,4 @@ cat runs/gpt2-diffusion/manifest.json # status completed; model_spec + convers Accepted with `uv` caching; an optional-extra dependency split is the documented upgrade path. 15. **Scope discipline: MoE-router-under-anneal + finetune (P4), BD3LM + schedulers (P5), the eval harness and eval-parity (P3), golden fixtures (candle-track), `head.py`, and the `datasets` dep are explicitly NOT built in P2.** - Flagged so they are not silently smuggled in; the MoE router monitor and eval parity in particular wait for their own phases because Phase 2 is the dense GPT-2 happy path only. \ No newline at end of file + Flagged so they are not silently smuggled in; the MoE router monitor and eval parity in particular wait for their own phases because Phase 2 is the dense GPT-2 happy path only. diff --git a/docs/SPEC-HANDOFF.md b/docs/SPEC-HANDOFF.md index 2a17cec..3209d16 100644 --- a/docs/SPEC-HANDOFF.md +++ b/docs/SPEC-HANDOFF.md @@ -216,9 +216,9 @@ a2d/ src/a2d_core/ ingest/ # ◄ EXTENSION POINT: format normalizers (copy-on-normalize) transform/ - attention.py # GPT-2 self.bias seam (attn.full) + AnnealState/schedule - gqa_attention.py # RoPE-family _update_causal_mask seam (attn.gqa) - swa_attention.py # Gemma 2/3 per-layer sliding-window seam (attn.swa) + attention.py # shared v5 mask-interface seam (all attn.*) + AnnealState/schedule + gqa_attention.py # RoPE-family gate on that shared seam (attn.gqa) + swa_attention.py # sliding-window-layer gate on that shared seam (attn.swa) apply.py # load model, resolve capabilities from its seam, apply handlers handlers/ # ◄ EXTENSION POINT: capability handlers (attn.full, attn.gqa, attn.swa, ffn.moe, …) objectives/ # ◄ EXTENSION POINT: mdlm.py, bd3lm.py (corrupt/loss iface) diff --git a/packages/a2d-worker-hf/src/a2d_core/sample/denoiser.py b/packages/a2d-worker-hf/src/a2d_core/sample/denoiser.py index 5ff73da..eec7d37 100644 --- a/packages/a2d-worker-hf/src/a2d_core/sample/denoiser.py +++ b/packages/a2d-worker-hf/src/a2d_core/sample/denoiser.py @@ -36,8 +36,9 @@ def denoise( ``prompt_ids`` is the untouched prefix; the ``canvas_len - len(prompt_ids)`` suffix positions start masked and are filled by iterative confidence reveal. - Installs the model's resolved attention transform (GPT-2 -> ``attn.full``, RoPE - family -> ``attn.gqa``) at ``alpha=1`` so attention is bidirectional. + Installs the model's resolved attention transform (Gemma 2/3 -> ``attn.swa``, RoPE + family -> ``attn.gqa``, GPT-2 -> ``attn.full``) at ``alpha=1`` so attention is + bidirectional. """ prompt_len = len(prompt_ids) if canvas_len < prompt_len: diff --git a/packages/a2d-worker-hf/src/a2d_core/transform/apply.py b/packages/a2d-worker-hf/src/a2d_core/transform/apply.py index 3c60f8f..dae7267 100644 --- a/packages/a2d-worker-hf/src/a2d_core/transform/apply.py +++ b/packages/a2d-worker-hf/src/a2d_core/transform/apply.py @@ -68,17 +68,21 @@ def resolve_mask_token(model: Any, tokenizer: Any, strategy: str = "grow") -> in def resolve_capabilities(model: Any) -> list[str]: - """The attention-handler capabilities a loaded model needs, chosen by its eager - causal seam (Decision 2). Three disjoint seams exist in this scope: - - - Gemma 2/3 route causality through ``_update_causal_mask`` AND add a per-layer - sliding window on their local decoder layers -> ``attn.swa``. Checked FIRST, - because these models also own ``_update_causal_mask`` (the ``attn.gqa`` signal); - the swa handler subsumes the gqa future-reveal and additionally opens the window. - - The RoPE family (Llama/Qwen2/Gemma 1) routes causality through the 4D mask that - ``_update_causal_mask`` builds -> ``attn.gqa`` (covers GQA, MQA, and full-attn - RoPE alike; the mask is family-independent). - - GPT-2 bakes causality into a per-layer ``self.bias`` buffer -> ``attn.full``. + """The attention-handler capabilities a loaded model needs, chosen by the structure + of its attention stack (Decision 2). Three disjoint shapes exist in this scope: + + - Gemma 2/3 give their local decoder layers a sliding-window mask, named by a + ``sliding_attention`` entry in ``config.layer_types`` -> ``attn.swa``. Checked + FIRST, because these models are RoPE-family too (the ``attn.gqa`` signal); the + swa handler subsumes the gqa future-reveal and additionally opens the window. + - The RoPE family (Llama/Qwen2/Gemma 1, and Mistral, which folds its window into + the one model-level mask) declares ``num_key_value_heads`` -> ``attn.gqa`` + (covers GQA, MQA, and full-attn RoPE alike; the mask is family-independent). + - GPT-2 has none of those and its own dense ``GPT2Attention`` -> ``attn.full``. + + All three now share ONE eager seam - the ``transformers>=5`` mask interface (see + ``transform/attention.py``) - but stay distinct capabilities because that is what + detect reports and what the handler registry is keyed on. The seam is read from the model itself, not from detect's tags: the ``ConversionJob`` does not carry the capability set, so the worker independently picks the correct, diff --git a/packages/a2d-worker-hf/src/a2d_core/transform/attention.py b/packages/a2d-worker-hf/src/a2d_core/transform/attention.py index c886561..af066e3 100644 --- a/packages/a2d-worker-hf/src/a2d_core/transform/attention.py +++ b/packages/a2d-worker-hf/src/a2d_core/transform/attention.py @@ -1,31 +1,39 @@ -"""Annealed causal->bidirectional attention patch for GPT-2's eager seam (Decision 2). - -``transformers==4.48.3`` GPT-2 bakes causality INSIDE the attention op: eager -``GPT2Attention`` masks scores with a lower-triangular ``self.bias`` buffer via -``torch.where(causal, scores, finfo(float32).min)``; the passed 4D -``attention_mask`` is padding-only. A model-level additive mask therefore does -NOT control causality, so this module patches the seam that does. - -The patch neutralizes ``self.bias`` (registers it all-True so ``torch.where`` -never masks) and re-supplies a single annealed additive mask driven by one shared -``AnnealState``: on/below the diagonal it is ``0``; strictly-future (``j>i``) it is -``clamp(log(alpha), finfo(float32).min)``. - -* At ``alpha=0`` a future entry is ``finfo.min``, and a finite pre-softmax score is - negligible against it in float32 (``score + finfo.min == finfo.min`` to the bit), - so patched-at-0 scores/logits are bit-identical to base. -* At ``alpha=1`` the penalty is ``log(1)=0`` and, with ``self.bias`` neutralized, - attention is fully bidirectional; intermediate ``alpha`` scales each future - position's pre-softmax mass by ``alpha`` (a smooth monotone reveal). - -This deliberately replaces the ML-recon's ``(1-alpha)*(-inf) + alpha*0`` blend, -which is ``-inf`` for every ``alpha<1`` and thus not an anneal at all (Risk 4). - -The eager path is patched at the module-global ``eager_attention_forward`` that -``GPT2Attention.forward`` resolves at call time. The replacement is inert for any -attention module without a ``_a2d_anneal`` tag, so an unpatched base model (the -identity gate's reference copy) keeps its original causal attention even though -the global is process-wide. +"""Annealed causal->bidirectional attention: the shared v5 mask seam (Decision 2). + +``transformers==5.x`` unified what used to be three distinct eager seams. GPT-2 no +longer bakes causality into a per-layer ``self.bias`` buffer, and the RoPE family no +longer builds its mask in a per-model ``_update_causal_mask`` method. Every family now +routes causality through the module-level ``create_causal_mask`` / +``create_sliding_window_causal_mask`` in ``transformers.masking_utils``, which build +the 4D additive mask by calling ``ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation]`` +and hand it to every layer; eager attention just does ``scores + attention_mask``. That +registry is a documented extension point, so a2d registers its own implementation +instead of monkeypatching library internals: one seam covers GPT-2, Llama, Qwen2, +Gemma 1/2/3 and Mistral, including Gemma 2/3's per-layer sliding window (a second +``sliding_attention`` mask built through the same interface) and Mistral's +single-mask window. + +The registered mask function calls HF's own ``eager_mask`` to get the exact base mask, +then re-reveals every cell that mask masked for a real (non-padded) key under one +shared ``AnnealState``: + +* At ``alpha=0`` the penalty is exactly ``finfo(dtype).min`` - the value the base mask + already carries at masked cells - so ``torch.where(reveal, penalty, base)`` returns a + tensor bit-identical to base and patched@0 logits equal base to the bit (the D13 + identity gate). +* At ``alpha=1`` the penalty is ``log(1)=0``, so masked cells become attendable and + attention is fully bidirectional AND unwindowed; intermediate ``alpha`` applies + ``log(alpha)``, a smooth monotone reveal. + +Deriving the result FROM the base mask (rather than rebuilding it) is what makes the +``alpha=0`` no-op bit-identical by construction, independent of dtype, cache offset or +padding. A genuine 2D padding mask is preserved: a masked cell is only revealed when +its key is a real token, so padding stays masked at every alpha. + +Isolation is per model, not per process: each install registers an implementation key +private to its ``AnnealState`` and points only that model's config at it, so a sibling +un-patched model in the same process - the identity gate's reference copy - keeps its +original causal attention. """ from __future__ import annotations @@ -116,9 +124,11 @@ def install_mask_anneal(model: Any, state: AnnealState) -> None: """Route ``model``'s mask construction through the annealed reveal. Requires ``attn_implementation="eager"`` (Decision 2); forces ``use_cache=False`` - (diffusion decodes the full canvas, ARCHITECTURE.md §7). The implementation key is - derived from ``state``, which the registry entry keeps alive, so the key is stable - and unique per anneal state: re-installing the same state is idempotent, while + (diffusion decodes the full canvas, ARCHITECTURE.md §7), but only AFTER the key has + taken, so a refused key raises with the model exactly as it was rather than leaving + it half-patched (cache off, seam still causal). The implementation key is derived + from ``state``, which the registry entry keeps alive, so the key is stable and + unique per anneal state: re-installing the same state is idempotent, while installing a fresh state swaps in a new key rather than being silently ignored. """ from transformers import AttentionInterface @@ -133,10 +143,10 @@ def install_mask_anneal(model: Any, state: AnnealState) -> None: key = f"{_KEY_PREFIX}{id(state)}" AttentionMaskInterface.register(key, lambda **kwargs: annealed_eager_mask(state, **kwargs)) AttentionInterface.register(key, _a2d_eager_attention) - model.config.use_cache = False model.set_attn_implementation(key) if model.config._attn_implementation != key: raise ValueError(f"transformers refused attn_implementation={key!r}: alpha stays causal") + model.config.use_cache = False def install_anneal_patch(model: Any, state: AnnealState) -> None: diff --git a/packages/a2d-worker-hf/src/a2d_core/transform/gqa_attention.py b/packages/a2d-worker-hf/src/a2d_core/transform/gqa_attention.py index e19840b..5887db7 100644 --- a/packages/a2d-worker-hf/src/a2d_core/transform/gqa_attention.py +++ b/packages/a2d-worker-hf/src/a2d_core/transform/gqa_attention.py @@ -1,40 +1,21 @@ """Annealed causal->bidirectional attention for the RoPE/GQA family (Decision 2). -Llama / Qwen2 / Gemma (``transformers==4.51.3``) do NOT bake causality into a -per-layer ``self.bias`` buffer the way GPT-2 does. Instead the decoder ``*Model`` -builds one 4D additive causal mask per forward in ``_update_causal_mask`` and hands -it down to every layer; eager attention just does ``scores + causal_mask``. Causality -therefore flows entirely through that mask, so annealing it (not a ``self.bias`` seam) -is what opens attention here. This is the ONE seam shared verbatim by all three -families, so a single handler covers Gemma (MQA), Qwen2/Llama (GQA), and even -full-attention RoPE models (Llama-2-7B): the mask is family-independent, and the GQA -group expansion (``repeat_kv``), RoPE, RMSNorm, and Gemma's sqrt(hidden) embedding -scaling all stay untouched in HF's own forward. - -The patch wraps the decoder's bound ``_update_causal_mask``: it calls the original -to get the exact base mask, then re-reveals every cell that mask masked for a real -(non-padded) key under one shared ``AnnealState``. For the full-attention RoPE -family those cells are exactly the strictly-future set; Mistral v0.1 and Qwen2 with -an active ``use_sliding_window`` fold their sliding window into this SAME model-level -mask (they have no per-layer window), so their far-past out-of-window cells reopen -through the identical anneal and ``alpha=1`` is fully non-causal AND unwindowed: - -* At ``alpha=0`` the penalty is exactly ``finfo(dtype).min`` - the same value the - base mask already carries at masked cells - so ``torch.where(reveal, penalty, - base)`` returns a tensor bit-identical to base, and patched@0 logits equal base to - the bit (the D13 identity gate). -* At ``alpha=1`` the penalty is ``log(1)=0`` so masked cells become attendable and - attention is fully bidirectional; intermediate ``alpha`` applies ``log(alpha)`` (a - smooth monotone reveal), mirroring the GPT-2 seam's semantics. - -Deriving the result FROM the base mask (rather than rebuilding it) is what makes the -``alpha=0`` no-op bit-identical by construction, independent of dtype, cache offset, -or padding. A genuine 2D padding mask is preserved: a masked cell is only revealed -when its key is a real (non-padded) token, so padding stays masked at every alpha. - -The override is installed on the model INSTANCE (it shadows the class method via the -instance ``__dict__``), so a sibling un-patched base model in the same process - the -identity gate's reference copy - keeps its original causal attention. +Llama / Qwen2 / Gemma 1 build no per-layer causal buffer of their own: the decoder +``*Model`` asks ``transformers.masking_utils`` for one 4D additive causal mask per +forward and hands it to every layer, so causality flows entirely through that mask. +Under ``transformers==4.x`` that mask came from a per-model ``_update_causal_mask`` +method, which v5 removed in favour of the shared ``create_causal_mask`` + +``ALL_MASK_ATTENTION_FUNCTIONS`` interface that GPT-2 now uses too. Annealing that +one interface is therefore the whole seam here - see ``attention.py`` for the reveal +semantics and the per-model isolation. + +The mask is family-independent, so this single handler covers Gemma 1 (MQA), +Qwen2/Llama (GQA) and full-attention RoPE models (Llama-2-7B) alike, with the GQA +group expansion (``repeat_kv``), RoPE, RMSNorm and Gemma's sqrt(hidden) embedding +scaling all left untouched in HF's own forward. Mistral v0.1 folds its sliding window +into this SAME single model-level mask (it has no per-layer window, hence no +``layer_types``), so its far-past out-of-window cells reopen through the identical +anneal and ``alpha=1`` is fully non-causal AND unwindowed. """ from __future__ import annotations diff --git a/packages/a2d-worker-hf/src/a2d_core/transform/handlers/full_attention.py b/packages/a2d-worker-hf/src/a2d_core/transform/handlers/full_attention.py index ac3e13e..003f60d 100644 --- a/packages/a2d-worker-hf/src/a2d_core/transform/handlers/full_attention.py +++ b/packages/a2d-worker-hf/src/a2d_core/transform/handlers/full_attention.py @@ -1,8 +1,12 @@ """``attn.full`` transform: install the annealed bidirectional attention patch. -GPT-2 dense causal attention -> annealed causal->bidirectional via the eager-seam -patch (Decision 2). The ``alpha=0`` identity gate and ``test_bidir`` together prove -the patch is both bit-identical to base and genuinely reaches GPT-2's causality. +GPT-2 no longer bakes causality into a per-layer ``self.bias`` buffer: under +``transformers>=5`` its 4D mask is built by the same shared mask interface the RoPE +family uses, so this handler installs that one mask-seam anneal (Decision 2), gated on +GPT-2's dense ``GPT2Attention`` module. ``reorder_and_upcast_attn=True`` is rejected by +name, because that path survives only a literal ``"eager"`` implementation key. The +``alpha=0`` identity gate and ``test_bidir`` together prove the patch is both +bit-identical to base and genuinely reaches past GPT-2's causality. """ from __future__ import annotations diff --git a/packages/a2d-worker-hf/src/a2d_core/transform/handlers/gqa_attention.py b/packages/a2d-worker-hf/src/a2d_core/transform/handlers/gqa_attention.py index 125c3fb..4af0e82 100644 --- a/packages/a2d-worker-hf/src/a2d_core/transform/handlers/gqa_attention.py +++ b/packages/a2d-worker-hf/src/a2d_core/transform/handlers/gqa_attention.py @@ -1,13 +1,15 @@ """``attn.gqa`` transform: annealed bidirectional attention for the RoPE/GQA family. -Llama / Qwen2 / Gemma route causality through the 4D mask that ``_update_causal_mask`` -builds (not GPT-2's per-layer ``self.bias``), so this handler installs the mask-seam -anneal patch. It covers GQA, MQA (Gemma's ``num_key_value_heads=1``), full-attention -RoPE models, and the single-mask windowed families (Mistral v0.1, Qwen2 with an active -sliding window) whose far-past window lives in that same mask, leaving RoPE, the -KV-group expansion, RMSNorm, and Gemma's embedding scaling to HF's own forward. The -``alpha=0`` identity gate and the ``alpha=1`` bidir test together prove the patch is -bit-identical to base yet genuinely reaches future tokens. +Llama / Qwen2 / Gemma route causality through the 4D mask the shared ``transformers>=5`` +mask interface builds, so this handler installs the mask-seam anneal patch. It covers +GQA, MQA (Gemma's ``num_key_value_heads=1``), full-attention RoPE models, and the +single-mask windowed family (Mistral v0.1) whose far-past window lives in that same +mask, leaving RoPE, the KV-group expansion, RMSNorm, and Gemma's embedding scaling to +HF's own forward. Families that name a per-layer ``sliding_attention`` in +``config.layer_types`` - Gemma 2/3 - request a second mask through the same interface +and so route to ``attn.swa`` instead. The ``alpha=0`` identity gate and the ``alpha=1`` +bidir test together prove the patch is bit-identical to base yet genuinely reaches +future tokens. """ from __future__ import annotations diff --git a/packages/a2d-worker-hf/src/a2d_core/transform/handlers/swa_attention.py b/packages/a2d-worker-hf/src/a2d_core/transform/handlers/swa_attention.py index 14391b8..fbc30a9 100644 --- a/packages/a2d-worker-hf/src/a2d_core/transform/handlers/swa_attention.py +++ b/packages/a2d-worker-hf/src/a2d_core/transform/handlers/swa_attention.py @@ -1,10 +1,10 @@ -"""``attn.swa`` transform: annealed bidirectional attention for the sliding-window -Gemma family (Gemma 2/3). +"""``attn.swa`` transform: annealed bidirectional attention for families that name a +per-layer sliding-window mask in ``config.layer_types`` (Gemma 2/3). -Gemma 2/3 route causality through the same model-level 4D mask as the RoPE/GQA family -(``_update_causal_mask``) but ADD a per-layer sliding window on their local layers. -This handler installs the two-part SWA anneal: the shared future-reveal on the full -causal mask plus a far-past reveal on every sliding decoder layer, so at ``alpha=0`` +Such a model asks the shared mask interface for a SECOND, sliding-window mask and hands +it to the local decoder layers ``config.layer_types`` names. Both masks reach the same +annealed seam, so one install reveals the full mask's future cells AND the sliding +mask's far-past cells, and at ``alpha=0`` the model is bit-identical to base (identity gate) and at ``alpha=1`` attention is fully non-causal AND unwindowed. RoPE, qk-norm, the KV-group layout, RMSNorm, and embedding scaling are left to HF's own forward. diff --git a/packages/a2d-worker-hf/src/a2d_core/transform/identity.py b/packages/a2d-worker-hf/src/a2d_core/transform/identity.py index 7711932..2a42dc5 100644 --- a/packages/a2d-worker-hf/src/a2d_core/transform/identity.py +++ b/packages/a2d-worker-hf/src/a2d_core/transform/identity.py @@ -33,8 +33,8 @@ def check_identity( """Compare an unpatched ``base`` to a patched model at ``alpha=0`` on ``probe``. ``base`` must be un-patched and un-grown; ``patched`` is the patched (possibly - grown) model whose ``_a2d_anneal`` is ``state``. Both are forced to CPU/eval so - the gate is deterministic float32 (Risk 2). + grown) model whose mask seam is the implementation key registered for ``state``. + Both are forced to CPU/eval so the gate is deterministic float32 (Risk 2). """ import torch diff --git a/packages/a2d-worker-hf/src/a2d_core/transform/swa_attention.py b/packages/a2d-worker-hf/src/a2d_core/transform/swa_attention.py index 6379736..9c1a6a5 100644 --- a/packages/a2d-worker-hf/src/a2d_core/transform/swa_attention.py +++ b/packages/a2d-worker-hf/src/a2d_core/transform/swa_attention.py @@ -1,43 +1,24 @@ -"""Annealed causal->bidirectional attention for the sliding-window Gemma family -(Gemma 2/3). - -Gemma 2 and Gemma 3 (``transformers==4.51.3``) are the RoPE/GQA seam PLUS a per-layer -sliding window. Their decoder ``*Model`` builds ONE full 4D additive causal mask in -``_update_causal_mask`` - exactly the Gemma 1 seam - and hands it to every layer. -Then each *sliding* decoder layer's ``forward`` further masks the strictly- -far-past (keys more than ``sliding_window`` behind the query) to ``finfo(dtype).min`` -with identical logic in both families; the remaining layers are *global* and skip -that step (Gemma 3 windows all but every ``sliding_window_pattern``-th layer, Gemma 2 -every other layer). Eager attention just does ``scores + attention_mask``, so BOTH -causality (future) AND the window (far-past) flow entirely through the additive mask, -layer by layer. (Mistral-style models instead fold their window into the model-level -mask and have no sliding layers; they are the ``attn.gqa`` seam, not this one.) - -Bidirectionalizing therefore needs TWO coordinated anneals over one shared -``AnnealState``: - -* the ``_update_causal_mask`` wrap (reused verbatim from the ``attn.gqa`` seam via - ``install_gqa_anneal_patch``) reveals strictly-future cells. This opens every - layer's future AND feeds the sliding layers their future-revealed base mask. -* a per-sliding-layer ``forward`` wrap reveals the strictly-far-past cells the layer - would otherwise window out. - -Both use the same penalty ramp as the GPT-2/GQA seams: ``finfo(dtype).min`` at -``alpha=0`` (so each reveal is bit-identical to base -> the D13 identity gate reads -``max_abs_diff == 0.0``) and ``clamp(log(alpha), finfo.min)``, reaching ``0`` at -``alpha=1`` (fully non-causal AND unwindowed). Global layers are untouched by the -second wrap; they open through the first alone. RoPE (Gemma 3's per-layer local vs -global theta), query-key norm, the GQA/MQA layout, RMSNorm, logit softcapping -(Gemma 2), and Gemma's sqrt(hidden) embedding scaling all stay in HF's own forward - -only the additive mask is patched, exactly as the Gemma 1 seam does. - -The far-past wrap temporarily flips the decoder layer's ``is_sliding`` to ``False`` -so HF's own (hard ``finfo.min``) window re-mask is skipped for that one call, and -supplies its annealed replacement instead. It does NOT touch ``self_attn.is_sliding``, -which selects the local RoPE embedding, so positions stay exactly as base computes -them. Overrides live on module INSTANCES (shadowing the class methods), so a sibling -un-patched base model in the same process - the identity gate's reference copy - -keeps its original causal+windowed attention. +"""Annealed causal->bidirectional attention for families whose ``config.layer_types`` +names a per-layer sliding-window mask (Gemma 2/3). + +Such a model is the RoPE/GQA seam PLUS a sliding window on its *local* decoder layers. +Under ``transformers==4.x`` that window was applied per decoder layer, inside each +sliding layer's ``forward``, on top of the one full mask ``_update_causal_mask`` built - +so bidirectionalizing needed two coordinated patches. v5 hoisted it: the decoder +``*Model`` now builds a ``{"full_attention": ..., "sliding_attention": ...}`` mask +mapping up front (both masks through the same ``ALL_MASK_ATTENTION_FUNCTIONS`` +interface) and hands each layer the mask its ``config.layer_types`` entry names. +Sliding layers hold no window logic of their own any more. + +So both anneals now happen at one seam: the shared mask patch reveals the +strictly-future cells of the full mask AND, in the sliding mask, the strictly-far-past +cells the window would otherwise close - both are just cells the base mask set to +``finfo(dtype).min`` (see ``attention.py``). At ``alpha=0`` each reveal is +bit-identical to base, so the D13 identity gate reads ``max_abs_diff == 0.0``; at +``alpha=1`` attention is fully non-causal AND unwindowed. Global layers open through +the full mask alone. RoPE (Gemma 3's per-layer local vs global theta), query-key norm, +the GQA/MQA layout, RMSNorm, logit softcapping (Gemma 2) and Gemma's sqrt(hidden) +embedding scaling all stay in HF's own forward - only the additive mask is patched. """ from __future__ import annotations @@ -66,6 +47,6 @@ def install_swa_anneal_patch(model: Any, state: AnnealState) -> None: if not has_sliding_window_seam(model): raise ValueError( "SWA anneal patch found no sliding-window decoder layers " - f"on {type(model).__name__!r} (not a Gemma 2/3 sliding-window model?)" + f"on {type(model).__name__!r} (no 'sliding_attention' in config.layer_types?)" ) install_mask_anneal(model, state) diff --git a/packages/a2d-worker-hf/src/a2d_core/worker.py b/packages/a2d-worker-hf/src/a2d_core/worker.py index 76c7b37..02a3ac2 100644 --- a/packages/a2d-worker-hf/src/a2d_core/worker.py +++ b/packages/a2d-worker-hf/src/a2d_core/worker.py @@ -171,11 +171,11 @@ def _convert(job: ConversionJob, emit: Callable[[dict[str, Any]], None]) -> int: mask_token_id = resolve_mask_token(model, tokenizer, cfg.mask_token) # 4. patch: install the annealed attention seam at alpha=0 (Decision 2). The seam - # is resolved from the model's own eager causal structure: GPT-2's self.bias - # (attn.full) vs the RoPE family's _update_causal_mask mask (attn.gqa - Gemma/ - # Qwen2/Llama) vs that same mask plus per-layer sliding windows (attn.swa - - # Gemma 2/3). The ConversionJob carries no capability set, so the worker picks - # the handler honestly from the model itself. + # is resolved from the model's own eager causal structure: GPT-2's dense + # attention (attn.full) vs the RoPE family (attn.gqa - Gemma/Qwen2/Llama) vs + # those plus a per-layer sliding-window mask (attn.swa - Gemma 2/3). The + # ConversionJob carries no capability set, so the worker picks the handler + # honestly from the model itself. # ponytail: attention seam only; feed the manifest's model_spec.capabilities # through the job when P6 adds sink handlers keyed off config-only fields. emit(progress("patch", 3, total)) diff --git a/packages/a2d-worker-hf/tests/conftest.py b/packages/a2d-worker-hf/tests/conftest.py index 393739e..0d88100 100644 --- a/packages/a2d-worker-hf/tests/conftest.py +++ b/packages/a2d-worker-hf/tests/conftest.py @@ -131,10 +131,11 @@ def tiny_gemma2() -> Callable[..., Any]: """Factory for a seeded tiny Gemma 2 (eager, eval, no download). Same ``(kwargs, seed)`` => bit-identical weights, so two calls give an aligned - base/patched pair. Exercises the Gemma 2 per-layer sliding-window seam: the window - logic is identical to Gemma 3's, but the decoder layer's ``forward`` takes ONE - ``position_embeddings`` pair where Gemma 3 takes a global/local pair - the - signature the swa wrapper must not hardcode away - and every EVEN layer slides. + base/patched pair. Exercises the Gemma 2 per-layer sliding-window seam: a second + decoder-layer shape (tanh-softcapped attention logits and no query/key norm, where + Gemma 3 softcaps neither and RMS-norms both) requesting the same + ``sliding_attention`` mask, so the shared reveal is proven family-independent, + and every EVEN layer slides. Defaults (2 layers) put both a local layer (0) and a global layer (1) in the stack; pass ``num_hidden_layers=1`` for an all-sliding model whose single layer's receptive field is exactly the window.""" @@ -169,10 +170,10 @@ def tiny_mistral() -> Callable[..., Any]: sliding window (default 2, well under the tests' seq_len 8). Same ``(kwargs, seed)`` => bit-identical weights, so two calls give an aligned - base/patched pair. Mistral folds its window into the single model-level 4D mask - ``_update_causal_mask`` builds and has NO sliding decoder layers, so it routes to - ``attn.gqa`` - the single-mask windowed flavor whose far-past cells the shared - reveal must open alongside the strictly-future ones.""" + base/patched pair. Mistral folds its window into ONE model-level 4D mask and names + no per-layer ``layer_types``, so it routes to ``attn.gqa`` - the single-mask + windowed flavor whose far-past cells the shared reveal must open alongside the + strictly-future ones.""" def _make(seed: int = 0, sliding_window: int = 2) -> Any: import torch diff --git a/packages/a2d-worker-hf/tests/test_bidir.py b/packages/a2d-worker-hf/tests/test_bidir.py index 68137ce..be12973 100644 --- a/packages/a2d-worker-hf/tests/test_bidir.py +++ b/packages/a2d-worker-hf/tests/test_bidir.py @@ -81,3 +81,17 @@ def test_install_rejects_reorder_and_upcast_attn(tiny_gpt2: Callable[..., Any]) with pytest.raises(ValueError, match="reorder_and_upcast_attn"): install_anneal_patch(model, AnnealState()) assert model.config._attn_implementation == "eager" + + +def test_install_leaves_the_model_untouched_when_transformers_refuses_the_key( + tiny_gpt2: Callable[..., Any], +) -> None: + """Atomic install: if the key does not take, the raise must leave no half-patched + model behind - cache off with the seam still causal is the worst of both.""" + model = tiny_gpt2(0) + model.set_attn_implementation = lambda *args, **kwargs: None + with pytest.raises(ValueError, match="transformers refused attn_implementation"): + install_anneal_patch(model, AnnealState()) + + assert model.config.use_cache is True + assert model.config._attn_implementation == "eager" diff --git a/packages/a2d-worker-hf/tests/test_gqa_attention.py b/packages/a2d-worker-hf/tests/test_gqa_attention.py index 51cad60..d81a62d 100644 --- a/packages/a2d-worker-hf/tests/test_gqa_attention.py +++ b/packages/a2d-worker-hf/tests/test_gqa_attention.py @@ -3,7 +3,7 @@ guard, plus the worker's structural handler dispatch. Mirrors ``test_identity``/``test_bidir`` (GPT-2) for the RoPE family, whose eager -seam is the 4D mask ``_update_causal_mask`` builds, not GPT-2's ``self.bias``. Also +seam is the 4D mask the shared ``transformers>=5`` mask interface builds. Also covers the single-mask windowed flavor (Mistral v0.1), whose sliding window is folded into that same model-level mask: the shared reveal must open its far-past window alongside the strictly-future cells, so a converted Mistral is never silently still @@ -132,7 +132,7 @@ def test_resolve_capabilities_picks_gqa_for_rope_family(tiny_gqa: Callable[..., def test_resolve_capabilities_keeps_gpt2_on_full(tiny_gpt2: Callable[..., Any]) -> None: - """GPT-2 must keep selecting its own attn.full self.bias seam, unchanged.""" + """GPT-2 must keep selecting attn.full, gated on its dense GPT2Attention module.""" assert resolve_capabilities(tiny_gpt2(0)) == ["attn.full"] @@ -209,9 +209,10 @@ def test_gqa_padding_stays_masked_and_identity_holds_with_padding( def test_gqa_reinstall_with_fresh_state_takes_effect( tiny_gqa: Callable[..., Any], family: str ) -> None: - """Re-installing on an already-patched model must swap in the NEW state (parity - with the GPT-2 seam's per-install re-tag) without double-wrapping: after a second - install at alpha=1, an earlier position sees a perturbed future token.""" + """Re-installing on an already-patched model must repoint it at the NEW state's + per-``id(state)`` registry key (the same seam GPT-2's ``attn.full`` install uses): + after a second install at alpha=1, an earlier position sees a perturbed future + token.""" model = tiny_gqa(family, 0) install_gqa_anneal_patch(model, AnnealState(alpha=0.0)) install_gqa_anneal_patch(model, AnnealState(alpha=1.0)) diff --git a/packages/a2d-worker-hf/tests/test_swa_attention.py b/packages/a2d-worker-hf/tests/test_swa_attention.py index 60fcbee..571ee7f 100644 --- a/packages/a2d-worker-hf/tests/test_swa_attention.py +++ b/packages/a2d-worker-hf/tests/test_swa_attention.py @@ -4,14 +4,13 @@ and the window-open (the sliding window's far-past opens) - plus the worker's structural handler dispatch. -The Gemma 2/3 eager seam is the single 4D ``_update_causal_mask`` (like Gemma 1) PLUS -a per-layer far-past window re-mask on the local (sliding) decoder layers. The window -test uses an all-sliding single-layer model so a query's receptive field is exactly -its window; the identity/future tests use a mixed 4-layer stack (local layers 0, 2 and -global layers 1, 3). Gemma 2 - whose decoder layer takes a single -``position_embeddings`` pair where Gemma 3 takes a global/local pair - guards the -wrapper's signature-agnostic re-bind. All hermetic: tiny random-weight configs on CPU -float32, no network. +The Gemma 2/3 eager seam is the shared ``transformers>=5`` mask interface, asked for +BOTH a full mask and a sliding-window mask; ``config.layer_types`` names which one each +decoder layer gets. The window test uses an all-sliding single-layer model so a query's +receptive field is exactly its window; the identity/future tests use a mixed 4-layer +stack (local layers 0, 2 and global layers 1, 3). Gemma 2 - a different decoder-layer +shape from Gemma 3 - guards that the seam is family-independent. All hermetic: tiny +random-weight configs on CPU float32, no network. """ from __future__ import annotations @@ -201,10 +200,9 @@ def test_swa_padding_stays_masked_and_identity_holds_with_padding( def test_swa_gemma2_patched_at_alpha0_is_bit_identical_to_base( tiny_gemma2: Callable[..., Any], ) -> None: - """Regression (review: swa-gemma2-signature-crash): Gemma 2's decoder layer takes a - single ``position_embeddings`` pair, so the wrapper must survive its keyword call - path without crashing, and patched@alpha=0 logits must equal base to 0.0 with both - a local and a global layer exercised.""" + """Gemma 2 has its own decoder-layer shape, so it guards that the shared mask seam + is family-independent: patched@alpha=0 logits must equal base to 0.0 with both a + local and a global layer exercised.""" base = tiny_gemma2() patched = tiny_gemma2() patched.load_state_dict(base.state_dict()) # guarantee identical weights @@ -225,7 +223,7 @@ def test_swa_gemma2_patched_at_alpha0_is_bit_identical_to_base( def test_swa_gemma2_opens_window_and_future_at_alpha1_not_alpha0( tiny_gemma2: Callable[..., Any], ) -> None: - """Gemma 2 end-to-end reveal through the signature-agnostic wrapper: on an + """Gemma 2 end-to-end reveal through the shared mask seam: on an all-sliding single-layer model an in-window key moves the query's logits even at alpha=0 (non-vacuity), while a far-past out-of-window key AND a strictly-future key move them at alpha=1 only.""" @@ -247,12 +245,11 @@ def test_swa_gemma2_opens_window_and_future_at_alpha1_not_alpha0( assert _shift(model, state, ids, q=3, k=6, alpha=1.0) > 1e-6 -def test_swa_wrapped_layer_survives_gradient_checkpointing( +def test_swa_annealed_mask_survives_gradient_checkpointing( tiny_gemma2: Callable[..., Any], tiny_gemma3: Callable[..., Any] ) -> None: - """The positional call path: under gradient checkpointing the model invokes the - decoder layer with every argument positional, which the wrapper must re-bind - against the original signature for both the Gemma 2 and Gemma 3 shapes.""" + """Gradient checkpointing re-enters the decoder layers under a no-grad/recompute + path; the annealed mask must survive it for both the Gemma 2 and Gemma 3 shapes.""" for model in (tiny_gemma2(), tiny_gemma3()): state = AnnealState(alpha=0.5) install_swa_anneal_patch(model, state) @@ -278,8 +275,8 @@ def test_swa_install_on_non_sliding_model_raises_and_leaves_model_unpatched( def test_swa_reinstall_with_fresh_state_takes_effect(tiny_gemma3: Callable[..., Any]) -> None: - """Re-installing on an already-patched model must swap in the NEW state without - double-wrapping: after a second install at alpha=1, an out-of-window token reaches the + """Re-installing on an already-patched model must repoint it at the NEW state's + registry key: after a second install at alpha=1, an out-of-window token reaches the query.""" model = tiny_gemma3(num_hidden_layers=1, sliding_window=2, sliding_window_pattern=6) install_swa_anneal_patch(model, AnnealState(alpha=0.0)) diff --git a/pyproject.toml b/pyproject.toml index e37c620..1bd17e2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,11 +38,13 @@ strict = true files = ["packages"] exclude = "packages/a2d-contracts/src/a2d_contracts/models" -# transformers 4.51.3 ships an (empty) py.typed marker, but the worker's transform -# layer patches HF's eager attention seams directly with runtime-only knowledge mypy -# cannot follow, so keep treating transformers as untyped (Any). `follow_imports=skip` -# is what makes that hold now that the marker exists - `ignore_missing_imports` alone -# no longer suffices. tokenizers is treated the same way. +# Still true on transformers 5.14.1: it ships an (empty) py.typed marker, so mypy reads +# its inline annotations, and those disagree with the runtime - `**kwargs` config fields +# read as unexpected keyword arguments, `Trainer.compute_loss` as an incompatible +# override - on top of the transform layer's runtime-only knowledge of HF's attention +# seams. Keep treating transformers as untyped (Any): `follow_imports=skip` is what makes +# that hold once the marker exists, `ignore_missing_imports` alone does not (dropping it +# reintroduces 68 errors). tokenizers ships no marker but is treated the same way. [[tool.mypy.overrides]] module = ["transformers", "transformers.*", "tokenizers", "tokenizers.*"] ignore_missing_imports = true