Upload Summer-0.5B-S0

.gitattributes

@@ -1,35 +1,4 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,75 @@
+---
+license: apache-2.0
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- qwen3
+- from-scratch
+- pretraining
+- bilingual
+- piece-tokenizer
+language:
+- en
+- zh
+---
+
+# Summer-0.5B-S0
+
+Summer-0.5B-S0 is a **from-scratch** bilingual (Chinese/English) base model:
+524,336,128 parameters, trained from random initialization on
+**12B tokens** with a self-trained 81,903-piece tokenizer.
+
+It is *not* a fine-tune or a tokenizer-swap of any existing model. The
+architecture follows `Qwen/Qwen3-0.6B-Base` (28 layers / hidden 1024 /
+GQA 16:8 / head_dim 128 / tied embeddings / RoPE theta 1e6), but every weight
+starts from `N(0, 0.02)`.
+
+## Stage
+
+**S0** — monolingual only. 12B tokens, Chinese/English 50:50, 45,149 steps.
+No parallel or instruction data at any point.
+
+## What to expect
+
+**This is a 13B-token model.** For scale: `Qwen3-0.6B-Base` saw 36T tokens —
+about 2,700x more. Treat the numbers below as what that budget buys, not as a
+competitive result.
+
+| WMT22 5-shot | BLEU | COMET |
+|---|---|---|
+| zh->en | 0.54 | 0.4638 |
+| en->zh | 3.97 | 0.5872 |
+
+**Few-shot translation is essentially zero, and that is the finding.** The
+model ignores the in-context examples entirely — on zh->en it does not even
+switch output language. Language modelling was learned; in-context learning
+was not. 12B monolingual tokens is not enough for ICL to emerge at 0.5B.
+
+Its value is as (a) the control for S1, and (b) a starting point for
+mid-training / SFT.
+
+## Tokenizer
+
+The tokenizer is a compiled C++ extension, **not** loadable by
+`AutoTokenizer`. The release ships `tokenizer.py` and `example_load.py`:
+
+```bash
+pip install git+https://github.com/Ismantic/PieceTokenizer
+python example_load.py
+```
+
+The model code (`model.py`, `checkpoint.py`) is bundled too — the package
+depends only on `torch` plus the tokenizer extension, not on `transformers`.
+
+## Training
+
+Full pipeline, data mixes and every design decision (including the mistakes)
+are documented in <https://github.com/Ismantic/Summer>. Notably
+`docs/WHY.md` records why fp32 master weights are mandatory, why the learning
+rate schedule is WSD rather than cosine, and what the vocabulary swap cost.
+
+## License
+
+Apache-2.0. Training corpora are public datasets (FineWeb-Edu, Cosmopedia,
+CCI3-HQ, SkyPile, WMT19, OPUS-100 and others; see `data/source.py` upstream).
+Please observe their respective licenses.

Summer-Tokenizer.pt

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

checkpoint.py

