"""Lab tổng hợp LLM Training: huấn luyện một LLM thu nhỏ từ dữ liệu đến checkpoint, bằng NumPy thuần.

Một bộ giải mã Transformer nhỏ (2 lớp, 1 head attention, RMSNorm, SwiGLU, mã hóa vị trí cố định)
được huấn luyện trên một kho chữ rất nhỏ, rồi đi tiếp qua SFT, LoRA, distillation, DPO và cuối cùng
được nạp lại từ checkpoint để suy luận. 13 giai đoạn [1]..[13] khớp các ô của sơ đồ kiến trúc ở Tab 01
(1A..1C, 2A..2D, 3A..3E, 4A..4G, 5A..5C).

Những gì là THẬT trong file này: mọi phép tính tensor, gradient giải tích (được kiểm bằng sai phân
hữu hạn), mọi giá trị loss, mọi số đếm byte và mọi sai số so sánh.
Những gì là GIẢ ĐỊNH hoặc ĐƠN GIẢN HÓA: mô hình và dữ liệu cực nhỏ (không có nghĩa ngôn ngữ thật);
"ZeRO" chỉ là phép chia mảng và đếm byte trong một tiến trình, không có GPU hay mạng; không có số
nào ở đây là tốc độ hoặc chất lượng của mô hình thật.
Mọi thứ chạy với seed cố định nên chạy lại cho đúng các số cũ.
"""
import io

import numpy as np

# NumPy trên macOS (Accelerate) có thể báo cảnh báo dấu phẩy động giả ở phép matmul với ma trận chuyển vị;
# kết quả vẫn đúng (gradient được kiểm bằng sai phân hữu hạn ở giai đoạn 5), nên tắt các cảnh báo này cho output gọn.
np.seterr(divide="ignore", over="ignore", invalid="ignore")

# ── Cấu hình (sửa các hằng số này ở phần bài tập rồi chạy lại) ───────────────
SEED = 7
D_MODEL = 16              # chiều vector ẩn d
N_LAYERS = 2              # số khối Transformer xếp chồng
FFN_HIDDEN = 32           # chiều trong của SwiGLU
POS_SCALE = 0.3           # độ lớn mã hóa vị trí
LR = 0.01                 # learning rate của huấn luyện nền (4C) và SFT
WEIGHT_DECAY = 0.01       # suy giảm trọng số của AdamW
BETA1, BETA2, EPS = 0.9, 0.999, 1e-8
BATCH_DOCS = 4            # số tài liệu mỗi batch
STEPS = 300               # số bước huấn luyện nền (giai đoạn 7)
EVAL_EVERY = 50
RESUME_AT = 100           # bước lưu checkpoint tiếp tục được
ZERO_WORKERS = 2          # số worker mô phỏng trong giai đoạn 8
SFT_STEPS = 150
MASK_RESPONSE_ONLY = True  # SFT chỉ tính loss trên token câu trả lời (1C)
LORA_RANK = 2
LORA_ALPHA = 4
LORA_LR = 0.02
LORA_STEPS = 150
STUDENT_D = 8             # chiều ẩn của student trong distillation
STUDENT_LAYERS = 1
DISTILL_T = 2.0           # nhiệt độ làm mềm phân phối của teacher
DISTILL_STEPS = 150
DPO_BETA = 0.5
DPO_LR = 0.003
DPO_STEPS = 60
MAX_NEW = 12              # số token sinh tối đa khi suy luận

PAD, EOS, USER, ASSISTANT = 0, 1, 2, 3
SPECIALS = ["<pad>", "<eos>", "<user>", "<assistant>"]

# Kho văn bản thô tự viết (có dòng trống, trùng lặp và chữ hoa để bước làm sạch có việc làm).
RAW_CORPUS = [
    "Meo an ca.", "cho an xuong.", "  meo uong sua. ", "cho uong nuoc.", "ca boi duoi nuoc.",
    "chim bay tren troi.", "meo ngu tren ghe.", "cho ngu duoi ban.", "meo an ca.", "Cho an xuong.", "",
]
VAL_CORPUS = ["cho an ca.", "meo uong nuoc."]   # tập kiểm định: câu chưa thấy khi huấn luyện
CHAT_PAIRS = [
    ("meo an gi", "ca."), ("cho an gi", "xuong."), ("meo uong gi", "sua."), ("cho uong gi", "nuoc."),
    ("ca o dau", "duoi nuoc."), ("chim o dau", "tren troi."), ("meo ngu o dau", "tren ghe."),
    ("cho ngu o dau", "duoi ban."),
]
UNSEEN_PROMPT = "meo bay o dau"   # câu hỏi chưa có trong dữ liệu SFT
PREF_PAIRS = [   # (prompt, chosen, rejected): rejected là câu trả lời "gần đúng" nhưng sai
    ("meo an gi", "ca.", "sua."), ("cho an gi", "xuong.", "nuoc."),
    ("meo uong gi", "sua.", "ca."), ("cho uong gi", "nuoc.", "xuong."),
    ("ca o dau", "duoi nuoc.", "tren troi."), ("chim o dau", "tren troi.", "duoi nuoc."),
]


# ── Tiện ích in ──────────────────────────────────────────────────────────────
def say(n, name, codes, lines):
    print(f"[{n} · {name}] ô {codes}")
    for line in lines:
        print("   " + line)
    print()


def vec(a, nd=3):
    return "[" + ", ".join(f"{x:.{nd}f}" for x in np.asarray(a, dtype=float).ravel()) + "]"


# ── [1] Dữ liệu → token IDs ──────────────────────────────────────────────────
def clean_corpus(raw):
    stripped = [t.strip().lower() for t in raw]
    non_empty = [t for t in stripped if t]
    unique = list(dict.fromkeys(non_empty))
    return unique, len(raw) - len(non_empty), len(non_empty) - len(unique)


CLEAN, N_EMPTY, N_DUP = clean_corpus(RAW_CORPUS)
CHARS = sorted(set("".join(CLEAN + VAL_CORPUS + [u + a for u, a in CHAT_PAIRS]
                           + [p + c + r for p, c, r in PREF_PAIRS] + [UNSEEN_PROMPT])))
ITOS = SPECIALS + CHARS
STOI = {t: i for i, t in enumerate(ITOS)}
V = len(ITOS)


def encode_text(text):
    return [STOI[c] for c in text]


def decode_ids(ids):
    return "".join(ITOS[i] for i in ids if i != EOS and i >= len(SPECIALS))


def doc_ids(text):
    return encode_text(text) + [EOS]


def make_batch(inputs, labels, doc, pos, mask):
    inputs, labels = np.asarray(inputs, dtype=int), np.asarray(labels, dtype=int)
    doc, pos, mask = np.asarray(doc), np.asarray(pos), np.asarray(mask, dtype=bool)
    allowed = (doc[:, None] == doc[None, :]) & (pos[None, :] <= pos[:, None])   # mặt nạ nhân quả theo từng tài liệu
    return dict(inputs=inputs, labels=labels, doc=doc, pos=pos, mask=mask, allowed=allowed)


def pack_lm(seqs):
    """Đóng gói nhiều tài liệu vào một batch: nhãn là token kế tiếp, mọi vị trí đều được chấm."""
    inputs, labels, doc, pos = [], [], [], []
    for i, ids in enumerate(seqs):
        inputs += ids[:-1]
        labels += ids[1:]
        doc += [i] * (len(ids) - 1)
        pos += list(range(len(ids) - 1))
    return make_batch(inputs, labels, doc, pos, np.ones(len(inputs), dtype=bool))


def chat_ids(user, answer):
    prompt = [USER] + encode_text(user) + [ASSISTANT]
    return prompt + encode_text(answer) + [EOS], len(prompt)


def pack_sft(pairs, response_only=True):
    inputs, labels, doc, pos, mask = [], [], [], [], []
    for i, (user, answer) in enumerate(pairs):
        ids, start = chat_ids(user, answer)
        n = len(ids) - 1
        inputs += ids[:-1]
        labels += ids[1:]
        doc += [i] * n
        pos += list(range(n))
        mask += [(t >= start - 1) or not response_only for t in range(n)]
    return make_batch(inputs, labels, doc, pos, mask)


# ── Mô hình: bộ giải mã Transformer nhỏ với gradient giải tích ───────────────
def rms(x):
    return x / np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + 1e-6)


def rms_backward(x, grad):
    r = 1 / np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + 1e-6)
    return grad * r - x * r ** 3 * np.mean(grad * x, axis=-1, keepdims=True)


def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)


def pos_enc(positions, d):
    return POS_SCALE * np.sin(np.asarray(positions)[:, None] / (100 ** (np.arange(d) / d)))


def init_params(seed, vocab, d, layers, ffn):
    rng = np.random.default_rng(seed)
    p = {"E": rng.normal(0, .15, (vocab, d)), "U": rng.normal(0, .15, (d, vocab))}
    for l in range(layers):
        for name, shape in (("Q", (d, d)), ("K", (d, d)), ("V", (d, d)), ("O", (d, d)),
                            ("G", (d, ffn)), ("Z", (d, ffn)), ("F", (ffn, d))):
            p[f"{name}{l}"] = rng.normal(0, .15, shape)
    return p


