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:
@@ -0,0 +1,52 @@
|
||||
from kda.training.data import IGNORE_INDEX, collate_sft, encode_sft_row, load_sft_rows
|
||||
from kda.training.prompts import instruction_prompt
|
||||
|
||||
|
||||
class _Tok:
|
||||
vocab_size = 32
|
||||
|
||||
def encode(self, text: str) -> list[int]:
|
||||
return [min((ord(c) % 30) + 1, 31) for c in text[:12]] or [1]
|
||||
|
||||
def decode(self, ids: list[int]) -> str:
|
||||
return "x" * len(ids)
|
||||
|
||||
|
||||
def test_instruction_matches_eval_template():
|
||||
assert instruction_prompt("你好", "en") == "Translate to English:\n你好"
|
||||
assert instruction_prompt("Hello", "zh") == "Translate to Chinese:\nHello"
|
||||
|
||||
|
||||
def test_prompt_tokens_are_ignored():
|
||||
tok = _Tok()
|
||||
src, tgt = "ab", "cd"
|
||||
ids, labels = encode_sft_row(tok, src, tgt, "en", max_len=64)
|
||||
prompt_n = len(tok.encode(instruction_prompt(src, "en")))
|
||||
assert labels[:prompt_n] == [IGNORE_INDEX] * prompt_n
|
||||
assert all(v != IGNORE_INDEX for v in labels[prompt_n:])
|
||||
assert ids[prompt_n:] == tok.encode(tgt)
|
||||
|
||||
|
||||
def test_collate_and_jsonl(tmp_path):
|
||||
path = tmp_path / "tiny.jsonl"
|
||||
path.write_text(
|
||||
'{"src": "a", "tgt": "b", "target_lang": "en"}\n'
|
||||
'{"src": "c", "tgt": "d", "target_lang": "zh"}\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
rows = load_sft_rows(path)
|
||||
assert len(rows) == 2
|
||||
x, y = collate_sft(rows, _Tok(), max_len=32)
|
||||
assert x.shape == y.shape
|
||||
assert x.size(0) == 2
|
||||
assert (y == IGNORE_INDEX).any()
|
||||
|
||||
|
||||
def test_toy_sft_file_parses():
|
||||
from pathlib import Path
|
||||
|
||||
path = Path(__file__).resolve().parents[2] / "data" / "sft" / "toy.jsonl"
|
||||
rows = load_sft_rows(path)
|
||||
assert len(rows) >= 20
|
||||
langs = {r["target_lang"] for r in rows}
|
||||
assert langs == {"en", "zh"}
|
||||
Reference in New Issue
Block a user