@@ -0,0 +1,123 @@
+"""safetensors 读写。**只依赖 torch。**
+
+格式很简单,自己实现比拉一个依赖划算:
+
+    [0:8]      u64 小端,头部 JSON 的字节数 N
+    [8:8+N]    JSON:{名字: {dtype, shape, data_offsets:[起, 止]}, "__metadata__": {...}}
+    [8+N:]     张量数据区,data_offsets 是相对这里的偏移
+
+## 为什么不用 safetensors 库
+
+`src/` 只依赖 torch —— 这条约束的意义是让读者能看见每一步在做什么。
+safetensors 的格式本身只有二十来行代码,藏进依赖里反而看不见了。
+
+## 不能改错的地方
+
+**读进来的 key 必须与 HF 的 Qwen3ForCausalLM 完全一致。** 已发布的
+`Ismantic/Qwen3-1.7B-Base-ReTok` 是标准 HF 命名,key 对不上不会报错 ——
+`load_state_dict(strict=False)` 会静默跳过,模型照样随机初始化跑起来、
+loss 照样降,只是那份权重白搭了。所以这里默认 `strict=True`。
+
+**tie_word_embeddings 的模型没有 `lm_head.weight`。** v18 的 checkpoint 就是
+这样:310 个张量,只有 `model.embed_tokens.weight` 和 `model.norm.weight` 加
+28 层 ×11。加载时要把 lm_head 绑回 embed,不能指望文件里有。
+"""
+import json
+import struct
+from pathlib import Path
+
+import torch
+
+# safetensors 的 dtype 名 → torch dtype
+_DTYPES = {
+    "BF16": torch.bfloat16, "F16": torch.float16, "F32": torch.float32,
+    "F64": torch.float64, "I64": torch.int64, "I32": torch.int32,
+    "I16": torch.int16, "I8": torch.int8, "U8": torch.uint8,
+    "BOOL": torch.bool,
+}
+_NAMES = {v: k for k, v in _DTYPES.items()}
+
+
+def read_header(path) -> tuple[dict, int]:
+    """返回 (头部 dict, 数据区起始偏移)。不读张量,几微秒。"""
+    with open(path, "rb") as f:
+        (n,) = struct.unpack("<Q", f.read(8))
+        hdr = json.loads(f.read(n))
+    return hdr, 8 + n
+
+
+def load_safetensors(path, device="cpu") -> dict[str, torch.Tensor]:
+    """读成 {名字: 张量}。
+
+    按 data_offsets 顺序 seek+read,每个张量一份可写内存。
+    不用 mmap:`torch.frombuffer` 不接受只读缓冲区,套一层 `bytearray` 反正
+    要拷一份,mmap 的零拷贝就没了 —— 还得处理 memoryview 挂着导出引用、
+    `mm.close()` 抛 BufferError 的麻烦。
+    """
+    path = Path(path)
+    hdr, base = read_header(path)
+    hdr.pop("__metadata__", None)
+
+    # 按文件内偏移排序后顺序读,避免来回 seek
+    items = sorted(hdr.items(), key=lambda kv: kv[1]["data_offsets"][0])
+    out = {}
+    with open(path, "rb") as f:
+        for name, spec in items:
+            lo, hi = spec["data_offsets"]
+            f.seek(base + lo)
+            buf = bytearray(hi - lo)
+            if f.readinto(buf) != hi - lo:
+                raise EOFError(f"{path}: 读 {name} 时文件提前结束")
+            t = torch.frombuffer(buf, dtype=_DTYPES[spec["dtype"]])
+            out[name] = t.reshape(spec["shape"]).to(device)
+    return out
+
+
+def save_safetensors(tensors: dict[str, torch.Tensor], path, metadata=None) -> None:
+    """写出。张量按名字排序,同样的输入产生同样的字节。
+
+    共享存储的张量(比如 tie 的 embed / lm_head)要**调用方自己去重** ——
+    这里不猜。写两份不会报错,只是文件大一倍、加载时可能对不上 tie 的预期。
+    """
+    hdr, blobs, off = {}, [], 0
+    for name in sorted(tensors):
+        t = tensors[name].detach().cpu().contiguous()
+        if t.dtype not in _NAMES:
+            raise TypeError(f"{name}: 不支持的 dtype {t.dtype}")
+        b = t.view(torch.uint8).reshape(-1).numpy().tobytes()
+        hdr[name] = {"dtype": _NAMES[t.dtype], "shape": list(t.shape),
+                     "data_offsets": [off, off + len(b)]}
+        blobs.append(b)
+        off += len(b)
+    if metadata:
+        hdr["__metadata__"] = {str(k): str(v) for k, v in metadata.items()}
+
+    raw = json.dumps(hdr, separators=(",", ":")).encode("utf-8")
+    # 头部按 8 字节对齐,张量数据区才能对齐
+    pad = (-len(raw)) % 8
+    raw += b" " * pad
+
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "wb") as f:
+        f.write(struct.pack("<Q", len(raw)))
+        f.write(raw)
+        for b in blobs:
+            f.write(b)
+
+
+def load_sharded(model_dir, device="cpu") -> dict[str, torch.Tensor]:
+    """加载一个模型目录:单文件 `model.safetensors`,或分片 + index.json。"""
+    model_dir = Path(model_dir)
+    single = model_dir / "model.safetensors"
+    if single.exists():
+        return load_safetensors(single, device)
+
+    index = model_dir / "model.safetensors.index.json"
+    if index.exists():
+        shards = sorted(set(json.loads(index.read_text())["weight_map"].values()))
+        out = {}
+        for s in shards:
+            out.update(load_safetensors(model_dir / s, device))
+        return out
+    raise FileNotFoundError(f"{model_dir} 下既没有 model.safetensors 也没有 index.json")

config.json

@@ -0,0 +1,19 @@
+{
+  "vocab_size": 81903,
+  "hidden_size": 1024,
+  "intermediate_size": 3072,
+  "num_hidden_layers": 28,
+  "num_attention_heads": 16,
+  "num_key_value_heads": 8,
+  "head_dim": 128,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000,
+  "max_position_embeddings": 32768,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "initializer_range": 0.02,
+  "model_type": "qwen3",
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ]
+}

example_load.py

@@ -0,0 +1,34 @@
+"""跑通这个发布包 —— 只需要 torch 和 piece_tokenizer。
+
+    pip install torch
+    pip install git+https://github.com/Ismantic/PieceTokenizer
+    python example_load.py
+"""
+import torch
+
+from model import Qwen3ForCausalLM
+from tokenizer import PieceTokenizerWrapper
+
+HERE = "."
+
+tok = PieceTokenizerWrapper(HERE)
+model = Qwen3ForCausalLM.from_pretrained(
+    HERE, device="cuda" if torch.cuda.is_available() else "cpu",
+    dtype=torch.bfloat16)
+
+prompt = "机器翻译的基本任务是"
+ids = tok.encode(prompt, add_special_tokens=False)
+print(f"{prompt!r} -> {len(ids)} tokens")
+
+# 贪心续写 40 步。没有 KV cache —— 每步重算前缀,短续写够用。
+x = torch.tensor([ids], device=next(model.parameters()).device)
+out = []
+with torch.no_grad():
+    for _ in range(40):
+        nxt = int(model(x)[0, -1].argmax())
+        if nxt == tok.eos_token_id:
+            break
+        out.append(nxt)
+        x = torch.cat([x, torch.tensor([[nxt]], device=x.device)], dim=1)
+
+print(prompt + tok.decode(out))

