lilo-100m: fp16 + group-wise int8 CPU release

README.md

@@ -0,0 +1,121 @@
+---
+license: apache-2.0
+language:
+  - id
+library_name: transformers
+tags:
+  - deepseek
+  - moe
+  - causal-lm
+  - chat
+  - cpu
+  - indonesian
+datasets:
+  - neosantara/IndoTalkSFT
+base_model: neosantara/wader-100m-base
+pipeline_tag: text-generation
+---
+
+# lilo-100m — Lightweight Local Intelligence & Optimizer
+
+A ~110M parameter language model implementing the DeepSeek-V4 architecture, fine-tuned for chat/instruction following in Indonesian. This is the **lilo** release of `neosantara/wader-100m`: same architecture and weights, re-packaged explicitly for **CPU-only, low-memory, low-footprint** use.
+
+- **Base**: [neosantara/wader-100m-base](https://huggingface.co/neosantara/wader-100m-base) · pretrained on [neosantara/IndoTalk](https://huggingface.co/datasets/neosantara/IndoTalk)
+- **Upstream SFT**: [neosantara/wader-100m](https://huggingface.co/neosantara/wader-100m) on [neosantara/IndoTalkSFT](https://huggingface.co/datasets/neosantara/IndoTalkSFT)
+- **Trained from scratch** — no weights from DeepSeek-V4 were used.
+
+## Why lilo
+
+| | wader-100m (original) | **lilo-100m** |
+|---|---|---|
+| Weights on disk | 421MB fp32 | **119MB** (int8) / **221MB** (fp16) |
+| CPU-friendly runtime | bf16 NaN bugs, no KV cache | fp32/fp16 stable, KV cache, inference_mode |
+| Runs on 2 vCPU / 2GB RAM VPS | untested | **measured** |
+
+## Variants
+
+| Variant | File | Disk | PPL delta vs fp32 | Greedy outputs |
+|---|---|---|---|---|
+| `fp16` (default) | `model.fp16.safetensors` | 221MB | +0.08% | 5/5 identical |
+| `int8` | `model.int8.safetensors` | 119MB | +0.32% | 5/5 identical |
+
+Notes:
+- `int8` is symmetric group-wise quantization (group=64) on all `nn.Linear` layers + embeddings; norms/router/HC params stay fp32. Scales cost ~7MB.
+- Both dequantize to fp32 compute on load — the **disk footprint is what shrinks**, because this CPU (Zen2, AVX2-only, no VNNI) does fp32 GEMM as fast as int8 anyway, and fp16 GEMM is *faster* still.
+- fp16 is the default because it is both the closest to the original quality **and** the fastest on CPU.
+
+## CPU speed (measured on 2 vCPU AMD EPYC 7K62 "Zen2", AVX2-only, box under load)
+
+| Variant | threads=1 | threads=2 |
+|---|---|---|
+| fp16 | **14–22 tok/s** | 11–12 tok/s |
+| int8 | 9–19 tok/s | 5–8 tok/s |
+
+On this tiny model (320 hidden, 8 layers) thread sync overhead dominates — `--threads 1` is fastest. On 8+ core machines try 2–4.
+
+## Memory safety (no OOM)
+
+Peak RSS measured on a 2 vCPU / 1.9GB RAM VPS: **~0.85GB (int8) / ~1.05GB (fp16)** including the Python/torch runtime. Rules baked into `run_cpu.sh`:
+
+- `OMP_NUM_THREADS` bounded, `MALLOC_ARENA_MAX=1` (glibc per-thread arenas can silently eat GBs)
+- `TMPDIR` forced to disk, never `/tmp` when `/tmp` is tmpfs
+- prefill chunked at 256 tokens so attention never allocates a big `[S,S]` matrix
+
+## Usage
+
+```bash
+git lfs install
+git clone https://huggingface.co/hzlo0/lilo-100m
+cd lilo-100m
+pip install torch --index-url https://download.pytorch.org/whl/cpu
+pip install transformers safetensors tokenizers accelerate
+
+bash run_cpu.sh                              # chat, default fp16
+CUDA_VISIBLE_DEVICES="" python chat_lilo.py --variant int8 --threads 1 "Halo, apa kabar?"
+```
+
+Chat template (DeepSeek-V4 style):
+
+```
+<|begin_of_sentence|><|User|>Halo<|Assistant|>
+```
+
+```python
+import torch
+from lilo_common import load_lilo, generate_lilo, chat
+
+torch.set_num_threads(1)
+tok, model = load_lilo(".", variant="fp16")          # or "int8"
+print(chat(tok, model, [{"role": "user", "content": "Halo, apa kabar?"}],
+           max_new_tokens=200, temperature=0.7, top_p=0.9))
+```
+
+Requires `trust_remote_code`-free loading: the model code ships as plain files (`modeling_deepseek_v4.py`) loaded by `lilo_common.py`. (The same code also works via `AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True)`.)
+
+## Architecture
+
+| Component | Details |
+|---|---|
+| Parameters | ~110M (41M embeddings, 69M non-embedding) |
+| Hidden size | 320 · Layers 8 |
+| Attention | 8 heads, 1 KV head (MQA), MLA q_lora_rank=160, grouped low-rank output |
+| MoE | 4 routed experts + 1 shared, top-2 routing |
+| Hyper-Connections | hc_mult=4, Sinkhorn routing (replaces residual stream) |
+| Vocab | 129,280 (DeepSeek-V4 tokenizer) · Context 2,048 |
+
+## What was fixed for CPU in this release
+
+1. **freqs_cis recomputed on demand** — original code left a buffer uninitialized after meta-device loading, corrupting RoPE beyond a few tokens.
+2. **KV cache honored** — original `forward` hardcoded `use_cache=False`, recomputing the whole sequence every decode step. This release caches and steps one token.
+3. **`return_dict` defaults to True** — transformers 5.x deprecated `config.use_return_dict`, which silently returned tuples.
+4. **Expert dtype bridge** — `w2(x.to(w2.weight.dtype))` silently cast activations to int8 under weight-only quantization.
+5. **bf16 NaN** — model runs fp32/fp16; no bf16 anywhere.
+
+## Limitations
+
+- Tiny model: 110M params, most capacity in embeddings. Generations are often incoherent or factually wrong.
+- Undertrained: 5K pretrain + 3K SFT steps (upstream). Production models train 100K+ steps on trillions of tokens.
+
+## License
+
+Apache-2.0 (inherited from upstream `neosantara/wader-100m`).

chat_lilo.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+"""LILO chat CLI — CPU-optimized, KV-cache generation.
+
+Usage:
+    python chat_lilo.py "prompt"                       # single prompt
+    python chat_lilo.py                                # interactive
+    python chat_lilo.py --temp 0.7 --top-p 0.9 --max 200 "prompt"
+"""
+import argparse, os, sys
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, HERE)
+
+from lilo_common import load_lilo, generate_lilo
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("prompt", nargs="*")
+    ap.add_argument("--variant", default="fp16", choices=["fp16", "int8"])
+    ap.add_argument("--max", type=int, default=200)
+    ap.add_argument("--temp", type=float, default=0.7)
+    ap.add_argument("--top-p", type=float, default=0.9)
+    ap.add_argument("--rep", type=float, default=1.1)
+    ap.add_argument("--threads", type=int, default=1)
+    ap.add_argument("--path", default=HERE)
+    a = ap.parse_args()
+
+    tok, model = load_lilo(a.path, variant=a.variant, threads=a.threads)
+
+    def run(prompt):
+        ids = tok.encode(prompt, add_special_tokens=False)
+        out = generate_lilo(model, tok, ids, max_new_tokens=a.max, temperature=a.temp,
+                            top_p=a.top_p, repetition_penalty=a.rep, stream=True)
+        return tok.decode(out, skip_special_tokens=True)
+
+    if a.prompt:
+        run(" ".join(a.prompt))
+        print()
+        return
+
+    print("LILO siap. Ketik pesan (Ctrl-C/Ctrl-D untuk keluar).")
+    while True:
+        try:
+            p = input("\n> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print()
+            break
+        if not p:
+            continue
+        run(p)
+        print()
+
+
+if __name__ == "__main__":
+    main()

chat_template.jinja

@@ -0,0 +1,33 @@
+{%- if not add_generation_prompt is defined -%}
+	{%- set add_generation_prompt = false -%}
+{%- endif -%}
+{%- set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt="", is_first_sp=true) -%}
+{%- for message in messages -%}
+	{%- if message["role"] == "system" -%}
+		{%- if ns.is_first_sp -%}
+			{%- set ns.system_prompt = ns.system_prompt + message["content"] -%}
+			{%- set ns.is_first_sp = false -%}
+		{%- else -%}
+			{%- set ns.system_prompt = ns.system_prompt + "\n\n" + message["content"] -%}
+		{%- endif -%}
+	{%- endif -%}
+{%- endfor -%}
+{{- bos_token -}}
+{{- ns.system_prompt -}}
+{%- for message in messages -%}
+	{%- if message["role"] == "user" -%}
+		{%- set ns.is_tool = false -%}
+		{{- "<｜User｜>" + message["content"] -}}
+	{%- endif -%}
+	{%- if message["role"] == "assistant" and message["content"] is not none -%}
+		{%- if ns.is_tool -%}
+			{{- "<｜tool▁outputs▁end｜>" + message["content"] + "<｜end▁of▁sentence｜>" -}}
+			{%- set ns.is_tool = false -%}
+		{%- else -%}
+			{{- "<｜Assistant｜>" + message["content"] + "<｜end▁of▁sentence｜>" -}}
+		{%- endif -%}
+	{%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+	{{- "<｜Assistant｜>" -}}
+{%- endif -%}

config.json

@@ -0,0 +1,76 @@
+{
+  "architectures": [
+    "DeepseekV4ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_deepseek_v4.DeepseekV4Config",
+    "AutoModelForCausalLM": "modeling_deepseek_v4.DeepseekV4ForCausalLM"
+  },
+  "bos_token_id": 0,
+  "compress_ratios": [
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0
+  ],
+  "compress_rope_theta": 160000.0,
+  "dtype": "float16",
+  "eos_token_id": 1,
+  "hc_eps": 1e-06,
+  "hc_mult": 4,
+  "hc_sinkhorn_iters": 2,
+  "head_dim": 96,
+  "hidden_act": "silu",
+  "hidden_size": 320,
+  "index_head_dim": 128,
+  "index_n_heads": 64,
+  "index_topk": 512,
+  "initializer_range": 0.02,
+  "max_position_embeddings": 2048,
+  "model_type": "deepseek_v4",
+  "moe_intermediate_size": 640,
+  "n_routed_experts": 4,
+  "n_shared_experts": 1,
+  "nope_head_dim": 64,
+  "norm_topk_prob": true,
+  "num_attention_heads": 8,
+  "num_experts_per_tok": 2,
+  "num_hash_layers": 0,
+  "num_hidden_layers": 8,
+  "num_key_value_heads": 1,
+  "num_nextn_predict_layers": 1,
+  "o_groups": 2,
+  "o_lora_rank": 80,
+  "pad_token_id": 1,
+  "q_lora_rank": 160,
+  "qk_rope_head_dim": 32,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "rope_theta": 10000.0,
+  "routed_scaling_factor": 1.5,
+  "scoring_func": "sqrtsoftplus",
+  "sliding_window": 128,
+  "swiglu_limit": 0.0,
+  "tie_word_embeddings": false,
+  "topk_method": "noaux_tc",
+  "transformers_version": "5.15.1",
+  "use_cache": false,
+  "vocab_size": 129280,
+  "name": "lilo-100m",
+  "quantization_config": {
+    "format": "int8-symmetric-groupwise",
+    "group_size": 64,
+    "modules": "linear+embeddings",
+    "file": "model.int8.safetensors"
+  }
+}

configuration_deepseek_v4.py

@@ -0,0 +1,140 @@
+"""DeepSeek-V4 model configuration.
+
+Adapted from the DeepSeek-V4 inference config (deepseek-ai/DeepSeek-V4-Pro)
+and the HF Transformers DeepSeek-V3 config for HF compatibility.
+
+Key V4-specific features vs V3:
+- Hyper-Connections (HC): multi-copy hidden states with Sinkhorn routing
+- Compressed Sparse Attention (CSA): compression + sliding window + sparse indexing
+- New MoE routing: sqrtsoftplus scoring, hash-based routing for first layers
+- Large head_dim (512), o_groups/o_lora_rank for grouped output projection
+- No kv_lora_rank (replaced by compress_ratios)
+- No v_head_dim/qk_nope_head_dim (replaced by head_dim)
+"""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class DeepseekV4Config(PretrainedConfig):
+    model_type = "deepseek_v4"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=129280,
+        hidden_size=4096,
+        num_hidden_layers=43,
+        num_attention_heads=64,
+        num_key_value_heads=1,
+        # MoE
+        moe_intermediate_size=2048,
+        n_routed_experts=256,
+        n_shared_experts=1,
+        num_experts_per_tok=6,
+        norm_topk_prob=True,
+        scoring_func="sqrtsoftplus",
+        routed_scaling_factor=1.5,
+        topk_method="noaux_tc",
+        num_hash_layers=3,
+        swiglu_limit=10.0,
+        # MLA / Attention
+        q_lora_rank=1024,
+        head_dim=512,
+        qk_rope_head_dim=64,
+        o_groups=8,
+        o_lora_rank=1024,
+        sliding_window=128,
+        # Compression
+        compress_ratios=None,
+        compress_rope_theta=160000.0,
+        # Index attention
+        index_n_heads=64,
+        index_head_dim=128,
+        index_topk=512,
+        # Hyper-Connections
+        hc_mult=4,
+        hc_sinkhorn_iters=20,
+        hc_eps=1e-6,
+        # MTP
+        num_nextn_predict_layers=1,
+        # Standard
+        hidden_act="silu",
+        max_position_embeddings=4096,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=0,
+        eos_token_id=1,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads or num_attention_heads
+
+        # MoE
+        self.moe_intermediate_size = moe_intermediate_size
+        self.n_routed_experts = n_routed_experts
+        self.n_shared_experts = n_shared_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.norm_topk_prob = norm_topk_prob
+        self.scoring_func = scoring_func
+        self.routed_scaling_factor = routed_scaling_factor
+        self.topk_method = topk_method
+        self.num_hash_layers = num_hash_layers
+        self.swiglu_limit = swiglu_limit
+
+        # Attention
+        self.q_lora_rank = q_lora_rank
+        self.head_dim = head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.nope_head_dim = head_dim - qk_rope_head_dim
+        self.o_groups = o_groups
+        self.o_lora_rank = o_lora_rank
+        self.sliding_window = sliding_window
+
+        # Compression
+        if compress_ratios is None:
+            # Default: no compression for small models
+            compress_ratios = [0] * (num_hidden_layers + 1)
+        self.compress_ratios = compress_ratios
+        self.compress_rope_theta = compress_rope_theta
+
+        # Index attention
+        self.index_n_heads = index_n_heads
+        self.index_head_dim = index_head_dim
+        self.index_topk = index_topk
+
+        # Hyper-Connections
+        self.hc_mult = hc_mult
+        self.hc_sinkhorn_iters = hc_sinkhorn_iters
+        self.hc_eps = hc_eps
+
+        # MTP
+        self.num_nextn_predict_layers = num_nextn_predict_layers
+
+        # Standard
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

generation_config.json

@@ -0,0 +1,12 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": [
+    1
+  ],
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 1,
+  "transformers_version": "5.15.1",
+  "use_cache": true
+}

