Upload AdamW-WSD attention 0.6B step 1000

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,43 @@
+---
+license: apache-2.0
+language:
+- en
+base_model: Qwen/Qwen3-0.6B
+tags:
+- parallax
+- qwen3
+- attention
+- wsd
+- 0.6b
+- pretraining-checkpoint
+---
+
+# Attention 0.6B AdamW-WSD - step 1000
+
+Intermediate checkpoint at optimizer step **1000** (of 20000) of the 0.6B
+Qwen3-style softmax-attention baseline used in the Parallax mechanism
+analysis.
+
+- **Architecture**: Qwen3-0.6B (28 layers, d_model=1024, 16 Q-heads, 8 KV-heads, GQA, RoPE theta 1e6)
+- **Optimizer**: AdamW (lr=3e-4, weight decay=0.1, betas=(0.9, 0.95))
+- **Scheduler**: WSD (warmup=0, last 20% linearly decayed)
+- **Tokens seen at this step**: ~3930.0M sequences * 4096 tokens = ~15.7B tokens
+- **Dataset**: Ultra-FineWeb (English split)
+- **Tokenizer**: Qwen3-0.6B (vocab 151936)
+- **Total steps in run**: 20000 (~80B tokens)
+
+The full training trajectory (every 500 steps) lives in the [Attention 0.6B
+AdamW-WSD training trajectory](https://huggingface.co/collections/YifeiZuo)
+collection.
+
+## Loading
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "YifeiZuo/attention-0.6b-adamw-wsd-step1000",
+    trust_remote_code=True,
+)
+tokenizer = AutoTokenizer.from_pretrained("YifeiZuo/attention-0.6b-adamw-wsd-step1000")
+```

configuration_parallax.py

@@ -0,0 +1,77 @@
+"""Parallax HF config, modular-style (inherits Qwen3Config).
+
+This is a second, separate HF interface for the Parallax model family. Unlike
+the self-contained ``torchtitan.hf_modeling.parallax`` package, this version
+inherits from ``transformers.models.qwen3.configuration_qwen3.Qwen3Config``
+so it picks up all of Qwen3's modern conventions for free (strict dataclass,
+``rope_parameters``, ``layer_types``, tp/pp plans, etc.).
+
+The only Parallax-specific additions are the four fields at the bottom that
+control the centered amortized local linear attention op. When
+``attn_type == "sdpa"`` this config describes a plain Qwen3-equivalent model.
+"""
+
+from __future__ import annotations
+
+from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+
+
+class ParallaxConfig(Qwen3Config):
+    r"""Config for the Parallax model (Qwen3-derived).
+
+    Args:
+        attn_type (`str`, *optional*, defaults to `"parallax"`):
+            Attention implementation. Must be one of ``"sdpa"`` or
+            ``"parallax"``. ``"sdpa"`` uses the plain Qwen3 causal attention
+            path (softmax attention). ``"parallax"`` uses the centered
+            amortized local linear attention op and requires the extra
+            ``r_proj`` / ``r_norm`` per-layer weights.
+        use_centering (`bool`, *optional*, defaults to `True`):
+            Whether the amortized-LLA math uses the centered variant
+            (``s + c*s - r_term``). Currently only the centered variant is
+            implemented; kept for config-roundtrip parity with the training
+            code, which also has a non-centered path.
+        amortized_rho_scale (`float`, *optional*, defaults to `1.0`):
+            Multiplicative scale applied to the ``r`` tensor after its
+            rotary embedding, mirroring the training-time
+            ``amortized_rho_scale`` hyperparameter.
+        learn_amortized_rho_scale (`bool`, *optional*, defaults to `False`):
+            Whether the model learns a per-head scalar ``a_proj`` modulating
+            ``r``. When enabled, adds an ``a_proj: Linear(hidden, n_heads)``
+            whose sigmoid output gates ``r`` per head, mirroring the
+            training-time ``wa`` module in torchtitan.
+        rope_r (`str`, *optional*, defaults to `"correct"`):
+            How to apply rotary positional embeddings to R.
+            ``"correct"`` — standard RoPE on R (matches fixed torchtitan code).
+            ``"none"`` — skip RoPE on R entirely.
+            ``"buggy"`` — reproduce the old torchtitan bug where transpose
+            was applied before RoPE, causing positions 0..H-1 to be applied
+            along the heads axis instead of the sequence axis. Use this to
+            faithfully evaluate checkpoints trained with the buggy code.
+    """
+
+    model_type = "parallax"
+
+    # ---- Parallax-specific extras ----
+    # "sdpa"    : plain causal self-attention (identical to Qwen3).
+    # "parallax": centered amortized local linear attention (adds r_proj + r_norm
+    #             and replaces softmax attention with the centered LLA op).
+    attn_type: str = "parallax"
+    use_centering: bool = True
+    amortized_rho_scale: float = 1.0
+    learn_amortized_rho_scale: bool = False
+    rope_r: str = "correct"
+
+    def __post_init__(self, **kwargs):
+        if self.attn_type not in ("sdpa", "parallax"):
+            raise ValueError(
+                f"attn_type must be 'sdpa' or 'parallax', got {self.attn_type!r}"
+            )
+        if self.rope_r not in ("correct", "none", "buggy"):
+            raise ValueError(
+                f"rope_r must be 'correct', 'none', or 'buggy', got {self.rope_r!r}"
+            )
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["ParallaxConfig"]

model-00001-of-00001.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b1a03efffcd26ef4ad15dff829444e3c5928fa163c40a55ba766b5d107f0001c
+size 1192137739

model.safetensors.index.json

@@ -0,0 +1,317 @@
+{
+  "metadata": {
+    "total_size": 1192099840
+  },
+  "weight_map": {
+    "model.embed_tokens.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.self_attn.q_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.self_attn.k_norm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00001-of-00001.safetensors",
+    "model.norm.weight": "model-00001-of-00001.safetensors"
+  }
+}

modeling_parallax.py

@@ -0,0 +1,367 @@
+"""Parallax HF modeling file, modular-style (inherits from Qwen3).
+
+This module is the second, separate HF interface for the Parallax family.
+It inherits almost everything from ``transformers.models.qwen3`` and only
+overrides the attention class to add the centered amortized local linear
+attention op (``attn_type="parallax"``).
+
+Design notes
+------------
+* All boilerplate — RMSNorm, MLP, rotary embeddings, decoder-layer scaffolding,
+  the base ``PreTrainedModel`` class, ``Qwen3Model``, ``Qwen3ForCausalLM``,
+  mask preparation, cache handling, ``GenerationMixin`` integration, etc. —
+  is inherited from ``transformers.models.qwen3``.
+* The only new code is:
+    - ``parallax_centered_attention_forward``: a pure-PyTorch attention kernel
+      matching the math of ``torchtitan/models/flashlla/dev_ops/reference.py``
+      (``centered_dense_reference``), but batched over heads and taking an
+      additive attention mask in the HF shape ``(B, 1, q_len, kv_len)``.
+    - ``ParallaxAttention``: inherits from ``Qwen3Attention``, adds an
+      ``r_proj`` / ``r_norm`` pair when ``config.attn_type == "parallax"``,
+      and overrides ``forward`` to dispatch to the parallax kernel in that
+      case. For ``attn_type == "sdpa"`` the forward simply delegates to
+      ``Qwen3Attention.forward`` unchanged, so all modern HF attention
+      backends (sdpa, flash_attention_2, flex_attention, eager) remain
+      available for that path.
+    - ``ParallaxDecoderLayer``, ``ParallaxModel``, ``ParallaxPreTrainedModel``,
+      ``ParallaxForCausalLM``: thin subclasses that rewire parent classes to
+      the ``ParallaxAttention`` / ``ParallaxConfig`` below.
+
+Why this file exists alongside ``torchtitan/hf_modeling/parallax/``
+-------------------------------------------------------------------
+The older ``torchtitan.hf_modeling.parallax`` package is fully self-contained
+(no imports from ``transformers.models.qwen3``). That makes it robust across
+transformers versions but also duplicates a lot of boilerplate that drifts
+from HF conventions. This ``parallax_qwen3`` package takes the opposite
+tradeoff: minimal code, surgical diffs from Qwen3, tracks HF updates for
+free, at the cost of depending on a stable ``transformers.models.qwen3`` API.
+
+Both produce numerically equivalent outputs on the same Parallax checkpoint
+weights; a parity test verifies this.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3Attention,
+    Qwen3DecoderLayer,
+    Qwen3ForCausalLM,
+    Qwen3Model,
+    Qwen3PreTrainedModel,
+    Qwen3RMSNorm,
+    apply_rotary_pos_emb,
+    repeat_kv,
+    rotate_half,
+)
+
+from .configuration_parallax import ParallaxConfig
+
+
+# ---------------------------------------------------------------------------
+# Parallax centered attention kernel (pure PyTorch reference)
+# ---------------------------------------------------------------------------
+
+
+def parallax_centered_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    r: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+) -> tuple[torch.Tensor, None]:
+    """Centered amortized local linear attention.
+
+    This is a batched generalization of ``centered_dense_reference`` from
+    ``torchtitan/models/flashlla/dev_ops/reference.py``:
+
+        p = softmax((Q Kᵀ) * scale + mask)
+        rk = R Kᵀ
+        s = p V
+        c = sum(p * rk, -1, keepdim=True)
+        r_term = (p * rk) V
+        out = s + c * s - r_term
+
+    Shapes
+    ------
+    query : (B, H, q_len, D)
+    r     : (B, H, q_len, D)          — Parallax-specific extra projection
+    key   : (B, H_kv, kv_len, D)      — GQA-compressed, expanded internally
+    value : (B, H_kv, kv_len, D)
+    attention_mask : (B, 1, q_len, kv_len) additive, 0 or -inf (HF convention)
+
+    Returns
+    -------
+    attn_output : (B, q_len, H, D)    — transposed to match Qwen3 eager impl
+    attn_weights : None               — not exposed (centered attention has no
+                                         meaningful "attention matrix" per head)
+    """
+    # Expand GQA-compressed k/v to match q heads
+    key = repeat_kv(key, module.num_key_value_groups)
+    value = repeat_kv(value, module.num_key_value_groups)
+
+    # Upcast to float32 for numerical stability (matches the training kernel,
+    # which uses bfloat16 for math but accumulates in fp32 within the softmax).
+    out_dtype = query.dtype
+    qf = query.float()
+    rf = r.float()
+    kf = key.float()
+    vf = value.float()
+
+    qk = torch.einsum("bhqd,bhkd->bhqk", qf, kf) * scaling
+    q_len = qk.shape[-2]
+    kv_len = qk.shape[-1]
+    if attention_mask is not None:
+        # HF causal masks may cover a larger key span than our current
+        # kv_len during decode; slice to match.
+        qk = qk + attention_mask[..., :q_len, :kv_len].to(qk.dtype)
+    elif q_len > 1:
+        # HF's `create_causal_mask` returns None when the model's default
+        # attention backend (e.g. sdpa) handles causality implicitly. Our
+        # parallax kernel is hand-written and needs the explicit additive
+        # mask, so build one here. Decode (q_len == 1) is naturally
+        # unmasked: the single query can attend to every cached key.
+        offset = kv_len - q_len
+        row = torch.arange(q_len, device=qk.device)[:, None] + offset
+        col = torch.arange(kv_len, device=qk.device)[None, :]
+        causal = row >= col
+        bias = torch.zeros(q_len, kv_len, dtype=qk.dtype, device=qk.device)
+        bias = bias.masked_fill(~causal, float("-inf"))
+        qk = qk + bias
+
+    p = torch.softmax(qk, dim=-1)
+    rk = torch.einsum("bhqd,bhkd->bhqk", rf, kf)
+
+    s = torch.einsum("bhqk,bhkd->bhqd", p, vf)
+    c = (p * rk).sum(dim=-1, keepdim=True)
+    r_term = torch.einsum("bhqk,bhkd->bhqd", p * rk, vf)
+
+    attn_output = (s + c * s - r_term).to(out_dtype)
+    # Match Qwen3 eager_attention_forward: return transposed (B, q_len, H, D)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+    return attn_output, None
+
+
+# ---------------------------------------------------------------------------
+# Attention / DecoderLayer / Model / ForCausalLM
+# ---------------------------------------------------------------------------
+
+
+class ParallaxAttention(Qwen3Attention):
+    """Qwen3 attention plus an optional Parallax (centered amortized LLA) path."""
+
+    def __init__(self, config: ParallaxConfig, layer_idx: int) -> None:
+        super().__init__(config, layer_idx)
+        self.attn_type = config.attn_type
+        self.amortized_rho_scale = config.amortized_rho_scale
+        self.learn_amortized_rho_scale = config.learn_amortized_rho_scale
+        self.rope_r = getattr(config, "rope_r", "correct")
+        self._buggy_rope_cos: torch.Tensor | None = None
+        self._buggy_rope_sin: torch.Tensor | None = None
+        if self.attn_type == "parallax":
+            # Extra projection + RMSNorm for the amortized local linear attn
+            self.r_proj = nn.Linear(
+                config.hidden_size,
+                config.num_attention_heads * self.head_dim,
+                bias=config.attention_bias,
+            )
+            self.r_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            # Learned per-head gate on r (sigmoid), mirroring torchtitan `wa`.
+            if self.learn_amortized_rho_scale:
+                self.a_proj = nn.Linear(
+                    config.hidden_size,
+                    config.num_attention_heads,
+                    bias=False,
+                )
+            else:
+                self.a_proj = None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if self.attn_type != "parallax":
+            # Plain Qwen3 attention — delegate to the parent, which supports
+            # sdpa / flash_attention_2 / flex_attention / eager backends.
+            return super().forward(
+                hidden_states=hidden_states,
+                position_embeddings=position_embeddings,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+                **kwargs,
+            )
+
+        # ---- Parallax path ----
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_norm(
+            self.q_proj(hidden_states).view(hidden_shape)
+        ).transpose(1, 2)
+        key_states = self.k_norm(
+            self.k_proj(hidden_states).view(hidden_shape)
+        ).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        # r: same num_heads as q (not num_key_value_heads), rope applied.
+        # Note hidden_shape uses -1 for the heads dim, so r_proj (which produces
+        # num_attention_heads*head_dim) will correctly view as (..., n_heads, D).
+        r_states = self.r_norm(
+            self.r_proj(hidden_states).view(hidden_shape)
+        ).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
+        # Apply rope to r according to rope_r mode.
+        if self.rope_r == "correct":
+            r_states, _ = apply_rotary_pos_emb(r_states, r_states, cos, sin)
+        elif self.rope_r == "buggy":
+            # Reproduce the old torchtitan bug: reshape_for_broadcast read
+            # (B, H, S, D) as (bz, seqlen=H, _, head_dim=D), so it applied
+            # positions 0..H-1 along the heads axis.  Every token in the
+            # sequence got the SAME rotation (the one for position h, where
+            # h is the head index).  This is a fixed per-head rotation that
+            # does not vary with sequence position.
+            #
+            # We lazily compute cos/sin for positions 0..H-1 and cache them.
+            # Shape: (1, H, 1, D) so head h gets rotation for position h,
+            # broadcast across all sequence positions.
+            if self._buggy_rope_cos is None or self._buggy_rope_cos.device != r_states.device:
+                H = r_states.shape[1]
+                # Reuse the same RoPE frequencies: extract from cos/sin for
+                # positions 0..H-1.  cos/sin are (B,1,T,D) during prefill
+                # (T >= H always since T >= 1 and H = num_heads).
+                # For robustness, recompute from config's rope_theta.
+                head_dim = r_states.shape[-1]
+                rope_theta = self.config.rope_parameters["rope_theta"]
+                inv_freq = 1.0 / (rope_theta ** (
+                    torch.arange(0, head_dim, 2, device=r_states.device, dtype=torch.float32) / head_dim
+                ))
+                pos = torch.arange(H, device=r_states.device, dtype=torch.float32)
+                freqs = torch.outer(pos, inv_freq)  # (H, D/2)
+                emb = torch.cat([freqs, freqs], dim=-1)  # (H, D)
+                self._buggy_rope_cos = emb.cos().view(1, H, 1, head_dim).to(r_states.dtype)
+                self._buggy_rope_sin = emb.sin().view(1, H, 1, head_dim).to(r_states.dtype)
+            r_states = (r_states * self._buggy_rope_cos) + (rotate_half(r_states) * self._buggy_rope_sin)
+        # else: rope_r == "none", skip RoPE on R entirely
+        if self.a_proj is not None:
+            # a: (B, n_heads, q_len), sigmoid per-head gate
+            a = torch.sigmoid(self.a_proj(hidden_states)).transpose(1, 2).to(dtype=r_states.dtype)
+            r_states = r_states * (self.amortized_rho_scale * a.unsqueeze(-1))
+        elif self.amortized_rho_scale != 1.0:
+            r_states = r_states * self.amortized_rho_scale
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx
+            )
+
+        attn_output, attn_weights = parallax_centered_attention_forward(
+            self,
+            query_states,
+            r_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class ParallaxDecoderLayer(Qwen3DecoderLayer):
+    def __init__(self, config: ParallaxConfig, layer_idx: int) -> None:
+        # Call the grandparent (nn.Module) init to avoid Qwen3DecoderLayer
+        # building a Qwen3Attention we'd immediately replace.
+        nn.Module.__init__(self)
+        self.hidden_size = config.hidden_size
+        self.attention_type = config.layer_types[layer_idx]
+        self.self_attn = ParallaxAttention(config=config, layer_idx=layer_idx)
+        # Reuse the MLP + norms from Qwen3 via its own constructor logic.
+        # Easiest is to rebuild them directly here to match Qwen3DecoderLayer.
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3MLP
+
+        self.mlp = Qwen3MLP(config)
+        self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+
+class ParallaxPreTrainedModel(Qwen3PreTrainedModel):
+    config: ParallaxConfig  # type: ignore[assignment]
+    config_class = ParallaxConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["ParallaxDecoderLayer"]
+    # Parallax path is hand-written and does not plug into ALL_ATTENTION_FUNCTIONS.
+    # The sdpa path (delegated to super) supports the standard backends.
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+
+class ParallaxModel(ParallaxPreTrainedModel, Qwen3Model):
+    def __init__(self, config: ParallaxConfig) -> None:
+        # Bypass Qwen3Model.__init__ (which builds Qwen3DecoderLayers) and
+        # rebuild the module tree with ParallaxDecoderLayers.
+        Qwen3PreTrainedModel.__init__(self, config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                ParallaxDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # Qwen3Model constructs its rotary embedding from the config; reuse.
+        from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding
+
+        self.rotary_emb = Qwen3RotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.has_sliding_layers = "sliding_attention" in (config.layer_types or [])
+        self.post_init()
+
+
+class ParallaxForCausalLM(ParallaxPreTrainedModel, Qwen3ForCausalLM):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config: ParallaxConfig) -> None:
+        # Bypass Qwen3ForCausalLM.__init__ (which builds a Qwen3Model) and
+        # rebuild with ParallaxModel.
+        Qwen3PreTrainedModel.__init__(self, config)
+        self.model = ParallaxModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+
+__all__ = [
+    "ParallaxAttention",
+    "ParallaxDecoderLayer",
+    "ParallaxForCausalLM",
+    "ParallaxModel",
+    "ParallaxPreTrainedModel",
+    "parallax_centered_attention_forward",
+]

sharded/shard-00001-model-00001-of-00001.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86cbf2242cae2269ee1f2faa0749f0fd7753232306b47668c30d8517183bddcf
+size 1192157480

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aeb13307a71acd8fe81861d94ad54ab689df773318809eed3cbe794b4492dae4
+size 11422654

tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": "<|im_start|>",
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if message.content is string %}\n        {%- set content = message.content %}\n    {%- else %}\n        {%- set content = '' %}\n    {%- endif %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is string %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in content %}\n                {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n                {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if loop.index0 > ns.last_query_index %}\n            {%- if loop.last or (not loop.last and reasoning_content) %}\n                {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n            {%- else %}\n                {{- '<|im_start|>' + message.role + '\\n' + content }}\n            {%- endif %}\n        {%- else %}\n            {{- '<|im_start|>' + message.role + '\\n' + content }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}