example_vllm.py

@@ -0,0 +1,46 @@
+"""用 vLLM 跑这个模型。
+
+    pip install vllm
+    pip install git+https://github.com/Ismantic/PieceTokenizer
+    python example_vllm.py
+
+## 为什么要 skip_tokenizer_init
+
+vLLM 能加载**权重** —— config.json 里 architectures 是 Qwen3ForCausalLM,
+它用自己那份 Qwen3 实现,按 state_dict 的 key 名灌进去。
+
+但它**认不了这个词表**:81903 的 piece 词表不是 HF 标准格式,AutoTokenizer
+走不通。所以要 skip_tokenizer_init=True,自己编码好 token id 用 TokensPrompt
+喂进去,拿回来的 id 自己解码。
+
+推论:`vllm serve` 那种 OpenAI 兼容服务**开箱用不了** —— 它要把文本转成
+token,而它转不了。调用方得自己传 token id。
+"""
+import torch
+from vllm import LLM, SamplingParams
+from vllm.inputs import TokensPrompt
+
+from tokenizer import PieceTokenizerWrapper
+
+HERE = "."
+
+tok = PieceTokenizerWrapper(HERE)
+llm = LLM(model=HERE, skip_tokenizer_init=True, dtype="bfloat16",
+          gpu_memory_utilization=0.85, max_model_len=4096,
+          trust_remote_code=True)
+
+prompts = [
+    "机器翻译的基本任务是",
+    "中文分词的目标是",
+]
+
+outputs = llm.generate(
+    [TokensPrompt(prompt_token_ids=tok.encode(p, add_special_tokens=False))
+     for p in prompts],
+    SamplingParams(temperature=0.0, max_tokens=64,
+                   stop_token_ids=[tok.eos_token_id]),
+)
+
+for prompt, out in zip(prompts, outputs):
+    print(f"\n{prompt}", end="")
+    print(tok.decode(list(out.outputs[0].token_ids)))

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 81899,
+  "max_new_tokens": 2048
+}

model.py