lilo_common.py

@@ -0,0 +1,177 @@
+"""LILO — Lightweight Local Intelligence & Optimizer
+
+Helpers for loading and running LILO-100M on CPU. Variants:
+  - lilo-100m (fp16) : fp16 weights (221MB on disk) — fastest, quality gate
+                       PPL delta 0.08% vs fp32, greedy outputs 5/5 identical
+  - lilo-100m (int8) : INT8 group-wise weights (group=64, 119MB on disk) —
+                       smallest/RAM-light, PPL delta 0.32%, greedy 5/5 identical
+
+Both ship patched modeling code (freqs_cis fix + working KV cache) so single-token
+decode is cheap. Generate with `generate_lilo()` (explicit KV cache).
+
+CPU speed (measured on 2 vCPU AMD EPYC 7K62 Zen2, AVX2-only, under load):
+  fp16  14-22 tok/s at --threads 1   (recommended default)
+  int8  9-19 tok/s at --threads 1
+More threads HURT on tiny models (sync overhead); scale up only on 8+ cores.
+"""
+
+import json
+import os
+from typing import List, Optional, Sequence
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import load_file
+
+from transformers import AutoConfig, AutoTokenizer
+
+from configuration_deepseek_v4 import DeepseekV4Config
+from modeling_deepseek_v4 import (
+    DeepseekV4ForCausalLM,
+    QuantizedLinear,
+    QuantizedEmbedding,
+)
+
+
+# submodule leaf names that get INT8-quantized in the int8 variant
+_LINEAR_TARGETS = (
+    "wq_a", "wq_b", "wkv", "wo_a", "wo_b",   # attention projections
+    "w1", "w2", "w3",                        # experts (gate/up/down)
+    "lm_head",
+)
+
+
+def _swap_quantized(model, quantize_embed=True):
+    """Swap nn.Linear -> QuantizedLinear and nn.Embedding -> QuantizedEmbedding.
+    Returns number of modules swapped."""
+    swapped = 0
+    for parent in list(dict(model.named_modules()).values()):
+        for name, child in list(parent.named_children()):
+            target = None
+            if name in _LINEAR_TARGETS and child.__class__.__name__ == "Linear":
+                target = QuantizedLinear(child.in_features, child.out_features)
+            elif quantize_embed and name == "embed_tokens" and child.__class__.__name__ == "Embedding":
+                target = QuantizedEmbedding(child.num_embeddings, child.embedding_dim)
+            if target is not None:
+                setattr(parent, name, target)
+                swapped += 1
+    return swapped
+
+
+def load_lilo(path_or_repo, variant="int8", device="cpu", threads=2, hf_token=None):
+    """Load a LILO model. variant: "int8" | "fp16" | "fp32"."""
+    path = path_or_repo
+    if not os.path.isdir(path):
+        from huggingface_hub import snapshot_download
+        path = snapshot_download(repo_id=path_or_repo, token=hf_token)
+
+    torch.set_num_threads(threads)
+
+    config = DeepseekV4Config.from_pretrained(path, trust_remote_code=True)
+    tok = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
+
+    if variant == "int8":
+        model = DeepseekV4ForCausalLM(config)
+        from safetensors.torch import load_file
+        sd = load_file(os.path.join(path, "model.int8.safetensors"))
+        with safe_open(os.path.join(path, "model.int8.safetensors"), framework="pt") as _sf:
+            emb_dtype = _sf.get_slice("model.embed_tokens.weight").get_dtype()
+        n = _swap_quantized(model, quantize_embed=(emb_dtype == "I8"))
+        missing, unexpected = model.load_state_dict(sd, strict=False)
+        bad_missing = [k for k in missing if not k.startswith("freqs_cis")]
+        assert not bad_missing and not unexpected, f"missing={bad_missing[:5]} unexpected={unexpected[:5]}"
+        model.eval()
+        model.model.get_freqs_cis()
+        print(f"[lilo] loaded INT8 variant from {path} ({n} modules swapped, "
+              f"{sum(t.numel() for t in sd.values())/1e6:.0f}M params)")
+    elif variant == "fp16":
+        model = DeepseekV4ForCausalLM(config)
+        from safetensors.torch import load_file
+        sd = load_file(os.path.join(path, "model.fp16.safetensors"))
+        missing, unexpected = model.load_state_dict(sd, strict=False)
+        assert not missing and not unexpected, f"missing={missing[:5]} unexpected={unexpected[:5]}"
+        model.eval()
+        model = model.to(torch.float16)
+        model.model.get_freqs_cis()
+        print(f"[lilo] loaded FP16 variant from {path}")
+    else:
+        model = DeepseekV4ForCausalLM.from_pretrained(
+            path, trust_remote_code=True, dtype=torch.float32, attn_implementation="eager"
+        )
+        model.eval()
+        model.model.get_freqs_cis()
+        print(f"[lilo] loaded FP32 variant from {path}")
+
+    model.to(device)
+    return tok, model
+
+
+@torch.inference_mode()
+def generate_lilo(
+    model,
+    tok,
+    prompt_ids: Sequence[int],
+    max_new_tokens: int = 200,
+    temperature: float = 0.7,
+    top_p: float = 0.9,
+    repetition_penalty: float = 1.0,
+    eos_token_id: int = 1,
+    stream: bool = False,
+):
+    """Greedy/sampling generation with an explicit KV cache.
+
+    Returns (generated_ids list). Repetition penalty is applied on the growing
+    sequence (standard formulation)."""
+    ids = torch.tensor([list(prompt_ids)], dtype=torch.long)
+    generated: List[int] = []
+    past = None
+
+    logits = None
+    for step in range(max_new_tokens):
+        if step == 0:
+            # chunked prefill to bound peak memory
+            seq = ids
+            chunk = 256
+            for s in range(0, seq.shape[1], chunk):
+                piece = seq[:, s : s + chunk]
+                out = model(input_ids=piece, past_key_values=past, use_cache=True)
+                past = out.past_key_values
+            logits = out.logits[0, -1].float()
+        else:
+            out = model(
+                input_ids=torch.tensor([[generated[-1]]], dtype=torch.long),
+                past_key_values=past,
+                use_cache=True,
+            )
+            past = out.past_key_values
+            logits = out.logits[0, -1].float()
+
+        if repetition_penalty != 1.0:
+            for t in set(generated):
+                logits[t] /= repetition_penalty if logits[t] > 0 else 1.0 / repetition_penalty
+
+        if temperature <= 0:
+            nxt = int(torch.argmax(logits))
+        else:
+            probs = torch.softmax(logits / max(temperature, 1e-5), dim=-1)
+            if top_p < 1.0:
+                sp, si = torch.sort(probs, descending=True)
+                cum = torch.cumsum(sp, dim=-1)
+                mask = cum - sp > top_p
+                sp[mask] = 0.0
+                probs = torch.zeros_like(probs).scatter_(0, si, sp)
+                probs = probs / probs.sum()
+            nxt = int(torch.multinomial(probs, 1))
+        generated.append(nxt)
+        if stream:
+            print(tok.decode([nxt]), end="", flush=True)
+        if nxt == eos_token_id:
+            break
+    return generated
+
+
+def chat(tok, model, messages, **kw):
+    prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    ids = tok.encode(prompt, add_special_tokens=False)
+    gen = generate_lilo(model, tok, ids, **kw)
+    return tok.decode(gen, skip_special_tokens=True)