def n_layers(p):
    return sum(1 for k in p if k[0] == "Q" and k[1:].isdigit())


def layer_w(p, l):
    return {n: p[f"{n}{l}"] for n in "QKVOGZF"}


def n_params(p):
    return int(sum(v.size for v in p.values()))


def block_forward(x, W, allowed, past=None, ad=None, l=0):
    """Một khối Transformer: RMSNorm → attention → cộng residual → RMSNorm → SwiGLU → cộng residual."""
    def proj(inp, name):
        y = inp @ W[name]
        if ad is not None and f"{name}{l}" in ad:
            A, B, s = ad[f"{name}{l}"]
            y = y + s * (inp @ A) @ B     # đường adapter chưa gộp: x·W + s·(x·A)·B
        return y
    d = x.shape[1]
    n = rms(x)
    q, kn, vn = proj(n, "Q"), proj(n, "K"), proj(n, "V")
    k = kn if past is None else np.concatenate([past[0], kn])
    v = vn if past is None else np.concatenate([past[1], vn])
    scores = np.where(allowed, q @ k.T / np.sqrt(d), -np.inf)
    a = softmax(scores)
    mixed = a @ v
    h = x + proj(mixed, "O")
    m = rms(h)
    gate, val = proj(m, "G"), proj(m, "Z")
    sig = 1 / (1 + np.exp(-gate))
    hidden = gate * sig * val
    f = h + proj(hidden, "F")
    cache = dict(x=x, n=n, q=q, k=k, v=v, a=a, mixed=mixed, h=h, m=m, gate=gate, val=val,
                 sig=sig, hidden=hidden, allowed=allowed)
    return f, cache, (k, v)


def block_backward(df, c, W):
    g = {"F": c["hidden"].T @ df}
    dhidden = df @ W["F"].T
    dgate = dhidden * c["val"] * (c["sig"] + c["gate"] * c["sig"] * (1 - c["sig"]))
    dval = dhidden * c["gate"] * c["sig"]
    g["G"], g["Z"] = c["m"].T @ dgate, c["m"].T @ dval
    dh = df + rms_backward(c["h"], dgate @ W["G"].T + dval @ W["Z"].T)
    g["O"] = c["mixed"].T @ dh
    dmixed = dh @ W["O"].T
    dv = c["a"].T @ dmixed
    da = dmixed @ c["v"].T
    ds = c["a"] * (da - (da * c["a"]).sum(axis=-1, keepdims=True)) * c["allowed"]
    d = c["x"].shape[1]
    dq, dk = ds @ c["k"] / np.sqrt(d), ds.T @ c["q"] / np.sqrt(d)
    g["Q"], g["K"], g["V"] = c["n"].T @ dq, c["n"].T @ dk, c["n"].T @ dv
    dn = dq @ W["Q"].T + dk @ W["K"].T + dv @ W["V"].T
    return dh + rms_backward(c["x"], dn), g


def forward(p, inputs, pos, allowed, ad=None):
    d = p["E"].shape[1]
    x = p["E"][inputs] + pos_enc(pos, d)
    caches = []
    for l in range(n_layers(p)):
        x, c, _ = block_forward(x, layer_w(p, l), allowed, ad=ad, l=l)
        caches.append(c)
    final = rms(x)
    return final @ p["U"], dict(inputs=inputs, caches=caches, xL=x, final=final)


def backward(p, saved, dlogits):
    grads = {"U": saved["final"].T @ dlogits}
    df = rms_backward(saved["xL"], dlogits @ p["U"].T)
    for l in reversed(range(n_layers(p))):
        df, g = block_backward(df, saved["caches"][l], layer_w(p, l))
        for name, val in g.items():
            grads[f"{name}{l}"] = val
    grads["E"] = np.zeros_like(p["E"])
    np.add.at(grads["E"], saved["inputs"], df)
    return grads


def cross_entropy(logits, labels, mask, denom=None):
    """Loss trung bình trên các vị trí có mask; trả thêm gradient theo logits = p - one-hot (chia theo denom)."""
    mask = np.asarray(mask, dtype=bool)
    denom = mask.sum() if denom is None else denom
    p = softmax(logits)
    idx = np.arange(len(labels))
    nll = -np.log(p[idx, labels])
    dlogits = p.copy()
    dlogits[idx, labels] -= 1
    dlogits *= mask[:, None] / denom
    return float((nll * mask).sum() / denom), dlogits, p


def loss_and_grads(p, batch, denom=None):
    logits, saved = forward(p, batch["inputs"], batch["pos"], batch["allowed"])
    loss, dlogits, _ = cross_entropy(logits, batch["labels"], batch["mask"], denom)
    return loss, backward(p, saved, dlogits), logits, saved


def loss_only(p, batch):
    logits, _ = forward(p, batch["inputs"], batch["pos"], batch["allowed"])
    return cross_entropy(logits, batch["labels"], batch["mask"])[0]


def adamw_update(w, g, m, v, t, lr):
    """Một bước AdamW trên mảng bất kỳ (phép tính theo từng phần tử)."""
    m = BETA1 * m + (1 - BETA1) * g
    v = BETA2 * v + (1 - BETA2) * g * g
    m_hat, v_hat = m / (1 - BETA1 ** t), v / (1 - BETA2 ** t)
    return w * (1 - lr * WEIGHT_DECAY) - lr * m_hat / (np.sqrt(v_hat) + EPS), m, v


def new_opt(p):
    return dict(t=0, m={k: np.zeros_like(x) for k, x in p.items()}, v={k: np.zeros_like(x) for k, x in p.items()})


def adamw_step(p, g, opt, lr):
    opt["t"] += 1
    out = {}
    for k in p:
        out[k], opt["m"][k], opt["v"][k] = adamw_update(p[k], g[k], opt["m"][k], opt["v"][k], opt["t"], lr)
    return out


def copy_params(p):
    return {k: x.copy() for k, x in p.items()}


def copy_opt(opt):
    return dict(t=opt["t"], m=copy_params(opt["m"]), v=copy_params(opt["v"]))


def predict_greedy(p, prompt, max_new=MAX_NEW):
    """Sinh bằng cách tính lại toàn bộ chuỗi ở mỗi bước (không dùng cache): chuẩn tham chiếu."""
    ids = list(prompt)
    for _ in range(max_new):
        n = len(ids)
        allowed = np.arange(n)[None, :] <= np.arange(n)[:, None]
        logits, _ = forward(p, np.array(ids), np.arange(n), allowed)
        nxt = int(np.argmax(logits[-1]))
        ids.append(nxt)
        if nxt == EOS:
            break
    return ids[len(prompt):]


def answer_of(p, user):
    prompt = [USER] + encode_text(user) + [ASSISTANT]
    return decode_ids(predict_greedy(p, prompt))


def exact_matches(p, pairs=CHAT_PAIRS):
    return sum(answer_of(p, u) == a for u, a in pairs)


def check_gradients(p, batch, grads, eps=1e-5):
    """So gradient giải tích với sai phân hữu hạn trung tâm ở phần tử có |gradient| lớn nhất của mỗi tensor."""
    worst, worst_name, show = 0.0, "", None
    for name, tensor in p.items():
        idx = np.unravel_index(np.abs(grads[name]).argmax(), tensor.shape)
        vals = []
        for sign in (1, -1):
            q = copy_params(p)
            q[name][idx] += sign * eps
            vals.append(loss_only(q, batch))
        numeric = (vals[0] - vals[1]) / (2 * eps)
        err = abs(numeric - grads[name][idx])
        if err > worst:
            worst, worst_name, show = err, name, (grads[name][idx], numeric)
    return worst, worst_name, show