@@ -0,0 +1,372 @@
+"""Qwen3 稠密模型,纯 torch 实现。**只依赖 torch。**
+
+替掉 `transformers.Qwen3ForCausalLM`。规模参考(v18 用的 1.7B):
+
+    28 层 / hidden 2048 / intermediate 6144 / 16 个 q 头 / 8 个 kv 头(GQA)
+    head_dim 128 / RMSNorm eps 1e-6 / RoPE θ=1e6 / SwiGLU / 全程无 bias
+    词表 81903 / 输入输出嵌入绑权重
+
+## 不能改错的地方
+
+**state_dict 的 key 必须与 HF 的 `Qwen3ForCausalLM` 完全一致。** 已发布的
+`Ismantic/Qwen3-1.7B-Base-ReTok` 用的是标准 HF 命名;改模块名或嵌套层级会让
+那份权重全部失配,**而模型照样能随机初始化跑起来、loss 照样降、不报错**。
+`test/test_model_equiv.py` 是这条的防线:同权重、与 transformers 逐层对齐。
+
+对应关系(每层 11 个张量,共 28×11+2=310):
+
+    model.embed_tokens.weight                      → embed_tokens
+    model.layers.{i}.input_layernorm.weight        → layers[i].input_layernorm
+    model.layers.{i}.self_attn.{q,k,v,o}_proj.weight
+    model.layers.{i}.self_attn.{q,k}_norm.weight   ← Qwen3 特有,逐头 RMSNorm
+    model.layers.{i}.post_attention_layernorm.weight
+    model.layers.{i}.mlp.{gate,up,down}_proj.weight
+    model.norm.weight                              → norm
+
+**没有 `lm_head.weight`。** tie_word_embeddings=true,权重与 embed 共享。
+
+**RoPE 用 half-split(NeoX)约定,不是 interleave。** 两种都"能跑",但结果
+完全不同 —— 用错了 loss 会明显偏高而不会报错。见 `_rotate_half`。
+
+**RMSNorm 的中间计算在 float32。** 直接用 bf16 算方差会有肉眼可见的偏差,
+同样不报错。
+"""
+import json
+import math
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+
+def _checkpoint_mod():
+    """`src.checkpoint`,或发布包里平铺的 `checkpoint`。
+
+    发布包(save/releases/<repo>/)是平铺的 —— model.py 和 checkpoint.py 并列,
+    没有包结构,相对 import 会 ImportError。所以两种都试。
+    """
+    try:
+        from . import checkpoint as _c            # 仓库内:src 是 package
+    except ImportError:                            # 发布包内:平铺
+        import checkpoint as _c                    # noqa: I001
+    return _c
+
+
+@dataclass
+class Qwen3Config:
+    vocab_size: int = 81903
+    hidden_size: int = 2048
+    intermediate_size: int = 6144
+    num_hidden_layers: int = 28
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 8
+    head_dim: int = 128
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 1_000_000.0
+    max_position_embeddings: int = 32768
+    tie_word_embeddings: bool = True
+    torch_dtype: str = "bfloat16"
+    # 随机初始化的标准差。只有从零训练(init_weights)会读它;从 checkpoint
+    # 加载时权重来自文件,这个值不参与。与 Qwen3 官方 config 的同名字段一致。
+    initializer_range: float = 0.02
+
+    @classmethod
+    def from_json(cls, path) -> "Qwen3Config":
+        raw = json.loads(Path(path).read_text())
+        if raw.get("model_type") not in (None, "qwen3"):
+            raise ValueError(f"不是 qwen3 架构:model_type={raw.get('model_type')}")
+        # head_dim 在有些 config 里缺省,按 hidden/heads 推
+        raw.setdefault("head_dim",
+                       raw["hidden_size"] // raw["num_attention_heads"])
+        raw.setdefault("torch_dtype", raw.get("dtype", "bfloat16"))
+        fields = {f for f in cls.__dataclass_fields__}
+        return cls(**{k: v for k, v in raw.items() if k in fields})
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return {"bfloat16": torch.bfloat16, "float16": torch.float16,
+                "float32": torch.float32}[self.torch_dtype]
+
+
+class RMSNorm(nn.Module):
+    """与 `Qwen3RMSNorm` 逐位对齐:方差在 float32 里算,再转回原 dtype 乘权重。"""
+
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dt = x.dtype
+        x = x.to(torch.float32)
+        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return self.weight * x.to(dt)
+
+
+class RotaryEmbedding(nn.Module):
+    """RoPE。inv_freq 存成 buffer(非持久化 —— 它是算出来的,不进 state_dict)。"""
+
+    def __init__(self, head_dim: int, theta: float):
+        super().__init__()
+        inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32)
+                               / head_dim))
+        self.register_buffer("inv_freq", inv, persistent=False)
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor, position_ids: torch.Tensor):
+        # position_ids [B,T] → freqs [B,T,head_dim/2] → emb [B,T,head_dim]
+        inv = self.inv_freq[None, :, None].to(x.device).float()      # [1,D/2,1]
+        pos = position_ids[:, None, :].float()                       # [B,1,T]
+
+        # **必须强制 float32,不能让 autocast 把这个 matmul 降到 bf16。**
+        # bf16 尾数只有 8 位,整数超过 256 就不精确:257→256、1023→1024。
+        # 位置编码一错,长序列上的注意力全歪 —— 而且不报错,短序列还测不出来。
+        # 实测代价:seq_len=1024 上 next-token loss 从 2.3331 涨到 2.7605。
+        # HF 的 Qwen3RotaryEmbedding 里也有这条防护(注释写着 Force float32)。
+        dev = x.device.type if x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=dev, enabled=False):
+            freqs = (inv.float() @ pos.float()).transpose(1, 2)      # [B,T,D/2]
+            emb = torch.cat((freqs, freqs), dim=-1)                  # [B,T,D]
+            cos, sin = emb.cos(), emb.sin()
+        return cos.to(x.dtype), sin.to(x.dtype)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """half-split(NeoX)约定:后半取负拼到前面。
+
+    **不是 interleave(GPT-J)约定。** 两种都能跑通、都不报错,但结果不同。
+    HF 的 Qwen3 用的是这个。
+    """
+    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rope(q, k, cos, sin):
+    cos, sin = cos.unsqueeze(1), sin.unsqueeze(1)                    # [B,1,T,D]
+    return (q * cos + _rotate_half(q) * sin,
+            k * cos + _rotate_half(k) * sin)
+
+
+def _repeat_kv(x: torch.Tensor, n: int) -> torch.Tensor:
+    """GQA:把 kv 头复制 n 份对齐 q 头数。n=1 时直接返回。"""
+    if n == 1:
+        return x
+    b, h, t, d = x.shape
+    return x[:, :, None].expand(b, h, n, t, d).reshape(b, h * n, t, d)
+
+
+class Attention(nn.Module):
+    def __init__(self, cfg: Qwen3Config):
+        super().__init__()
+        self.n_head = cfg.num_attention_heads
+        self.n_kv = cfg.num_key_value_heads
+        self.head_dim = cfg.head_dim
+        self.n_rep = self.n_head // self.n_kv
+        self.scaling = self.head_dim ** -0.5
+
+        h, d = cfg.hidden_size, cfg.head_dim
+        self.q_proj = nn.Linear(h, self.n_head * d, bias=False)
+        self.k_proj = nn.Linear(h, self.n_kv * d, bias=False)
+        self.v_proj = nn.Linear(h, self.n_kv * d, bias=False)
+        self.o_proj = nn.Linear(self.n_head * d, h, bias=False)
+        # Qwen3 特有:对每个头的 d 维做 RMSNorm,在 RoPE 之前
+        self.q_norm = RMSNorm(d, cfg.rms_norm_eps)
+        self.k_norm = RMSNorm(d, cfg.rms_norm_eps)
+
+    def forward(self, x, cos, sin):
+        b, t, _ = x.shape
+        # 顺序要紧:proj → view 成 [B,T,H,D] → 逐头 norm → transpose → RoPE
+        q = self.q_norm(self.q_proj(x).view(b, t, self.n_head, self.head_dim)).transpose(1, 2)
+        k = self.k_norm(self.k_proj(x).view(b, t, self.n_kv, self.head_dim)).transpose(1, 2)
+        v = self.v_proj(x).view(b, t, self.n_kv, self.head_dim).transpose(1, 2)
+
+        q, k = _apply_rope(q, k, cos, sin)
+        k, v = _repeat_kv(k, self.n_rep), _repeat_kv(v, self.n_rep)
+
+        o = F.scaled_dot_product_attention(q, k, v, is_causal=True,
+                                           scale=self.scaling)
+        return self.o_proj(o.transpose(1, 2).reshape(b, t, -1))
+
+
+class MLP(nn.Module):
+    def __init__(self, cfg: Qwen3Config):
+        super().__init__()
+        h, i = cfg.hidden_size, cfg.intermediate_size
+        self.gate_proj = nn.Linear(h, i, bias=False)
+        self.up_proj = nn.Linear(h, i, bias=False)
+        self.down_proj = nn.Linear(i, h, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, cfg: Qwen3Config):
+        super().__init__()
+        self.self_attn = Attention(cfg)
+        self.mlp = MLP(cfg)
+        self.input_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+
+    def forward(self, x, cos, sin):
+        x = x + self.self_attn(self.input_layernorm(x), cos, sin)
+        return x + self.mlp(self.post_attention_layernorm(x))
+
+
+class Qwen3Model(nn.Module):
+    """对应 HF 的 `model.*` 那一层前缀。"""
+
+    def __init__(self, cfg: Qwen3Config):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
+        self.layers = nn.ModuleList(DecoderLayer(cfg) for _ in range(cfg.num_hidden_layers))
+        self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.rotary_emb = RotaryEmbedding(cfg.head_dim, cfg.rope_theta)
+        self.gradient_checkpointing = False
+
+    def forward(self, input_ids, position_ids=None):
+        x = self.embed_tokens(input_ids)
+        if position_ids is None:
+            position_ids = torch.arange(input_ids.shape[1],
+                                        device=input_ids.device)[None]
+        cos, sin = self.rotary_emb(x, position_ids)
+        for layer in self.layers:
+            if self.gradient_checkpointing and self.training:
+                # use_reentrant=False:新式实现,能和 requires_grad=False 的
+                # 冻结参数共存(Phase 1 冻结 transformer 时正是这个场景)
+                x = torch.utils.checkpoint.checkpoint(
+                    layer, x, cos, sin, use_reentrant=False)
+            else:
+                x = layer(x, cos, sin)
+        return self.norm(x)
+
+
+class Qwen3ForCausalLM(nn.Module):
+    def __init__(self, cfg: Qwen3Config):
+        super().__init__()
+        self.config = cfg
+        self.model = Qwen3Model(cfg)
+        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
+        if cfg.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+
+    def forward(self, input_ids, position_ids=None):
+        return self.lm_head(self.model(input_ids, position_ids))
+
+    def gradient_checkpointing_enable(self, enable: bool = True) -> None:
+        self.model.gradient_checkpointing = enable
+
+    @torch.no_grad()
+    def init_weights(self, seed: int | None = None) -> "Qwen3ForCausalLM":
+        """从零训练用的随机初始化。**没有这一步就只有 torch 的默认值。**
+
+        torch 的默认值对 transformer 是错的,而且错得不报错:
+
+            nn.Embedding  → N(0, 1)              std 比该有的 0.02 大 **50 倍**
+            nn.Linear     → U(±1/√fan_in)        均匀分布,且 fan_in 一变标准差就变
+
+        嵌入 std=1 意味着第一层的 RMSNorm 之前激活就是 O(1)*50,而 tie 之后
+        lm_head 也是同一份权重 —— 初始 logits 的量级直接决定第一步的 loss 和
+        梯度。这类错**照样能训**,只是慢、抖、有时候崩,回头查不出原因。
+
+        用的是 Llama / GPT-NeoX 那套:
+
+            所有权重       N(0, initializer_range)
+            残差出口       N(0, initializer_range / √(2L))  ← o_proj 和 down_proj
+            RMSNorm        全 1
+
+        残差出口缩放是为了让 28 层累加之后的方差不随深度线性增长。少了它
+        深层的激活会越叠越大,表现为训练早期的 loss spike。
+
+        tie 的模型里 `lm_head.weight` 就是 `embed_tokens.weight` 同一个张量,
+        `named_parameters()` 默认去重,所以只会被初始化一次 —— 不要改成
+        `remove_duplicate=False`,那会按两套 std 初始化同一块内存。
+        """
+        if seed is not None:
+            torch.manual_seed(seed)
+        std = self.config.initializer_range
+        resid_std = std / math.sqrt(2 * self.config.num_hidden_layers)
+        for name, p in self.named_parameters():
+            if name.endswith("o_proj.weight") or name.endswith("down_proj.weight"):
+                nn.init.normal_(p, mean=0.0, std=resid_std)
+            elif "norm" in name or "layernorm" in name:
+                nn.init.ones_(p)
+            else:
+                nn.init.normal_(p, mean=0.0, std=std)
+        return self
+
+    @classmethod
+    def from_config(cls, config_path, device="cpu", dtype=None,
+                    seed: int | None = None) -> "Qwen3ForCausalLM":
+        """只读 config.json,随机初始化权重 —— 从零预训练的起点。
+
+        与 `from_pretrained` 的区别就是不加载 checkpoint。刻意做成两个入口:
+        「加载失败于是静默随机初始化」是这个项目最怕的一类错。
+        """
+        cfg = Qwen3Config.from_json(Path(config_path) / "config.json")
+        model = cls(cfg).to(dtype or torch.float32)
+        model.init_weights(seed=seed)
+        if cfg.tie_word_embeddings:
+            model.lm_head.weight = model.model.embed_tokens.weight
+        return model.to(device)
+
+    def save_pretrained(self, out_dir) -> None:
+        """存成 HF 布局:`config.json` + `model.safetensors`。
+
+        tie 的模型不写 `lm_head.weight` —— 与上游 checkpoint 一致(310 个张量)。
+        分词器产物(piece.model / dict.txt / token_mapping.json)由调用方负责
+        拷过来,这里不碰文本相关的东西。
+        """
+        import dataclasses
+
+        save_safetensors = _checkpoint_mod().save_safetensors
+
+        out_dir = Path(out_dir)
+        out_dir.mkdir(parents=True, exist_ok=True)
+        sd = self.state_dict_for_save()
+        cfg = dataclasses.asdict(self.config)
+        cfg["model_type"] = "qwen3"
+        cfg["architectures"] = ["Qwen3ForCausalLM"]
+        # torch_dtype 写实际存下来的精度,不写 config 里那个。
+        # 从零预训练的参数是 float32,而 config 从模板继承来的可能还是
+        # bfloat16 —— 对不上时 vLLM 按 config 说的加载,权重被悄悄转一次,
+        # 不报错。以文件里的张量为准才是唯一不会骗人的口径。
+        cfg["torch_dtype"] = str(next(iter(sd.values())).dtype).removeprefix("torch.")
+        (out_dir / "config.json").write_text(json.dumps(cfg, indent=2))
+        save_safetensors(sd, out_dir / "model.safetensors", metadata={"format": "pt"})
+
+    @classmethod
+    def from_pretrained(cls, model_dir, device="cpu", dtype=None):
+        """从 HF 格式的目录加载。key 对不上就报错,不静默跳过。"""
+        load_sharded = _checkpoint_mod().load_sharded
+
+        model_dir = Path(model_dir)
+        cfg = Qwen3Config.from_json(model_dir / "config.json")
+        model = cls(cfg).to(dtype or cfg.dtype)
+
+        sd = load_sharded(model_dir, device="cpu")
+        # tie 的 checkpoint 里没有 lm_head.weight —— 补上指向 embed 的引用,
+        # 这样 strict=True 才不会误报缺失
+        if cfg.tie_word_embeddings and "lm_head.weight" not in sd:
+            sd["lm_head.weight"] = sd["model.embed_tokens.weight"]
+        missing, unexpected = model.load_state_dict(sd, strict=False)
+        # 手动检查而不是 strict=True:好给出能看懂的报错
+        if missing or unexpected:
+            raise RuntimeError(
+                f"state_dict 对不上 {model_dir}:\n"
+                f"  缺少 {len(missing)} 个:{sorted(missing)[:6]}\n"
+                f"  多余 {len(unexpected)} 个:{sorted(unexpected)[:6]}\n"
+                f"  这类错**不会自己暴露** —— 权重没加载上模型照样跑。")
+        if cfg.tie_word_embeddings:
+            model.lm_head.weight = model.model.embed_tokens.weight
+        return model.to(device).eval()
+
+    def state_dict_for_save(self) -> dict[str, torch.Tensor]:
+        """导出成 HF 格式:tie 的话去掉 `lm_head.weight`,与上游 checkpoint 一致。"""
+        sd = self.state_dict()
+        if self.config.tie_word_embeddings:
+            sd.pop("lm_head.weight", None)
+        return sd

