Keep attnres on resume, fix final chunk_index, default 0.5b to Yi-6B

CLI default attnres=off was overwriting block checkpoints on resume
so later loads hit Unexpected key(s). Only apply flags the user
passed. Track next_chunk so a budget-exit save does not skip the
untrained yield. 0.5b now uses 01-ai/Yi-6B (64k); refuse resume
when the ckpt tokenizer does not match.

README.md

@@ -138,7 +138,7 @@ PYTHONPATH=. python train.py
 
 ### 复现训练
 
-目标是 **~0.5B zh↔en 指令翻译模型**（`K3Config.preset("0.5b")` = 482M，tied Qwen3 词表）。本机 RTX 3060 6GB 只跑 8M 全流程孪生；0.5B 预训练需要 **32–40GB Ampere bf16**。
+目标是 **~0.5B zh↔en 指令翻译模型**（`K3Config.preset("0.5b")` ≈ 415M，tied Yi-6B 64k 词表）。本机 RTX 3060 6GB 只跑 8M 全流程孪生；0.5B 预训练需要 **32–40GB Ampere bf16**。
 
 成功标准是冻结集上的 `translation_success()`，**不是** wiki train loss。wiki 预训练没见过 `Translate to English:\n...`，预训练阶段 `eval_mt` 的 success_rate 预期 ≈0。
 