def batch_indices(step, n_docs, per_batch):
    per_epoch = max(1, n_docs // per_batch)
    epoch, k = divmod(step, per_epoch)
    order = np.random.default_rng(SEED + 100 + epoch).permutation(n_docs)
    return order[k * per_batch:(k + 1) * per_batch]


def train_lm(p, opt, seqs, start_step, end_step, lr, hist=None, val_batch=None, full_batch=None):
    """Vòng lặp: batch → forward → loss → backward → AdamW; định kỳ đo loss trên cả tập huấn luyện và kiểm định."""
    for step in range(start_step, end_step):
        idx = batch_indices(step, len(seqs), BATCH_DOCS)
        batch = pack_lm([seqs[i] for i in idx])
        loss, grads, _, _ = loss_and_grads(p, batch)
        p = adamw_step(p, grads, opt, lr)
        if hist is not None and (step + 1) % EVAL_EVERY == 0:
            hist.append((step + 1, loss, loss_only(p, full_batch), loss_only(p, val_batch)))
    return p


def save_checkpoint(p, opt=None, extra=None):
    """Ghi checkpoint thành byte (npz). Bản weights-only không có m, v; bản tiếp tục được có thêm m, v và bước."""
    arrays = {f"w.{k}": x for k, x in p.items()}
    arrays["vocab"] = np.array(ITOS)
    if opt is not None:
        arrays.update({f"m.{k}": x for k, x in opt["m"].items()})
        arrays.update({f"v.{k}": x for k, x in opt["v"].items()})
        arrays["step"] = np.array(opt["t"])
    buf = io.BytesIO()
    np.savez(buf, **arrays)
    return buf.getvalue()


def load_checkpoint(blob):
    z = np.load(io.BytesIO(blob), allow_pickle=False)
    p = {k[2:]: z[k] for k in z.files if k.startswith("w.")}
    opt = None
    if "step" in z.files:
        opt = dict(t=int(z["step"]), m={k[2:]: z[k] for k in z.files if k.startswith("m.")},
                   v={k[2:]: z[k] for k in z.files if k.startswith("v.")})
    return p, opt, [str(t) for t in z["vocab"]]


# ── Chuẩn bị chung ───────────────────────────────────────────────────────────
TRAIN_SEQS = [doc_ids(t) for t in CLEAN]
VAL_SEQS = [doc_ids(t) for t in VAL_CORPUS]
TRAIN_ALL = pack_lm(TRAIN_SEQS)
VAL_ALL = pack_lm(VAL_SEQS)
P0 = init_params(SEED, V, D_MODEL, N_LAYERS, FFN_HIDDEN)
N_P = n_params(P0)
LN_V = float(np.log(V))


def main():
    print("=== LAB TỔNG HỢP LLM TRAINING: TỪ DỮ LIỆU ĐẾN CHECKPOINT VÀ SUY LUẬN ===")
    print("MỤC TIÊU: tự tay chạy cả vòng huấn luyện một LLM thu nhỏ, đọc số thật ở từng ô của sơ đồ rồi đưa checkpoint sang suy luận.")
    print(f"MÔ HÌNH: bộ giải mã Transformer {N_LAYERS} lớp, 1 head attention, chiều ẩn d = {D_MODEL}, SwiGLU {FFN_HIDDEN}, "
          f"từ vựng {V} token, tổng {N_P} tham số, trọng số khởi tạo ngẫu nhiên seed {SEED}.")
    print(f"DỮ LIỆU: {len(RAW_CORPUS)} dòng thô, {len(CHAT_PAIRS)} cặp hỏi đáp cho SFT, {len(PREF_PAIRS)} cặp ưu tiên cho DPO; tất cả do tác giả tự viết.")
    print("CÁCH ĐỌC: mỗi khối [n · Tên] mang tên giai đoạn và mã ô trên sơ đồ Tab 01. Đọc các số, tự nói điều chúng cho thấy, rồi đối chiếu dòng Kết luận và Không chứng minh.")
    print("LƯU Ý: mọi mô hình ở đây là đồ chơi để quan sát cơ chế; không số nào là tốc độ hay chất lượng của mô hình thật.\n")

    # ── [1] Dữ liệu → token IDs (1A) ────────────────────────────────────────
    first = CLEAN[0]
    ids_first = doc_ids(first)
    batch_demo_idx = batch_indices(0, len(TRAIN_SEQS), BATCH_DOCS)
    n_batch_tokens = sum(len(TRAIN_SEQS[i]) - 1 for i in batch_demo_idx)
    say(1, "Dữ liệu → token IDs", "1A", [
        f"Quan sát: kho thô có {len(RAW_CORPUS)} dòng; bỏ {N_EMPTY} dòng trống và {N_DUP} dòng trùng (sau khi bỏ khoảng trắng thừa và đưa về chữ thường) còn {len(CLEAN)} tài liệu huấn luyện; tách riêng {len(VAL_CORPUS)} tài liệu kiểm định.",
        f"Tokenizer ký tự: {len(CHARS)} ký tự (lấy từ mọi văn bản dùng trong lab để không có ký tự lạ) + 4 token đặc biệt {SPECIALS} = {V} token. Bảng từ vựng: {ITOS}",
        f"Tài liệu đầu \"{first}\" -> token IDs {ids_first} (cuối là <eos> có ID {EOS}).",
        f"Một batch gồm {BATCH_DOCS} tài liệu (chỉ số {batch_demo_idx.tolist()}) được đóng gói thành {n_batch_tokens} vị trí; mỗi bước huấn luyện lấy một batch như vậy, {len(TRAIN_SEQS) // BATCH_DOCS} bước là hết một epoch (một lượt qua dữ liệu).",
        "Kết luận: mô hình không bao giờ thấy chữ, chỉ thấy các số nguyên này; cùng văn bản phải luôn cho cùng ID, nên tokenizer phải được lưu cùng checkpoint.",
        "Không chứng minh: ID là số thứ tự trong bảng từ vựng, chưa mang ý nghĩa; kho chỉ có vài chục ký tự nên không nói gì về chất lượng dữ liệu thật hay tokenizer BPE.",
    ])

    # ── [2] Next-token targets (1B) ─────────────────────────────────────────
    b = pack_lm([ids_first])
    pairs = [f"{ITOS[i]!r}->{ITOS[j]!r}" for i, j in zip(b["inputs"][:6], b["labels"][:6])]
    two = pack_lm(TRAIN_SEQS[:2])
    cross_doc = bool(two["allowed"][len(TRAIN_SEQS[0]) - 1, 0])
    say(2, "Next-token targets", "1B", [
        f"Quan sát: input = ids[:-1] có {len(b['inputs'])} vị trí, label = ids[1:] dịch đi một vị trí. Sáu cặp đầu (input -> label): {', '.join(pairs)}.",
        f"Ma trận mặt nạ nhân quả có shape {b['allowed'].shape} (hàng = vị trí đang tính, cột = vị trí được nhìn); vị trí 3 nhìn được {int(b['allowed'][3].sum())} vị trí (0..3), không nhìn thấy tương lai.",
        f"Khi đóng gói 2 tài liệu, vị trí đầu của tài liệu thứ hai có nhìn thấy tài liệu thứ nhất không? {cross_doc} (mặt nạ cắt theo ranh giới tài liệu).",
        f"Mọi vị trí đều được chấm ở bước này: {int(b['mask'].sum())} trên {len(b['mask'])}. (Giai đoạn 9 sẽ cho thấy mặt nạ loss chỉ chấm một phần.)",
        "Kết luận: một tài liệu n token tự sinh ra n-1 bài toán dự đoán, mỗi bài có đáp án là token đứng ngay sau; không cần ai gán nhãn.",
        "Không chứng minh: việc dịch nhãn chỉ chuẩn bị đề bài; chưa nói mô hình giải đúng hay sai bài nào.",
    ])

    # ── [3] Forward (2A–2D) ─────────────────────────────────────────────────
    logits, saved = forward(P0, b["inputs"], b["pos"], b["allowed"])
    c0 = saved["caches"][0]
    row = c0["a"][3]
    x0 = P0["E"][b["inputs"]] + pos_enc(b["pos"], D_MODEL)
    alt = b["inputs"].copy()
    alt[-1] = (alt[-1] + 1) % V
    logits_alt, _ = forward(P0, alt, b["pos"], b["allowed"])
    causal_diff = float(np.abs(logits[:-1] - logits_alt[:-1]).max())
    rms_before, rms_after = float(np.sqrt((x0[0] ** 2).mean())), float(np.sqrt((rms(x0)[0] ** 2).mean()))
    zeroed = copy_params(P0)
    zeroed["O0"], zeroed["F0"] = np.zeros_like(P0["O0"]), np.zeros_like(P0["F0"])
    x_in = P0["E"][b["inputs"]] + pos_enc(b["pos"], D_MODEL)
    out_zero = block_forward(x_in, layer_w(zeroed, 0), b["allowed"])[0]
    residual_gap = float(np.abs(out_zero - x_in).max())
    p_row = softmax(logits[3])
    say(3, "Forward: embedding, attention, MLP, norm → logits", "2A–2D", [
        f"Quan sát: {len(b['inputs'])} token ID đi vào; bảng embedding E có shape {P0['E'].shape} tra ra x có shape {x0.shape} = (số vị trí, chiều ẩn d).",
        f"Attention (lớp 0): hàng của vị trí 3 = {vec(row)} có tổng {row.sum():.3f} và chỉ có {int((row > 0).sum())} phần tử khác 0 (các vị trí 0..3); đây là trọng số trên VỊ TRÍ, không phải xác suất của token.",
        f"Mặt nạ nhân quả: đổi token CUỐI rồi chạy lại, logits ở các vị trí trước đó đổi tối đa {causal_diff:.1e} (đúng 0 nghĩa là tương lai không ảnh hưởng quá khứ).",
        f"RMSNorm: độ lớn trung bình (RMS) của một vector x từ {rms_before:.3f} về {rms_after:.3f}; norm giữ biên độ ổn định giữa các lớp.",
        f"Residual: đầu ra của khối = đầu vào + phần attention + phần MLP (SwiGLU). Thử đặt ma trận ra O và F của lớp 0 bằng 0 (hai nhánh không đóng góp gì): đầu ra khối lệch đầu vào tối đa {residual_gap:.1e}, tức đường tắt cho thông tin đi thẳng qua; nhờ đó thông tin và gradient không bị bóp nghẹt khi xếp nhiều lớp.",
        f"Đầu ra logits có shape {logits.shape} = (số vị trí, {V} token); hàng vị trí 3 sau softmax có tổng {p_row.sum():.3f}, token cao nhất là {ITOS[int(p_row.argmax())]!r} với xác suất {p_row.max():.3f}. Mô hình chưa học nên chưa ưu tiên token nào rõ rệt (đoán đều sẽ cho 1/{V} = {1 / V:.3f} mỗi token).",
        "Kết luận: forward biến từng ID thành một bảng điểm cho mọi token trong từ vựng ở mọi vị trí cùng lúc, qua 4 bước: attention lấy thông tin từ quá khứ, MLP xử lý riêng từng vị trí, norm giữ biên độ, residual giữ đường tắt.",
        "Không chứng minh: trọng số là ngẫu nhiên, nên attention và logits ở đây chưa mang ý nghĩa ngôn ngữ; ví dụ chỉ cho thấy luồng tensor và tính nhân quả.",
    ])

    # ── [4] Cross-entropy loss (3A) ─────────────────────────────────────────
    loss0, dl0, p_all = cross_entropy(logits, b["labels"], b["mask"])
    pc = p_all[np.arange(len(b["labels"])), b["labels"]]
    rows = [f"vị trí {t}: đáp án {ITOS[b['labels'][t]]!r}, p(đáp án) = {pc[t]:.4f}, -ln p = {-np.log(pc[t]):.3f}" for t in range(4)]
    say(4, "Cross-entropy loss", "3A", [
        f"Quan sát: loss = -ln p(token đúng), trung bình trên các vị trí có mặt nạ. Bốn vị trí đầu: " + " | ".join(rows) + ".",
        f"Trung bình trên {int(b['mask'].sum())} vị trí = {loss0:.4f}. So với đoán đều: ln {V} = {LN_V:.4f}; mô hình chưa học nên loss gần mức này.",
        f"Nếu mô hình chắc chắn đúng (p = 1) thì loss = 0; nếu gán p = 0,01 cho đáp án thì loss = {-np.log(0.01):.3f}. Loss càng thấp = đáp án nhận xác suất càng cao.",
        "Kết luận: loss gói cả bảng điểm vào MỘT con số dương duy nhất; chính con số này (và gradient của nó) là thứ vòng huấn luyện cố làm nhỏ đi.",
        "Không chứng minh: loss thấp không có nghĩa câu trả lời đúng, hữu ích hay an toàn; nó chỉ đo xác suất gán cho token kế tiếp trên dữ liệu kiểu này.",
    ])

    # ── [5] Backward & gradient từng lớp (3B–3C) ────────────────────────────
    loss_b, grads, _, _ = loss_and_grads(P0, TRAIN_ALL)
    logits_t, saved_t = forward(P0, TRAIN_ALL["inputs"], TRAIN_ALL["pos"], TRAIN_ALL["allowed"])
    _, dl_t, _ = cross_entropy(logits_t, TRAIN_ALL["labels"], TRAIN_ALL["mask"])
    row_sum = float(np.abs(dl_t.sum(axis=1)).max())
    layer_norm = {l: float(np.sqrt(sum((grads[f"{n}{l}"] ** 2).sum() for n in "QKVOGZF"))) for l in range(N_LAYERS)}
    worst, worst_name, show = check_gradients(P0, TRAIN_ALL, grads)
    assert worst < 1e-6, f"gradient giải tích lệch sai phân hữu hạn {worst}"
    say(5, "Backward & gradient từng lớp", "3B–3C", [
        f"Quan sát: gradient của loss theo logits là p - one-hot; tổng mỗi hàng bằng {row_sum:.1e} (xác suất tổng 1, one-hot tổng 1).",
        f"Độ lớn gradient (chuẩn L2): embedding E = {np.linalg.norm(grads['E']):.4f}, unembedding U = {np.linalg.norm(grads['U']):.4f}; gom theo lớp (Q, K, V, O, G, Z, F): " + ", ".join(f"lớp {l} = {layer_norm[l]:.4f}" for l in range(N_LAYERS)) + ". Mỗi lớp có gradient riêng vì gradient đi ngược qua từng khối.",
        f"Kiểm tra bằng sai phân hữu hạn (dịch một phần tử ±1e-5 rồi tính lại loss), ở phần tử có gradient lớn nhất của mỗi trong {len(P0)} tensor: sai lệch lớn nhất {worst:.2e} (tensor {worst_name}; giải tích {show[0]:.8f} so với số {show[1]:.8f}).",
        "Kết luận: backward đi từ loss ngược qua unembedding, các lớp (từng lớp tách gradient thành nhánh residual và nhánh attention/MLP) tới embedding; cho mỗi tham số biết tăng hay giảm bao nhiêu thì loss giảm; kết quả khớp phép đo số trực tiếp.",
        "Không chứng minh: khớp ở các phần tử lớn nhất của từng tensor là bằng chứng mạnh nhưng không phải kiểm tra mọi phần tử; gradient chỉ nói hướng dốc xuống cục bộ, chưa nói sửa bao nhiêu (đó là việc của optimizer).",
    ])

    # ── [6] AdamW (3D) ──────────────────────────────────────────────────────
    name, shape_idx = "Q0", np.unravel_index(np.abs(grads["Q0"]).argmax(), P0["Q0"].shape)
    g1, w1 = grads["Q0"][shape_idx], P0["Q0"][shape_idx]
    m1, v1 = (1 - BETA1) * g1, (1 - BETA2) * g1 * g1
    mh, vh = m1 / (1 - BETA1), v1 / (1 - BETA2)
    by_hand = w1 * (1 - LR * WEIGHT_DECAY) - LR * mh / (np.sqrt(vh) + EPS)
    opt_demo = new_opt(P0)
    p_after = adamw_step(P0, grads, opt_demo, LR)
    assert abs(by_hand - p_after["Q0"][shape_idx]) < 1e-12
    steps_adam = np.concatenate([np.abs(p_after[k] - P0[k]).ravel() for k in P0])
    steps_sgd = np.concatenate([np.abs(LR * grads[k]).ravel() for k in P0])
    nz = np.concatenate([np.abs(grads[k]).ravel() for k in P0]) > 1e-12
    loss_after = loss_only(p_after, TRAIN_ALL)
    say(6, "AdamW", "3D", [
        f"Quan sát (bước 1, phần tử {name}{tuple(int(i) for i in shape_idx)}): gradient g = {g1:.6f}; m = (1-β1)·g = {m1:.6f}; v = (1-β2)·g² = {v1:.3e}; sau hiệu chỉnh m̂ = {mh:.6f}, v̂ = {vh:.3e}.",
        f"Cập nhật = lr·m̂/(√v̂+ε) = {LR * mh / (np.sqrt(vh) + EPS):.6f} (ở bước đầu gần bằng lr = {LR}, đúng dấu của g); suy giảm trọng số lr·wd·w = {LR * WEIGHT_DECAY * w1:.2e} tách riêng. Trọng số {w1:.6f} -> {p_after['Q0'][shape_idx]:.6f}; tính tay khớp mã (sai lệch < 1e-12).",
        f"Toàn mô hình, bước đầu: AdamW đổi mỗi tham số có gradient khác 0 một lượng trong [{steps_adam[nz].min():.4f}, {steps_adam[nz].max():.4f}], còn SGD với cùng lr đổi từ {steps_sgd[nz].min():.1e} đến {steps_sgd[nz].max():.1e}.",
        f"Một bước trên cả tập huấn luyện: loss {loss_b:.4f} -> {loss_after:.4f}.",
        "Kết luận: gradient chỉ hướng, AdamW quyết định bước dài bao nhiêu bằng trung bình trượt của gradient (m) và của gradient bình phương (v); mỗi tham số có độ dài bước riêng.",
        "Không chứng minh: loss giảm sau MỘT bước trên đúng batch đó chưa nói gì về bước sau hay dữ liệu chưa thấy; lr quá lớn vẫn làm loss tăng.",
    ])

    # ── [7] Vòng lặp, eval, checkpoint (3E) ─────────────────────────────────
    hist = []
    resume_at = min(RESUME_AT, STEPS)
    p_run = copy_params(P0)
    opt_run = new_opt(P0)
    p_run = train_lm(p_run, opt_run, TRAIN_SEQS, 0, resume_at, LR, hist, VAL_ALL, TRAIN_ALL)
    ckpt_resume = save_checkpoint(p_run, opt_run)
    ckpt_weights = save_checkpoint(p_run)
    p_mid, opt_mid = copy_params(p_run), copy_opt(opt_run)
    p_run = train_lm(p_run, opt_run, TRAIN_SEQS, resume_at, STEPS, LR, hist, VAL_ALL, TRAIN_ALL)
    PRE = copy_params(p_run)
    val_pre = loss_only(PRE, VAL_ALL)
    # Chạy tiếp từ checkpoint đầy đủ so với chỉ nạp trọng số.
    p_r, opt_r, _ = load_checkpoint(ckpt_resume)
    p_r = train_lm(p_r, opt_r, TRAIN_SEQS, resume_at, resume_at + 10, LR)
    p_u = copy_params(p_mid)
    opt_u = new_opt(p_mid)
    p_u = train_lm(p_u, opt_u, TRAIN_SEQS, resume_at, resume_at + 10, LR)
    p_ref = train_lm(copy_params(p_mid), copy_opt(opt_mid), TRAIN_SEQS, resume_at, resume_at + 10, LR)
    gap_resume = max(float(np.abs(p_r[k] - p_ref[k]).max()) for k in p_ref)
    gap_weights_only = max(float(np.abs(p_u[k] - p_ref[k]).max()) for k in p_ref)
    prompt_demo = [STOI[c] for c in "meo an"]
    greedy_demo = decode_ids(predict_greedy(PRE, prompt_demo))
    best_val = min(hist, key=lambda h: h[3])
    table = ["bước | loss batch | loss cả tập huấn luyện | loss kiểm định"] + \
            [f"{s:4d} | {a:.4f} | {t:.4f} | {v:.4f}" for s, a, t, v in hist]
    say(7, "Vòng lặp, eval, checkpoint", "3E, 4C", [
        f"Quan sát: {STEPS} bước AdamW (lr {LR}, batch {BATCH_DOCS} tài liệu, {len(TRAIN_SEQS) // BATCH_DOCS} bước mỗi epoch) trên {len(TRAIN_SEQS)} tài liệu. Đây là huấn luyện nền (pretraining) thu nhỏ: nhãn là token kế tiếp (ô 4C).",
        *["   " + t for t in table],
        f"Loss khởi đầu ≈ ln {V} = {LN_V:.3f}; sau {STEPS} bước loss trên tập huấn luyện = {hist[-1][2]:.4f} (giảm {LN_V - hist[-1][2]:.2f} nat), loss kiểm định = {hist[-1][3]:.4f}.",
        f"Thử sinh: prompt \"meo an\" -> \"{greedy_demo}\" (chọn token điểm cao nhất mỗi bước).",
        f"Checkpoint tại bước {resume_at}: bản weights-only {len(ckpt_weights)} byte; bản tiếp tục được (thêm m, v và số bước) {len(ckpt_resume)} byte, gấp {len(ckpt_resume) / len(ckpt_weights):.2f} lần.",
        f"Chạy tiếp 10 bước từ bản đầy đủ: trọng số lệch tối đa {gap_resume:.1e} so với chạy liền không ngắt; nạp chỉ trọng số (m = v = 0, bộ đếm bước về 0): lệch {gap_weights_only:.1e}.",
        "Kết luận: loop lặp đi lặp lại chu trình ở cột 1-3, mỗi lần chỉ sửa một chút, và định kỳ đo loss trên dữ liệu không dùng để cập nhật; checkpoint đầy đủ khôi phục đúng đường huấn luyện, còn bản weights-only đủ cho suy luận nhưng không tiếp tục đúng.",
        f"Cách đọc loss kiểm định: thấp nhất {best_val[3]:.4f} ở bước {best_val[0]}, ở bước cuối {val_pre:.4f}. " + (f"Loss huấn luyện vẫn giảm còn loss kiểm định đi lên là dấu hiệu mô hình bắt đầu học thuộc {len(TRAIN_SEQS)} tài liệu thay vì học điều tổng quát." if best_val[0] < hist[-1][0] else "Ở các mốc đã đo, loss kiểm định chưa đi lên; chưa thấy dấu hiệu học thuộc."),
        f"Không chứng minh: kho chỉ có {len(TRAIN_SEQS)} tài liệu và kiểm định chỉ {len(VAL_SEQS)} câu, nên kết luận về học thuộc hay tổng quát hóa ở đây chỉ mang tính minh họa và không nói gì về mô hình thật.",
    ])

    # ── [8] Bộ nhớ: ZeRO & recompute (4A–4B) ────────────────────────────────
    def zero_bytes(P, N):
        return {"data parallel": 16 * P, "ZeRO-1": 4 * P + 12 * P / N, "ZeRO-2": 2 * P + 14 * P / N, "ZeRO-3": 16 * P / N}
    lab_b = zero_bytes(N_P, ZERO_WORKERS)
    big = {k: v / 1e9 for k, v in zero_bytes(7e9, 8).items()}
    # Data parallel: mỗi worker tính gradient trên nửa batch, chuẩn hóa theo tổng token của cả batch.
    shares = [[TRAIN_SEQS[i] for i in idx_] for idx_ in np.array_split(np.arange(len(TRAIN_SEQS)), ZERO_WORKERS)]
    total_tokens = TRAIN_ALL["mask"].sum()
    g_parts = [loss_and_grads(PRE, pack_lm(s), denom=total_tokens)[1] for s in shares if s]
    g_full = loss_and_grads(PRE, TRAIN_ALL)[1]
    dp_gap = max(float(np.abs(sum(g[k] for g in g_parts) - g_full[k]).max()) for k in g_full)
    # ZeRO-1: mỗi worker giữ một đoạn m, v và cập nhật đoạn đó.
    keys = list(PRE.keys())
    flat = lambda d: np.concatenate([d[k].ravel() for k in keys])
    p_pre_next = adamw_step(PRE, g_full, new_opt(PRE), LR)
    flat_w, flat_g = flat(PRE), flat(g_full)
    m0, v0 = np.zeros_like(flat_w), np.zeros_like(flat_w)
    shards = np.array_split(np.arange(flat_w.size), ZERO_WORKERS)
    new_parts = [adamw_update(flat_w[s], flat_g[s], m0[s], v0[s], 1, LR)[0] for s in shards]
    zero_gap = float(np.abs(np.concatenate(new_parts) - flat(p_pre_next)).max())
    # Recompute: chỉ giữ biên của từng lớp, tính lại khi backward cần.
    logits_f, saved_f = forward(PRE, TRAIN_ALL["inputs"], TRAIN_ALL["pos"], TRAIN_ALL["allowed"])
    _, dlog, _ = cross_entropy(logits_f, TRAIN_ALL["labels"], TRAIN_ALL["mask"])
    boundaries = [PRE["E"][TRAIN_ALL["inputs"]] + pos_enc(TRAIN_ALL["pos"], D_MODEL)]
    for l in range(N_LAYERS):
        boundaries.append(block_forward(boundaries[l], layer_w(PRE, l), TRAIN_ALL["allowed"])[0])
    grads_rc = {"U": saved_f["final"].T @ dlog}
    df = rms_backward(boundaries[-1], dlog @ PRE["U"].T)
    for l in reversed(range(N_LAYERS)):
        _, cache_l, _ = block_forward(boundaries[l], layer_w(PRE, l), TRAIN_ALL["allowed"])   # tính lại activation của lớp l
        df, g = block_backward(df, cache_l, layer_w(PRE, l))
        for nm, val in g.items():
            grads_rc[f"{nm}{l}"] = val
    grads_rc["E"] = np.zeros_like(PRE["E"])
    np.add.at(grads_rc["E"], TRAIN_ALL["inputs"], df)
    rc_gap = max(float(np.abs(grads_rc[k] - g_full[k]).max()) for k in g_full)
    stored_full = sum(v.size for c in saved_f["caches"] for k, v in c.items() if isinstance(v, np.ndarray) and k != "allowed")
    stored_rc = sum(x.size for x in boundaries[:-1]) + sum(v.size for k, v in saved_f["caches"][0].items() if isinstance(v, np.ndarray) and k != "allowed")
    L_demo, act, bound = 32, 10, 1
    say(8, "Bộ nhớ: ZeRO & recompute", "4A–4B", [
        f"Quan sát (ZeRO, đếm byte): mô hình lab có {N_P} tham số; mixed precision với AdamW tốn 2 (trọng số) + 2 (gradient) + 12 (bản chính 32 bit, m, v) = 16 byte mỗi tham số.",
        f"Với {ZERO_WORKERS} worker, byte mỗi worker: data parallel {lab_b['data parallel']:.0f}; ZeRO-1 {lab_b['ZeRO-1']:.0f}; ZeRO-2 {lab_b['ZeRO-2']:.0f}; ZeRO-3 {lab_b['ZeRO-3']:.0f}.",
        f"Cùng công thức cho mô hình 7 tỉ tham số trên 8 GPU (tính tay, không đo): data parallel {big['data parallel']:.0f} GB, ZeRO-1 {big['ZeRO-1']:.1f} GB, ZeRO-2 {big['ZeRO-2']:.2f} GB, ZeRO-3 {big['ZeRO-3']:.0f} GB mỗi GPU, chưa tính activation.",
        f"Mô phỏng: {ZERO_WORKERS} worker, mỗi worker tính gradient trên phần dữ liệu riêng (chia theo tổng token của cả batch), cộng lại khớp gradient cả batch (lệch {dp_gap:.1e}); với ZeRO-1 mỗi worker giữ {[len(s_) for s_ in shards]} phần tử của m và của v (tổng {flat_w.size} = số tham số), cập nhật phần của mình rồi gom lại: trọng số mới lệch {zero_gap:.1e} so với cập nhật không chia.",
        f"Quan sát (recompute): lưu mọi activation của {N_LAYERS} lớp tốn {stored_full} số; chỉ lưu biên (đầu vào) của từng lớp cộng activation của một lớp đang được tính lại tốn khoảng {stored_rc} số. Gradient của hai đường lệch tối đa {rc_gap:.1e}.",
        f"Với mô hình đơn vị tự đặt {L_demo} lớp, mỗi lớp {act} đơn vị activation và biên {bound}: không checkpoint {L_demo * act}, checkpoint từng lớp {L_demo * bound + act} đơn vị ({(L_demo * bound + act) / (L_demo * act):.0%}); đổi lại thêm một lượt forward (bước từ 3F lên 4F, tăng khoảng 33% theo mô hình đơn giản).",
        "Kết luận: ZeRO và recompute đổi một thứ lấy một thứ: bộ nhớ lấy giao tiếp hoặc phép tính, và cả hai KHÔNG đổi toán học của huấn luyện (cập nhật và gradient giống hệt).",
        "Không chứng minh: đây là phép đếm byte và chia mảng trong một tiến trình; chưa đo bộ nhớ đỉnh, thời gian, băng thông mạng hay tốc độ huấn luyện phân tán.",
    ])
    assert dp_gap < 1e-12 and zero_gap == 0.0 and rc_gap < 1e-12

    # ── [9] SFT: chat template + response mask (1C, 4D) ─────────────────────
    user0, ans0 = CHAT_PAIRS[0]
    ids0, start0 = chat_ids(user0, ans0)
    sft_full = pack_sft(CHAT_PAIRS, response_only=False)
    sft_resp = pack_sft(CHAT_PAIRS, response_only=True)
    sft_batch = pack_sft(CHAT_PAIRS, response_only=MASK_RESPONSE_ONLY)
    flags0 = [int((t >= start0 - 1) or not MASK_RESPONSE_ONLY) for t in range(len(ids0) - 1)]
    loss_full_mask, loss_resp_mask = loss_only(PRE, sft_full), loss_only(PRE, sft_resp)
    lg, _ = forward(PRE, sft_resp["inputs"], sft_resp["pos"], sft_resp["allowed"])
    _, dl_resp, _ = cross_entropy(lg, sft_resp["labels"], sft_resp["mask"])
    prompt_rows = float(np.abs(dl_resp[~sft_resp["mask"]]).sum())
    before = exact_matches(PRE)
    sft_prompt = dict(sft_resp, mask=~sft_resp["mask"])      # chỉ các vị trí thuộc prompt
    ans_before, prompt_before = loss_only(PRE, sft_resp), loss_only(PRE, sft_prompt)
    p_sft, opt_sft = copy_params(PRE), new_opt(PRE)
    sft_hist = []
    for step in range(SFT_STEPS):
        loss, g, _, _ = loss_and_grads(p_sft, sft_batch)
        p_sft = adamw_step(p_sft, g, opt_sft, LR)
        if step == 0:
            loss_start = loss
    SFT = copy_params(p_sft)
    after = exact_matches(SFT)
    ans_after, prompt_after = loss_only(SFT, sft_resp), loss_only(SFT, sft_prompt)
    sample = [(u, answer_of(SFT, u)) for u, _ in CHAT_PAIRS[:3]]
    unseen = answer_of(SFT, UNSEEN_PROMPT)
    ckpt_sft = save_checkpoint(SFT)
    say(9, "SFT: chat template + response mask", "1C, 4D", [
        f"Quan sát: mẫu hội thoại ghép <user> câu hỏi <assistant> câu trả lời <eos> thành một chuỗi. Cặp đầu: \"{user0}\" -> \"{ans0}\" có token IDs {ids0}.",
        f"Mặt nạ loss của cặp đầu (1 = được chấm, nhãn dịch một vị trí): {flags0}; chỉ phần câu trả lời và <eos> được chấm khi MASK_RESPONSE_ONLY = {MASK_RESPONSE_ONLY}.",
        f"Hàng gradient ở vị trí bị che có độ lớn tổng {prompt_rows:.1e} (đúng 0: prompt không sinh gradient). Trên checkpoint nền, loss tính cả prompt = {loss_full_mask:.4f}, chỉ tính câu trả lời = {loss_resp_mask:.4f}.",
        f"Trước SFT (checkpoint nền, bước 7): {before}/{len(CHAT_PAIRS)} câu hỏi cho đúng câu trả lời mẫu. Sau {SFT_STEPS} bước AdamW (lr {LR}) từ checkpoint nền: {after}/{len(CHAT_PAIRS)}; loss SFT {loss_start:.4f} -> {loss_only(SFT, sft_batch):.4f}.",
        f"Loss tách theo vùng (trước -> sau SFT): trên token câu trả lời {ans_before:.4f} -> {ans_after:.4f}; trên token prompt {prompt_before:.4f} -> {prompt_after:.4f}. " + ("Phần prompt không được chấm thì loss của nó không được cải thiện (ở đây còn tăng)." if MASK_RESPONSE_ONLY else "Prompt cũng được chấm nên loss trên prompt giảm: một phần công sức học đi vào việc chép lại câu hỏi."),
        "Ví dụ sau SFT: " + "; ".join(f"\"{u}\" -> \"{a}\"" for u, a in sample) + f". Câu hỏi chưa có trong dữ liệu \"{UNSEEN_PROMPT}\" -> \"{unseen}\".",
        f"Checkpoint SFT được lưu ({len(ckpt_sft)} byte) để dùng ở giai đoạn 10-13.",
        "Kết luận: SFT là cùng vòng lặp huấn luyện với hai thay đổi: dữ liệu là cặp hỏi đáp trong một mẫu cố định và loss chỉ chấm phần câu trả lời, để mô hình học cách đáp chứ không học chép câu hỏi.",
        f"Không chứng minh: {len(CHAT_PAIRS)} cặp có thể bị học thuộc; trả lời đúng cặp đã thấy chưa phải khả năng trả lời câu hỏi mới (xem dòng chưa thấy), và không đo mức hữu ích hay an toàn.",
    ])
    assert prompt_rows == 0.0

    # ── [10] LoRA (4E) ──────────────────────────────────────────────────────
    s_lora = LORA_ALPHA / LORA_RANK
    targets = [f"{n}{l}" for l in range(N_LAYERS) for n in "QKVOGZF"]
    rng = np.random.default_rng(SEED + 11)
    adapters = {}
    for t in targets:
        din, dout = PRE[t].shape
        adapters[t + ".A"] = rng.normal(0, .1, (din, LORA_RANK))
        adapters[t + ".B"] = np.zeros((LORA_RANK, dout))
    n_ad = sum(a.size for a in adapters.values())
    n_target_full = sum(PRE[t].size for t in targets)

    def merged(base, ad):
        out = dict(base)
        for t in targets:
            out[t] = base[t] + s_lora * ad[t + ".A"] @ ad[t + ".B"]
        return out
    ad_unmerged = {t: (adapters[t + ".A"], adapters[t + ".B"], s_lora) for t in targets}
    lg_base, _ = forward(PRE, sft_batch["inputs"], sft_batch["pos"], sft_batch["allowed"])
    lg_ad0, _ = forward(PRE, sft_batch["inputs"], sft_batch["pos"], sft_batch["allowed"], ad=ad_unmerged)
    step0_gap = float(np.abs(lg_base - lg_ad0).max())
    base_before = copy_params(PRE)
    opt_lora = dict(t=0, m={k: np.zeros_like(a) for k, a in adapters.items()}, v={k: np.zeros_like(a) for k, a in adapters.items()})
    ad_cur = {k: a.copy() for k, a in adapters.items()}
    for step in range(LORA_STEPS):
        eff = merged(PRE, ad_cur)
        loss, g, _, _ = loss_and_grads(eff, sft_batch)
        if step == 0:
            loss_lora_start = loss
        opt_lora["t"] += 1
        for t in targets:
            gA = s_lora * g[t] @ ad_cur[t + ".B"].T
            gB = s_lora * ad_cur[t + ".A"].T @ g[t]
            for key, grad in ((t + ".A", gA), (t + ".B", gB)):
                ad_cur[key], opt_lora["m"][key], opt_lora["v"][key] = adamw_update(
                    ad_cur[key], grad, opt_lora["m"][key], opt_lora["v"][key], opt_lora["t"], LORA_LR)
    base_gap = max(float(np.abs(PRE[k] - base_before[k]).max()) for k in PRE)
    ad_unmerged = {t: (ad_cur[t + ".A"], ad_cur[t + ".B"], s_lora) for t in targets}
    LORA_MERGED = merged(PRE, ad_cur)
    lg_unm, _ = forward(PRE, sft_batch["inputs"], sft_batch["pos"], sft_batch["allowed"], ad=ad_unmerged)
    lg_mrg, _ = forward(LORA_MERGED, sft_batch["inputs"], sft_batch["pos"], sft_batch["allowed"])
    merge_gap = float(np.abs(lg_unm - lg_mrg).max())
    lora_match = exact_matches(LORA_MERGED)
    delta = float(np.abs(LORA_MERGED[targets[0]] - PRE[targets[0]]).max())
    say(10, "LoRA", "4E", [
        f"Quan sát: đóng băng {N_P} trọng số gốc; gắn cặp adapter A (d×r) và B (r×d), r = {LORA_RANK}, hệ số alpha/r = {s_lora:g}, vào {len(targets)} ma trận (Q, K, V, O, G, Z, F của {N_LAYERS} lớp). Tham số học được: {n_ad}, bằng {n_ad / N_P:.1%} số tham số gốc (và {n_ad / n_target_full:.1%} số tham số của chính các ma trận được gắn).",
        f"Khởi tạo B = 0 nên mô hình có adapter y hệt mô hình gốc ở bước 0: logits lệch {step0_gap:.1e}.",
        f"Sau {LORA_STEPS} bước AdamW (lr {LORA_LR}) trên cùng dữ liệu SFT: loss {loss_lora_start:.4f} -> {loss_only(LORA_MERGED, sft_batch):.4f}; trả lời đúng {lora_match}/{len(CHAT_PAIRS)} cặp (SFT đủ ở giai đoạn 9: {after}/{len(CHAT_PAIRS)}).",
        f"Trọng số gốc sau huấn luyện lệch tối đa {base_gap:.1e} so với trước (đúng 0 nghĩa là không bị sửa); độ đổi lớn nhất của một ma trận sau khi gộp {delta:.4f}.",
        f"Gộp W' = W + (alpha/r)·A·B rồi chạy, so với chạy đường tách x·W + (alpha/r)·(x·A)·B: logits lệch tối đa {merge_gap:.1e}.",
        "Kết luận: LoRA chỉ huấn luyện một phần rất nhỏ; ghi ra chỉ cần cặp A, B, còn khi phục vụ có thể gộp vào W để không tốn thêm phép tính (kết quả không đổi).",
        "Không chứng minh: kết quả trên 8 cặp không nói LoRA ngang hay kém tinh chỉnh toàn bộ ở mô hình thật; chọn hạng r, các ma trận được gắn và lượng tử hóa (QLoRA) không được thử ở đây.",
    ])
    assert base_gap == 0.0 and step0_gap < 1e-12 and merge_gap < 1e-9

    # ── [11] Distillation (4F) ──────────────────────────────────────────────
    stu0 = init_params(SEED + 5, V, STUDENT_D, STUDENT_LAYERS, STUDENT_D * 2)
    n_stu = n_params(stu0)
    t_logits, _ = forward(SFT, sft_resp["inputs"], sft_resp["pos"], sft_resp["allowed"])
    mask_r = sft_resp["mask"]
    pt1, pt2 = softmax(t_logits), softmax(t_logits / DISTILL_T)
    first_pos = int(np.flatnonzero(mask_r)[0])
    top = np.argsort(-pt1[first_pos])[:3]

    def kl_rows(p_t, p_s):
        return (p_t * (np.log(p_t + 1e-12) - np.log(p_s + 1e-12))).sum(axis=1)

    def distill_loss_grad(student, T):
        sl, sv = forward(student, sft_resp["inputs"], sft_resp["pos"], sft_resp["allowed"])
        ps, pt = softmax(sl / T), softmax(t_logits / T)
        n = mask_r.sum()
        loss = float((kl_rows(pt, ps) * mask_r).sum() / n) * T * T
        dl = T * (ps - pt) * mask_r[:, None] / n      # T² · d(KL)/d logits = T · (p_s - p_t)
        return loss, backward(student, sv, dl), sl

    def kl_to_teacher(student):
        sl, _ = forward(student, sft_resp["inputs"], sft_resp["pos"], sft_resp["allowed"])
        return float((kl_rows(softmax(t_logits), softmax(sl)) * mask_r).sum() / mask_r.sum())
    # Kiểm tra gradient KL theo sai phân hữu hạn.
    _, g_kl, _ = distill_loss_grad(stu0, DISTILL_T)
    nm = "U"
    idx = np.unravel_index(np.abs(g_kl[nm]).argmax(), stu0[nm].shape)
    vals = []
    for sg in (1, -1):
        q = copy_params(stu0)
        q[nm][idx] += sg * 1e-5
        vals.append(distill_loss_grad(q, DISTILL_T)[0])
    kl_fd_err = abs((vals[0] - vals[1]) / 2e-5 - g_kl[nm][idx])
    kl_start = kl_to_teacher(stu0)
    # Student A: học từ soft target của teacher. Student B: chỉ học nhãn cứng như SFT.
    sa, oa = copy_params(stu0), new_opt(stu0)
    sb, ob = copy_params(stu0), new_opt(stu0)
    for _ in range(DISTILL_STEPS):
        sa = adamw_step(sa, distill_loss_grad(sa, DISTILL_T)[1], oa, LR * 3)
        sb = adamw_step(sb, loss_and_grads(sb, sft_resp)[1], ob, LR * 3)
    kl_a, kl_b = kl_to_teacher(sa), kl_to_teacher(sb)
    acc_a, acc_b = exact_matches(sa), exact_matches(sb)
    STUDENT = sa
    say(11, "Distillation", "4F", [
        f"Quan sát: teacher = mô hình SFT ({n_params(SFT)} tham số, d = {D_MODEL}, {N_LAYERS} lớp); student = mô hình nhỏ hơn ({n_stu} tham số, d = {STUDENT_D}, {STUDENT_LAYERS} lớp), khởi tạo ngẫu nhiên, bằng {n_stu / n_params(SFT):.0%} kích thước teacher.",
        f"Mục tiêu mềm tại vị trí đáp án đầu của cặp đầu (nhãn cứng chỉ là 'c' với xác suất 1): teacher cho " + ", ".join(f"{ITOS[t]!r} {pt1[first_pos][t]:.3f}" for t in top) + f" ở T = 1, và {', '.join(f'{ITOS[t]!r} {pt2[first_pos][t]:.3f}' for t in top)} ở T = {DISTILL_T:g} (nhiệt độ làm phân phối phẳng hơn, để lộ thêm thứ tự ưu tiên).",
        f"Loss = T²·KL(teacher || student) trên các vị trí đáp án; kiểm tra gradient bằng sai phân hữu hạn: lệch {kl_fd_err:.1e}. KL(teacher || student) ở T = 1 trước khi học: {kl_start:.4f}.",
        f"Sau {DISTILL_STEPS} bước: student học từ soft target có KL {kl_a:.4f} và trả lời đúng {acc_a}/{len(CHAT_PAIRS)}; student cùng cấu hình chỉ học nhãn cứng có KL {kl_b:.4f} và đúng {acc_b}/{len(CHAT_PAIRS)}.",
        "Kết luận: distillation dạy student bắt chước cả phân phối xác suất của teacher (không chỉ đáp án cao nhất), nên student nhỏ có thể nhận nhiều tín hiệu hơn từ mỗi vị trí so với nhãn cứng.",
        "Không chứng minh: teacher và dữ liệu đều rất nhỏ, kết quả này không nói student nhỏ giữ được chất lượng của teacher thật; chênh lệch KL giữa hai student có thể đổi theo seed và nhiệt độ.",
    ])
    assert kl_fd_err < 1e-6 and kl_a >= 0

    # ── [12] Preference (DPO) (4G) ──────────────────────────────────────────
    def seq_logp_and_dlogits(p, user, answer):
        ids, start = chat_ids(user, answer)
        inputs, labels = np.array(ids[:-1]), np.array(ids[1:])
        pos = np.arange(len(inputs))
        allowed = pos[None, :] <= pos[:, None]
        lg, sv = forward(p, inputs, pos, allowed)
        pr = softmax(lg)
        mask = pos >= start - 1
        idx = np.arange(len(labels))
        logp = float(np.log(pr[idx, labels])[mask].sum())
        onehot_minus_p = -pr
        onehot_minus_p[idx, labels] += 1                # d(logp)/d(logits) = one-hot - p
        return logp, onehot_minus_p * mask[:, None], sv

    ref = SFT
    ref_lp = [(seq_logp_and_dlogits(ref, u, c)[0], seq_logp_and_dlogits(ref, u, r)[0]) for u, c, r in PREF_PAIRS]

    def dpo_eval(p):
        zs, losses = [], []
        for (u, c, r), (rc, rr) in zip(PREF_PAIRS, ref_lp):
            lc_, lr_ = seq_logp_and_dlogits(p, u, c)[0], seq_logp_and_dlogits(p, u, r)[0]
            z = DPO_BETA * ((lc_ - rc) - (lr_ - rr))
            zs.append(z)
            losses.append(float(np.log1p(np.exp(-z))))
        return float(np.mean(zs)), float(np.mean(losses)), float(np.mean([z > 0 for z in zs]))

    def dpo_grads(p):
        total = {k: np.zeros_like(x) for k, x in p.items()}
        for (u, c, r), (rc, rr) in zip(PREF_PAIRS, ref_lp):
            lc_, dc, sc = seq_logp_and_dlogits(p, u, c)
            lr_, dr, sr = seq_logp_and_dlogits(p, u, r)
            z = DPO_BETA * ((lc_ - rc) - (lr_ - rr))
            coef = DPO_BETA / (1 + np.exp(z)) / len(PREF_PAIRS)        # -dL/dz · β / số cặp, với dL/dz = -σ(-z)
            gc, gr = backward(p, sc, -coef * dc), backward(p, sr, coef * dr)
            for k in total:
                total[k] += gc[k] + gr[k]
        return total
    z0, l0, acc0 = dpo_eval(ref)
    g_dpo = dpo_grads(ref)
    nm = "U"
    idx = np.unravel_index(np.abs(g_dpo[nm]).argmax(), ref[nm].shape)
    vals = []
    for sg in (1, -1):
        q = copy_params(ref)
        q[nm][idx] += sg * 1e-5
        vals.append(dpo_eval(q)[1])
    dpo_fd_err = abs((vals[0] - vals[1]) / 2e-5 - g_dpo[nm][idx])
    pol, opt_dpo = copy_params(ref), new_opt(ref)
    for _ in range(DPO_STEPS):
        pol = adamw_step(pol, dpo_grads(pol), opt_dpo, DPO_LR)
    z1, l1, acc1 = dpo_eval(pol)
    u0, c0_, r0_ = PREF_PAIRS[0]
    lp_c0, lp_r0 = seq_logp_and_dlogits(pol, u0, c0_)[0], seq_logp_and_dlogits(pol, u0, r0_)[0]
    say(12, "Preference (DPO)", "4G", [
        f"Quan sát: {len(PREF_PAIRS)} cặp (prompt, chosen, rejected); mô hình tham chiếu = checkpoint SFT đóng băng, policy bắt đầu từ chính nó. Cặp đầu: \"{u0}\" -> chosen \"{c0_}\" so với rejected \"{r0_}\".",
        f"DPO loss = -ln σ(β·[(ln π(chosen) - ln π_ref(chosen)) - (ln π(rejected) - ln π_ref(rejected))]), β = {DPO_BETA:g}; biên z là biểu thức trong ngoặc nhân β. Khi policy = tham chiếu thì z = {z0:.4f} và loss = {l0:.4f} = ln 2 = {np.log(2):.4f}.",
        f"Kiểm tra gradient DPO bằng sai phân hữu hạn: lệch {dpo_fd_err:.1e}.",
        f"Sau {DPO_STEPS} bước AdamW (lr {DPO_LR}): biên z trung bình {z0:.3f} -> {z1:.3f}; loss {l0:.4f} -> {l1:.4f}; tỉ lệ cặp có biên dương {acc0:.0%} -> {acc1:.0%}.",
        f"Cặp đầu: xác suất cả câu chosen {np.exp(ref_lp[0][0]):.3f} (tham chiếu) -> {np.exp(lp_c0):.3f} (sau DPO), {'giảm' if lp_c0 < ref_lp[0][0] else 'tăng'}; xác suất cả câu rejected {np.exp(ref_lp[0][1]):.2e} -> {np.exp(lp_r0):.2e}. Biên chỉ đo HIỆU của hai thay đổi, nên nó có thể tăng dù chosen không tăng.",
        "Kết luận: DPO không cần mô hình phần thưởng riêng: nó đẩy policy tăng xác suất tương đối của câu được chọn so với câu bị loại, đo so với một mô hình tham chiếu.",
        "Không chứng minh: biên tăng chỉ nói policy phân biệt tốt hơn 6 cặp huấn luyện; không nói câu trả lời tốt hơn theo con người, không có dữ liệu ưu tiên thật, không có RLHF/PPO hay kiểm tra phần thưởng.",
    ])
    assert dpo_fd_err < 1e-6 and abs(z0) < 1e-12

    # ── [13] Từ checkpoint sang suy luận (5A–5C) ────────────────────────────
    p_srv, opt_srv, vocab_srv = load_checkpoint(ckpt_sft)
    load_gap = max(float(np.abs(p_srv[k] - SFT[k]).max()) for k in SFT)
    user_q = CHAT_PAIRS[2][0]
    prompt_ids = [USER] + encode_text(user_q) + [ASSISTANT]            # 5A: mẫu hội thoại + tokenizer
    # 5B: prefill rồi decode với KV cache
    cache = [None] * N_LAYERS
    positions_processed = 0

    def step_forward(new_ids, start):
        nonlocal positions_processed
        pos = np.arange(start, start + len(new_ids))
        x = p_srv["E"][np.array(new_ids)] + pos_enc(pos, D_MODEL)
        for l in range(N_LAYERS):
            past = cache[l]
            total = (0 if past is None else len(past[0])) + len(new_ids)
            allowed = np.arange(total)[None, :] <= pos[:, None]
            x, _, kv = block_forward(x, layer_w(p_srv, l), allowed, past=past)
            cache[l] = kv
        positions_processed += len(new_ids)
        return rms(x) @ p_srv["U"]
    lg_pre = step_forward(prompt_ids, 0)
    kv_shape = tuple(cache[0][0].shape)
    gen, steps_gap = [int(np.argmax(lg_pre[-1]))], 0.0
    full_ids = list(prompt_ids)
    while gen[-1] != EOS and len(gen) < MAX_NEW:
        full_ids = prompt_ids + gen
        lg_step = step_forward([gen[-1]], len(full_ids) - 1)[-1]
        n = len(full_ids)
        ref_logits, _ = forward(p_srv, np.array(full_ids), np.arange(n), np.arange(n)[None, :] <= np.arange(n)[:, None])
        steps_gap = max(steps_gap, float(np.abs(lg_step - ref_logits[-1]).max()))
        gen.append(int(np.argmax(lg_step)))
    text_out = decode_ids(gen)
    wrong_p, _, _ = load_checkpoint(save_checkpoint(init_params(SEED + 1, V, D_MODEL, N_LAYERS, FFN_HIDDEN)))
    wrong_text = decode_ids(predict_greedy(wrong_p, prompt_ids))
    no_cache = sum(len(prompt_ids) + k for k in range(len(gen)))
    ref_text = answer_of(SFT, user_q)
    cache_len = cache[0][0].shape[0]
    w_bytes = len(ckpt_sft)
    say(13, "Từ checkpoint sang suy luận", "5A–5C", [
        f"Quan sát: nạp checkpoint SFT ({w_bytes} byte, chỉ trọng số và bảng từ vựng, không có m, v): trọng số nạp lại khớp bản đang huấn luyện (lệch {load_gap:.1e}); không có optimizer ({'không' if opt_srv is None else 'có'} m, v) và không tính gradient.",
        f"5A · Prompt → token IDs: \"{user_q}\" qua mẫu hội thoại và đúng tokenizer đó -> {prompt_ids} ({len(prompt_ids)} token).",
        f"5B · Prefill: cả {len(prompt_ids)} token đi qua mô hình trong một lượt, KV cache của lớp 0 có K shape {kv_shape} = (số vị trí, chiều ẩn); từ hàng logits cuối chọn token đầu tiên {ITOS[gen[0]]!r}.",
        f"5B · Decode: mỗi bước chỉ đưa MỘT token mới, đọc K, V cũ từ cache rồi thêm K, V mới (cache dài {len(prompt_ids)} -> {cache_len}); ở mọi bước logits có cache và logits tính lại cả chuỗi lệch tối đa {steps_gap:.1e}.",
        f"Số vị trí phải xử lý: có cache {positions_processed} (prefill {len(prompt_ids)} + {len(gen) - 1} bước decode), tính lại cả chuỗi mỗi bước {no_cache}; đây là số vị trí, chưa phải thời gian.",
        f"5C · Output: dừng khi gặp <eos> hoặc đủ {MAX_NEW} token; token {gen} -> văn bản \"{text_out}\" (đường tính lại không cache cho \"{ref_text}\").",
        f"Nạp nhầm: một checkpoint khác cùng cấu hình (khởi tạo seed {SEED + 1}, chưa huấn luyện) nạp vào không báo lỗi vì cùng shape, nhưng cho \"{wrong_text}\" thay vì \"{text_out}\"; code không tự biết checkpoint nào là đúng, nên cần ghi tên, phiên bản hoặc mã băm cùng checkpoint.",
        "Kết luận: suy luận dùng lại đúng checkpoint với trọng số đóng băng và cùng phép forward; prefill đọc cả prompt một lượt, decode đi từng token với KV cache cho kết quả giống hệt tính lại.",
        "Không chứng minh: khớp giữa cache và tính lại chỉ xác nhận việc tái dùng K, V cũ là đúng, không chứng minh tốc độ hay chất lượng câu trả lời; đây không phải engine phục vụ thật (không có batching, lập lịch hay sampling).",
    ])
    assert load_gap == 0.0 and steps_gap < 1e-9 and text_out == ref_text and opt_srv is None

    print("TỔNG KẾT: Bạn vừa quan sát được gì?")
    print(f"1. Loss từ {LN_V:.2f} (đoán đều) xuống {hist[-1][2]:.2f} trên tập huấn luyện sau {STEPS} bước; gradient giải tích khớp sai phân hữu hạn.")
    print("2. ZeRO và recompute đổi bộ nhớ lấy giao tiếp hoặc phép tính mà không đổi gradient hay cập nhật (lệch bằng 0 hoặc cỡ 1e-16).")
    print(f"3. SFT cho {after}/{len(CHAT_PAIRS)} câu đúng, LoRA {lora_match}/{len(CHAT_PAIRS)} với {n_ad / N_P:.0%} tham số học được, distillation đưa student {n_stu / n_params(SFT):.0%} kích thước teacher tới KL {kl_a:.3f}, DPO tăng biên z lên {z1:.2f}.")
    print("4. Checkpoint weights-only đủ để suy luận; prefill + decode với KV cache cho đúng cùng logits như tính lại.")
    print("5. Mọi số ở đây là kết quả của mô hình đồ chơi: không phải tốc độ GPU, không phải chất lượng ngôn ngữ thật.")
    print("PASS: các phép kiểm tra cơ chế đạt (gradient khớp sai phân hữu hạn, ZeRO và recompute không đổi toán học, LoRA gộp = không gộp, suy luận có cache = tính lại).")


main()