model.safetensors

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

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.6
+# 分词器(C++,同时提供 81903 词表):
+#   pip install git+https://github.com/Ismantic/PieceTokenizer

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "unk_token": "<unk>",
+  "pad_token": "<pad>"
+}

token_mapping.json

@@ -0,0 +1,17 @@
+{
+  "base_vocab_size": 81903,
+  "total_vocab_size": 81903,
+  "special_tokens": {
+    "<pad>": 81899,
+    "<user>": 81900,
+    "<assistant>": 81901,
+    "<system>": 81902
+  },
+  "bos_id": 1,
+  "eos_id": 2,
+  "unk_id": 0,
+  "pad_id": 81899,
+  "user_id": 81900,
+  "assistant_id": 81901,
+  "system_id": 81902
+}

tokenizer.py

@@ -0,0 +1,230 @@
+"""PieceTokenizer 的 HuggingFace 风格外壳。
+
+## dict.txt 是必需的,不是可选的
+
+没有中文分词词典时,**中文的 token id 会变** —— 不只是慢。2026-07-27 实测
+(同一 piece.model,带/不带 dict):
+
+    中文 22 字     0.169ms → 1.060ms   6.3×    id 序列不同(长度都是 22)
+    中文 ×20       3.289ms → 21.125ms  6.4×    id 序列不同(长度都是 440)
+    英文           0.110ms → 0.110ms   1.0×    完全相同(dict 只管中文切分)
+
+**round-trip 正确会掩盖这个问题**:两种模式 decode 都能还原原文,看起来没事,
+但喂给模型的 id 已经不是训练时那套了。所以这里缺 dict 直接报错,不静默降级
+(旧版本是 `else: load(model_file)`,不报错 —— 那是个真隐患)。
+
+确实不需要 dict 的场合(比如只查 vocab_size)传 `require_dict=False`。
+
+## 词表从哪来
+
+本仓库**不存词表副本**。规范位置是 clone 的 PieceTokenizer:
+
+    deps/PieceTokenizer/save/Summer-Tokenizer.pt
+    deps/PieceTokenizer/save/Summer-Tokenizer.dict.txt
+
+用 `resolve_assets()` 反查(经 `piece_tokenizer.__file__`)。checkpoint 目录里
+另有一份 piece.model + dict.txt,与上面逐字节相同(sha256 已核),训练/评测时
+直接从 checkpoint 目录加载即可。
+"""
+import os
+import json
+import piece_tokenizer as pt
+
+PIECE_MODEL_NAME = "Summer-Tokenizer.pt"
+PIECE_DICT_NAME = "Summer-Tokenizer.dict.txt"
+
+
+def resolve_assets():
+    """从已安装的 piece_tokenizer 反查词表文件,返回 (piece_model, cn_dict)。
+
+    本仓库不留副本,所以要经 `piece_tokenizer.__file__` 找到 clone 的仓库根,
+    再取 `save/` 下那两个文件。找不到就报错 —— 报错好过静默用错词表。
+
+    `SUMMER_PIECE_MODEL` / `SUMMER_PIECE_DICT` 可以覆盖(与 BERTc 的
+    `BERTC_PIECE_MODEL` 同一套做法)。
+    """
+    env_m = os.environ.get("SUMMER_PIECE_MODEL")
+    env_d = os.environ.get("SUMMER_PIECE_DICT")
+    if env_m and env_d:
+        return env_m, env_d
+
+    root = os.path.dirname(os.path.abspath(pt.__file__))
+    tried = []
+    for cand in (root, os.path.dirname(root)):
+        save = os.path.join(cand, "save")
+        model, cn_dict = (os.path.join(save, PIECE_MODEL_NAME),
+                          os.path.join(save, PIECE_DICT_NAME))
+        tried.append(save)
+        if os.path.exists(model) and os.path.exists(cn_dict):
+            return model, cn_dict
+    raise FileNotFoundError(
+        f"找不到 {PIECE_MODEL_NAME} / {PIECE_DICT_NAME}。找过:{tried}\n"
+        f"  先跑 `bash prepare/install_deps.sh` clone 并安装 PieceTokenizer。")
+
+
+# checkpoint 目录里词表文件的名字。**上游名优先。**
+#
+# `Summer-Tokenizer.pt` / `.dict.txt` 与 PieceTokenizer 仓库 `save/` 下的文件
+# 同名 —— 这样从任何一个 checkpoint 都能一眼看出词表出自哪里,不用靠 sha256
+# 去反推(BERTc 一直是这么做的,它的发布包里就叫 `BERTc-Tokenizer.pt`)。
+#
+# `piece.model` / `dict.txt` 是改造之前的名字。已发布的
+# `Ismantic/Qwen3-1.7B-Base-ReTok` 和 v18 的 checkpoint 用的都是旧名,所以
+# **保留为回退**,不然那些目录一个都加载不了。新产出统一用上游名。
+_MODEL_NAMES = (PIECE_MODEL_NAME, "piece.model", "piece_mt.model")
+_DICT_NAMES = (PIECE_DICT_NAME, "dict.txt")
+
+
+def _first_in(model_dir, names):
+    for n in names:
+        p = os.path.join(model_dir, n)
+        if os.path.exists(p):
+            return p
+    return None
+
+
+def has_piece_vocab(model_dir) -> bool:
+    """这个目录是 piece 词表的模型吗?
+
+    评测入口靠它决定用 PieceTokenizerWrapper 还是 AutoTokenizer。**不要各自
+    硬编码文件名** —— 新产出用上游名 `Summer-Tokenizer.pt`,旧的用
+    `piece.model`,漏一个就会静默退回 AutoTokenizer,而 AutoTokenizer 对这个
+    词表走不通,结果是错的。
+    """
+    return _first_in(model_dir, _MODEL_NAMES) is not None
+
+
+class PieceTokenizerWrapper:
+    def __init__(self, model_dir, require_dict=True):
+        """从模型目录加载。
+
+        词表文件按 `_MODEL_NAMES` / `_DICT_NAMES` 的顺序找 —— 上游名优先,
+        旧名回退。
+        """
+        self._tok = pt.Tokenizer()
+
+        model_file = _first_in(model_dir, _MODEL_NAMES)
+        if model_file is None:
+            raise FileNotFoundError(
+                f"{model_dir} 里找不到词表。试过:{list(_MODEL_NAMES)}")
+
+        cn_dict = _first_in(model_dir, _DICT_NAMES)
+        if cn_dict is not None:
+            self._tok.load(model_file, cn_dict)
+        elif require_dict:
+            raise FileNotFoundError(
+                f"{model_dir} 里没有中文分词词典(试过 {list(_DICT_NAMES)})。"
+                f"缺了它中文的 token id 会变(不只是慢),而且 decode 照样能"
+                f"还原原文、不会报错。\n"
+                f"  从 checkpoint 或 PieceTokenizer 的 save/{PIECE_DICT_NAME} "
+                f"拷���份过来;确实不需要就传 require_dict=False。")
+        else:
+            self._tok.load(model_file)
+        self.piece_model_path = model_file
+        self.cn_dict_path = cn_dict
+
+        # Load token mapping
+        mapping_file = os.path.join(model_dir, "token_mapping.json")
+        if os.path.exists(mapping_file):
+            with open(mapping_file) as f:
+                mapping = json.load(f)
+            self.pad_token_id = mapping["pad_id"]
+            self.bos_token_id = mapping["bos_id"]
+            self.eos_token_id = mapping["eos_id"]
+            self.user_token_id = mapping.get("user_id")
+            self.assistant_token_id = mapping.get("assistant_id")
+            self.system_token_id = mapping.get("system_id")
+        else:
+            # Fallback to piece_to_id lookups
+            self.bos_token_id = self._tok.piece_to_id("<s>")
+            self.eos_token_id = self._tok.piece_to_id("</s>")
+            self.pad_token_id = self._tok.piece_to_id("<pad>")
+            self.user_token_id = self._tok.piece_to_id("<user>")
+            self.assistant_token_id = self._tok.piece_to_id("<assistant>")
+            self.system_token_id = self._tok.piece_to_id("<system>")
+            if self.pad_token_id < 0:
+                self.pad_token_id = 0
+
+    @property
+    def vocab_size(self):
+        return self._tok.vocab_size()
+
+    def encode(self, text, add_special_tokens=False):
+        ids = self._tok.encode_as_ids(text)
+        if add_special_tokens:
+            ids = [self.bos_token_id] + ids + [self.eos_token_id]
+        return ids
+
+    def decode(self, ids, skip_special_tokens=True):
+        if skip_special_tokens:
+            special = {self.bos_token_id, self.eos_token_id, self.pad_token_id,
+                       self.user_token_id, self.assistant_token_id, self.system_token_id}
+            ids = [i for i in ids if i not in special]
+        try:
+            return self._tok.decode(ids)
+        except UnicodeDecodeError:
+            # Model emitted byte-fallback piece(s) that don't form valid UTF-8.
+            # Per-piece fallback: keep ids that decode cleanly, drop the rest.
+            parts = []
+            for i in ids:
+                try:
+                    parts.append(self._tok.id_to_piece(i))
+                except UnicodeDecodeError:
+                    continue
+            return "".join(parts).replace("▁", " ")
+
+    def apply_chat_template(self, messages, tokenize=True, add_generation_prompt=False, **kwargs):
+        """Build chat-formatted token sequence from messages."""
+        ids = []
+
+        # Check for system message
+        start = 0
+        if messages and messages[0]["role"] == "system":
+            ids.append(self.bos_token_id)
+            ids.extend(self._tok.encode_as_ids(messages[0]["content"]))
+            ids.append(self.system_token_id)
+            start = 1
+        else:
+            ids.append(self.bos_token_id)
+
+        for msg in messages[start:]:
+            if msg["role"] == "user":
+                ids.append(self.user_token_id)
+                ids.extend(self._tok.encode_as_ids(msg["content"]))
+            elif msg["role"] == "assistant":
+                ids.append(self.assistant_token_id)
+                ids.extend(self._tok.encode_as_ids(msg["content"]))
+                ids.append(self.eos_token_id)
+
+        if add_generation_prompt:
+            ids.append(self.assistant_token_id)
+
+        if tokenize:
+            return ids
+        else:
+            # Return as string (rarely needed)
+            return self._tok.decode(ids)
+
+    def save_pretrained(self, output_dir):
+        """Save tokenizer files to directory (for checkpoint saving)."""
+        import shutil
+        os.makedirs(output_dir, exist_ok=True)
+        # Copy piece.model
+        # 从自己加载时的实际路径拷,不猜文件名
+        if self.piece_model_path:
+            shutil.copy2(self.piece_model_path,
+                         os.path.join(output_dir, PIECE_MODEL_NAME))
+        if self.cn_dict_path:
+            shutil.copy2(self.cn_dict_path,
+                         os.path.join(output_dir, PIECE_DICT_NAME))
+        # Save mapping
+        mapping = {
+            "bos_id": self.bos_token_id,
+            "eos_id": self.eos_token_id,
+            "pad_id": self.pad_token_id,
+            "user_id": self.user_token_id,
+            "assistant_id": self.assistant_token_id,
+            "system_id": self.system_token_id,
+        }
+        with open(os.path.join(output_dir, "token_mapping.json"), "w") as f:
+            json.dump(mapping, f, indent=2)

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "model_type": "qwen3",
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "unk_token": "<unk>",
+  "pad_token": "<pad>",
+  "clean_up_tokenization_spaces": false,
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}