Upload folder using huggingface_hub

.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

.meta.json

@@ -0,0 +1,69 @@
+{
+  "source_model": {
+    "name": "nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16",
+    "commit_id": "b366b87f762ea334ba3eb55912031baba1515810"
+  },
+  "env": {
+    "python": "3.12.13",
+    "packages": {
+      "huggingface_hub": "2.3.0",
+      "safetensors": "0.8.0",
+      "tokenizers": "0.23.3",
+      "torch": "2.13.0+cu126",
+      "transformers": "5.20.2"
+    },
+    "cuda": "12.6"
+  },
+  "environments": {
+    "creation": {
+      "python": "3.12.13",
+      "packages": {
+        "huggingface_hub": "2.3.0",
+        "safetensors": "0.8.0",
+        "tokenizers": "0.23.3",
+        "torch": "2.13.0+cu126",
+        "transformers": "5.20.2"
+      },
+      "cuda": "12.6"
+    },
+    "vllm-offline": {
+      "python": "3.12.13",
+      "packages": {
+        "flashinfer-python": "0.7.0.post1",
+        "huggingface_hub": "2.3.0",
+        "safetensors": "0.8.0",
+        "tokenizers": "0.23.3",
+        "torch": "2.13.0+cu130",
+        "transformers": "5.19.0",
+        "vllm": "0.32.0"
+      },
+      "cuda": "13.0"
+    },
+    "vllm-serving": {
+      "python": "3.12.13",
+      "packages": {
+        "flashinfer-python": "0.7.0.post1",
+        "huggingface_hub": "2.3.0",
+        "safetensors": "0.8.0",
+        "tokenizers": "0.23.3",
+        "torch": "2.13.0+cu130",
+        "transformers": "5.19.0",
+        "vllm": "0.32.0"
+      },
+      "cuda": "13.0"
+    },
+    "vllm-verify": {
+      "python": "3.12.13",
+      "packages": {
+        "flashinfer-python": "0.7.0.post1",
+        "huggingface_hub": "2.3.0",
+        "safetensors": "0.8.0",
+        "tokenizers": "0.23.3",
+        "torch": "2.13.0+cu130",
+        "transformers": "5.19.0",
+        "vllm": "0.32.0"
+      },
+      "cuda": "13.0"
+    }
+  }
+}

README.md

