Upload Monostich-2 SFT EMA (smoltalk+hermes 1 epoch, step 10619, 21.70h)

README.md

@@ -0,0 +1,175 @@
+---
+license: apache-2.0
+language:
+  - en
+tags:
+  - text-generation
+  - causal-lm
+  - pytorch
+  - sft
+  - instruction-tuned
+  - chat
+  - hybrid
+  - gated-deltanet
+  - gqa
+  - monostich
+pipeline_tag: text-generation
+library_name: tiny_gdn
+datasets:
+  - HuggingFaceTB/smoltalk
+  - NousResearch/Hermes-3-Dataset
+base_model: kerzgrr/Monostich-2-base
+model-index:
+  - name: Monostich-2
+    results: []
+---
+
+<div align="center">
+
+# Monostich-2
+
+### Instruction-tuned chat model (~150M) — Monostich-2 family
+
+[![Model](https://img.shields.io/badge/Model-~150M_params-blue)](.)
+[![Stage](https://img.shields.io/badge/Stage-SFT_(chat)-green.svg)](.)
+[![License](https://img.shields.io/badge/License-Apache_2.0-green.svg)](LICENSE)
+[![Base](https://img.shields.io/badge/Base-Monostich--2--base-orange.svg)](https://huggingface.co/kerzgrr/Monostich-2-base)
+
+*Second-generation Monostich-class model — hybrid GDN-2 + GQA, chat-tuned*
+
+</div>
+
+---
+
+## What this is
+
+**Monostich-2** is the **supervised fine-tuned (SFT) chat checkpoint** for the Monostich-2 family.
+
+- Base (pretrain): [`kerzgrr/Monostich-2-base`](https://huggingface.co/kerzgrr/Monostich-2-base)
+- Successor to [`kerzgrr/Monostich`](https://huggingface.co/kerzgrr/Monostich) (~100M LLaMA-style)
+- Architecture: hybrid **Gated DeltaNet-2** + **gated GQA** (not plain LLaMA)
+- This repo is **chat / instruction** (ChatML) — use the base repo for raw continuation
+
+---
+
+## Training
+
+### Pretrain → SFT
+
+| Stage | Details |
+|-------|---------|
+| **Base** | FineWeb-Edu pretrain → [`Monostich-2-base`](https://huggingface.co/kerzgrr/Monostich-2-base) |
+| **SFT mix** | [HuggingFaceTB/smoltalk](https://huggingface.co/datasets/HuggingFaceTB/smoltalk) + [NousResearch/Hermes-3-Dataset](https://huggingface.co/datasets/NousResearch/Hermes-3-Dataset) |
+| **Epochs** | 1 full epoch |
+| **Wall time** | **21.70 hours** |
+| **Final step** | optimizer step **10619** |
+| **Weights** | EMA (Hub `model.safetensors` is EMA @ bfloat16) |
+| **Seq length** | 8192 (packed SFT) |
+| **Peak LR** | 1 × 10⁻⁴ AdamW, cosine → 10% min |
+| **Final val loss (EMA)** | **1.532** (ppl ≈ 4.63) |
+
+### Chat template (ChatML)
+
+```
+<|begin_of_text|><|im_start|>system
+{system}<|im_end|>
+<|im_start|>user
+{user}<|im_end|>
+<|im_start|>assistant
+{assistant}<|im_end|>
+```
+
+Generation prompt ends at `<|im_start|>assistant\n`.
+
+---
+
+## Model Architecture
+
+Same TinyGDN hybrid as the base (~149.1M params):
+
+| | |
+|--|--|
+| **Layers** | 32 (GDN-2 ×3 + GQA every 4th) |
+| **Hidden** | 512 |
+| **MLP** | SwiGLU 1472 |
+| **Attention** | 4 Q / 1 KV, head dim 128, partial RoPE |
+| **Linear** | Gated DeltaNet-2, 4 heads × 128 |
+| **Vocab** | 49,152 BPE |
+
+---
+
+## Install & run
+
+```bash
+pip install torch safetensors tokenizers huggingface_hub
+hf download kerzgrr/Monostich-2 inference.py --local-dir .
+python inference.py --prompt "What is the capital of France?"
+```
+
+`inference.py` auto-downloads weights/tokenizer/`tiny_gdn/` and **auto-installs** pinned `flash-linear-attention` (Windows applies Hub patches). Git required on `PATH`.
+
+**Interactive chat:**
+
+```bash
+python inference.py
+```
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--prompt` | — | One-shot user message |
+| `--system` | — | Optional system prompt |
+| `--temperature` | `0.7` | Sampling temperature |
+| `--top-p` | `0.9` | Nucleus sampling |
+| `--top-k` | `50` | Top-k |
+| `--max-new-tokens` | `256` | Max generation length |
+| `--device` | `cuda` if available | `cuda` / `cpu` |
+
+---
+
+## Limitations
+
+- **~150M** research / edge chat model — not frontier quality
+- Can **hallucinate**, repeat, fail simple arithmetic, lose multi-turn facts
+- **Safety is incomplete** — do not deploy without your own filters
+- Requires `flash-linear-attention` (GDN-2); **not GGUF / llama.cpp** compatible today
+
+---
+
+## Model family
+
+| Model | Stage | Hub |
+|-------|-------|-----|
+| Monostich | SFT (~100M LLaMA) | [`kerzgrr/Monostich`](https://huggingface.co/kerzgrr/Monostich) |
+| Monostich-2-base | Pretrain (~150M hybrid) | [`kerzgrr/Monostich-2-base`](https://huggingface.co/kerzgrr/Monostich-2-base) |
+| **Monostich-2** | **SFT (~150M hybrid)** | **this repo** |
+
+---
+
+## Citation
+
+```bibtex
+@misc{monostich22026,
+  title={Monostich-2: A Hybrid GDN-2 + GQA Chat Model},
+  author={kerzgrr},
+  year={2026},
+  url={https://huggingface.co/kerzgrr/Monostich-2}
+}
+```
+
+---
+
+## Acknowledgments
+
+- [flash-linear-attention](https://github.com/fla-org/flash-linear-attention) (Gated DeltaNet-2)
+- [HuggingFaceTB/smoltalk](https://huggingface.co/datasets/HuggingFaceTB/smoltalk)
+- [NousResearch/Hermes-3-Dataset](https://huggingface.co/datasets/NousResearch/Hermes-3-Dataset)
+- Base: [`kerzgrr/Monostich-2-base`](https://huggingface.co/kerzgrr/Monostich-2-base)
+
+---
+
+<div align="center">
+
+*A monostich is a poem of a single line — small, but complete.*  
+*Monostich-2 renews the form.*
+
+</div>

chat_template.jinja

@@ -0,0 +1,17 @@
+{%- for message in messages -%}
+{%- if loop.first -%}{{- bos_token -}}{%- endif -%}
+{{- '<|im_start|>' + message['role'] + '\n' -}}
+{%- if message['role'] == 'tool' -%}{{- '<|tool_response|>\n' -}}{%- endif -%}
+{%- if message['content'] is string -%}
+{{- message['content'] -}}
+{%- elif message['content'] is iterable -%}
+{%- for item in message['content'] -%}
+{%- if item['type'] == 'text' -%}{{- item['text'] -}}{%- endif -%}
+{%- endfor -%}
+{%- endif -%}
+{%- if message['tool_calls'] is defined and message['tool_calls'] -%}
+{{- '\n<|tool_call|>\n' + (message['tool_calls'] | tojson) -}}
+{%- endif -%}
+{{- '<|im_end|>\n' -}}
+{%- endfor -%}
+{%- if add_generation_prompt -%}{{- '<|im_start|>assistant\n' -}}{%- endif -%}

config.json

@@ -0,0 +1,82 @@
+{
+  "allow_negative_eigenvalues": false,
+  "architecture": "TinyGDNForCausalLM",
+  "architectures": [
+    "TinyGDNForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "attention_head_dim": 128,
+  "base_model": "kerzgrr/Monostich-2-base",
+  "bos_token_id": 0,
+  "checkpoint_step": 10619,
+  "eos_token_id": 1,
+  "full_attention_interval": 4,
+  "hidden_size": 512,
+  "initializer_range": 0.02,
+  "intermediate_size": 1472,
+  "layer_types": [
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention"
+  ],
+  "linear_conv_kernel_dim": 4,
+  "linear_expand_v": 1.0,
+  "linear_head_dim": 128,
+  "linear_num_heads": 4,
+  "linear_num_value_heads": 4,
+  "max_position_embeddings": 32768,
+  "model_family": "Monostich-2",
+  "model_name": "Monostich-2",
+  "model_type": "tiny_gdn",
+  "mtp_adapter_rank": 128,
+  "mtp_loss_weight": 0.0,
+  "mtp_num_heads": 0,
+  "num_attention_heads": 4,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 1,
+  "pad_token_id": 2,
+  "partial_rotary_factor": 0.5,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000.0,
+  "sft_hours": 21.7,
+  "sft_val_loss_ema": 1.5321617420986058,
+  "sft_val_ppl_ema": 4.628170928002891,
+  "shared_layer_indices": [],
+  "stage": "sft",
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "training_sequence_length": 8192,
+  "transformers_version": "4.45.0",
+  "unk_token_id": 3,
+  "vocab_size": 49152,
+  "weights": "ema"
+}

inference.py

@@ -0,0 +1,473 @@
+#!/usr/bin/env python3
+"""Standalone chat inference for kerzgrr/Monostich-2 (SFT).
+
+Downloads model assets from the Hub (cached after first run), auto-installs
+flash-linear-attention when needed, and streams ChatML assistant replies.
+
+Examples:
+  python inference.py --prompt "What is the capital of France?"
+  python inference.py
+  python inference.py --prompt "Write a haiku about GPUs" --temperature 0.7
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import platform
+import shutil
+import subprocess
+import sys
+import time
+import warnings
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+from tokenizers import Tokenizer
+
+
+def _silence_runtime_warnings() -> None:
+    patterns = (
+        r"tl\.make_block_ptr is deprecated",
+        r"Memory efficient kernel not used because",
+        r"Memory Efficient attention has been runtime disabled",
+        r"Flash attention kernel not used because",
+        r"Torch was not compiled with flash attention",
+        r"cuDNN attention kernel not used because",
+        r"cuDNN attention has been runtime disabled",
+    )
+    for pattern in patterns:
+        warnings.filterwarnings("ignore", message=pattern)
+
+
+REPO_ID = "kerzgrr/Monostich-2"
+FLA_COMMIT = "cbb0a72efb55c18ca0ef4f298298317573ad2cb3"
+FLA_REPO = "https://github.com/fla-org/flash-linear-attention.git"
+PATCH_FILES = (
+    "fla/__init__.py",
+    "fla/ops/__init__.py",
+    "fla/layers/__init__.py",
+    "fla/ops/simple_gla/__init__.py",
+)
+
+
+def _download(filename: str, local_dir: Path | None) -> Path:
+    from huggingface_hub import hf_hub_download
+
+    return Path(
+        hf_hub_download(
+            repo_id=REPO_ID,
+            filename=filename,
+            local_dir=str(local_dir) if local_dir else None,
+        )
+    )
+
+
+def _run(cmd: list[str], *, cwd: Path | None = None, env: dict | None = None) -> None:
+    print("+", " ".join(cmd), flush=True)
+    merged = os.environ.copy()
+    if env:
+        merged.update(env)
+    merged.setdefault("PYTHONUTF8", "1")
+    merged.setdefault("PYTHONIOENCODING", "utf-8")
+    subprocess.check_call(cmd, cwd=str(cwd) if cwd else None, env=merged)
+
+
+def _pip_install(*args: str) -> None:
+    _run([sys.executable, "-m", "pip", "install", *args])
+
+
+def _fla_importable() -> tuple[bool, str]:
+    try:
+        from fla.layers.gdn2 import GatedDeltaNet2  # noqa: F401
+    except Exception as error:  # noqa: BLE001
+        return False, str(error)
+    return True, ""
+
+
+def _cache_root() -> Path:
+    override = os.environ.get("MONOSTICH_CACHE")
+    if override:
+        path = Path(override).expanduser().resolve()
+    else:
+        path = Path.home() / ".cache" / "monostich-2"
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def _ensure_git() -> None:
+    if shutil.which("git") is None:
+        raise RuntimeError(
+            "git is required to auto-install flash-linear-attention. "
+            "Install Git and ensure it is on PATH."
+        )
+
+
+def _apply_windows_fla_patches(fla_root: Path, local_dir: Path | None) -> None:
+    print("Applying Windows FLA import patches from the Hub …", flush=True)
+    for relative in PATCH_FILES:
+        source = _download(f"windows_fla_patches/{relative}", local_dir)
+        target = fla_root / relative
+        target.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(source, target)
+        print(f"  patched {relative}", flush=True)
+
+
+def _install_fla(local_dir: Path | None) -> None:
+    print("flash-linear-attention missing/broken — installing automatically …", flush=True)
+    _pip_install("einops", "numpy")
+    if platform.system() != "Windows":
+        _pip_install("--no-deps", f"git+{FLA_REPO}@{FLA_COMMIT}")
+        return
+
+    _ensure_git()
+    fla_root = _cache_root() / "flash-linear-attention"
+    if (fla_root / ".git").is_dir():
+        _run(["git", "fetch", "--depth", "1", "origin", FLA_COMMIT], cwd=fla_root)
+        _run(["git", "checkout", "--force", FLA_COMMIT], cwd=fla_root)
+    else:
+        if fla_root.exists():
+            shutil.rmtree(fla_root)
+        _run(["git", "clone", "--filter=blob:none", FLA_REPO, str(fla_root)])
+        _run(["git", "fetch", "--depth", "1", "origin", FLA_COMMIT], cwd=fla_root)
+        _run(["git", "checkout", "--force", FLA_COMMIT], cwd=fla_root)
+
+    _apply_windows_fla_patches(fla_root, local_dir)
+    _pip_install("--no-build-isolation", "--no-deps", "-e", str(fla_root))
+
+
+def _ensure_fla(local_dir: Path | None) -> None:
+    ok, error = _fla_importable()
+    if ok:
+        return
+    print(f"FLA not ready ({error})", flush=True)
+    try:
+        _install_fla(local_dir)
+    except Exception as install_error:  # noqa: BLE001
+        raise RuntimeError(
+            "Automatic flash-linear-attention install failed.\n"
+            f"Original import error: {error}\n"
+            f"Install error: {install_error}"
+        ) from install_error
+
+    for name in list(sys.modules):
+        if name == "fla" or name.startswith("fla."):
+            del sys.modules[name]
+
+    ok, error = _fla_importable()
+    if not ok:
+        raise RuntimeError(
+            "flash-linear-attention installed but still failed to import "
+            f"GatedDeltaNet2: {error}"
+        )
+    print("flash-linear-attention ready.", flush=True)
+
+
+def _ensure_tiny_gdn(local_dir: Path | None) -> Path:
+    here = Path(__file__).resolve().parent
+    if (here / "tiny_gdn" / "__init__.py").is_file():
+        return here
+    if local_dir and (local_dir / "tiny_gdn" / "__init__.py").is_file():
+        return local_dir
+    for name in ("tiny_gdn/__init__.py", "tiny_gdn/config.py", "tiny_gdn/model.py"):
+        _download(name, local_dir)
+    return _download("tiny_gdn/__init__.py", local_dir).parent.parent
+
+
+def _sample(
+    logits: torch.Tensor,
+    *,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    generator: torch.Generator,
+) -> int:
+    logits = logits.float()
+    if temperature <= 1e-5:
+        return int(torch.argmax(logits).item())
+    logits = logits / temperature
+    if 0 < top_k < logits.shape[-1]:
+        threshold = torch.topk(logits, top_k).values[-1]
+        logits = logits.masked_fill(logits < threshold, -torch.inf)
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        probs = torch.softmax(sorted_logits, dim=-1)
+        remove = torch.cumsum(probs, dim=-1) > top_p
+        remove[1:] = remove[:-1].clone()
+        remove[0] = False
+        sorted_logits = sorted_logits.masked_fill(remove, -torch.inf)
+        logits = torch.full_like(logits, -torch.inf)
+        logits.scatter_(0, sorted_indices, sorted_logits)
+    probs = torch.softmax(logits, dim=-1)
+    return int(torch.multinomial(probs, 1, generator=generator).item())
+
+
+def _apply_repetition_penalty(
+    logits: torch.Tensor,
+    token_ids: list[int],
+    penalty: float,
+    window: int,
+) -> torch.Tensor:
+    if penalty == 1.0 or not token_ids:
+        return logits
+    recent = token_ids[-window:] if window > 0 else token_ids
+    unique = torch.tensor(list(set(recent)), dtype=torch.long, device=logits.device)
+    score = logits[unique]
+    logits[unique] = torch.where(score > 0, score / penalty, score * penalty)
+    return logits
+
+
+def encode_chat(tokenizer: Tokenizer, messages: list[dict[str, str]]) -> list[int]:
+    bos = tokenizer.token_to_id("<|begin_of_text|>")
+    im_start = tokenizer.token_to_id("<|im_start|>")
+    im_end = tokenizer.token_to_id("<|im_end|>")
+    if bos is None or im_start is None or im_end is None:
+        raise RuntimeError("Tokenizer missing ChatML specials")
+    ids = [bos]
+    newline = tokenizer.encode("\n", add_special_tokens=False).ids
+    for message in messages:
+        role = message["role"]
+        content = message["content"]
+        ids.append(im_start)
+        ids.extend(tokenizer.encode(f"{role}\n", add_special_tokens=False).ids)
+        ids.extend(tokenizer.encode(content, add_special_tokens=False).ids)
+        ids.append(im_end)
+        ids.extend(newline)
+    ids.append(im_start)
+    ids.extend(tokenizer.encode("assistant\n", add_special_tokens=False).ids)
+    return ids
+
+
+@torch.inference_mode()
+def generate(
+    model,
+    tokenizer: Tokenizer,
+    prompt_ids: list[int],
+    *,
+    max_new_tokens: int,
+    context_length: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    repetition_penalty: float,
+    repetition_window: int,
+    seed: int,
+    stream: bool,
+    device: torch.device,
+) -> tuple[str, int, str]:
+    eos_id = int(model.config.eos_token_id)
+    im_end = tokenizer.token_to_id("<|im_end|>")
+    stop = {eos_id}
+    if im_end is not None:
+        stop.add(im_end)
+
+    token_ids = list(prompt_ids[-context_length:])
+    generated: list[int] = []
+    decoded = ""
+    stop_reason = "max_new_tokens"
+    generator = torch.Generator(device=device)
+    generator.manual_seed(seed)
+    started = time.perf_counter()
+
+    for _ in range(max_new_tokens):
+        context = token_ids[-context_length:]
+        input_ids = torch.tensor([context], dtype=torch.long, device=device)
+        output = model(input_ids, return_logits=True, logits_to_keep=1)
+        if output.logits is None:
+            raise RuntimeError("Model returned no logits")
+        next_logits = _apply_repetition_penalty(
+            output.logits[0, -1],
+            token_ids,
+            repetition_penalty,
+            repetition_window,
+        )
+        next_id = _sample(
+            next_logits,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            generator=generator,
+        )
+        if next_id in stop:
+            stop_reason = "stop"
+            break
+        token_ids.append(next_id)
+        generated.append(next_id)
+        current = tokenizer.decode(generated, skip_special_tokens=True)
+        delta = current[len(decoded) :] if current.startswith(decoded) else current
+        decoded = current
+        if stream and delta:
+            print(delta, end="", flush=True)
+
+    if stream:
+        print(flush=True)
+    elapsed = time.perf_counter() - started
+    tps = len(generated) / max(elapsed, 1e-9)
+    if stream:
+        print(
+            f"[done] tokens={len(generated)} stop={stop_reason} {tps:.1f} tok/s",
+            file=sys.stderr,
+            flush=True,
+        )
+    return decoded, len(generated), stop_reason
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Monostich-2 SFT chat inference")
+    parser.add_argument("--prompt", default=None, help="Single user prompt")
+    parser.add_argument("--system", default="", help="Optional system prompt")
+    parser.add_argument("--max-new-tokens", type=int, default=256)
+    parser.add_argument("--temperature", type=float, default=0.7)
+    parser.add_argument("--top-p", type=float, default=0.9)
+    parser.add_argument("--top-k", type=int, default=50)
+    parser.add_argument("--repetition-penalty", type=float, default=1.08)
+    parser.add_argument("--repetition-window", type=int, default=256)
+    parser.add_argument("--context-length", type=int, default=2048)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--device",
+        default="cuda" if torch.cuda.is_available() else "cpu",
+        choices=["cuda", "cpu"],
+    )
+    parser.add_argument("--no-stream", action="store_true")
+    parser.add_argument("--local-dir", default=None)
+    parser.add_argument("--repo-id", default=REPO_ID)
+    return parser.parse_args()
+
+
+def main() -> int:
+    _silence_runtime_warnings()
+    args = parse_args()
+    global REPO_ID
+    REPO_ID = args.repo_id
+    local_dir = Path(args.local_dir).resolve() if args.local_dir else None
+
+    print(f"Loading Monostich-2 from huggingface.co/{REPO_ID} …", flush=True)
+    try:
+        package_root = _ensure_tiny_gdn(local_dir)
+    except Exception as error:  # noqa: BLE001
+        print(f"Failed to resolve tiny_gdn package: {error}", file=sys.stderr)
+        return 1
+
+    if str(package_root) not in sys.path:
+        sys.path.insert(0, str(package_root))
+
+    try:
+        _ensure_fla(local_dir)
+    except Exception as error:  # noqa: BLE001
+        print(str(error), file=sys.stderr)
+        return 1
+
+    try:
+        from tiny_gdn import TinyGDNConfig, TinyGDNForCausalLM
+    except ImportError as error:
+        print(f"Could not import tiny_gdn: {error}", file=sys.stderr)
+        return 1
+
+    weights_path = _download("model.safetensors", local_dir)
+    tok_path = _download("tokenizer.json", local_dir)
+    cfg_path = _download("config.json", local_dir)
+
+    raw = json.loads(cfg_path.read_text(encoding="utf-8"))
+    from dataclasses import fields
+
+    allowed = {item.name for item in fields(TinyGDNConfig)}
+    payload = {key: value for key, value in raw.items() if key in allowed}
+    if "shared_layer_indices" in payload:
+        payload["shared_layer_indices"] = tuple(payload["shared_layer_indices"])
+    config = TinyGDNConfig(**payload)
+
+    device = torch.device(args.device)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        print("CUDA requested but unavailable; falling back to CPU.", flush=True)
+        device = torch.device("cpu")
+    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
+
+    print(
+        f"Building TinyGDN ({config.num_hidden_layers}L / {config.hidden_size}d) "
+        f"on {device} …",
+        flush=True,
+    )
+    model = TinyGDNForCausalLM(config)
+    state = load_file(str(weights_path), device="cpu")
+    model.load_state_dict(state, strict=True)
+    del state
+    model = model.to(device=device, dtype=dtype)
+    model.eval()
+    model.requires_grad_(False)
+
+    tokenizer = Tokenizer.from_file(str(tok_path))
+    context_length = min(args.context_length, config.max_position_embeddings)
+    stream = not args.no_stream
+
+    def run_chat(messages: list[dict[str, str]]) -> str:
+        prompt_ids = encode_chat(tokenizer, messages)
+        text, _, _ = generate(
+            model,
+            tokenizer,
+            prompt_ids,
+            max_new_tokens=args.max_new_tokens,
+            context_length=context_length,
+            temperature=args.temperature,
+            top_p=args.top_p,
+            top_k=args.top_k,
+            repetition_penalty=args.repetition_penalty,
+            repetition_window=args.repetition_window,
+            seed=args.seed,
+            stream=stream,
+            device=device,
+        )
+        return text
+
+    if args.prompt is not None:
+        messages: list[dict[str, str]] = []
+        if args.system.strip():
+            messages.append({"role": "system", "content": args.system.strip()})
+        messages.append({"role": "user", "content": args.prompt})
+        if stream:
+            print("assistant> ", end="", flush=True)
+        text = run_chat(messages)
+        if not stream:
+            print(text)
+        return 0
+
+    print(
+        "Interactive chat. Commands: /exit  /quit  /reset\n"
+        "This is the SFT chat model (ChatML).",
+        flush=True,
+    )
+    history: list[dict[str, str]] = []
+    if args.system.strip():
+        history.append({"role": "system", "content": args.system.strip()})
+    while True:
+        try:
+            user_input = input("user> ")
+        except (EOFError, KeyboardInterrupt):
+            print()
+            break
+        text = user_input.strip()
+        if not text:
+            continue
+        if text.lower() in {"/exit", "/quit"}:
+            break
+        if text.lower() == "/reset":
+            history = []
+            if args.system.strip():
+                history.append({"role": "system", "content": args.system.strip()})
+            print("(history cleared)", flush=True)
+            continue
+        turn = history + [{"role": "user", "content": text}]
+        if stream:
+            print("assistant> ", end="", flush=True)
+        reply = run_chat(turn)
+        history = turn + [{"role": "assistant", "content": reply}]
+        if not stream:
+            print(f"assistant> {reply}")
+        print(flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

model.safetensors

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

requirements.txt

@@ -0,0 +1,8 @@
+torch>=2.4.0
+safetensors>=0.4.0
+tokenizers>=0.20.0
+huggingface_hub>=0.26.0
+einops>=0.8.0
+numpy>=1.26.0
+
+# flash-linear-attention is auto-installed by inference.py on first run.

sft_config.json

@@ -0,0 +1,37 @@
+{
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.95,
+  "adam_epsilon": 1e-08,
+  "cpu_prefetch_factor": 2,
+  "cpu_workers": 2,
+  "dataset_manifest": "data/tokenized/smoltalk-hermes-sft/manifest.json",
+  "ema_inv_gamma": 1.0,
+  "ema_max_decay": 0.9999,
+  "ema_power": 0.75,
+  "enable_gradient_checkpointing": true,
+  "full_shuffle": true,
+  "gradient_accumulation_steps": 16,
+  "gradient_clip_norm": 1.0,
+  "initial_run_dir": "runs/tiny-gdn-150m-smoltalk-sft",
+  "initial_weights": "ema",
+  "learning_rate": 0.0001,
+  "log_every_optimizer_steps": 1,
+  "logit_z_loss_coefficient": 0.0001,
+  "maximum_checkpoints": 5,
+  "maximum_optimizer_steps": null,
+  "micro_batch_size": 1,
+  "minimum_learning_rate_ratio": 0.1,
+  "muon_learning_rate": 0.01,
+  "muon_momentum": 0.95,
+  "optimizer": "adamw",
+  "output_dir": "runs/tiny-gdn-150m-smoltalk-hermes-sft",
+  "require_complete_pretraining": false,
+  "require_complete_sft_source": false,
+  "save_every_optimizer_steps": 100,
+  "seed": 20260720,
+  "sequence_length": 8192,
+  "shuffle_block_sequences": 512,
+  "validation_batches": 16,
+  "warmup_ratio": 0.01,
+  "weight_decay": 0.0
+}

special_token_ids.json

@@ -0,0 +1,12 @@
+{
+  "bos_token": "<|begin_of_text|>",
+  "eos_token": "<|end_of_text|>",
+  "pad_token": "<|padding|>",
+  "unk_token": "<|unknown|>",
+  "im_start": "<|im_start|>",
+  "im_end": "<|im_end|>",
+  "bos_token_id": 0,
+  "eos_token_id": 1,
+  "pad_token_id": 2,
+  "unk_token_id": 3
+}

special_tokens_map.json

@@ -0,0 +1,68 @@
+{
+  "bos_token": "<|begin_of_text|>",
+  "eos_token": "<|end_of_text|>",
+  "pad_token": "<|padding|>",
+  "unk_token": "<|unknown|>",
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|tool_call|>",
+    "<|tool_response|>",
+    "<|think_start|>",
+    "<|think_end|>",
+    "<|fim_prefix|>",
+    "<|fim_middle|>",
+    "<|fim_suffix|>",
+    "<|fim_pad|>",
+    "<|reserved_000|>",
+    "<|reserved_001|>",
+    "<|reserved_002|>",
+    "<|reserved_003|>",
+    "<|reserved_004|>",
+    "<|reserved_005|>",
+    "<|reserved_006|>",
+    "<|reserved_007|>",
+    "<|reserved_008|>",
+    "<|reserved_009|>",
+    "<|reserved_010|>",
+    "<|reserved_011|>",
+    "<|reserved_012|>",
+    "<|reserved_013|>",
+    "<|reserved_014|>",
+    "<|reserved_015|>",
+    "<|reserved_016|>",
+    "<|reserved_017|>",
+    "<|reserved_018|>",
+    "<|reserved_019|>",
+    "<|reserved_020|>",
+    "<|reserved_021|>",
+    "<|reserved_022|>",
+    "<|reserved_023|>",
+    "<|reserved_024|>",
+    "<|reserved_025|>",
+    "<|reserved_026|>",
+    "<|reserved_027|>",
+    "<|reserved_028|>",
+    "<|reserved_029|>",
+    "<|reserved_030|>",
+    "<|reserved_031|>",
+    "<|reserved_032|>",
+    "<|reserved_033|>",
+    "<|reserved_034|>",
+    "<|reserved_035|>",
+    "<|reserved_036|>",
+    "<|reserved_037|>",
+    "<|reserved_038|>",
+    "<|reserved_039|>",
+    "<|reserved_040|>",
+    "<|reserved_041|>",
+    "<|reserved_042|>",
+    "<|reserved_043|>",
+    "<|reserved_044|>",
+    "<|reserved_045|>",
+    "<|reserved_046|>",
+    "<|reserved_047|>",
+    "<|reserved_048|>",
+    "<|reserved_049|>"
+  ]
+}

tiny_gdn/__init__.py

@@ -0,0 +1,8 @@
+from tiny_gdn.config import TinyGDNConfig
+from tiny_gdn.model import TinyGDNForCausalLM, TinyGDNOutput
+
+__all__ = [
+    "TinyGDNConfig",
+    "TinyGDNForCausalLM",
+    "TinyGDNOutput",
+]

tiny_gdn/config.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any
+
+
+@dataclass(frozen=True)
+class TinyGDNConfig:
+    architecture: str = "TinyGDNForCausalLM"
+    model_type: str = "tiny_gdn"
+
+    vocab_size: int = 49_152
+    # Deep-thin sizing is deliberate: controlled sub-billion studies find
+    # depth materially more valuable than width around the 125M-150M scale.
+    hidden_size: int = 512
+    intermediate_size: int = 1_472
+    num_hidden_layers: int = 32
+
+    num_attention_heads: int = 4
+    num_key_value_heads: int = 1
+    attention_head_dim: int = 128
+    full_attention_interval: int = 4
+    attention_dropout: float = 0.0
+    partial_rotary_factor: float = 0.5
+    rope_theta: float = 1_000_000.0
+
+    linear_num_heads: int = 4
+    linear_num_value_heads: int = 4
+    linear_head_dim: int = 128
+    linear_expand_v: float = 1.0
+    linear_conv_kernel_dim: int = 4
+    allow_negative_eigenvalues: bool = False
+
+    max_position_embeddings: int = 32_768
+    training_sequence_length: int = 2_048
+    rms_norm_eps: float = 1e-6
+    initializer_range: float = 0.02
+    tie_word_embeddings: bool = True
+    shared_layer_indices: tuple[int, ...] = ()
+
+    # MTP is an opt-in ablation at this scale; static MTP is not assumed to
+    # improve a 150M model without a controlled pilot.
+    mtp_num_heads: int = 0
+    mtp_adapter_rank: int = 128
+    mtp_loss_weight: float = 0.0
+
+    bos_token_id: int = 0
+    eos_token_id: int = 1
+    pad_token_id: int = 2
+    unk_token_id: int = 3
+
+    def __post_init__(self) -> None:
+        if self.vocab_size <= 0 or self.vocab_size > 65_536:
+            raise ValueError("vocab_size must fit the uint16 token dataset")
+        if self.hidden_size != self.num_attention_heads * self.attention_head_dim:
+            raise ValueError("hidden_size must equal num_attention_heads * attention_head_dim")
+        if self.hidden_size != self.linear_num_heads * self.linear_head_dim:
+            raise ValueError("hidden_size must equal linear_num_heads * linear_head_dim")
+        if self.linear_num_value_heads < self.linear_num_heads:
+            raise ValueError("linear_num_value_heads must be at least linear_num_heads")
+        if self.linear_num_value_heads % self.linear_num_heads != 0:
+            raise ValueError("linear_num_value_heads must be divisible by linear_num_heads")
+        if self.num_attention_heads % self.num_key_value_heads != 0:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        if self.num_hidden_layers % self.full_attention_interval != 0:
+            raise ValueError("num_hidden_layers must be divisible by full_attention_interval")
+        if not 0.0 < self.partial_rotary_factor <= 1.0:
+            raise ValueError("partial_rotary_factor must be in (0, 1]")
+        rotary_dim = int(self.attention_head_dim * self.partial_rotary_factor)
+        if rotary_dim <= 0 or rotary_dim % 2:
+            raise ValueError("The partial rotary dimension must be positive and even")
+        if self.training_sequence_length > self.max_position_embeddings:
+            raise ValueError("training_sequence_length exceeds max_position_embeddings")
+        if len(set(self.shared_layer_indices)) != len(self.shared_layer_indices):
+            raise ValueError("shared_layer_indices must be unique")
+        if any(
+            index < 0 or index >= self.num_hidden_layers
+            for index in self.shared_layer_indices
+        ):
+            raise ValueError("shared_layer_indices contains an invalid layer")
+        if self.mtp_num_heads < 0:
+            raise ValueError("mtp_num_heads cannot be negative")
+        if self.mtp_num_heads and self.mtp_adapter_rank <= 0:
+            raise ValueError("mtp_adapter_rank must be positive when MTP is enabled")
+        if not 0.0 <= self.mtp_loss_weight <= 1.0:
+            raise ValueError("mtp_loss_weight must be between zero and one")
+        for token_id in (
+            self.bos_token_id,
+            self.eos_token_id,
+            self.pad_token_id,
+            self.unk_token_id,
+        ):
+            if not 0 <= token_id < self.vocab_size:
+                raise ValueError(f"Special token ID {token_id} is outside the vocabulary")
+
+    @property
+    def layer_types(self) -> tuple[str, ...]:
+        return tuple(
+            "full_attention" if (index + 1) % self.full_attention_interval == 0 else "gdn2"
+            for index in range(self.num_hidden_layers)
+        )
+
+    @property
+    def rotary_dim(self) -> int:
+        return int(self.attention_head_dim * self.partial_rotary_factor)
+
+    @property
+    def effective_num_layers(self) -> int:
+        return self.num_hidden_layers + len(self.shared_layer_indices)
+
+    def to_dict(self) -> dict[str, Any]:
+        payload = asdict(self)
+        payload["layer_types"] = list(self.layer_types)
+        return payload
+
+    def save_json(self, path: Path) -> None:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(
+            json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+
+    @classmethod
+    def from_json(cls, path: Path) -> TinyGDNConfig:
+        payload = json.loads(path.read_text(encoding="utf-8"))
+        payload.pop("layer_types", None)
+        if "shared_layer_indices" in payload:
+            payload["shared_layer_indices"] = tuple(payload["shared_layer_indices"])
+        return cls(**payload)

tiny_gdn/model.py

@@ -0,0 +1,587 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from safetensors.torch import load_model, save_model
+from torch import nn
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from torch.utils.checkpoint import checkpoint
+
+from tiny_gdn.config import TinyGDNConfig
+
+try:
+    # Import the module directly — `from fla.layers import GatedDeltaNet2`
+    # executes layers/__init__.py and eagerly loads every attention kernel.
+    from fla.layers.gdn2 import GatedDeltaNet2
+except ImportError as import_error:
+    GatedDeltaNet2 = None
+    FLA_IMPORT_ERROR: ImportError | None = import_error
+else:
+    FLA_IMPORT_ERROR = None
+
+
+@dataclass
+class TinyGDNOutput:
+    loss: torch.Tensor | None
+    logits: torch.Tensor | None
+    main_loss: torch.Tensor | None
+    mtp_loss: torch.Tensor | None
+    z_loss: torch.Tensor | None
+    hidden_states: torch.Tensor | None = None
+
+
+class RMSNorm(nn.Module):
+    """Zero-centered RMSNorm as used by Qwen3-Next."""
+
+    def __init__(self, hidden_size: int, eps: float) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(hidden_size))
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        normalized = hidden_states.float()
+        normalized = normalized * torch.rsqrt(normalized.square().mean(dim=-1, keepdim=True) + self.eps)
+        normalized = normalized * (1.0 + self.weight.float())
+        return normalized.to(dtype=input_dtype)
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, rotary_dim: int, rope_theta: float) -> None:
+        super().__init__()
+        inverse_frequency = 1.0 / (
+            rope_theta
+            ** (
+                torch.arange(0, rotary_dim, 2, dtype=torch.float32)
+                / rotary_dim
+            )
+        )
+        self.rotary_dim = rotary_dim
+        self.register_buffer("inverse_frequency", inverse_frequency, persistent=False)
+
+    def forward(
+        self,
+        sequence_length: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        position_offset: int = 0,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        positions = torch.arange(
+            position_offset,
+            position_offset + sequence_length,
+            device=device,
+            dtype=torch.float32,
+        )
+        frequencies = torch.outer(positions, self.inverse_frequency.float())
+        embeddings = torch.cat((frequencies, frequencies), dim=-1)
+        return embeddings.cos().to(dtype=dtype), embeddings.sin().to(dtype=dtype)
+
+
+def rotate_half(hidden_states: torch.Tensor) -> torch.Tensor:
+    first, second = hidden_states.chunk(2, dim=-1)
+    return torch.cat((-second, first), dim=-1)
+
+
+def apply_rotary_embedding(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    cosine: torch.Tensor,
+    sine: torch.Tensor,
+    rotary_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    cosine = cosine[None, None, :, :]
+    sine = sine[None, None, :, :]
+    query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:]
+    key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:]
+    query_rotary = query_rotary * cosine + rotate_half(query_rotary) * sine
+    key_rotary = key_rotary * cosine + rotate_half(key_rotary) * sine
+    return (
+        torch.cat((query_rotary, query_pass), dim=-1),
+        torch.cat((key_rotary, key_pass), dim=-1),
+    )
+
+
+class GatedGroupedQueryAttention(nn.Module):
+    """QK-normalized, partially rotary GQA with a learned sigmoid output gate."""
+
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.head_dim = config.attention_head_dim
+        self.rotary_dim = config.rotary_dim
+        self.dropout = config.attention_dropout
+
+        query_size = self.num_heads * self.head_dim
+        key_value_size = self.num_key_value_heads * self.head_dim
+        self.q_gate_proj = nn.Linear(config.hidden_size, query_size * 2, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, key_value_size, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, key_value_size, bias=False)
+        self.o_proj = nn.Linear(query_size, config.hidden_size, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+        self.rotary = RotaryEmbedding(self.rotary_dim, config.rope_theta)
+
+    def _attention_mask(
+        self,
+        attention_mask: torch.Tensor | None,
+        sequence_length: int,
+        device: torch.device,
+    ) -> torch.Tensor | None:
+        if attention_mask is None:
+            return None
+        if attention_mask.ndim != 2:
+            raise ValueError("attention_mask must have shape [batch, sequence]")
+        if attention_mask.shape[1] != sequence_length:
+            raise ValueError("attention_mask sequence length does not match input")
+
+        causal = torch.ones(
+            sequence_length,
+            sequence_length,
+            dtype=torch.bool,
+            device=device,
+        ).tril()
+        valid_keys = attention_mask[:, None, None, :].to(dtype=torch.bool, device=device)
+        return causal[None, None, :, :] & valid_keys
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        batch_size, sequence_length, _ = hidden_states.shape
+        query_and_gate = self.q_gate_proj(hidden_states)
+        query, output_gate = query_and_gate.chunk(2, dim=-1)
+
+        query = query.view(batch_size, sequence_length, self.num_heads, self.head_dim)
+        key = self.k_proj(hidden_states).view(
+            batch_size,
+            sequence_length,
+            self.num_key_value_heads,
+            self.head_dim,
+        )
+        value = self.v_proj(hidden_states).view(
+            batch_size,
+            sequence_length,
+            self.num_key_value_heads,
+            self.head_dim,
+        )
+
+        query = self.q_norm(query).transpose(1, 2)
+        key = self.k_norm(key).transpose(1, 2)
+        value = value.transpose(1, 2)
+
+        cosine, sine = self.rotary(
+            sequence_length,
+            device=hidden_states.device,
+            dtype=query.dtype,
+        )
+        query, key = apply_rotary_embedding(
+            query,
+            key,
+            cosine,
+            sine,
+            rotary_dim=self.rotary_dim,
+        )
+
+        sdpa_mask = self._attention_mask(
+            attention_mask,
+            sequence_length,
+            hidden_states.device,
+        )
+        sdpa_options = {
+            "attn_mask": sdpa_mask,
+            "dropout_p": self.dropout if self.training else 0.0,
+            "is_causal": sdpa_mask is None,
+            "enable_gqa": True,
+        }
+        # Prefer Flash / mem-efficient when available; fall back to MATH for
+        # Windows PyTorch builds that ship without FlashAttention kernels.
+        backends = (
+            [
+                SDPBackend.FLASH_ATTENTION,
+                SDPBackend.EFFICIENT_ATTENTION,
+                SDPBackend.CUDNN_ATTENTION,
+                SDPBackend.MATH,
+            ]
+            if query.is_cuda
+            else [SDPBackend.MATH]
+        )
+        with sdpa_kernel(backends):
+            attention_output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                **sdpa_options,
+            )
+        attention_output = attention_output.transpose(1, 2).reshape(
+            batch_size,
+            sequence_length,
+            -1,
+        )
+        attention_output = attention_output * torch.sigmoid(output_gate)
+        return self.o_proj(attention_output)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.gate_up_proj = nn.Linear(
+            config.hidden_size,
+            config.intermediate_size * 2,
+            bias=False,
+        )
+        self.down_proj = nn.Linear(
+            config.intermediate_size,
+            config.hidden_size,
+            bias=False,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        gate, up = self.gate_up_proj(hidden_states).chunk(2, dim=-1)
+        return self.down_proj(F.silu(gate) * up)
+
+
+class TinyGDNBlock(nn.Module):
+    def __init__(self, config: TinyGDNConfig, layer_index: int) -> None:
+        super().__init__()
+        layer_type = config.layer_types[layer_index]
+        self.layer_type = layer_type
+        self.token_mixer_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+
+        if layer_type == "gdn2":
+            if GatedDeltaNet2 is None:
+                raise ImportError(
+                    "Gated DeltaNet-2 requires the pinned flash-linear-attention dependency"
+                ) from FLA_IMPORT_ERROR
+            self.token_mixer = GatedDeltaNet2(
+                hidden_size=config.hidden_size,
+                expand_v=config.linear_expand_v,
+                head_dim=config.linear_head_dim,
+                num_heads=config.linear_num_heads,
+                num_v_heads=config.linear_num_value_heads,
+                mode="chunk",
+                use_short_conv=True,
+                allow_neg_eigval=config.allow_negative_eigenvalues,
+                conv_size=config.linear_conv_kernel_dim,
+                conv_bias=False,
+                layer_idx=layer_index,
+                norm_eps=config.rms_norm_eps,
+            )
+        elif layer_type == "full_attention":
+            self.token_mixer = GatedGroupedQueryAttention(config)
+        else:
+            raise ValueError(f"Unsupported layer type: {layer_type}")
+
+        self.mlp = SwiGLU(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        normalized = self.token_mixer_norm(hidden_states)
+        if self.layer_type == "gdn2":
+            mixed, _, _ = self.token_mixer(
+                normalized,
+                attention_mask=attention_mask,
+                use_cache=False,
+            )
+        else:
+            mixed = self.token_mixer(normalized, attention_mask=attention_mask)
+        hidden_states = residual + mixed
+        hidden_states = hidden_states + self.mlp(self.mlp_norm(hidden_states))
+        return hidden_states
+
+
+class MultiTokenPredictionAdapter(nn.Module):
+    """A lightweight residual adapter for one additional prediction horizon."""
+
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.down_proj = nn.Linear(
+            config.hidden_size,
+            config.mtp_adapter_rank,
+            bias=False,
+        )
+        self.up_proj = nn.Linear(
+            config.mtp_adapter_rank,
+            config.hidden_size,
+            bias=False,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        adapted = self.up_proj(F.silu(self.down_proj(self.norm(hidden_states))))
+        return hidden_states + adapted
+
+
+class TinyGDNForCausalLM(nn.Module):
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            TinyGDNBlock(config, layer_index)
+            for layer_index in range(config.num_hidden_layers)
+        )
+        self.final_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.mtp_adapters = nn.ModuleList(
+            MultiTokenPredictionAdapter(config)
+            for _ in range(config.mtp_num_heads)
+        )
+        self.gradient_checkpointing = False
+
+        self.apply(self._initialize_module)
+        self._initialize_residual_projections()
+
+    def _initialize_module(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(
+                module.weight,
+                mean=0.0,
+                std=self.config.initializer_range,
+            )
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(
+                module.weight,
+                mean=0.0,
+                std=self.config.initializer_range,
+            )
+
+    def _initialize_residual_projections(self) -> None:
+        residual_std = self.config.initializer_range / math.sqrt(
+            2 * self.config.num_hidden_layers
+        )
+        for layer in self.layers:
+            nn.init.normal_(
+                layer.token_mixer.o_proj.weight,
+                mean=0.0,
+                std=residual_std,
+            )
+            nn.init.normal_(
+                layer.mlp.down_proj.weight,
+                mean=0.0,
+                std=residual_std,
+            )
+        for adapter in self.mtp_adapters:
+            nn.init.normal_(adapter.up_proj.weight, mean=0.0, std=residual_std)
+
+    def enable_gradient_checkpointing(self, enabled: bool = True) -> None:
+        self.gradient_checkpointing = enabled
+
+    def project_to_vocabulary(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return F.linear(hidden_states, self.embed_tokens.weight)
+
+    def _run_layer(
+        self,
+        layer: TinyGDNBlock,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if self.gradient_checkpointing and self.training:
+            return checkpoint(
+                layer,
+                hidden_states,
+                attention_mask,
+                use_reentrant=False,
+            )
+        return layer(hidden_states, attention_mask)
+
+    def _causal_loss(
+        self,
+        hidden_states: torch.Tensor,
+        labels: torch.Tensor,
+        target_offset: int,
+        adapter: nn.Module | None = None,
+        compute_z_loss: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if target_offset < 0:
+            raise ValueError("target_offset cannot be negative")
+        if target_offset and hidden_states.shape[1] <= target_offset:
+            raise ValueError(
+                f"Sequence length must exceed target offset {target_offset}"
+            )
+        if target_offset:
+            prediction_states = hidden_states[:, :-target_offset, :]
+            targets = labels[:, target_offset:].contiguous()
+        else:
+            prediction_states = hidden_states
+            targets = labels.contiguous()
+        if adapter is not None:
+            prediction_states = adapter(prediction_states)
+        logits = self.project_to_vocabulary(prediction_states)
+        cross_entropy = F.cross_entropy(
+            logits.reshape(-1, self.config.vocab_size),
+            targets.reshape(-1),
+            ignore_index=-100,
+        )
+        z_loss = None
+        if compute_z_loss:
+            valid_targets = targets.ne(-100)
+            log_partition = torch.logsumexp(logits.float(), dim=-1)
+            z_loss = log_partition.square()[valid_targets].mean()
+        return cross_entropy, z_loss
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        *,
+        return_logits: bool = True,
+        return_hidden_states: bool = False,
+        labels_are_shifted: bool = False,
+        include_mtp_loss: bool = True,
+        mtp_loss_weight: float | None = None,
+        z_loss_coefficient: float = 0.0,
+        logits_to_keep: int | None = None,
+    ) -> TinyGDNOutput:
+        if input_ids.ndim != 2:
+            raise ValueError("input_ids must have shape [batch, sequence]")
+        if input_ids.shape[1] > self.config.max_position_embeddings:
+            raise ValueError("Input exceeds max_position_embeddings")
+        if labels is not None and labels.shape != input_ids.shape:
+            raise ValueError("labels must have the same shape as input_ids")
+        if z_loss_coefficient < 0.0:
+            raise ValueError("z_loss_coefficient cannot be negative")
+        if logits_to_keep is not None and logits_to_keep <= 0:
+            raise ValueError("logits_to_keep must be positive")
+        effective_mtp_weight = (
+            self.config.mtp_loss_weight
+            if mtp_loss_weight is None
+            else mtp_loss_weight
+        )
+        if not 0.0 <= effective_mtp_weight <= 1.0:
+            raise ValueError("mtp_loss_weight must be between zero and one")
+
+        hidden_states = self.embed_tokens(input_ids)
+        shared_layer_indices = set(self.config.shared_layer_indices)
+        for layer_index, layer in enumerate(self.layers):
+            hidden_states = self._run_layer(
+                layer,
+                hidden_states,
+                attention_mask,
+            )
+            if layer_index in shared_layer_indices:
+                hidden_states = self._run_layer(
+                    layer,
+                    hidden_states,
+                    attention_mask,
+                )
+        hidden_states = self.final_norm(hidden_states)
+
+        main_loss = None
+        mtp_loss = None
+        z_loss = None
+        total_loss = None
+        if labels is not None:
+            main_target_offset = 0 if labels_are_shifted else 1
+            main_loss, z_loss = self._causal_loss(
+                hidden_states,
+                labels,
+                target_offset=main_target_offset,
+                compute_z_loss=z_loss_coefficient > 0.0,
+            )
+            if self.mtp_adapters and include_mtp_loss:
+                auxiliary_losses = [
+                    self._causal_loss(
+                        hidden_states,
+                        labels,
+                        target_offset=(
+                            head_index + 1
+                            if labels_are_shifted
+                            else head_index + 2
+                        ),
+                        adapter=adapter,
+                    )[0]
+                    for head_index, adapter in enumerate(self.mtp_adapters)
+                ]
+                mtp_loss = torch.stack(auxiliary_losses).mean()
+                total_loss = main_loss + effective_mtp_weight * mtp_loss
+            else:
+                total_loss = main_loss
+            if z_loss is not None:
+                total_loss = total_loss + z_loss_coefficient * z_loss
+
+        output_states = (
+            hidden_states
+            if logits_to_keep is None
+            else hidden_states[:, -logits_to_keep:, :]
+        )
+        logits = self.project_to_vocabulary(output_states) if return_logits else None
+        return TinyGDNOutput(
+            loss=total_loss,
+            logits=logits,
+            main_loss=main_loss,
+            mtp_loss=mtp_loss,
+            z_loss=z_loss,
+            hidden_states=hidden_states if return_hidden_states else None,
+        )
+
+    def parameter_report(self) -> dict[str, int]:
+        total = sum(parameter.numel() for parameter in self.parameters())
+        mtp = sum(parameter.numel() for parameter in self.mtp_adapters.parameters())
+        embeddings = self.embed_tokens.weight.numel()
+        return {
+            "deployable_core": total - mtp,
+            "training_total": total,
+            "embedding": embeddings,
+            "mtp_auxiliary": mtp,
+            "non_embedding_core": total - mtp - embeddings,
+        }
+
+    def save_checkpoint(self, output_dir: Path) -> None:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        self.config.save_json(output_dir / "config.json")
+        save_model(self, output_dir / "model.safetensors")
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        checkpoint_dir: Path,
+        *,
+        device: str | torch.device = "cpu",
+        dtype: torch.dtype | None = None,
+    ) -> TinyGDNForCausalLM:
+        config = TinyGDNConfig.from_json(checkpoint_dir / "config.json")
+        model = cls(config).to(device=device, dtype=dtype)
+        load_model(model, checkpoint_dir / "model.safetensors", device=str(device))
+        return model
+
+    def extra_repr(self) -> str:
+        report = self.parameter_report()
+        return (
+            f"core_parameters={report['deployable_core']:,}, "
+            f"training_parameters={report['training_total']:,}"
+        )
+
+    def get_architecture_metadata(self) -> dict[str, Any]:
+        return {
+            "architecture": self.config.architecture,
+            "layer_types": list(self.config.layer_types),
+            "effective_num_layers": self.config.effective_num_layers,
+            "shared_layer_indices": list(self.config.shared_layer_indices),
+            "parameter_report": self.parameter_report(),
+            "features": [
+                "32-layer deep-thin parameter allocation",
+                "Gated DeltaNet-2 recurrent memory",
+                "3:1 recurrent-to-full-attention hybrid",
+                "gated grouped-query attention",
+                "QK normalization",
+                "partial rotary embeddings",
+                "zero-centered RMSNorm",
+                "SwiGLU",
+                "tied input-output embeddings",
+                "optional multi-token prediction auxiliaries",
+            ],
+        }

tokenizer_config.json

@@ -0,0 +1,74 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "bos_token": "<|begin_of_text|>",
+  "eos_token": "<|end_of_text|>",
+  "pad_token": "<|padding|>",
+  "unk_token": "<|unknown|>",
+  "model_max_length": 2048,
+  "clean_up_tokenization_spaces": false,
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "chat_template": "{%- for message in messages -%}\n{%- if loop.first -%}{{- bos_token -}}{%- endif -%}\n{{- '<|im_start|>' + message['role'] + '\\n' -}}\n{%- if message['role'] == 'tool' -%}{{- '<|tool_response|>\\n' -}}{%- endif -%}\n{%- if message['content'] is string -%}\n{{- message['content'] -}}\n{%- elif message['content'] is iterable -%}\n{%- for item in message['content'] -%}\n{%- if item['type'] == 'text' -%}{{- item['text'] -}}{%- endif -%}\n{%- endfor -%}\n{%- endif -%}\n{%- if message['tool_calls'] is defined and message['tool_calls'] -%}\n{{- '\\n<|tool_call|>\\n' + (message['tool_calls'] | tojson) -}}\n{%- endif -%}\n{{- '<|im_end|>\\n' -}}\n{%- endfor -%}\n{%- if add_generation_prompt -%}{{- '<|im_start|>assistant\\n' -}}{%- endif -%}",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|tool_call|>",
+    "<|tool_response|>",
+    "<|think_start|>",
+    "<|think_end|>",
+    "<|fim_prefix|>",
+    "<|fim_middle|>",
+    "<|fim_suffix|>",
+    "<|fim_pad|>",
+    "<|reserved_000|>",
+    "<|reserved_001|>",
+    "<|reserved_002|>",
+    "<|reserved_003|>",
+    "<|reserved_004|>",
+    "<|reserved_005|>",
+    "<|reserved_006|>",
+    "<|reserved_007|>",
+    "<|reserved_008|>",
+    "<|reserved_009|>",
+    "<|reserved_010|>",
+    "<|reserved_011|>",
+    "<|reserved_012|>",
+    "<|reserved_013|>",
+    "<|reserved_014|>",
+    "<|reserved_015|>",
+    "<|reserved_016|>",
+    "<|reserved_017|>",
+    "<|reserved_018|>",
+    "<|reserved_019|>",
+    "<|reserved_020|>",
+    "<|reserved_021|>",
+    "<|reserved_022|>",
+    "<|reserved_023|>",
+    "<|reserved_024|>",
+    "<|reserved_025|>",
+    "<|reserved_026|>",
+    "<|reserved_027|>",
+    "<|reserved_028|>",
+    "<|reserved_029|>",
+    "<|reserved_030|>",
+    "<|reserved_031|>",
+    "<|reserved_032|>",
+    "<|reserved_033|>",
+    "<|reserved_034|>",
+    "<|reserved_035|>",
+    "<|reserved_036|>",
+    "<|reserved_037|>",
+    "<|reserved_038|>",
+    "<|reserved_039|>",
+    "<|reserved_040|>",
+    "<|reserved_041|>",
+    "<|reserved_042|>",
+    "<|reserved_043|>",
+    "<|reserved_044|>",
+    "<|reserved_045|>",
+    "<|reserved_046|>",
+    "<|reserved_047|>",
+    "<|reserved_048|>",
+    "<|reserved_049|>"
+  ]
+}

validation.json

@@ -0,0 +1,21 @@
+{
+  "checkpoint": "checkpoint-00010619",
+  "ema": {
+    "assistant_targets": 91313,
+    "batches": 16,
+    "elapsed_seconds": 1.9623924390034517,
+    "loss": 1.5321617420986058,
+    "perplexity": 4.628170928002891,
+    "targets_per_second": 46531.46750115428
+  },
+  "normal": {
+    "assistant_targets": 91313,
+    "batches": 16,
+    "elapsed_seconds": 1.9647464910085546,
+    "loss": 1.5321617420986058,
+    "perplexity": 4.628170928002891,
+    "targets_per_second": 46475.716036589896
+  },
+  "optimizer_step": 10619,
+  "type": "sft_validation"
+}

windows_fla_patches/fla/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# Keep package import light. Eagerly importing fla.layers pulls every Triton
+# kernel and breaks on Windows + Triton 3.7.
+
+from pkgutil import extend_path
+
+__path__ = extend_path(__path__, __name__)
+__version__ = "0.5.2"
+__all__: list[str] = []

windows_fla_patches/fla/layers/__init__.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# Lazy layer exports — avoid compiling every Triton kernel at import time.
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+_EXPORTS: dict[str, tuple[str, str]] = {
+    "ABCAttention": (".abc", "ABCAttention"),
+    "Attention": (".attn", "Attention"),
+    "BasedLinearAttention": (".based", "BasedLinearAttention"),
+    "BitAttention": (".bitattn", "BitAttention"),
+    "Comba": (".comba", "Comba"),
+    "DeltaNet": (".delta_net", "DeltaNet"),
+    "DeltaFormerAttention": (".deltaformer", "DeltaFormerAttention"),
+    "ForgettingAttention": (".forgetting_attn", "ForgettingAttention"),
+    "GatedDeltaNet": (".gated_deltanet", "GatedDeltaNet"),
+    "GatedDeltaProduct": (".gated_deltaproduct", "GatedDeltaProduct"),
+    "GatedDeltaNet2": (".gdn2", "GatedDeltaNet2"),
+    "GatedLinearAttention": (".gla", "GatedLinearAttention"),
+    "GatedSlotAttention": (".gsa", "GatedSlotAttention"),
+    "HGRNAttention": (".hgrn", "HGRNAttention"),
+    "HGRN2Attention": (".hgrn2", "HGRN2Attention"),
+    "KimiDeltaAttention": (".kda", "KimiDeltaAttention"),
+    "LightNetAttention": (".lightnet", "LightNetAttention"),
+    "LinearAttention": (".linear_attn", "LinearAttention"),
+    "LogLinearMamba2": (".log_linear_mamba2", "LogLinearMamba2"),
+    "Mamba": (".mamba", "Mamba"),
+    "Mamba2": (".mamba2", "Mamba2"),
+    "Mamba3": (".mamba3", "Mamba3"),
+    "MesaNet": (".mesa_net", "MesaNet"),
+    "MultiheadLatentAttention": (".mla", "MultiheadLatentAttention"),
+    "MoBA": (".moba", "MoBA"),
+    "MomAttention": (".mom", "MomAttention"),
+    "MultiScaleRetention": (".multiscale_retention", "MultiScaleRetention"),
+    "NativeSparseAttention": (".nsa", "NativeSparseAttention"),
+    "Parallax": (".parallax", "Parallax"),
+    "PaTHAttention": (".path_attn", "PaTHAttention"),
+    "Raven": (".raven", "Raven"),
+    "ReBasedLinearAttention": (".rebased", "ReBasedLinearAttention"),
+    "RodimusAttention": (".rodimus", "RodimusAttention"),
+    "SlidingWindowSharedKeyAttention": (".rodimus", "SlidingWindowSharedKeyAttention"),
+    "RWKV6Attention": (".rwkv6", "RWKV6Attention"),
+    "RWKV7Attention": (".rwkv7", "RWKV7Attention"),
+    "WallAttention": (".wall_attn", "WallAttention"),
+    "YOCOCrossAttention": (".yoco", "YOCOCrossAttention"),
+    "YOCOGatedRetention": (".yoco", "YOCOGatedRetention"),
+    "YOCOSharedKVBuilder": (".yoco", "YOCOSharedKVBuilder"),
+}
+
+__all__ = list(_EXPORTS)
+
+
+def __getattr__(name: str) -> Any:
+    spec = _EXPORTS.get(name)
+    if spec is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module_name, attr = spec
+    value = getattr(importlib.import_module(module_name, __name__), attr)
+    globals()[name] = value
+    return value

windows_fla_patches/fla/ops/__init__.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# Lazy public exports so importing fla.ops.utils / fla.ops.gdn2 does not
+# eagerly compile every Triton kernel (needed on Windows + Triton 3.7).
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+_EXPORTS: dict[str, str] = {
+    "chunk_abc": "fla.ops.abc",
+    "parallel_attn": "fla.ops.attn",
+    "fused_attnres": "fla.ops.attnres",
+    "fused_chunk_based": "fla.ops.based",
+    "parallel_based": "fla.ops.based",
+    "chunk_comba": "fla.ops.comba",
+    "fused_recurrent_comba": "fla.ops.comba",
+    "chunk_delta_rule": "fla.ops.delta_rule",
+    "fused_chunk_delta_rule": "fla.ops.delta_rule",
+    "fused_recurrent_delta_rule": "fla.ops.delta_rule",
+    "parallel_forgetting_attn": "fla.ops.forgetting_attn",
+    "chunk_gated_delta_rule": "fla.ops.gated_delta_rule",
+    "chunk_gdn": "fla.ops.gated_delta_rule",
+    "fused_recurrent_gated_delta_rule": "fla.ops.gated_delta_rule",
+    "fused_recurrent_gdn": "fla.ops.gated_delta_rule",
+    "chunk_dplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "chunk_iplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "fused_recurrent_dplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "fused_recurrent_iplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "chunk_gla": "fla.ops.gla",
+    "fused_chunk_gla": "fla.ops.gla",
+    "fused_recurrent_gla": "fla.ops.gla",
+    "chunk_gsa": "fla.ops.gsa",
+    "fused_recurrent_gsa": "fla.ops.gsa",
+    "fused_recurrent_hgrn": "fla.ops.hgrn",
+    "chunk_kda": "fla.ops.kda",
+    "fused_recurrent_kda": "fla.ops.kda",
+    "chunk_lightning_attn": "fla.ops.lightning_attn",
+    "fused_recurrent_lightning_attn": "fla.ops.lightning_attn",
+    "chunk_linear_attn": "fla.ops.linear_attn",
+    "fused_chunk_linear_attn": "fla.ops.linear_attn",
+    "fused_recurrent_linear_attn": "fla.ops.linear_attn",
+    "chunk_log_linear_attn": "fla.ops.log_linear_attn",
+    "chunk_mesa_net": "fla.ops.mesa_net",
+    "parallel_nsa": "fla.ops.nsa",
+    "parallel_parallax": "fla.ops.parallax",
+    "parallel_path_attn": "fla.ops.path_attn",
+    "chunk_retention": "fla.ops.retention",
+    "fused_chunk_retention": "fla.ops.retention",
+    "fused_recurrent_retention": "fla.ops.retention",
+    "parallel_retention": "fla.ops.retention",
+    "chunk_rwkv6": "fla.ops.rwkv6",
+    "fused_recurrent_rwkv6": "fla.ops.rwkv6",
+    "chunk_rwkv7": "fla.ops.rwkv7",
+    "fused_recurrent_rwkv7": "fla.ops.rwkv7",
+    "chunk_simple_gla": "fla.ops.simple_gla",
+    "fused_chunk_simple_gla": "fla.ops.simple_gla",
+    "fused_recurrent_simple_gla": "fla.ops.simple_gla",
+    "parallel_simple_gla": "fla.ops.simple_gla",
+    "parallel_wall_attn": "fla.ops.wall_attn",
+    "parallel_wall_attn_decode": "fla.ops.wall_attn",
+    "chunk_gdn2": "fla.ops.gdn2",
+    "fused_recurrent_gdn2": "fla.ops.gdn2",
+}
+
+__all__ = list(_EXPORTS)
+
+
+def __getattr__(name: str) -> Any:
+    module_name = _EXPORTS.get(name)
+    if module_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    value = getattr(importlib.import_module(module_name), name)
+    globals()[name] = value
+    return value

windows_fla_patches/fla/ops/simple_gla/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+from .chunk import chunk_simple_gla
+from .fused_chunk import fused_chunk_simple_gla
+from .fused_recurrent import fused_recurrent_simple_gla
+
+# Triton 3.7 on Windows can fail while decorating parallel kernels at import time.
+try:
+    from .parallel import parallel_simple_gla
+except Exception:  # noqa: BLE001
+    parallel_simple_gla = None
+
+__all__ = [
+    'chunk_simple_gla',
+    'fused_chunk_simple_gla',
+    'fused_recurrent_simple_gla',
+    'parallel_simple_gla',
+]