william-0g/MSA-Qwen3.5

MSA × Qwen3.5 — Memory Sparse Attention (Generative Retrieval)

Reference implementation and training recipe for adding Memory Sparse Attention (MSA) — a Generative-Retrieval memory module — on top of a Qwen3.5-4B hybrid (linear + full attention) backbone.

MSA equips the back layers of the model with a cross-document memory: at query time a per-layer router retrieves the top-k documents from a chunk-pooled K̄/V̄ memory bank via cosine routing, and only those chunks participate in sparse attention. The LM autoregressively emits document IDs ([k]<|object_ref_end|>), a controller injects the referenced doc text, and generation continues to <End-of-Retrieve> + the answer — so the LM performs both retrieval (slot emission) and answer extraction.

Contents

code/
  model_msa_qwen3_5.py        # MSA overlay ported to the Qwen3.5-4B hybrid backbone
                              #   (24 linear DGA layers + 8 full-attn layers; MSA on back 16)
  train_qwen3_5_msa_cpt.py    # Continual-pre-training (CPT) trainer, two-phase paper schedule
docs/
  CPT_TRAINING_PLAYBOOK.md    # Launch commands, mandatory/forbidden flags, tuning notes
  MSA_CPT_SFT_REPORT.md       # Full CPT + two-stage SFT reproduction report & results
sft/                          # Two-stage SFT (post-training) — see sft/SFT_TRAINING_METHOD.md
  SFT_TRAINING_METHOD.md      #   S1 (~8K) + S2 (~64K) recipe, the late-June run, S2 OOM record
  code/                       #   train_qwen3_5_msa_sft.py, dataset_msa_interleave.py
  scripts/                    #   sft_v11_weakfix3.sh (late-June S1-only launch)
  logs/                       #   verbatim S2 OOM traceback

Architecture summary

• Backbone: Qwen3.5-4B, 32 layers — every 4th layer is full-attention, the rest are Delta-Gated linear-attention (DGA). Vocab 248320.
• MSA layers: back-16 (indices 16–31): 12 linear + 4 full.
  • Linear (DGA) layers — paper-faithful initial_state injection: encode each doc through DGA, route on chunk-pooled K̄_r vs Q_r, sum top-k docs' recurrent states into S_ext, and pass it as initial_state to chunk_gated_delta_rule on the query path.
  • Full (softmax) layers — standard MSA: concat top-k chunks' K/V with local KV in sparse attention.
• Only LoRA (or optional OFT) adapters + the MSA overlay/router params are trainable; the base backbone stays frozen.

Training (CPT)

Two-phase schedule (paper §3.3.1):

Phase	lm_coef	aux_coef	base LR
warmup	0.1	1.0	1e-4
main	1.0	0.1	1e-5 → cosine to lr_min

See docs/CPT_TRAINING_PLAYBOOK.md for the canonical launch command and the mandatory/forbidden flags established over the v8→v20 iteration.

Notes

The reproduction report (docs/MSA_CPT_SFT_REPORT.md) documents a scaled-down single-GPU reproduction (originally on a Qwen3-1.7B backbone, ~0.16B CPT tokens vs the paper's 158.95B) and the ablations that isolated the correct CPT learning-rate schedule. The code here is the later Qwen3.5-4B port of the same method.