@@ -0,0 +1,899 @@
+---
+library_name: vllm
+pipeline_tag: text-generation
+base_model:
+- nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16
+---
+
+This tiny model is randomly initialized for debugging using [nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16).
+
+Dory preserves two recurrent passes, input reinjection, nGPT residuals,
+SWA/full attention, and untied BF16 token matrices. The pinned official vLLM
+implementation is included with tensor-parallel vector loading and distributed
+nGPT normalization. Transformers has no Dory model implementation or source
+auto_map; checkpoint generation follows the official source tensor schema.
+
+| Feature | Source | Tiny |
+|---|---|---|
+| Checkpoint tensor elements | 3,975,229,440 BF16 | 689,248 BF16 |
+| Safetensors weight files | 7,950,515,472 bytes (model.safetensors) | 1,412,480 bytes (model.safetensors) |
+| Input / recurrent / output sublayers | 16 / 40 / 16 | 16 / 16 / 16 |
+| Recurrent passes | 2 | 2 |
+| SWA/full cycle | 3 SWA+MLP pairs, then full+MLP | Same; two cycles per partition |
+| Vocabulary / embedding tying | 131,072 / false | Same |
+| Hidden width / MLP width | 2,560 / 7,168 | 2 / 32 |
+| Q heads / KV heads / head dimension | 16 / 4 / 256 | 4 / 1 / 32 |
+| Checkpoint precision | BF16 | Same |
+
+Install the included runtime once with `pip install --no-deps <snapshot directory>`
+before running the offline example. It registers Dory in every vLLM worker.
+The native vLLM sampler avoids FlashInfer sampling JIT; attention uses FlashAttention.
+Width 1 was rejected because scalar nGPT normalization collapses the latent state:
+one and two recurrent passes produced identical final states. Width 2 retains
+continuous recurrent updates; both untied vocabulary matrices keep their source rows.
+TP=1 and TP=4 were validated with English, Chinese and a 1,100-token prompt,
+generating ten tokens each. Both per-loop and last-loop caches and CUDA Graph
+were exercised; TP=1/TP=4 returned identical token IDs with at most 0.000644
+logprob difference. The independent creation and offline snippets passed outside
+the checkout with repository imports blocked.
+The current offline and HTTP serving examples were verified on four NVIDIA H20
+GPUs with TP=4, FlashAttention and CUDA Graph. Both chat requests (English and
+Chinese) and a 1,100-token completion request generated ten tokens with finite
+log probabilities. The CUDA 13.0 test environment used NVIDIA's CUDA compatibility
+libraries on a CUDA 12.6 driver; use a compatible driver/runtime combination.
+
+| File path | Size |
+|------|------|
+| model.safetensors | 1.4MB |
+
+### Example usage:
+
+- vLLM offline (✅ verified)
+
+```python
+def run_vllm(model_id: str):
+    import os
+    from pathlib import Path
+    from huggingface_hub import snapshot_download
+    from transformers import AutoTokenizer
+
+    # Use vLLM's native sampler; FlashAttention and CUDA Graph remain enabled.
+    os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
+    os.environ.setdefault("VLLM_ALLREDUCE_USE_FLASHINFER", "0")
+    os.environ.setdefault("VLLM_ALLREDUCE_USE_SYMM_MEM", "0")
+    from vllm import LLM, SamplingParams
+
+    model_id = str(model_id)
+    folder = Path(model_id) if Path(model_id).is_dir() else Path(snapshot_download(model_id))
+    tokenizer = AutoTokenizer.from_pretrained(folder, trust_remote_code=True)
+    prompt = tokenizer.apply_chat_template(
+        [{"role": "user", "content": "Hello"}], tokenize=False, add_generation_prompt=True
+    )
+    llm = LLM(
+        model=str(folder),
+        trust_remote_code=True,
+        tensor_parallel_size=4,
+        max_model_len=2048,
+        max_num_batched_tokens=256,
+        max_num_seqs=4,
+        gpu_memory_utilization=0.2,
+        disable_custom_all_reduce=True,
+        compilation_config={"cudagraph_capture_sizes": [1, 2, 4, 8]},
+    )
+    try:
+        print(llm)
+        print(llm.generate([prompt], SamplingParams(temperature=0, max_tokens=10, ignore_eos=True))[0].outputs[0].text)
+    finally:
+        llm.llm_engine.engine_core.shutdown(timeout=30)
+
+
+if __name__ == "__main__":
+    run_vllm(model_id="tiny-random/dory-1.0")
+```
+
+- vLLM serving (✅ verified)
+
+Download `tiny-random/dory-1.0` and install its included runtime before starting the server.
+
+```bash
+export VLLM_USE_FLASHINFER_SAMPLER=0 VLLM_ALLREDUCE_USE_FLASHINFER=0 VLLM_ALLREDUCE_USE_SYMM_MEM=0
+python - <<'PY'
+from huggingface_hub import snapshot_download
+snapshot_download("tiny-random/dory-1.0", local_dir="dory-1.0")
+PY
+pip install --no-deps ./dory-1.0
+
+vllm serve ./dory-1.0 \
+  --trust-remote-code \
+  --tensor-parallel-size 4 \
+  --hf-overrides '{"n_recurrent_loops": 2, "recurrent_kv_cache_mode": "per_loop"}' \
+  --max-model-len 2048 \
+  --max-num-batched-tokens 256 \
+  --max-num-seqs 4 \
+  --gpu-memory-utilization 0.2 \
+  --disable-custom-all-reduce \
+  --compilation-config '{"cudagraph_capture_sizes": [1, 2, 4, 8]}'
+```
+
+### Codes to create this repo:
+
+<details><summary>Click to expand</summary>
+
+```python
+def create_model(source_model_id: str, source_revision: str, save_folder: str):
+    import json
+    import urllib.request
+    from pathlib import Path
+
+    import torch
+    from huggingface_hub import hf_hub_download
+    from safetensors.torch import save_file
+
+    runtime_revision = "b5007e42b09c3990685fe564a31ae6917394a80b"
+    folder = Path(save_folder)
+    folder.mkdir(parents=True, exist_ok=True)
+    config = json.loads(Path(hf_hub_download(source_model_id, "config.json", revision=source_revision)).read_text())
+    # Two complete SWA/SWA/SWA/full cycles in every physical partition.
+    config.update(
+        hidden_size=2,
+        intermediate_size=32,
+        num_attention_heads=4,
+        num_key_value_heads=1,
+        head_dim=32,
+        num_hidden_layers=48,
+        n_input_layers=16,
+        n_recurrent_layers=16,
+        n_output_layers=16,
+    )
+    config["layer_types"] = config["layer_types"][:48]
+    config["auto_map"] = {"AutoConfig": "configuration_dory.DoryConfig"}
+    (folder / "config.json").write_text(json.dumps(config, indent=2) + "\n")
+    for name in (
+        "tokenizer.json",
+        "tokenizer_config.json",
+        "special_tokens_map.json",
+        "chat_template.jinja",
+        "generation_config.json",
+    ):
+        (folder / name).write_bytes(Path(hf_hub_download(source_model_id, name, revision=source_revision)).read_bytes())
+
+    # No Transformers model class or source auto_map exists. Preserve the pinned
+    # official checkpoint schema and ship its vLLM runtime with TP adaptations.
+    base = f"https://raw.githubusercontent.com/vllm-project/vllm/{runtime_revision}/"
+    for remote, local in (
+        ("vllm/transformers_utils/configs/dory.py", "configuration_dory.py"),
+        ("vllm/model_executor/models/dory.py", "vllm_dory.py"),
+    ):
+        with urllib.request.urlopen(base + remote) as response:
+            code = response.read().decode()
+        if local == "vllm_dory.py":
+            code = code.replace(
+                "from .utils import maybe_prefix", "from vllm.model_executor.models.utils import maybe_prefix"
+            )
+            code = code.replace(
+                "from vllm.distributed import get_tensor_model_parallel_world_size",
+                "from vllm.distributed import (get_tensor_model_parallel_world_size, "
+                "get_tensor_model_parallel_rank, tensor_model_parallel_all_reduce)",
+            )
+            code = code.replace(
+                "EPS = 1e-6",
+                "EPS = 1e-6\n\n"
+                "def partition_norm(x):\n"
+                "    total = tensor_model_parallel_all_reduce(x.float().square().sum(dim=-1, keepdim=True))\n"
+                "    return (x.float() * torch.rsqrt(total + EPS)).to(x.dtype)\n\n"
+                "def load_vector(param, weight):\n"
+                "    rank = get_tensor_model_parallel_rank()\n"
+                "    param.data.copy_(weight.narrow(0, rank * param.numel(), param.numel()))\n",
+            )
+            code = code.replace(
+                "self.num_heads = config.num_attention_heads",
+                "self.num_heads = config.num_attention_heads // get_tensor_model_parallel_world_size()",
+            )
+            code = code.replace(
+                "self.num_kv_heads = config.num_key_value_heads",
+                "self.num_kv_heads = max(1, config.num_key_value_heads // get_tensor_model_parallel_world_size())",
+            )
+            code = code.replace(
+                "            self.num_heads,\n            self.num_kv_heads,",
+                "            config.num_attention_heads,\n            config.num_key_value_heads,",
+            )
+            code = code.replace(
+                "            self.q_size,\n            bias=False,",
+                "            config.num_attention_heads * self.head_dim,\n            bias=False,",
+            )
+            code = code.replace(
+                "            self.q_size,\n            self.hidden_size,",
+                "            config.num_attention_heads * self.head_dim,\n            self.hidden_size,",
+            )
+            code = code.replace(
+                "self.s_gate = nn.Parameter(torch.empty(self.q_size))",
+                "self.s_gate = nn.Parameter(torch.empty(self.q_size))\n"
+                "        self.sqk.weight_loader = self.s_gate.weight_loader = load_vector",
+            )
+            code = code.replace(
+                "self.ffn = config.intermediate_size",
+                "self.ffn = config.intermediate_size // get_tensor_model_parallel_world_size()",
+            )
+            code = code.replace("[self.ffn] * 2", "[config.intermediate_size] * 2")
+            code = code.replace(
+                "            self.ffn,\n            config.hidden_size,",
+                "            config.intermediate_size,\n            config.hidden_size,",
+            )
+            code = code.replace(
+                "self.suv_up = nn.Parameter(torch.empty(self.ffn))",
+                "self.suv_up = nn.Parameter(torch.empty(self.ffn))\n"
+                "        self.suv_gate.weight_loader = self.suv_up.weight_loader = load_vector",
+            )
+            code = code.replace("out = justnorm(out)", "out = partition_norm(out)")
+            code = code.replace("m = justnorm(m)", "m = partition_norm(m)")
+            code = code.replace(
+                "        if get_tensor_model_parallel_world_size() != 1:\n"
+                '            raise ValueError("Dory currently supports tensor_parallel_size=1 only")\n',
+                "",
+            )
+            code = code.replace(
+                "            self._check_weight_shape(name, w, shape)\n" "            getattr(param,",
+                '            if name.endswith((".sqk", ".s_gate")):\n'
+                "                shape = (self.config.num_attention_heads * self.config.head_dim,)\n"
+                '            elif name.endswith((".suv_gate", ".suv_up")):\n'
+                "                shape = (self.config.intermediate_size,)\n"
+                '            elif name.endswith(".g_proj.weight"):\n'
+                "                shape = (self.config.num_attention_heads * self.config.head_dim, self.config.hidden_size)\n"
+                '            elif name.endswith(".o_proj.weight"):\n'
+                "                shape = (self.config.hidden_size, self.config.num_attention_heads * self.config.head_dim)\n"
+                '            elif name.endswith(".down_proj.weight"):\n'
+                "                shape = (self.config.hidden_size, self.config.intermediate_size)\n"
+                "            self._check_weight_shape(name, w, shape)\n"
+                "            getattr(param,",
+            )
+            assert "tensor_parallel_size=1 only" not in code
+            assert code.count("partition_norm(") == 3
+        (folder / local).write_text(code)
+    # vLLM spawns fresh interpreters. A general plugin registers the model in
+    # every process; parent-process ModelRegistry changes alone do not survive.
+    (folder / "dory_plugin.py").write_text(
+        "def register():\n"
+        "    from vllm import ModelRegistry\n"
+        "    ModelRegistry.register_model('DoryForCausalLM', 'vllm_dory:DoryForCausalLM')\n"
+    )
+    (folder / "pyproject.toml").write_text(
+        '[build-system]\nrequires = ["setuptools>=61"]\nbuild-backend = "setuptools.build_meta"\n'
+        '[project]\nname = "dory-inference"\nversion = "0.1.0"\n'
+        '[project.entry-points."vllm.general_plugins"]\ndory = "dory_plugin:register"\n'
+        '[tool.setuptools]\npy-modules = ["dory_plugin", "vllm_dory", "configuration_dory"]\n'
+    )
+
+    torch.manual_seed(42)
+    h, f = config["hidden_size"], config["intermediate_size"]
+    q, kv = config["num_attention_heads"] * config["head_dim"], config["num_key_value_heads"] * config["head_dim"]
+    tensors = {}
+
+    class CheckpointModule(torch.nn.Module):
+        def extra_repr(self):
+            return "\n".join(
+                f"{name}: shape={tuple(p.shape)}, dtype={p.dtype}" for name, p in self.named_parameters(recurse=False)
+            )
+
+    graph = CheckpointModule()
+
+    def add(name, shape):
+        tensor = torch.randn(shape, dtype=torch.bfloat16) * config["initializer_range"]
+        tensors[name] = tensor
+        module = graph
+        parts = name.split(".")
+        for part in parts[:-1]:
+            if part not in module._modules:
+                module.add_module(part, CheckpointModule())
+            module = module._modules[part]
+        module.register_parameter(parts[-1], torch.nn.Parameter(tensor))
+
+    add("backbone.embeddings.weight", (config["vocab_size"], h))
+    add("lm_head.weight", (config["vocab_size"], h))
+    add("lm_head.logit_scale", (config["vocab_size"],))
+    for i, kind in enumerate(config["layer_types"]):
+        prefix = f"backbone.layers.{i}"
+        if kind == "mlp":
+            add(prefix + ".mlp_alpha", (h,))
+            for name in ("gate_proj", "up_proj"):
+                add(prefix + f".mixer.{name}.weight", (f, h))
+            add(prefix + ".mixer.down_proj.weight", (h, f))
+            for name in ("suv_gate", "suv_up"):
+                add(prefix + f".mixer.{name}", (f,))
+        else:
+            add(prefix + ".attn_alpha", (h,))
+            for name, rows in (("q_proj", q), ("k_proj", kv), ("v_proj", kv), ("g_proj", q)):
+                add(prefix + f".mixer.{name}.weight", (rows, h))
+            add(prefix + ".mixer.o_proj.weight", (h, q))
+            for name in ("sqk", "s_gate"):
+                add(prefix + f".mixer.{name}", (q,))
+    save_file(tensors, str(folder / "model.safetensors"), metadata={"format": "pt"})
+    return str(graph)
+
+
+create_model(
+    source_model_id="nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16",
+    source_revision="b366b87f762ea334ba3eb55912031baba1515810",
+    save_folder="/mlx_devbox/users/yujie.pan/playground/tiny-random/tmp/scratch/dory-builder-v1",
+)
+```
+
+</details>
+
+### Printing the model:
+
+<details><summary>Click to expand</summary>
+
+```text
+CheckpointModule(
+  (backbone): CheckpointModule(
+    (embeddings): CheckpointModule(weight: shape=(131072, 2), dtype=torch.bfloat16)
+    (layers): CheckpointModule(
+      (0): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (1): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (2): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (3): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (4): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (5): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (6): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (7): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (8): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (9): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (10): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (11): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (12): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (13): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (14): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (15): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (16): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (17): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (18): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (19): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (20): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (21): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (22): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (23): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (24): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (25): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (26): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (27): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (28): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (29): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (30): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (31): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (32): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (33): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (34): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (35): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (36): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (37): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (38): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (39): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (40): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (41): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (42): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (43): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (44): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (45): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+      (46): CheckpointModule(
+        attn_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          sqk: shape=(128,), dtype=torch.bfloat16
+          s_gate: shape=(128,), dtype=torch.bfloat16
+          (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16)
+          (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16)
+        )
+      )
+      (47): CheckpointModule(
+        mlp_alpha: shape=(2,), dtype=torch.bfloat16
+        (mixer): CheckpointModule(
+          suv_gate: shape=(32,), dtype=torch.bfloat16
+          suv_up: shape=(32,), dtype=torch.bfloat16
+          (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16)
+          (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16)
+        )
+      )
+    )
+  )
+  (lm_head): CheckpointModule(
+    weight: shape=(131072, 2), dtype=torch.bfloat16
+    logit_scale: shape=(131072,), dtype=torch.bfloat16
+  )
+)
+```
+
+</details>
+
+### Creation environment:
+
+- python: 3.12.13
+- cuda: 12.6
+- huggingface_hub: 2.3.0
+- safetensors: 0.8.0
+- tokenizers: 0.23.3
+- torch: 2.13.0+cu126
+- transformers: 5.20.2
+
+### Test environment:
+
+**vLLM offline / vLLM serving**
+
+- python: 3.12.13
+- cuda: 13.0
+- flashinfer-python: 0.7.0.post1
+- huggingface_hub: 2.3.0
+- safetensors: 0.8.0
+- tokenizers: 0.23.3
+- torch: 2.13.0+cu130
+- transformers: 5.19.0
+- vllm: 0.32.0
+
+### Change log:
+
+- 2026-10-25: Created BF16 Dory with a two-dimensional latent and verified TP=4 inference.
+- 2026-10-26: Verified the current standalone offline and HTTP serving examples on four H20 GPUs.
+
+This repo is generated by GPT-6.1-Sol