@@ -160,7 +160,7 @@ uv run python train_k3.py --preset 0.5b --attnres block \
   --max-tokens 1000000000 --warmup 2000
 ```
 
-`0.5b` 预设：`d=768`，`L=24`（6×3 KDA + 1 MLA），`H=12`，`head=64`，`chunk=64`，LatentMoE `ℓ=384` / 16 routed / Top-2 / shared 2，tied Qwen3 embedding，activation checkpoint 默认开。训练默认 seq 2048、micro-batch 2、grad-acc 8、lr 3e-4、warmup **64 optimizer steps**。checkpoint：`ckpts/k3_0.5b.pt`，另写 `_last` / `_best`（best 按 held-out CE）。
+`0.5b` 预设：`d=768`，`L=24`（6×3 KDA + 1 MLA），`H=12`，`head=64`，`chunk=64`，LatentMoE `ℓ=384` / 16 routed / Top-2 / shared 2，tied Yi-6B embedding（64k，有 EOS），activation checkpoint 默认开。训练默认 seq 2048、micro-batch 2、grad-acc 8、lr 3e-4、warmup **64 optimizer steps**。checkpoint：`ckpts/k3_0.5b.pt`，另写 `_last` / `_best`（best 按 held-out CE）。**不能**从 Qwen3 词表的旧 ckpt `--resume`。
 
 Token 会计：`step` 仍是 micro-batch；有效 token = `batch × seq_len × micro_steps`。默认 0.5b 冒烟是 **8.2M token ≈ 0.017 tok/param**。翻译前置 LM 的最低有意义预算是 **1B token**（`--max-tokens`），不是 2000 step。
 

kda/models/k3_config.py

@@ -9,13 +9,15 @@ Hybrid Attention (K3): 每 4 层 1 次 Gated MLA, 末层强制 MLA.
 
 Presets:
     toy  — ~8M, 自训 8k SP, 本地过拟合
-    0.5b — ~482M, Qwen3 词表, 32–40GB bf16；默认 step 是冒烟，翻译前置用 --max-tokens
+    0.5b — ~415M, Yi-6B 词表 (64k), 32–40GB bf16；默认 step 是冒烟，翻译前置用 --max-tokens
 """
 from __future__ import annotations
 
 from dataclasses import dataclass
 
-# Qwen3 config.json; train_k3 overrides with len(tokenizer).
+# 01-ai/Yi-6B config.json; train_k3 overrides with len(tokenizer).
+YI6B_VOCAB_SIZE = 64000
+# Kept for old Qwen3 checkpoints / docs.
 QWEN3_VOCAB_SIZE = 151936
 
 
@@ -24,7 +26,7 @@ class K3Config:
     # 主干
     hidden_size: int = 256
     num_hidden_layers: int = 4
-    vocab_size: int = 8192               # toy: data/spm_4k; 0.5b: Qwen3
+    vocab_size: int = 8192               # toy: data/spm_4k; 0.5b: Yi-6B 64k
     initializer_range: float = 0.02
     norm_eps: float = 1e-6
     tie_word_embeddings: bool = False
@@ -76,11 +78,11 @@ class K3Config:
             return cls()
         if name in {"0.5b", "500m"}:
             # H * head_dim == hidden. Routed 16 Top-2; LatentMoE padded bmm.
-            # ~482M with tied Qwen3 embeddings. 6×(3 KDA + 1 MLA).
+            # ~415M with tied Yi-6B embeddings. 6×(3 KDA + 1 MLA).
             return cls(
                 hidden_size=768,
                 num_hidden_layers=24,
-                vocab_size=QWEN3_VOCAB_SIZE,
+                vocab_size=YI6B_VOCAB_SIZE,
                 tie_word_embeddings=True,
                 max_position_embeddings=2048,
                 num_heads=12,

tests/integration/test_k3_arch.py

@@ -78,6 +78,7 @@ def test_preset_0_5b_schedule():
     assert cfg.hidden_size == 768
     assert cfg.num_heads * cfg.head_dim == cfg.hidden_size
     assert cfg.num_hidden_layers == 24
+    assert cfg.vocab_size == 64000
     assert cfg.tie_word_embeddings
     assert cfg.chunk_size == 64
     assert cfg.gradient_checkpointing is True

tests/integration/test_train_k3_resume.py

@@ -0,0 +1,81 @@
+"""Resume must not clobber attnres or skip a chunk at budget exit."""
+from argparse import Namespace
+from dataclasses import asdict
+
+import pytest
+import torch
+
+from kda.models.causal_lm import CausalLM
+from kda.models.k3_config import K3Config
+from kda.training.toy import load_ckpt
+from train_k3 import _apply_cli_overrides
+
+
+def _tiny_block():
+    return K3Config(
+        hidden_size=32,
+        num_hidden_layers=4,
+        num_heads=4,
+        head_dim=8,
+        chunk_size=4,
+        vocab_size=64,
+        moe_latent_size=16,
+        moe_d_ff=16,
+        n_routed=4,
+        top_k=2,
+        n_shared=1,
+        kv_lora_rank=8,
+        q_lora_rank=16,
+        qk_nope_head_dim=8,
+        v_head_dim=8,
+        attnres="block",
+        attnres_block_size=2,
+    )
+
+
+def _cli(**kwargs):
+    base = dict(
+        attnres=None,
+        attnres_block_size=None,
+        grad_checkpoint=None,
+        moe_aux_coef=None,
+        moe_z_coef=None,
+    )
+    base.update(kwargs)
+    return Namespace(**base)
+
+
+def test_resume_without_attnres_flag_keeps_block():
+    cfg = _tiny_block()
+    _apply_cli_overrides(cfg, _cli())
+    assert cfg.attnres == "block"
+    assert cfg.attnres_block_size == 2
+
+
+def test_explicit_attnres_overrides_resume():
+    cfg = _tiny_block()
+    _apply_cli_overrides(cfg, _cli(attnres="full", attnres_block_size=1))
+    assert cfg.attnres == "full"
+    assert cfg.attnres_block_size == 1
+
+
+def test_block_state_with_off_config_cannot_load(tmp_path):
+    cfg = _tiny_block()
+    model = CausalLM(cfg)
+    payload = {"config": asdict(cfg), "model_state": model.state_dict()}
+    payload["config"]["attnres"] = "off"
+    payload["config"]["attnres_block_size"] = None
+    path = str(tmp_path / "polluted.pt")
+    torch.save(payload, path)
+    with pytest.raises(RuntimeError, match="Unexpected key"):
+        load_ckpt(path)
+
+
+def test_budget_break_does_not_skip_yielded_chunk():
+    next_chunk = 10
+    for chunk_index in (10, 11, 12):
+        tokens = 100
+        if tokens >= 100:
+            break
+        next_chunk = chunk_index + 1
+    assert next_chunk == 10

train_k3.py

@@ -41,7 +41,7 @@ _TOY_TRAIN = {
     "gen_every": 200,
 }
 _B500M_TRAIN = {
-    "tokenizer": "Qwen/Qwen3-8B",
+    "tokenizer": "01-ai/Yi-6B",
     "out": "ckpts/k3_0.5b.pt",
     "limit": 20000,
     "batch": 2,
@@ -184,6 +184,20 @@ def _apply_moe_coefs(model, cfg: K3Config) -> None:
             module.z_loss_coef = cfg.moe_z_loss_coef
 
 
+def _apply_cli_overrides(cfg: K3Config, args: argparse.Namespace) -> None:
+    """Copy only flags the user actually passed. CLI defaults must not clobber a resume."""
+    if args.attnres is not None:
+        cfg.attnres = args.attnres
+    if args.attnres_block_size is not None:
+        cfg.attnres_block_size = args.attnres_block_size
+    if args.grad_checkpoint is not None:
+        cfg.gradient_checkpointing = args.grad_checkpoint
+    if args.moe_aux_coef is not None:
+        cfg.moe_aux_loss_coef = args.moe_aux_coef
+    if args.moe_z_coef is not None:
+        cfg.moe_z_loss_coef = args.moe_z_coef
+
+
 def _moe_log(model) -> dict:
     frac = moe_route_frac(model)
     if frac is None:
@@ -267,9 +281,10 @@ def main() -> None:
     p.add_argument("--device", default="auto")
     p.add_argument(
         "--attnres",
-        default="off",
+        default=None,
         choices=["off", "full", "block"],
-        help="depth mixer: off=standard residual, block=K3 AttnRes, full=per-layer AttnRes",
+        help="depth mixer: off=standard residual (preset default), block=K3 AttnRes, "
+        "full=per-layer AttnRes. Omit on --resume to keep the checkpoint value",
     )
     p.add_argument(
         "--attnres-block-size",
@@ -329,14 +344,7 @@ def main() -> None:
     tok = load_tokenizer(args.tokenizer)
     cfg = K3Config.preset(args.preset)
     cfg.vocab_size = tok.vocab_size
-    cfg.attnres = args.attnres
+    _apply_cli_overrides(cfg, args)
-    cfg.attnres_block_size = args.attnres_block_size
-    if args.grad_checkpoint is not None:
-        cfg.gradient_checkpointing = args.grad_checkpoint
-    if args.moe_aux_coef is not None:
-        cfg.moe_aux_loss_coef = args.moe_aux_coef
-    if args.moe_z_coef is not None:
-        cfg.moe_z_loss_coef = args.moe_z_coef
 
     langs = [part.strip() for part in args.langs.split(",") if part.strip()]
     tpm = tokens_per_micro(args.batch, args.seq_len)
@@ -364,19 +372,16 @@ def main() -> None:
                 f"{type(loaded_cfg).__name__}"
             )
         cfg = loaded_cfg
-        cfg.attnres = args.attnres
+        _apply_cli_overrides(cfg, args)
-        cfg.attnres_block_size = args.attnres_block_size
-        if args.grad_checkpoint is not None:
-            cfg.gradient_checkpointing = args.grad_checkpoint
-        if args.moe_aux_coef is not None:
-            cfg.moe_aux_loss_coef = args.moe_aux_coef
-        if args.moe_z_coef is not None:
-            cfg.moe_z_loss_coef = args.moe_z_coef
         model.gradient_checkpointing = cfg.gradient_checkpointing
         model.to(device)
         payload = torch.load(args.resume, map_location="cpu", weights_only=False)
         if payload.get("tokenizer") and payload["tokenizer"] != args.tokenizer:
-            print(f"warning: ckpt tokenizer {payload['tokenizer']} != {args.tokenizer}")
+            raise SystemExit(
+                f"tokenizer mismatch: ckpt {payload['tokenizer']!r} vs "
+                f"CLI {args.tokenizer!r}; embeddings are not interchangeable "
+                f"(do not resume a Qwen ckpt with Yi)"
+            )
         micro_step = int(payload.get("micro_step", 0))
         opt_step = int(payload.get("opt_step", 0))
         tokens = int(payload.get("tokens", 0))
@@ -469,6 +474,10 @@ def main() -> None:
     model.train()
     t0 = time.perf_counter()
     tokens_at_t0 = tokens
+    # Index of the next untrained chunk. Mid-loop saves use last_trained+1.
+    # The final save must NOT +1 again: the loop may break on a yielded chunk
+    # that was never trained (budget check is at the top).
+    next_chunk = chunk_index
     for chunk_index, x, y in iter_indexed(train_chunks, start=chunk_index):
         if args.max_tokens is not None:
             if tokens >= args.max_tokens:
@@ -497,6 +506,7 @@ def main() -> None:
         lr_now = optim.param_groups[0]["lr"]
         if raw_loss < best_train:
             best_train = raw_loss
+        next_chunk = chunk_index + 1
 
         metrics = {
             "train/loss": raw_loss,
@@ -542,7 +552,7 @@ def main() -> None:
                     micro_step=micro_step,
                     opt_step=opt_step,
                     tokens=tokens,
-                    chunk_index=chunk_index + 1,
+                    chunk_index=next_chunk,
                     best_heldout=best_heldout,
                 )
                 _save(_sibling(args.out, "_best"), payload)
@@ -570,7 +580,7 @@ def main() -> None:
                 micro_step=micro_step,
                 opt_step=opt_step,
                 tokens=tokens,
-                chunk_index=chunk_index + 1,
+                chunk_index=next_chunk,
                 best_heldout=best_heldout,
             )
             _save(_sibling(args.out, "_last"), payload)
@@ -586,7 +596,7 @@ def main() -> None:
         micro_step=micro_step,
         opt_step=opt_step,
         tokens=tokens,
-        chunk_index=chunk_index + 1,
+        chunk_index=next_chunk,
         best_heldout=best_heldout,
     )
     _save(args.out, payload)
@@ -596,6 +606,12 @@ def main() -> None:
     )
     if tracker is not None:
         tracker.finish()
+    if device == "cuda":
+        try:
+            torch.cuda.synchronize()
+            torch.cuda.empty_cache()
+        except Exception:
+            pass
 
 
 if __name__ == "__main__":