Initial K3 snapshot: 0.5B KDA/MLA/MoE train path

Standalone tree split from LLMRL/projects/kda. Includes Triton dt_bias
backward fix, train_k3 --preset 0.5b, SFT, Docker runtime, and tests.
This commit is contained in:
dela
2026-08-25 14:43:17 +08:00
commit 584f7e9e73
140 changed files with 21592 additions and 0 deletions
+51
View File
@@ -0,0 +1,51 @@
"""Train a SentencePiece tokenizer on a Chinese Wikipedia subset.
用法:
uv run python kda/training/train_tokenizer.py \
--out data/spm_4k --vocab-size 4096 --limit 20000
产出:
data/spm_4k.model / data/spm_4k.vocab (BPE/unigram, 中文小语料)
"""
from __future__ import annotations
import argparse
import sentencepiece as spm
from .data import fetch_wiki_texts
def main() -> None:
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--out", default="data/spm_4k", help="输出前缀 (model/vocab 文件)")
p.add_argument("--vocab-size", type=int, default=8192)
p.add_argument("--limit", type=int, default=20000, help="用于训练的 wiki 文章数")
p.add_argument("--model-type", default="unigram", choices=["unigram", "bpe"])
p.add_argument("--character-coverage", type=float, default=0.9995)
args = p.parse_args()
texts = fetch_wiki_texts(args.limit)
corpus = "".join(texts)
tmp = args.out + ".corpus.txt"
with open(tmp, "w", encoding="utf-8") as f:
f.write(corpus)
print(f"corpus: {len(corpus):,} chars from {len(texts)} articles")
spm.SentencePieceTrainer.train(
input=tmp,
model_prefix=args.out,
vocab_size=args.vocab_size,
model_type=args.model_type,
character_coverage=args.character_coverage,
unk_id=0,
pad_id=1,
bos_id=-1,
eos_id=-1,
num_threads=4,
)
print(f"tokenizer saved: {args.out}.model / {args.out}.vocab")
if __name__ == "__main__":
main()