model.fp16.safetensors

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

model.int8.safetensors

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

modeling_deepseek_v4.py

@@ -0,0 +1,842 @@
+"""DeepSeek-V4 model implementation for HuggingFace Transformers.
+
+Ported from deepseek-ai/DeepSeek-V4-Pro inference/model.py to be compatible
+with HF Trainer, SFTTrainer, and AutoModelForCausalLM.
+
+Key V4 architecture features implemented:
+- Hyper-Connections (HC): multi-copy hidden states with Sinkhorn routing
+- Compressed Sparse Attention (CSA) with sliding window
+- MoE with sqrtsoftplus scoring and hash-based routing
+- Grouped low-rank output projection (o_groups + o_lora_rank)
+- Multi-Token Prediction (MTP) layers (disabled for small models)
+
+Custom kernels (tilelang) are NOT required — all ops are pure PyTorch.
+For training from scratch in bf16, this is sufficient and simpler.
+
+LILO patches (CPU-optimized, behavior-preserving):
+- freqs_cis buffer is recomputed on demand (fixes uninitialized-memory bug
+  after transformers>=5 meta-device loading which corrupted RoPE and caused
+  NaN logits on sequences longer than a few tokens)
+- KV cache is honored in attention (original code hardcoded use_cache=False,
+  which recomputed the full sequence every decode step)
+- Optional symmetric INT8 weight-only quantization for nn.Linear / embedding
+  layers (dequantized per-layer during forward, negligible compute overhead)
+"""
+
+import math
+from typing import Optional, Tuple, List
+from functools import lru_cache
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.generation import GenerationMixin
+
+try:
+    from .configuration_deepseek_v4 import DeepseekV4Config
+except ImportError:
+    from configuration_deepseek_v4 import DeepseekV4Config
+
+
+# ---------------------------------------------------------------------------
+# Utility functions
+# ---------------------------------------------------------------------------
+
+class DeepseekV4RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x = x.float()
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps)
+        return (self.weight * x).to(dtype)
+
+
+def precompute_freqs_cis(dim, seqlen, base=10000.0):
+    """Precompute cos/sin for rotary embeddings (real-valued, compile-friendly)."""
+    freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    t = torch.arange(seqlen, dtype=torch.float32)
+    freqs = torch.outer(t, freqs)  # [S, D//2]
+    cos = freqs.cos()
+    sin = freqs.sin()
+    return torch.stack([cos, sin], dim=0)  # [2, S, D//2]
+
+
+def apply_rotary_emb(x: torch.Tensor, cos_sin: torch.Tensor) -> torch.Tensor:
+    """Apply rotary positional embeddings (real-valued, no complex ops).
+    
+    x: [..., D] where D is even
+    cos_sin: [2, S, D//2] - precomputed cos and sin
+    """
+    cos, sin = cos_sin[0], cos_sin[1]  # each [S, D//2]
+    d = x.shape[-1] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    # Broadcast cos/sin to match x shape
+    while cos.ndim < x1.ndim:
+        cos = cos.unsqueeze(0)
+        sin = sin.unsqueeze(0)
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], dim=-1).to(x.dtype)
+
+
+# ---------------------------------------------------------------------------
+# INT8 quantization (LILO): symmetric per-output-channel weights.
+# Checkpoint stores int8 tensors + fp32 per-row/per-output-channel scales.
+# ---------------------------------------------------------------------------
+
+def quantize_per_channel(w: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Symmetric INT8 quantization along dim 0 (output channels / rows)."""
+    wf = w.detach().float()
+    scale = wf.abs().amax(dim=1, keepdim=True) / 127.0
+    scale = torch.clamp(scale, min=1e-12)
+    q = torch.round(wf / scale).clamp(-127, 127).to(torch.int8)
+    return q, scale.to(torch.float32)
+
+
+class QuantizedLinear(nn.Module):
+    """Drop-in for nn.Linear(bias=False) holding INT8 weights + scales.
+
+    Supports two scale layouts:
+    - per-output-channel: weight_scale shape [out, 1]
+    - group-wise (LILO g64): weight_scale shape [out, in/G] where G is inferred
+    """
+
+    def __init__(self, in_features: int, out_features: int):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        g = 64 if in_features % 64 == 0 else 32
+        self.register_buffer("weight", torch.zeros(out_features, in_features, dtype=torch.int8), persistent=True)
+        self.register_buffer("weight_scale", torch.ones(out_features, in_features // g, dtype=torch.float32), persistent=True)
+        self._w_cache = None
+
+    def dequantized_weight(self, dtype: torch.dtype = torch.float32) -> torch.Tensor:
+        if self._w_cache is None or self._w_cache.dtype != dtype:
+            q = self.weight.to(torch.float32)
+            s = self.weight_scale
+            if s.dim() == 2 and s.shape[1] > 1:  # group-wise
+                g = s.shape[1]
+                self._w_cache = (q.view(s.shape[0], g, q.shape[1] // g) * s.unsqueeze(-1)).view(q.shape).to(dtype)
+            else:  # per-output-channel
+                self._w_cache = (q * s).to(dtype)
+        return self._w_cache
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.linear(x, self.dequantized_weight(x.dtype))
+
+
+class QuantizedEmbedding(nn.Module):
+    """Drop-in for nn.Embedding holding INT8 weights + per-row scales.
+
+    Supports per-output-channel [N, 1] and group-wise [N, D/G] scale layouts.
+    """
+
+    def __init__(self, num_embeddings: int, embedding_dim: int):
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        g = 64 if embedding_dim % 64 == 0 else 32
+        self.register_buffer("weight", torch.zeros(num_embeddings, embedding_dim, dtype=torch.int8), persistent=True)
+        self.register_buffer("weight_scale", torch.ones(num_embeddings, embedding_dim // g, dtype=torch.float32), persistent=True)
+
+    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+        q = self.weight[ids].to(torch.float32)
+        s = self.weight_scale[ids]
+        if s.shape[-1] > 1:  # group-wise
+            g = s.shape[-1]
+            in_f = q.shape[-1]
+            out = (q.reshape(-1, g, in_f // g) * s.reshape(-1, g, 1)).reshape(-1, in_f)
+            return out.reshape(*q.shape[:-1], in_f)
+        return (q * s).to(torch.float32)
+
+
+
+# ---------------------------------------------------------------------------
+# Hyper-Connections (HC)
+# ---------------------------------------------------------------------------
+
+def hc_split_sinkhorn(mixes, hc_scale, hc_base, hc_mult=4, sinkhorn_iters=20, eps=1e-6):
+    """Pure PyTorch implementation of HC split + Sinkhorn normalization.
+    
+    Args:
+        mixes: [B, S, (2+hc_mult)*hc_mult] - mixed scores from linear projection
+        hc_scale: [3] - scale parameters
+        hc_base: [(2+hc_mult)*hc_mult] - bias parameters
+        hc_mult: number of HC copies
+        sinkhorn_iters: number of Sinkhorn normalization iterations
+        eps: numerical stability epsilon
+    
+    Returns:
+        pre: [B, S, hc_mult] - pre-connection weights
+        post: [B, S, hc_mult] - post-connection weights
+        comb: [B, S, hc_mult, hc_mult] - combination matrix
+    """
+    # Split into pre, post, and combination parts
+    pre_raw = mixes[..., :hc_mult]
+    post_raw = mixes[..., hc_mult:2*hc_mult]
+    comb_raw = mixes[..., 2*hc_mult:].reshape(*mixes.shape[:-1], hc_mult, hc_mult)
+    
+    # Apply scale and base
+    pre = torch.sigmoid(pre_raw * hc_scale[0] + hc_base[:hc_mult]) + eps
+    post = 2 * torch.sigmoid(post_raw * hc_scale[1] + hc_base[hc_mult:2*hc_mult])
+    
+    # Combination matrix with Sinkhorn normalization
+    comb = comb_raw * hc_scale[2] + hc_base[2*hc_mult:].reshape(hc_mult, hc_mult)
+    
+    # Initial softmax along last dim + eps
+    comb = F.softmax(comb, dim=-1) + eps
+    # Normalize along dim=-2
+    comb = comb / (comb.sum(dim=-2, keepdim=True) + eps)
+    
+    # Sinkhorn iterations
+    for _ in range(sinkhorn_iters - 1):
+        comb = comb / (comb.sum(dim=-1, keepdim=True) + eps)
+        comb = comb / (comb.sum(dim=-2, keepdim=True) + eps)
+    
+    return pre, post, comb
+
+
+# ---------------------------------------------------------------------------
+# Attention
+# ---------------------------------------------------------------------------
+
+class DeepseekV4Attention(nn.Module):
+    """Multi-head Latent Attention (MLA) with sliding window.
+    
+    V4 attention uses:
+    - Low-rank Q projection (wq_a -> q_norm -> wq_b)
+    - Direct KV projection (wkv -> kv_norm) - no kv_lora_rank
+    - Grouped low-rank O projection (wo_a -> wo_b)
+    - Sliding window attention
+    - RoPE on last qk_rope_head_dim dims
+    """
+    
+    def __init__(self, config: DeepseekV4Config, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.head_dim
+        self.qk_rope_head_dim = config.qk_rope_head_dim
+        self.nope_head_dim = config.head_dim - config.qk_rope_head_dim
+        self.q_lora_rank = config.q_lora_rank
+        self.o_groups = config.o_groups
+        self.o_lora_rank = config.o_lora_rank
+        self.scaling = config.head_dim ** -0.5
+        
+        # Q projection: low-rank
+        self.wq_a = nn.Linear(self.hidden_size, self.q_lora_rank, bias=False)
+        self.q_norm = DeepseekV4RMSNorm(self.q_lora_rank, config.rms_norm_eps)
+        self.wq_b = nn.Linear(self.q_lora_rank, self.num_heads * self.head_dim, bias=False)
+        
+        # KV projection: direct (no lora, single head)
+        self.wkv = nn.Linear(self.hidden_size, self.head_dim, bias=False)
+        self.kv_norm = DeepseekV4RMSNorm(self.head_dim, config.rms_norm_eps)
+        
+        # O projection: grouped low-rank  
+        # wo_a: [num_heads * head_dim / o_groups] -> [o_groups * o_lora_rank]
+        group_head_dim = self.num_heads * self.head_dim // self.o_groups
+        self.wo_a = nn.Linear(group_head_dim, self.o_groups * self.o_lora_rank, bias=False)
+        self.wo_b = nn.Linear(self.o_groups * self.o_lora_rank, self.hidden_size, bias=False)
+        
+        # Learnable attention sink bias
+        self.attn_sink = nn.Parameter(torch.zeros(self.num_heads))
+    
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        freqs_cis: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        """LILO-patched: honors KV cache (original ignored it and recomputed
+        the whole sequence every decode step). past_key_value = (kv, kv) with
+        kv shaped [B, 1, T, head_dim]."""
+        bsz, seqlen, _ = hidden_states.shape
+        past_len = 0
+        if past_key_value is not None and past_key_value[0] is not None:
+            past_len = int(past_key_value[0].shape[2])
+
+        # Q: low-rank projection
+        q = self.q_norm(self.wq_a(hidden_states))
+        q = self.wq_b(q)
+        q = q.view(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2)
+        # RMSNorm on q per-head
+        q = q * torch.rsqrt(q.float().pow(2).mean(-1, keepdim=True) + self.config.rms_norm_eps)
+        q = q.to(hidden_states.dtype)
+
+        # KV: direct projection (single KV head, shared across all Q heads)
+        kv = self.kv_norm(self.wkv(hidden_states))
+        kv = kv.unsqueeze(1)  # [B, 1, S, head_dim]
+
+        # Apply RoPE to last qk_rope_head_dim dims of q and kv
+        if freqs_cis is not None:
+            q_rope = q[..., -self.qk_rope_head_dim:]
+            kv_rope = kv[..., -self.qk_rope_head_dim:]
+            q_rope = apply_rotary_emb(q_rope, freqs_cis)
+            kv_rope = apply_rotary_emb(kv_rope, freqs_cis)
+            q = torch.cat([q[..., :-self.qk_rope_head_dim], q_rope], dim=-1)
+            kv = torch.cat([kv[..., :-self.qk_rope_head_dim], kv_rope], dim=-1)
+
+        # KV cache
+        if past_key_value is not None and past_key_value[0] is not None:
+            past_k, _ = past_key_value
+            kv = torch.cat([past_k, kv], dim=2)
+        new_cache = (kv, kv) if use_cache else None
+
+        total_len = kv.shape[2]
+
+        # Attention mask: causal, absolute positions (works for prefill and decode)
+        if past_len == 0 and seqlen > 1:
+            # fresh causal block: SDPA fast path
+            attn_mask = None
+            is_causal = True
+        elif seqlen == 1 and past_len == 0:
+            attn_mask = None
+            is_causal = False
+        else:
+            q_pos = torch.arange(past_len, past_len + seqlen, device=kv.device).unsqueeze(1)
+            k_pos = torch.arange(total_len, device=kv.device).unsqueeze(0)
+            allowed = q_pos >= k_pos  # [S, T]
+            if bool(allowed.all()):
+                attn_mask = None
+                is_causal = False
+            else:
+                attn_mask = torch.zeros(seqlen, total_len, device=kv.device, dtype=kv.dtype)
+                attn_mask = attn_mask.masked_fill(~allowed, float("-inf")).unsqueeze(0).unsqueeze(0)
+                is_causal = False
+
+        # Use PyTorch SDPA (fused kernel, memory-efficient)
+        # q: [B, H, S, D], kv: [B, 1, T, D] (MQA via enable_gqa)
+        # Note: attn_sink bias is small and omitted in SDPA path for speed.
+        # It's a learnable per-head scalar — its effect is minimal and the model
+        # will learn to compensate through other parameters.
+        attn_output = F.scaled_dot_product_attention(
+            q, kv, kv,
+            attn_mask=attn_mask,
+            is_causal=is_causal,
+            scale=self.scaling,
+            enable_gqa=True,
+        )
+
+        # De-rotate RoPE on output (inverse rotation = negate sin)
+        if freqs_cis is not None:
+            cos, sin = freqs_cis[0], freqs_cis[1]  # [S, D//2]
+            cos_inv = cos.unsqueeze(0).unsqueeze(0)  # [1, 1, S, D//2]
+            sin_inv = -sin.unsqueeze(0).unsqueeze(0)  # negate for inverse
+            out_rope = attn_output[..., -self.qk_rope_head_dim:]
+            d = out_rope.shape[-1] // 2
+            o1, o2 = out_rope[..., :d], out_rope[..., d:]
+            out_rope = torch.cat([o1 * cos_inv + o2 * sin_inv, o1 * (-sin_inv) + o2 * cos_inv], dim=-1)
+            attn_output = torch.cat([attn_output[..., :-self.qk_rope_head_dim], out_rope.to(attn_output.dtype)], dim=-1)
+
+        # Grouped output projection
+        attn_output = attn_output.transpose(1, 2)  # [B, S, H, D]
+        attn_output = attn_output.reshape(bsz, seqlen, self.o_groups, -1)
+
+        # wo_a applied per group: [B, S, G, H*D/G] -> [B, S, G, o_lora_rank]
+        wo_a_w = self._wo_a_weight().view(self.o_groups, self.o_lora_rank, -1)
+        attn_output = torch.einsum("bsgd,grd->bsgr", attn_output, wo_a_w)
+        attn_output = attn_output.flatten(2)  # [B, S, G*o_lora_rank]
+        attn_output = self.wo_b(attn_output)
+
+        return attn_output, new_cache
+
+    def _wo_a_weight(self) -> torch.Tensor:
+        if isinstance(self.wo_a, QuantizedLinear):
+            return self.wo_a.dequantized_weight()
+        return self.wo_a.weight
+
+
+
+# ---------------------------------------------------------------------------
+# MoE
+# ---------------------------------------------------------------------------
+
+class DeepseekV4Expert(nn.Module):
+    """Single MoE expert with SwiGLU activation."""
+    
+    def __init__(self, hidden_size: int, intermediate_size: int, swiglu_limit: float = 0.0):
+        super().__init__()
+        self.w1 = nn.Linear(hidden_size, intermediate_size, bias=False)  # gate
+        self.w2 = nn.Linear(intermediate_size, hidden_size, bias=False)  # down
+        self.w3 = nn.Linear(hidden_size, intermediate_size, bias=False)  # up
+        self.swiglu_limit = swiglu_limit
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate = self.w1(x).float()
+        up = self.w3(x).float()
+        if self.swiglu_limit > 0:
+            up = up.clamp(-self.swiglu_limit, self.swiglu_limit)
+            gate = gate.clamp(max=self.swiglu_limit)
+        x = F.silu(gate) * up
+        w2_is_q = isinstance(self.w2, QuantizedLinear)
+        return self.w2(x.to(self.w2.dequantized_weight().dtype if w2_is_q else self.w2.weight.dtype))
+
+
+class DeepseekV4Gate(nn.Module):
+    """MoE gating with sqrtsoftplus scoring."""
+    
+    def __init__(self, config: DeepseekV4Config, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.topk = config.num_experts_per_tok
+        self.scoring_func = config.scoring_func
+        self.route_scale = config.routed_scaling_factor
+        self.is_hash_layer = layer_idx < config.num_hash_layers
+        
+        self.weight = nn.Parameter(torch.empty(config.n_routed_experts, config.hidden_size))
+        if not self.is_hash_layer:
+            self.bias = nn.Parameter(torch.zeros(config.n_routed_experts))
+        else:
+            self.register_parameter("bias", None)
+    
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        scores = F.linear(x.float(), self.weight.float())
+        
+        if self.scoring_func == "softmax":
+            scores = scores.softmax(dim=-1)
+        elif self.scoring_func == "sigmoid":
+            scores = scores.sigmoid()
+        elif self.scoring_func == "sqrtsoftplus":
+            scores = F.softplus(scores).sqrt()
+        
+        original_scores = scores
+        
+        if self.bias is not None:
+            scores = scores + self.bias
+        
+        # Top-k selection
+        indices = scores.topk(self.topk, dim=-1)[1]
+        weights = original_scores.gather(1, indices)
+        
+        if self.scoring_func != "softmax":
+            weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-20)
+        
+        weights = weights * self.route_scale
+        return weights.to(x.dtype), indices
+
+
+class DeepseekV4MoE(nn.Module):
+    """Mixture of Experts layer."""
+    
+    def __init__(self, config: DeepseekV4Config, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.n_routed_experts = config.n_routed_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+        
+        self.gate = DeepseekV4Gate(config, layer_idx)
+        self.experts = nn.ModuleList([
+            DeepseekV4Expert(config.hidden_size, config.moe_intermediate_size, config.swiglu_limit)
+            for _ in range(config.n_routed_experts)
+        ])
+        self.shared_expert = DeepseekV4Expert(config.hidden_size, config.moe_intermediate_size)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = x.shape
+        x_flat = x.view(-1, self.hidden_size)
+        
+        weights, indices = self.gate(x_flat)
+        
+        y = torch.zeros_like(x_flat, dtype=torch.float32)
+        
+        # Route tokens to experts
+        counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts)
+        for i in range(self.n_routed_experts):
+            if counts[i] == 0:
+                continue
+            idx, top = torch.where(indices == i)
+            expert_out = self.experts[i](x_flat[idx])
+            y[idx] += (weights[idx, top].unsqueeze(-1) * expert_out.float())
+        
+        # Add shared expert
+        y = y + self.shared_expert(x_flat).float()
+        
+        return y.to(x.dtype).view(shape)
+
+
+# ---------------------------------------------------------------------------
+# Transformer Block
+# ---------------------------------------------------------------------------
+
+class DeepseekV4Block(nn.Module):
+    """Transformer block with Hyper-Connections.
+    
+    Instead of simple residuals, HC maintains hc_mult copies of the hidden state.
+    hc_pre: reduces hc copies -> 1 via learned weighted sum.
+    hc_post: expands 1 -> hc copies via learned post-weights + combination matrix.
+    """
+    
+    def __init__(self, config: DeepseekV4Config, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hc_mult = config.hc_mult
+        self.norm_eps = config.rms_norm_eps
+        self.hc_eps = config.hc_eps
+        self.hc_sinkhorn_iters = config.hc_sinkhorn_iters
+        
+        self.attn = DeepseekV4Attention(config, layer_idx)
+        self.ffn = DeepseekV4MoE(config, layer_idx)
+        self.attn_norm = DeepseekV4RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.ffn_norm = DeepseekV4RMSNorm(config.hidden_size, config.rms_norm_eps)
+        
+        # HC parameters for attention and FFN sub-layers
+        mix_hc = (2 + config.hc_mult) * config.hc_mult
+        hc_dim = config.hc_mult * config.hidden_size
+        
+        self.hc_attn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim))
+        self.hc_ffn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim))
+        self.hc_attn_base = nn.Parameter(torch.empty(mix_hc))
+        self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc))
+        self.hc_attn_scale = nn.Parameter(torch.empty(3))
+        self.hc_ffn_scale = nn.Parameter(torch.empty(3))
+    
+    def hc_pre(self, x, hc_fn, hc_scale, hc_base):
+        """Reduce hc_mult copies to 1 via learned weighted sum.
+        
+        x: [B, S, hc_mult, D]
+        Returns: y [B, S, D], post [B, S, hc_mult], comb [B, S, hc_mult, hc_mult]
+        """
+        shape = x.size()
+        dtype = x.dtype
+        x_flat = x.flatten(2).float()  # [B, S, hc_mult*D]
+        
+        rsqrt = torch.rsqrt(x_flat.pow(2).mean(-1, keepdim=True) + self.norm_eps)
+        mixes = F.linear(x_flat, hc_fn.float()) * rsqrt  # [B, S, mix_hc]
+        
+        pre, post, comb = hc_split_sinkhorn(
+            mixes, hc_scale, hc_base,
+            self.hc_mult, self.hc_sinkhorn_iters, self.hc_eps
+        )
+        
+        # Weighted sum: pre [B, S, hc] * x [B, S, hc, D] -> y [B, S, D]
+        y = (pre.unsqueeze(-1) * x.float()).sum(dim=2)
+        return y.to(dtype), post, comb
+    
+    def hc_post(self, x, residual, post, comb):
+        """Expand 1 -> hc_mult copies.
+        
+        x: [B, S, D] - output from sub-layer
+        residual: [B, S, hc_mult, D] - input HC state
+        post: [B, S, hc_mult]
+        comb: [B, S, hc_mult, hc_mult]
+        """
+        # post * x + comb * residual
+        y = (post.unsqueeze(-1) * x.unsqueeze(2).float() + 
+             torch.einsum("bsij,bsjd->bsid", comb.float(), residual.float()))
+        return y.to(x.dtype)
+    
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        freqs_cis: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        """
+        x: [B, S, hc_mult, D] - HC state
+        """
+        # Attention with HC
+        residual = x
+        y, post, comb = self.hc_pre(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base)
+        y = self.attn_norm(y)
+        y, new_cache = self.attn(y, attention_mask=attention_mask, position_ids=position_ids,
+                                  freqs_cis=freqs_cis, past_key_value=past_key_value, use_cache=use_cache)
+        x = self.hc_post(y, residual, post, comb)
+        
+        # FFN with HC
+        residual = x
+        y, post, comb = self.hc_pre(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base)
+        y = self.ffn_norm(y)
+        y = self.ffn(y)
+        x = self.hc_post(y, residual, post, comb)
+        
+        return x, new_cache
+
+
+# ---------------------------------------------------------------------------
+# Full Model
+# ---------------------------------------------------------------------------
+
+class DeepseekV4PreTrainedModel(PreTrainedModel):
+    config_class = DeepseekV4Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["DeepseekV4Block"]
+    _skip_keys_device_placement = ["past_key_values"]
+    
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+        elif isinstance(module, DeepseekV4RMSNorm):
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, DeepseekV4Gate):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, DeepseekV4Block):
+            # Initialize HC parameters
+            nn.init.normal_(module.hc_attn_fn, std=0.01)
+            nn.init.normal_(module.hc_ffn_fn, std=0.01)
+            nn.init.zeros_(module.hc_attn_base)
+            nn.init.zeros_(module.hc_ffn_base)
+            nn.init.ones_(module.hc_attn_scale)
+            nn.init.ones_(module.hc_ffn_scale)
+        elif isinstance(module, DeepseekV4Attention):
+            nn.init.zeros_(module.attn_sink)
+
+
+class DeepseekV4Model(DeepseekV4PreTrainedModel):
+    def __init__(self, config: DeepseekV4Config):
+        super().__init__(config)
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([
+            DeepseekV4Block(config, layer_idx)
+            for layer_idx in range(config.num_hidden_layers)
+        ])
+        self.norm = DeepseekV4RMSNorm(config.hidden_size, config.rms_norm_eps)
+        
+        # HC head parameters (for contracting hc_mult -> 1 at output)
+        hc_dim = config.hc_mult * config.hidden_size
+        self.hc_head_fn = nn.Parameter(torch.empty(config.hc_mult, hc_dim))
+        self.hc_head_base = nn.Parameter(torch.empty(config.hc_mult))
+        self.hc_head_scale = nn.Parameter(torch.empty(1))
+        
+        # Precomputed RoPE frequencies (non-persistent so it never enters
+        # checkpoints; recomputed on demand — see get_freqs_cis).
+        self.register_buffer(
+            "freqs_cis",
+            precompute_freqs_cis(config.qk_rope_head_dim, config.max_position_embeddings, config.rope_theta),
+            persistent=False,
+        )
+
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_freqs_cis(self) -> torch.Tensor:
+        """Return cos/sin table [2, max_pos, D//2], recomputing if the buffer
+        was left as uninitialized memory by meta-device loading (transformers>=5)."""
+        buf = getattr(self, "freqs_cis", None)
+        ok = (
+            isinstance(buf, torch.Tensor)
+            and buf.dim() == 3
+            and int(buf.shape[1]) >= int(self.config.max_position_embeddings)
+            and bool(torch.isfinite(buf).all())
+            and float(buf.abs().max()) <= 2.0
+        )
+        if not ok:
+            buf = precompute_freqs_cis(
+                self.config.qk_rope_head_dim,
+                self.config.max_position_embeddings,
+                self.config.rope_theta,
+            ).to(self.device if hasattr(self, "device") else "cpu")
+            self.freqs_cis = buf
+        return self.freqs_cis
+
+    
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        # HC head initialization
+        if module is self:
+            nn.init.normal_(self.hc_head_fn, std=0.01)
+            nn.init.zeros_(self.hc_head_base)
+            nn.init.ones_(self.hc_head_scale)
+    
+    def hc_head(self, x):
+        """Contract hc_mult copies to 1 for final output.
+        
+        x: [B, S, hc_mult, D] -> [B, S, D]
+        """
+        shape = x.size()
+        dtype = x.dtype
+        x_flat = x.flatten(2).float()  # [B, S, hc_mult*D]
+        
+        rsqrt = torch.rsqrt(x_flat.pow(2).mean(-1, keepdim=True) + self.config.rms_norm_eps)
+        mixes = F.linear(x_flat, self.hc_head_fn.float()) * rsqrt  # [B, S, hc_mult]
+        
+        pre = torch.sigmoid(mixes * self.hc_head_scale.float() + self.hc_head_base.float()) + self.config.hc_eps
+        y = (pre.unsqueeze(-1) * x.float()).sum(dim=2)
+        return y.to(dtype)
+    
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> BaseModelOutputWithPast:
+        return_dict = return_dict if return_dict is not None else True  # LILO: return_dict safe default (transformers 5.x)
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("Cannot specify both input_ids and inputs_embeds")
+        
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        
+        bsz, seqlen = inputs_embeds.shape[:2]
+
+        # LILO: honor KV cache (original hardcoded use_cache=False which
+        # recomputed the full sequence on every decode step)
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        if use_cache and past_key_values is not None and len(past_key_values) > 0 and past_key_values[0] is not None:
+            past_len = int(past_key_values[0][0].shape[2])
+        else:
+            past_key_values = None
+            past_len = 0
+
+        if position_ids is None:
+            position_ids = torch.arange(past_len, past_len + seqlen, device=inputs_embeds.device).unsqueeze(0)
+
+        # Get freqs for RoPE
+        # freqs_cis is [2, max_seq, D//2], index by position
+        pos = position_ids.squeeze(0)
+        freqs_cis = self.get_freqs_cis()[:, pos].to(inputs_embeds.device)  # [2, seqlen, D//2]
+        
+        # The attention layer builds its own causal mask from absolute
+        # positions (KV-cache aware), so we do not pass a mask down.
+        causal_mask = None
+
+        
+        # Expand to hc_mult copies
+        hidden_states = inputs_embeds.unsqueeze(2).expand(-1, -1, self.config.hc_mult, -1)
+        hidden_states = hidden_states.contiguous()
+        
+        new_past_key_values = [] if use_cache else None
+        
+        for i, layer in enumerate(self.layers):
+            past_kv = past_key_values[i] if past_key_values is not None and i < len(past_key_values) else None
+            
+            if self.gradient_checkpointing and self.training:
+                hidden_states, new_cache = torch.utils.checkpoint.checkpoint(
+                    layer, hidden_states, causal_mask, position_ids, freqs_cis, past_kv, use_cache,
+                    use_reentrant=False,
+                )
+            else:
+                hidden_states, new_cache = layer(
+                    hidden_states, attention_mask=causal_mask, position_ids=position_ids,
+                    freqs_cis=freqs_cis, past_key_value=past_kv, use_cache=use_cache,
+                )
+            
+            if use_cache:
+                new_past_key_values.append(new_cache)
+        
+        # Contract HC copies -> single hidden state
+        hidden_states = self.hc_head(hidden_states)
+        hidden_states = self.norm(hidden_states)
+        
+        if not return_dict:
+            return (hidden_states, new_past_key_values)
+        
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=new_past_key_values,
+        )
+
+
+class DeepseekV4ForCausalLM(DeepseekV4PreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    
+    def __init__(self, config: DeepseekV4Config):
+        super().__init__(config)
+        self.model = DeepseekV4Model(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+    
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+    
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+    
+    def get_output_embeddings(self):
+        return self.lm_head
+    
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+    
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        return_dict = return_dict if return_dict is not None else True  # LILO: return_dict safe default (transformers 5.x)
+        
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_hidden_states=output_hidden_states,
+            return_dict=False,  # always tuple for compile compatibility
+        )
+        
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+        
+        past_kv = outputs[1] if len(outputs) > 1 else None
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=past_kv,
+        )
+    
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+        
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": True,
+        }

run_cpu.sh

@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+# LILO CPU runner — memory-safe: never touches /tmp (tmpfs), bounded threads.
+set -euo pipefail
+cd "$(dirname "$0")"
+
+export TMPDIR="$PWD/.tmp"          # MUST be disk-backed, not tmpfs
+export HF_HOME="${HF_HOME:-$PWD/.hf}"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-2}"
+export MALLOC_ARENA_MAX=1          # keep glibc from reserving RAM per thread
+export TOKENIZERS_PARALLELISM=false
+
+mkdir -p "$TMPDIR" "$HF_HOME"
+
+PY="${PYTHON:-python3}"
+if [ -x "$PWD/.venv/bin/python" ]; then
+    PY="$PWD/.venv/bin/python"
+fi
+
+exec "$PY" chat_lilo.py "$@"

tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<｜begin▁of▁sentence｜>",
+  "eos_token": "<｜end▁of▁sentence｜>",
+  "is_local": true,
+  "local_files_only": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<｜end▁of▁sentence｜>",
+  "tokenizer_class": "TokenizersBackend"
+}