chat_template.jinja

@@ -0,0 +1,195 @@
+{% macro render_extra_keys(json_dict, handled_keys) %}
+    {%- if json_dict is mapping %}
+        {%- for json_key in json_dict if json_key not in handled_keys %}
+            {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
+                {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '</' ~ json_key ~ '>' }}
+            {%- else %}
+                {{-'\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}
+            {%- endif %}
+        {%- endfor %}
+    {%- endif %}
+{% endmacro %}
+{%- set enable_thinking = enable_thinking if enable_thinking is defined else True %}
+{%- set medium_effort = medium_effort if medium_effort is defined else False %}
+{%- set truncate_history_thinking = truncate_history_thinking if truncate_history_thinking is defined else True %}
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- set loop_messages = messages %}
+{%- for m in loop_messages %}
+  {%- if m["role"] == "user" %}
+    {%- set ns.last_user_idx = loop.index0 %}
+  {%- endif %}
+{%- endfor %}
+{%- if messages[0]["role"] == "system" %}
+    {%- set system_message = messages[0]["content"] %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set system_message = "" %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+{%- if not tools is defined %}
+    {%- set tools = [] %}
+{%- endif %}
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- for m in loop_messages %}
+  {%- if m["role"] == "user" %}
+    {%- set ns.last_user_idx = loop.index0 %}
+  {%- endif %}
+{%- endfor %}
+{%- if system_message is defined %}
+    {{- "<|im_start|>system\n" + system_message }}
+{%- else %}
+    {%- if tools is iterable and tools | length > 0 %}
+        {{- "<|im_start|>system\n" }}
+    {%- endif %}
+{%- endif %}
+{%- if tools is iterable and tools | length > 0 %}
+    {%- if system_message is defined and system_message | length > 0 %}
+        {{- "\n\n" }}
+    {%- endif %}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n" }}
+    {{- "<tools>" }}
+    {%- for tool in tools %}
+        {%- if tool.function is defined %}
+            {%- set tool = tool.function %}
+        {%- endif %}
+        {{- "\n<function>\n<name>" ~ tool.name ~ "</name>" }}
+        {%- if tool.description is defined %}
+            {{- '\n<description>' ~ (tool.description | trim) ~ '</description>' }}
+        {%- endif %}
+        {{- '\n<parameters>' }}
+        {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
+            {%- for param_name, param_fields in tool.parameters.properties|items %}
+                {{- '\n<parameter>' }}
+                {{- '\n<name>' ~ param_name ~ '</name>' }}
+                {%- if param_fields.type is defined %}
+                    {{- '\n<type>' ~ (param_fields.type | string) ~ '</type>' }}
+                {%- endif %}
+                {%- if param_fields.description is defined %}
+                    {{- '\n<description>' ~ (param_fields.description | trim) ~ '</description>' }}
+                {%- endif %}
+                {%- if param_fields.enum is defined %}
+                    {{- '\n<enum>' ~ (param_fields.enum | tojson | safe) ~ '</enum>' }}
+                {%- endif %}
+                {%- set handled_keys = ['name', 'type', 'description', 'enum'] %}
+                {{- render_extra_keys(param_fields, handled_keys) }}
+                {{- '\n</parameter>' }}
+            {%- endfor %}
+        {%- endif %}
+        {% set handled_keys = ['type', 'properties', 'required'] %}
+        {{- render_extra_keys(tool.parameters, handled_keys) }}
+        {%- if tool.parameters is defined and tool.parameters.required is defined %}
+            {{- '\n<required>' ~ (tool.parameters.required | tojson | safe) ~ '</required>' }}
+        {%- endif %}
+        {{- '\n</parameters>' }}
+        {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
+        {{- render_extra_keys(tool, handled_keys) }}
+        {{- '\n</function>' }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+{%- endif %}
+{%- if system_message is defined %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if tools is iterable and tools | length > 0 %}
+        {{- '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in loop_messages %}
+    {%- if message.role == "assistant" %}
+        {%- if message.reasoning_content is defined and message.reasoning_content is string and message.reasoning_content | trim | length > 0 %}
+            {%- set content = "<think>\n" ~ message.reasoning_content ~ "</think>" ~ (message.content | default('', true)) %}
+        {%- else %}
+            {%- set content = message.content | default('', true) %}
+            {%- if content is string -%}
+                {%- if '<think>' not in content and '</think>' not in content -%}
+                    {%- set content = "<think></think>" ~ content -%}
+                {%- endif -%}
+            {%- else -%}
+                {%- set content = content -%}
+            {%- endif -%}
+        {%- endif %}
+        {%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}
+            {{- '<|im_start|>assistant\n' }}
+                {%- set include_content = not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+                {%- if content is string and content | trim | length > 0 %}
+                    {%- if include_content %}
+                        {{- (content | trim) ~ '\n' -}}
+                    {%- else %}
+                        {%- set c = (content | string) %}
+                        {%- if '</think>' in c %}
+                            {%- set c = c.split('</think>')[-1] %}
+                        {%- elif '<think>' in c %}
+                            {%- set c = c.split('<think>')[0] %}
+                        {%- endif %}
+                        {%- set c = "<think></think>" ~ c %}
+                        {%- if c | length > 0 %}
+                            {{- c ~ '\n' -}}
+                        {%- endif %}
+                    {%- endif %}
+                {%- else %}
+                    {{- "<think></think>" -}}
+                {%- endif %}
+                {%- for tool_call in message.tool_calls %}
+                    {%- if tool_call.function is defined %}
+                        {%- set tool_call = tool_call.function %}
+                    {%- endif %}
+                    {{- '<tool_call>\n<function=' ~ tool_call.name ~ '>\n' -}}
+                        {%- if tool_call.arguments is defined %}
+                            {%- for args_name, args_value in tool_call.arguments|items %}
+                                {{- '<parameter=' ~ args_name ~ '>\n' -}}
+                                    {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                                {{- args_value ~ '\n</parameter>\n' -}}
+                            {%- endfor %}
+                        {%- endif %}
+                    {{- '</function>\n</tool_call>\n' -}}
+                {%- endfor %}
+                {{- '<|im_end|>\n' }}
+        {%- else %}
+            {%- if not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+                {{- '<|im_start|>assistant\n' ~ (content | default('', true) | string | trim) ~ '<|im_end|>\n' }}
+            {%- else %}
+                {%- set c = (content | default('', true) | string) %}
+                {%- if '<think>' in c and '</think>' in c %}
+                    {%- set c = "<think></think>" ~ c.split('</think>')[-1] %}
+                {%- endif %}
+                {%- set c = c | trim %}
+                {%- if c | length > 0 %}
+                    {{- '<|im_start|>assistant\n' ~ c ~ '<|im_end|>\n' }}
+                {%- else %}
+                    {{- '<|im_start|>assistant\n<|im_end|>\n' }}
+                {%- endif %}
+            {%- endif %}
+        {%- endif %}
+    {%- elif message.role == "user" or message.role == "system" %}
+        {{- '<|im_start|>' + message.role + '\n' }}
+        {%- set content = message.content | string %}
+        {%- if message.role == "user" and loop.index0 == ns.last_user_idx and medium_effort %}
+            {{- content + '\n\n{reasoning effort: efficient}' }}
+        {%- else %}
+            {{- content }}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user\n' }}
+        {%- endif %}
+        {{- '<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>\n' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {%- if enable_thinking %}
+        {{- '<|im_start|>assistant\n<think>\n' }}
+    {%- else %}
+        {{- '<|im_start|>assistant\n<think></think>' }}
+    {%- endif %}
+{%- endif %}

config.json

@@ -0,0 +1,92 @@
+{
+  "architectures": [
+    "DoryForCausalLM"
+  ],
+  "model_type": "dory",
+  "vocab_size": 131072,
+  "hidden_size": 2,
+  "num_hidden_layers": 48,
+  "intermediate_size": 32,
+  "num_attention_heads": 4,
+  "num_key_value_heads": 1,
+  "max_position_embeddings": 262144,
+  "tie_word_embeddings": false,
+  "ngpt_alpha_init": 0.05,
+  "attention_output_gate_logit_scale_init": null,
+  "mlp_bias": false,
+  "attention_bias": false,
+  "n_input_layers": 16,
+  "n_recurrent_layers": 16,
+  "n_output_layers": 16,
+  "n_recurrent_loops": 2,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "initializer_range": 0.0177,
+  "dtype": "bfloat16",
+  "sliding_window": 1024,
+  "rope_parameters": {
+    "sliding_attention": {
+      "rope_type": "default",
+      "rope_theta": 10000,
+      "partial_rotary_factor": 1.0
+    },
+    "full_attention": {
+      "rope_type": "default",
+      "rope_theta": 10000000,
+      "partial_rotary_factor": 1.0
+    }
+  },
+  "layer_types": [
+    "sliding_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "full_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "full_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "full_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "full_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "full_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "sliding_attention",
+    "mlp",
+    "full_attention",
+    "mlp"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_dory.DoryConfig"
+  }
+}

configuration_dory.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Configuration for the normalized latent recurrent Dory architecture."""
+
+from huggingface_hub.dataclasses import strict
+from transformers.configuration_utils import PreTrainedConfig
+
+
+@strict
+class DoryConfig(PreTrainedConfig):
+    model_type = "dory"
+    default_theta = 10_000.0
+
+    vocab_size: int = 131072
+    hidden_size: int = 2560
+    num_hidden_layers: int = 72
+    intermediate_size: int = 7168
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 4
+    head_dim: int = 256
+    max_position_embeddings: int = 8192
+    tie_word_embeddings: bool = False
+    hidden_act: str = "silu"
+    # Also normalizes the learned nGPT scales at inference time.
+    initializer_range: float = 0.02
+    sliding_window: int | None = 1024
+    rope_parameters: dict | None = None
+    layer_types: list[str] | None = None
+    ngpt_alpha_init: float = 0.05
+    attention_output_gate_logit_scale_init: float | None = None
+    mlp_bias: bool = False
+    attention_bias: bool = False
+    n_input_layers: int = 0
+    n_recurrent_layers: int | None = None
+    n_output_layers: int = 0
+    n_recurrent_loops: int = 1
+    # "last_loop" reuses one cache, overwritten by each recurrent loop.
+    recurrent_kv_cache_mode: str = "per_loop"
+
+    def __post_init__(self, **kwargs):
+        if self.layer_types is None:
+            self.layer_types = [
+                "mlp"
+                if i % 2
+                else "full_attention"
+                if i % 8 == 6
+                else "sliding_attention"
+                for i in range(self.num_hidden_layers)
+            ]
+        if self.rope_parameters is None:
+            self.rope_parameters = {
+                kind: {
+                    "rope_type": "default",
+                    "partial_rotary_factor": 1.0,
+                }
+                for kind in ("sliding_attention", "full_attention")
+                if kind in self.layer_types
+            }
+        if self.n_recurrent_layers is None:
+            self.n_recurrent_layers = (
+                self.num_hidden_layers - self.n_input_layers - self.n_output_layers
+            )
+        super().__post_init__(**kwargs)
+
+    def validate_layer_type(self):
+        # Each physical block is attention-only or MLP-only.
+        valid_types = {"sliding_attention", "full_attention", "mlp"}
+        if self.layer_types is None or any(
+            kind not in valid_types for kind in self.layer_types
+        ):
+            raise ValueError(
+                "layer_types must be a list of sliding_attention, "
+                "full_attention, or mlp"
+            )
+        if len(self.layer_types) != self.num_hidden_layers:
+            raise ValueError("layer_types must have num_hidden_layers entries")
+
+    def validate_architecture(self):
+        if self.layer_types is None or self.n_recurrent_layers is None:
+            raise ValueError("Dory layer layout must be initialized")
+        partitions = (
+            self.n_input_layers,
+            self.n_recurrent_layers,
+            self.n_output_layers,
+        )
+        if any(n < 0 for n in partitions) or sum(partitions) != self.num_hidden_layers:
+            raise ValueError(
+                "Dory input, recurrent, and output layers must sum to num_hidden_layers"
+            )
+        if self.n_recurrent_loops < 1:
+            raise ValueError("n_recurrent_loops must be a positive integer")
+        if self.recurrent_kv_cache_mode not in ("per_loop", "last_loop"):
+            raise ValueError("recurrent_kv_cache_mode must be per_loop or last_loop")
+        if any(
+            n <= 0
+            for n in (self.head_dim, self.num_attention_heads, self.num_key_value_heads)
+        ):
+            raise ValueError("Attention dimensions must be positive integers")
+        if self.num_attention_heads % self.num_key_value_heads:
+            raise ValueError(
+                "num_attention_heads must be divisible by num_key_value_heads"
+            )
+        if self.initializer_range <= 0:
+            raise ValueError("initializer_range must be positive")
+        if "sliding_attention" in self.layer_types and (
+            self.sliding_window is None or self.sliding_window <= 0
+        ):
+            raise ValueError("sliding_window must be a positive integer")
+        if self.attention_bias or self.mlp_bias:
+            raise ValueError("Dory requires bias-free projections")
+        if self.hidden_act != "silu" or self.tie_word_embeddings:
+            raise ValueError("Dory requires SiLU and untied embeddings")
+
+    def validate_rope(self):
+        if self.layer_types is None:
+            raise ValueError("Dory layer layout must be initialized")
+        attention_types = set(self.layer_types) - {"mlp"}
+        if self.rope_parameters is None or (
+            set(self.rope_parameters) - {"sliding_attention", "full_attention"}
+            or attention_types - self.rope_parameters.keys()
+        ):
+            raise ValueError("rope_parameters must define each attention block type")
+        for rope in self.rope_parameters.values():
+            if (
+                not isinstance(rope, dict)
+                or rope.get("rope_type") != "default"
+                or rope.get("rope_theta", 0) <= 0
+            ):
+                raise ValueError(
+                    "Dory requires default RoPE with a positive rope_theta"
+                )
+            fraction = rope.get("partial_rotary_factor", 1.0)
+            rotary_dim = int(self.head_dim * fraction)
+            if not 0 < fraction <= 1 or rotary_dim <= 0 or rotary_dim % 2:
+                raise ValueError(
+                    "Rotary dimensions must be positive, even, and <= head_dim"
+                )

dory_plugin.py

@@ -0,0 +1,3 @@
+def register():
+    from vllm import ModelRegistry
+    ModelRegistry.register_model('DoryForCausalLM', 'vllm_dory:DoryForCausalLM')

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "do_sample": true,
+  "eos_token_id": [
+    2,
+    11
+  ],
+  "pad_token_id": 11,
+  "top_p": 0.95,
+  "temperature": 1.0,
+  "transformers_version": "4.57.6"
+}

model.safetensors

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

pyproject.toml

@@ -0,0 +1,10 @@
+[build-system]
+requires = ["setuptools>=61"]
+build-backend = "setuptools.build_meta"
+[project]
+name = "dory-inference"
+version = "0.1.0"
+[project.entry-points."vllm.general_plugins"]
+dory = "dory_plugin:register"
+[tool.setuptools]
+py-modules = ["dory_plugin", "vllm_dory", "configuration_dory"]

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

vllm_dory.py

@@ -0,0 +1,410 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Dory: normalized latent recurrent transformer."""
+
+import math
+from collections.abc import Iterable
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.distributed import (get_tensor_model_parallel_world_size, get_tensor_model_parallel_rank, tensor_model_parallel_all_reduce)
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.parameter import PerTensorScaleParameter
+from vllm.sequence import IntermediateTensors
+
+from vllm.model_executor.models.utils import maybe_prefix
+
+EPS = 1e-6
+
+def partition_norm(x):
+    total = tensor_model_parallel_all_reduce(x.float().square().sum(dim=-1, keepdim=True))
+    return (x.float() * torch.rsqrt(total + EPS)).to(x.dtype)
+
+def load_vector(param, weight):
+    rank = get_tensor_model_parallel_rank()
+    param.data.copy_(weight.narrow(0, rank * param.numel(), param.numel()))
+
+
+
+def justnorm(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
+    """NGPT normalization: x * rsqrt(sum(x^2) + eps), fp32 accumulation."""
+    xf = x.float()
+    y = xf * torch.rsqrt(xf.pow(2).sum(dim=dim, keepdim=True) + EPS)
+    return y.to(x.dtype)
+
+
+def slerp_norm(a: torch.Tensor, b: torch.Tensor, lr: torch.Tensor) -> torch.Tensor:
+    ya, yb = justnorm(a).float(), justnorm(b).float()
+    return justnorm(ya + lr.float() * (yb - ya)).to(a.dtype)
+
+
+def slerp_alpha(a, b, alpha, alpha_scale):
+    return slerp_norm(a, b, (alpha.float() * alpha_scale).abs())
+
+
+class DoryAttention(nn.Module):
+    """GQA attention with nGPT qk-norm, output gate, and fngpt justnorm."""
+
+    def __init__(
+        self, config, layer_idx, cache_config, quant_config, prefix, num_loops=1
+    ):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads // get_tensor_model_parallel_world_size()
+        self.num_kv_heads = max(1, config.num_key_value_heads // get_tensor_model_parallel_world_size())
+        self.head_dim = config.head_dim
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        std = config.initializer_range
+
+        self.qkv_proj = QKVParallelLinear(
+            self.hidden_size,
+            self.head_dim,
+            config.num_attention_heads,
+            config.num_key_value_heads,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+        self.g_proj = ColumnParallelLinear(
+            self.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.g_proj",
+        )
+        self.o_proj = RowParallelLinear(
+            config.num_attention_heads * self.head_dim,
+            self.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        self.sqk = nn.Parameter(torch.empty(self.q_size))
+        self.s_gate = nn.Parameter(torch.empty(self.q_size))
+        self.sqk.weight_loader = self.s_gate.weight_loader = load_vector
+        self._sqk_factor = self.head_dim**0.5 / std
+        g0 = config.attention_output_gate_logit_scale_init
+        if g0 is None:
+            g0 = config.hidden_size**0.5
+        self._s_gate_factor = g0**0.5 / std
+
+        layer_type = config.layer_types[layer_idx]
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=config.max_position_embeddings,
+            rope_parameters=config.rope_parameters[layer_type],
+            is_neox_style=True,
+        )
+
+        self.per_loop_kv = config.recurrent_kv_cache_mode == "per_loop"
+        # Match LoopCoder's virtual layer indices for independent loop caches.
+        # Projection prefixes keep the physical layer index to share weights.
+        self.attn = nn.ModuleList()
+        for loop_idx in range(num_loops if self.per_loop_kv else 1):
+            unique_layer_idx = loop_idx * config.num_hidden_layers + layer_idx
+            unique_prefix = prefix.replace(
+                f".layers.{layer_idx}.", f".layers.{unique_layer_idx}."
+            )
+            self.attn.append(
+                Attention(
+                    self.num_heads,
+                    self.head_dim,
+                    scale=1.0 / math.sqrt(self.head_dim),
+                    num_kv_heads=self.num_kv_heads,
+                    cache_config=cache_config,
+                    quant_config=quant_config,
+                    per_layer_sliding_window=(
+                        config.sliding_window
+                        if layer_type == "sliding_attention"
+                        else None
+                    ),
+                    prefix=f"{unique_prefix}.attn",
+                )
+            )
+
+    def forward(self, positions, hidden_states, loop_idx=0):
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        gate, _ = self.g_proj(hidden_states)
+
+        q, k = self.rotary_emb(positions, q, k)
+
+        # nGPT qk-norm AFTER rope: q scaled by (sqk*f)^2 post-justnorm, k plain justnorm
+        t = q.shape[0]
+        q = q.view(t, self.num_heads, self.head_dim)
+        sqk = (self.sqk.float() * self._sqk_factor).view(
+            1, self.num_heads, self.head_dim
+        )
+        q = (sqk.square() * justnorm(q).float()).to(q.dtype).view(t, self.q_size)
+        k = k.view(t, self.num_kv_heads, self.head_dim)
+        k = justnorm(k).view(t, self.kv_size)
+
+        # In last-loop mode every loop writes this step's KV to the same cache,
+        # overwriting the previous loop's. Loop i thus attends to loop-i KV for
+        # this step's tokens and last-loop KV for earlier tokens.
+        out = self.attn[loop_idx if self.per_loop_kv else 0](q, k, v)
+
+        # output gate (sigmoid, pre-o_proj), then fngpt justnorm over np*hn
+        sg = (self.s_gate.float() * self._s_gate_factor).square()
+        out = (out.float() * torch.sigmoid(gate.float() * sg)).to(out.dtype)
+        out = partition_norm(out)
+        out, _ = self.o_proj(out)
+        return out
+
+
+class DoryMLP(nn.Module):
+    """SwiGLU with nGPT suv scaling and fngpt justnorm."""
+
+    def __init__(self, config, quant_config, prefix):
+        super().__init__()
+        self.ffn = config.intermediate_size // get_tensor_model_parallel_world_size()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            config.hidden_size,
+            [config.intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            config.intermediate_size,
+            config.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        self.suv_gate = nn.Parameter(torch.empty(self.ffn))
+        self.suv_up = nn.Parameter(torch.empty(self.ffn))
+        self.suv_gate.weight_loader = self.suv_up.weight_loader = load_vector
+        self._suv_factor = config.hidden_size**0.5
+
+    def forward(self, hidden_states):
+        m, _ = self.gate_up_proj(hidden_states)
+        g, u = m.chunk(2, dim=-1)
+        g = (g.float() * (self.suv_gate.float() * self._suv_factor)).to(g.dtype)
+        u = (u.float() * (self.suv_up.float() * self._suv_factor)).to(u.dtype)
+        m = (F.silu(g.float()) * u.float()).to(m.dtype)
+        m = partition_norm(m)
+        out, _ = self.down_proj(m)
+        return out
+
+
+class DoryDecoderLayer(nn.Module):
+    """One physical layer: a single sublayer (attention OR mlp) + slerp residual."""
+
+    def __init__(
+        self, config, layer_idx, cache_config, quant_config, prefix, num_loops=1
+    ):
+        super().__init__()
+        self.is_attention = config.layer_types[layer_idx] != "mlp"
+        self._alpha_scale = config.ngpt_alpha_init / config.initializer_range
+        if self.is_attention:
+            self.mixer = DoryAttention(
+                config,
+                layer_idx,
+                cache_config,
+                quant_config,
+                f"{prefix}.mixer",
+                num_loops=num_loops,
+            )
+            self.attn_alpha = nn.Parameter(torch.empty(config.hidden_size))
+        else:
+            self.mixer = DoryMLP(config, quant_config, f"{prefix}.mixer")
+            self.mlp_alpha = nn.Parameter(torch.empty(config.hidden_size))
+
+    def forward(self, positions, hidden_states, loop_idx=0):
+        residual = hidden_states
+        if self.is_attention:
+            out = self.mixer(positions, hidden_states, loop_idx=loop_idx)
+            return slerp_alpha(residual, out, self.attn_alpha, self._alpha_scale)
+        out = self.mixer(hidden_states)
+        return slerp_alpha(residual, out, self.mlp_alpha, self._alpha_scale)
+
+
+@support_torch_compile
+class DoryBackbone(nn.Module):
+    def __init__(self, config, cache_config, quant_config, prefix):
+        super().__init__()
+        self.config = config
+        self.num_loops = config.n_recurrent_loops
+        n_in = config.n_input_layers
+        n_rec = config.n_recurrent_layers
+        self.embeddings = VocabParallelEmbedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            [
+                DoryDecoderLayer(
+                    config,
+                    i,
+                    cache_config,
+                    quant_config,
+                    f"{prefix}.layers.{i}",
+                    num_loops=self.num_loops if n_in <= i < n_in + n_rec else 1,
+                )
+                for i in range(config.num_hidden_layers)
+            ]
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        h = inputs_embeds if inputs_embeds is not None else self.embeddings(input_ids)
+        c = self.config
+        n_in, n_rec = c.n_input_layers, c.n_recurrent_layers
+
+        for layer in self.layers[:n_in]:
+            h = layer(positions, h)
+        input_h = h
+
+        state = input_h
+        for i in range(self.num_loops):
+            combined = justnorm(state + input_h)
+            h = combined
+            for layer in self.layers[n_in : n_in + n_rec]:
+                h = layer(positions, h, loop_idx=i)
+            state = justnorm(h)
+
+        h = state
+        for layer in self.layers[n_in + n_rec :]:
+            h = layer(positions, h)
+        return h
+
+
+class DoryForCausalLM(nn.Module):
+    packed_modules_mapping = {
+        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        self.config = config
+        # hf_overrides can change architecture fields after config initialization.
+        config.validate()
+        self.backbone = DoryBackbone(
+            config,
+            vllm_config.cache_config,
+            vllm_config.quant_config,
+            prefix=maybe_prefix(prefix, "backbone"),
+        )
+        self.lm_head = ParallelLMHead(
+            config.vocab_size,
+            config.hidden_size,
+            bias=False,
+            prefix=maybe_prefix(prefix, "lm_head"),
+        )
+        self.lm_head.logit_scale = nn.Parameter(torch.empty(config.vocab_size))
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+
+    def embed_input_ids(self, input_ids):
+        return self.backbone.embeddings(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ):
+        return self.backbone(input_ids, positions, inputs_embeds=inputs_embeds)
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        if logits is not None:
+            logits = logits.float() * (
+                self.lm_head.logit_scale.float() / self.config.initializer_range
+            )
+        return logits
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked = [
+            # (target, artifact shard name, shard id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+        params = dict(self.named_parameters())
+        loaded: set[str] = set()
+        for name, w in weights:
+            mapped = False
+            for target, shard_name, shard_id in stacked:
+                if f".{shard_name}." not in name:
+                    continue
+                pname = name.replace(f".{shard_name}.", f".{target}.")
+                param = params[pname]
+                if target == "qkv_proj":
+                    heads = (
+                        self.config.num_attention_heads
+                        if shard_id == "q"
+                        else self.config.num_key_value_heads
+                    )
+                    rows = heads * self.config.head_dim
+                else:
+                    rows = self.config.intermediate_size
+                if isinstance(param, PerTensorScaleParameter):
+                    shape = () if w.ndim == 0 else (1,)
+                else:
+                    # Preserve packed input columns and per-block scale columns.
+                    shape = (rows, *param.shape[1:])
+                self._check_weight_shape(name, w, shape)
+                param.weight_loader(param, w, shard_id)
+                loaded.add(pname)  # report the MODEL param name, not the artifact name
+                mapped = True
+                break
+            if mapped:
+                continue
+
+            param = params[name]
+            shape = (
+                (self.config.vocab_size, self.config.hidden_size)
+                if name in ("backbone.embeddings.weight", "lm_head.weight")
+                else tuple(param.shape)
+            )
+            if isinstance(param, PerTensorScaleParameter) and w.ndim == 0:
+                shape = ()
+            if name.endswith((".sqk", ".s_gate")):
+                shape = (self.config.num_attention_heads * self.config.head_dim,)
+            elif name.endswith((".suv_gate", ".suv_up")):
+                shape = (self.config.intermediate_size,)
+            elif name.endswith(".g_proj.weight"):
+                shape = (self.config.num_attention_heads * self.config.head_dim, self.config.hidden_size)
+            elif name.endswith(".o_proj.weight"):
+                shape = (self.config.hidden_size, self.config.num_attention_heads * self.config.head_dim)
+            elif name.endswith(".down_proj.weight"):
+                shape = (self.config.hidden_size, self.config.intermediate_size)
+            self._check_weight_shape(name, w, shape)
+            getattr(param, "weight_loader", default_weight_loader)(param, w)
+            loaded.add(name)
+        return loaded
+
+    @staticmethod
+    def _check_weight_shape(name, weight, expected):
+        # TP loaders may narrow oversized tensors; reject incompatible exports.
+        if tuple(weight.shape) != expected:
+            raise ValueError(
+                f"Dory weight {name}: expected shape {expected}, "
+                f"got {tuple(weight.shape)}"
+            )