Lab tổng hợp · 13 giai đoạn

Lab tổng hợp LLM Training · huấn luyện một LLM thu nhỏ từ dữ liệu đến checkpoint và suy luận bằng NumPy

Chạy bằng NumPy một LLM thu nhỏ đi hết vòng huấn luyện: dữ liệu, forward, loss, backward, AdamW, ZeRO, SFT, LoRA, distillation, DPO, rồi nạp checkpoint để suy luận.

LESSON ID training-end-to-end-lab GROUP Bài Lab tổng hợp

1Định nghĩa

Bài Lab tổng hợp này cho bạn chạy bằng NumPy (thư viện tính toán số của Python) ngay trong trình duyệt toàn bộ vòng đời huấn luyện của một mô hình ngôn ngữ lớn (LLM, large language model) thu nhỏ: chuẩn bị dữ liệu, tính điểm số bằng forward, đo sai lầm bằng loss, tính hướng sửa bằng backward, cập nhật trọng số, rồi đưa kết quả sang suy luận. Mỗi giai đoạn trong output mang đúng tên các ô trên sơ đồ kiến trúc ở Tab 01, nên bạn nhìn thấy số thật chảy qua từng ô thay vì chỉ đọc mô tả. Đây là cách ôn lại vòng lặp huấn luyện bằng giá trị cụ thể.

Cơ chế chạy như sau: một mô hình 2 lớp với 5856 tham số học cách đoán token kế tiếp qua forward, cross-entropy loss, backward và AdamW; sau đó cùng khung vòng lặp ấy được dùng lại cho bốn phương pháp huấn luyện khác nhau là SFT, LoRA, distillation và DPO, còn bộ nhớ của cả vòng được tính bằng ZeRO và tính lại activation. Cuối cùng checkpoint được nạp lại để tiếp nhận một yêu cầu suy luận. Mọi loss, đếm số, đếm byte và sai số là kết quả tính thật trên mô hình đồ chơi, còn lab không đo thời gian hay tốc độ GPU thật.

Ví dụ đời thường: Lab giống một xưởng dạy nghề thu nhỏ có đủ mọi công đoạn. Người học đọc đề (token ID), làm bài (forward), chấm điểm (loss), tìm chỗ sai (backward), sửa cách làm theo từng mức độ khác nhau (AdamW), thỉnh thoảng làm bài thi thử chưa gặp (kiểm định) và cất vở vào tủ (checkpoint). Sau đó xưởng thử bốn cách dạy tiếp: học theo mẫu hỏi đáp, chỉ chỉnh một vài thao tác, bắt chước thầy giỏi và chọn giữa hai cách trả lời. Phép so sánh dừng lại ở quy mô: xưởng này chỉ có tám bài tập nên học thuộc chứ chưa phải thành thạo.

2Luồng

Sơ đồ kiến trúc huấn luyện · 22 ô (1A → 5C)

Kiến trúc huấn luyện LLM · Ánh xạ với Bài Lab tổng hợp

Thứ tự đọc: 1A → 1B ở cột 1, 2A → 2D ở cột 2, 3A → 3E ở cột 3; mũi tên next batch từ 3E quay về cột 1 và mũi tên load weights đi xuống băng 5 (5A → 5C). Cột 4 không phải bước kế tiếp: 4A, 4B tiết kiệm bộ nhớ, còn 4C–4G là các phương pháp huấn luyện thay thế nhau. Di chuột hoặc bấm vào từng ô để đọc pop-over.

Sơ đồ kiến trúc huấn luyện LLM gồm 22 ô từ 1A đến 5C

1A1 · Data & Targets

văn bản thô → clean/dedup → tokenizer → [12, 87, 5, 40]

1A · Corpus → Token IDs

Điểm bắt đầu: biến kho văn bản thô thành các dãy số nguyên (token ID) mà mô hình đọc được.

Nhận gì: Văn bản thô từ kho dữ liệu (corpus: tập hợp mọi văn bản dùng để dạy mô hình). Đây là ô đầu tiên của sơ đồ, chưa nhận gì từ ô khác.
Làm gì, vì sao: Làm sạch và bỏ các bản trùng để mô hình không học thuộc dữ liệu lặp. Sau đó tokenizer (bộ cắt văn bản thành các mảnh nhỏ gọi là token) đổi mỗi mảnh thành một số nguyên gọi là token ID, vì mô hình chỉ tính trên số.
Ra gì: Các dãy token ID, đóng gói thành batch (nhóm tài liệu xử lý cùng lúc) và đi xuống ô 1B bằng mũi tên đi xuống.
Ví dụ, cách đọc ô: Đọc nhãn trong ô từ trái sang phải: corpus → clean/dedup → tokenizer → [12, 87, 5, 40]. Bốn số cuối là token ID của một đoạn văn; ID 87 chỉ là số thứ tự của một mảnh trong bảng từ vựng.
Lưu ý: Token ID chưa mang ý nghĩa, chỉ là số thứ tự. Đổi tokenizer thì mọi ID đổi theo, nên tokenizer phải được lưu cùng checkpoint.

Đọc bài 1A · Corpus → Token IDs Giai đoạn 1 trong Bài Lab: Dữ liệu → token IDs

Sơ đồ luồng (13 bước)

1 · Dữ liệu → token IDs
11 dòng thô thành 8 tài liệu sạch, từ vựng 23 token; tài liệu đầu thành 11 ID; một batch 4 tài liệu có 63 vị trí.
2 · Next-token targets
Tài liệu 11 token cho 10 cặp (input, nhãn dịch một vị trí); mặt nạ cắt theo ranh giới tài liệu.
3 · Forward: embedding, attention, MLP, norm → logits
Attention vị trí 3 chia trọng số cho đúng 4 vị trí 0..3; đổi token cuối không làm đổi logits phía trước (0.0); logits có shape (10, 23).
4 · Cross-entropy loss
Loss ban đầu 3.2481 gần ln 23 = 3.1355 vì mô hình chưa học; p(đáp án) từ 0.0226 đến 0.1285 ở bốn vị trí đầu.
5 · Backward & gradient từng lớp
Gradient theo logits là p trừ one-hot, tổng mỗi hàng bằng 0; gradient giải tích khớp sai phân hữu hạn (sai lệch cỡ 5e-10) ở 16 tensor.
6 · AdamW
Bước đầu mỗi tham số đổi gần đúng lr = 0.01 theo dấu của gradient (khoảng 0.0098 đến 0.0101); loss một bước 3.2746 xuống 2.8802.
7 · Vòng lặp, eval, checkpoint
300 bước: loss huấn luyện 0.1014, loss kiểm định thấp nhất 0.8048 ở bước 50 rồi lên 1.0469; checkpoint tiếp tục được gấp 2.96 lần bản weights-only.
8 · Bộ nhớ: ZeRO & recompute
5856 tham số cần 16 byte mỗi tham số; ZeRO-3 với 2 worker còn 46848 byte mỗi worker; cập nhật chia theo ZeRO-1 và gradient tính lại lệch 0.0.
9 · SFT: chat template + response mask
Mặt nạ chỉ chấm 4 nhãn cuối của cặp đầu; từ 0 trên 8 lên 8 trên 8 câu đúng, loss 4.1347 xuống 0.0036.
10 · LoRA
Chỉ 1088 tham số (18.6%) được học; trọng số gốc lệch 0.0; mô hình gộp và chưa gộp cho cùng logits; 8 trên 8 câu đúng.
11 · Distillation
Student 1008 tham số (17% teacher): học từ soft target có KL 0.0068 và 8 trên 8 câu đúng, học nhãn cứng có KL 0.1955 và 5 trên 8.
12 · Preference (DPO)
Khi policy bằng tham chiếu loss đúng ln 2 = 0.6931; sau 60 bước biên 8.317, xác suất câu chosen của cặp đầu lại giảm từ 0.982 xuống 0.520.
13 · Từ checkpoint sang suy luận
Prompt 13 token, prefill một lượt rồi decode với KV cache cho đúng văn bản sua. như tính lại; 17 vị trí so với 75.

Mã Mermaid

flowchart TD
  step1["1 · Dữ liệu → token IDs"]
  step2["2 · Next-token targets"]
  step3["3 · Forward: embedding, attention, MLP, norm → logits"]
  step4["4 · Cross-entropy loss"]
  step5["5 · Backward #amp; gradient từng lớp"]
  step6["6 · AdamW"]
  step7["7 · Vòng lặp, eval, checkpoint"]
  step8["8 · Bộ nhớ: ZeRO #amp; recompute"]
  step9["9 · SFT: chat template + response mask"]
  step10["10 · LoRA"]
  step11["11 · Distillation"]
  step12["12 · Preference (DPO)"]
  step13["13 · Từ checkpoint sang suy luận"]
  step1 --> step2
  step2 --> step3
  step3 --> step4
  step4 --> step5
  step5 --> step6
  step6 --> step7
  step7 --> step8
  step8 --> step9
  step9 --> step10
  step10 --> step11
  step11 --> step12
  step12 --> step13

3Giải pháp

Mỗi bài trong khóa này giải thích một ô của sơ đồ, nhưng đọc riêng từng bài rất khó thấy các ô nối nhau thế nào: dữ liệu ra khỏi ô nào, vào ô nào, và con số nào ở một ô thay đổi khi ô khác thay đổi. Lab gom 13 giai đoạn vào một file để bạn theo dõi cùng một mô hình nhỏ từ lúc chưa biết gì (loss ban đầu 3.2481, gần mức đoán đều ln 23 = 3.1355) đến lúc trả lời được câu hỏi và được nạp lại cho suy luận.

Đầu vào gồm dữ liệu tự viết và một cấu hình nhỏ. Kho nền có 11 dòng thô, sau khi bỏ dòng trống, hai dòng trùng và chuẩn hóa chữ thường còn 8 tài liệu; tách riêng 2 câu kiểm định. Dữ liệu SFT là 8 cặp hỏi đáp, dữ liệu ưu tiên là 6 bộ (câu hỏi, câu được chọn, câu bị loại). Tokenizer tách theo ký tự nên từ vựng chỉ có 23 token. Mô hình là bộ giải mã Transformer 2 lớp, chiều ẩn 16, một head attention, khởi tạo ngẫu nhiên với seed 7; chạy lại cho đúng các số cũ.

Giai đoạn 1 (ô 1A) biến văn bản thành số: tài liệu đầu "meo an ca." thành 11 token ID, cuối là eos (token báo hết tài liệu); một batch (nhóm tài liệu xử lý cùng lúc) gồm 4 tài liệu có 63 vị trí, nên một epoch (một lượt qua toàn bộ dữ liệu) là 2 bước. Giai đoạn 2 (ô 1B) dịch nhãn đi một vị trí: tài liệu 11 token cho 10 cặp (input, token đứng ngay sau), và mặt nạ nhân quả (che các vị trí phía sau) được cắt theo ranh giới tài liệu, nên khi đóng gói hai tài liệu, tài liệu sau không nhìn thấy tài liệu trước.

Giai đoạn 3 (ô 2A đến 2D) đi qua forward: bảng embedding (bảng tra ID ra vector) cho x shape (10, 16), tức 10 vị trí với 16 số mỗi vị trí; attention ở vị trí 3 chia trọng số tổng bằng 1 cho đúng 4 vị trí 0 đến 3 (các vị trí sau có trọng số 0), và đó là trọng số trên vị trí chứ chưa phải xác suất của token. Đổi token cuối rồi chạy lại, logits (điểm số của mọi token trong từ vựng) ở các vị trí trước đó đổi 0.0. Norm đưa độ lớn trung bình của vector từ 0.136 về 1.000; nếu đặt ma trận ra của attention và MLP bằng 0 thì đầu ra khối bằng đầu vào, minh họa đường tắt residual. Giai đoạn 4 (ô 3A) chấm điểm: loss là trung bình của trừ ln p(token đúng), ví dụ vị trí 3 có p = 0.1285 cho 2.052 nat (nat là đơn vị của loss khi dùng logarit tự nhiên), trung bình 3.2481.

Giai đoạn 5 (ô 3B, 3C) tính backward. Gradient theo logits là p trừ one-hot nên tổng mỗi hàng bằng 0; từ đó gradient đi ngược qua unembedding, từng lớp (lớp 0 có chuẩn gradient 0.9951, lớp 1 có 0.6791) và embedding. Để chắc công thức đúng, lab dịch phần tử có gradient lớn nhất của mỗi trong 16 tensor đi ±1e-5, đo loss thay đổi bao nhiêu (sai phân hữu hạn) và thấy khớp, sai lệch lớn nhất cỡ 5e-10. Giai đoạn 6 (ô 3D) theo một tham số qua AdamW: gradient -0.017946, m = (1 - β1) nhân g = -0.001795, v = (1 - β2) nhân g² = 3.221e-07; sau hiệu chỉnh ở bước đầu m̂ trên căn v̂ bằng đúng dấu của g, nên bước cập nhật gần đúng bằng lr = 0.01. Vì thế ở bước đầu AdamW đổi mỗi tham số một lượng 0.0098 đến 0.0101, còn SGD cùng lr đổi từ 6e-09 đến 2.6e-03; một bước giảm loss 3.2746 xuống 2.8802.

Giai đoạn 7 (ô 3E và 4C) chạy 300 bước. Loss huấn luyện đo trên cả tập đi 0.2657 (bước 50), 0.1259, 0.1128, 0.1070, 0.1031 rồi 0.1014 (bước 300); loss kiểm định thấp nhất 0.8048 ở bước 50 rồi lên 1.0469. Hai đường tách nhau: đó là dấu hiệu bắt đầu học thuộc chỉ 8 tài liệu, chưa phải kết luận chắc chắn vì kiểm định chỉ có 2 câu. Checkpoint đầy đủ ở bước 100 nặng khoảng 153,6 nghìn byte, gấp 2.96 lần bản weights-only 51960 byte vì mang thêm m và v. Chạy tiếp 10 bước từ bản đầy đủ lệch 0.0 so với chạy liền; chỉ nạp trọng số (m, v về 0 và bộ đếm bước về 0) lệch 9.0e-02.

Giai đoạn 8 (ô 4A, 4B) tính bộ nhớ mà không đổi toán học. Mỗi tham số cần 2 byte trọng số, 2 byte gradient và 12 byte trạng thái bộ tối ưu (bản chính 32 bit cộng m và v), tổng 16 byte; với 5856 tham số và 2 worker, mỗi worker cần 93696 byte (data parallel), 58560 (ZeRO-1), 52704 (ZeRO-2) hoặc 46848 (ZeRO-3). Cùng công thức cho mô hình 7 tỉ tham số trên 8 GPU cho 112, 38.5, 26.25 và 14 GB, tính tay chưa gồm activation. Lab còn chia dữ liệu cho 2 worker (gradient cộng lại lệch cỡ 1e-17), chia m và v thành hai đoạn 2928 phần tử (cập nhật chia lệch 0.0) và tính lại activation từng lớp khi backward (gradient lệch 0.0, số lưu giảm từ 90240 còn 48960).

Giai đoạn 9 đến 12 là bốn phương pháp khác nhau, mỗi phương pháp đổi một phần của vòng lặp. SFT (ô 1C, 4D) ghép mỗi cặp thành chuỗi "user câu hỏi assistant câu trả lời eos" và chỉ chấm các nhãn của câu trả lời: từ 0 trên 8 lên 8 trên 8 câu đúng, loss 4.1347 xuống 0.0036; loss trên token prompt (không được chấm) đi từ 3.8980 lên 5.0391, và câu hỏi chưa gặp "meo bay o dau" vẫn được trả lời bằng "tren ghe." nên không có bằng chứng về khả năng tổng quát. LoRA (ô 4E) đóng băng 5856 trọng số gốc và chỉ học cặp ma trận hạng thấp A, B (1088 tham số, 18.6%): B bắt đầu bằng 0 nên mô hình y hệt gốc ở bước 0, trọng số gốc lệch 0.0 sau huấn luyện và mô hình gộp bằng mô hình chưa gộp.

Distillation (ô 4F): student 1008 tham số (17% teacher) học khớp phân phối xác suất của teacher bằng KL (đo hai phân phối khác nhau bao nhiêu): KL 0.0068 và 8 trên 8 câu đúng, so với 0.1955 và 5 trên 8 khi cùng student chỉ học nhãn cứng. DPO (ô 4G) so câu được chọn và câu bị loại với một mô hình tham chiếu: khi policy còn bằng tham chiếu loss đúng bằng ln 2 = 0.6931; sau 60 bước biên z là 8.317, nhưng xác suất cả câu được chọn của cặp đầu lại giảm từ 0.982 xuống 0.520 còn câu bị loại giảm từ 1.02e-03 xuống 2.70e-10. Biên đo hiệu của hai thay đổi nên có thể tăng dù câu được chọn không tăng.

Giai đoạn 13 (ô 5A đến 5C) nạp checkpoint SFT (51960 byte, không có m, v) và chạy suy luận. Prompt "meo uong gi" qua mẫu hội thoại thành 13 token ID; prefill (xử lý cả prompt một lượt) ghi KV cache của lớp 0 shape (13, 16) và chọn token đầu "s"; mỗi bước decode (sinh từng token) chỉ đưa một token mới, đọc K, V cũ từ cache rồi thêm K, V mới, cache dài 13 lên 17. Logits có cache khớp logits tính lại cả chuỗi (sai lệch cỡ 1e-15) và văn bản sinh ra là "sua.". Số vị trí phải xử lý là 17 với cache so với 75 nếu mỗi bước tính lại cả chuỗi; đây là số vị trí, chưa phải thời gian. Nạp nhầm một checkpoint khác cùng cấu hình không báo lỗi nhưng cho đầu ra khác hẳn.

Giới hạn của kết luận: mô hình và dữ liệu cực nhỏ, nên các số thuộc về lab này chứ không đại diện cho mô hình thật; số câu đúng đo trên chính các cặp đã học; ZeRO và recompute chỉ là đếm byte và so gradient; LoRA, distillation và DPO mỗi phương pháp chỉ chạy một seed và vài chục ví dụ. Dòng PASS chỉ xác nhận các phép kiểm tra cơ chế đã đạt, không xác nhận mô hình học tốt.

4Profiling

Kiểm tra tính đúng trước khi đọc các con số học được
Trước khi so loss hay số câu đúng, hãy xem các phép kiểm tra cơ chế: gradient giải tích khớp sai phân hữu hạn (sai lệch cỡ 5e-10) ở 16 tensor; đổi token cuối không đổi logits phía trước; gradient cộng từ các worker bằng gradient cả batch; cập nhật chia theo ZeRO-1 và gradient tính lại cho sai lệch 0.0; LoRA gộp bằng chưa gộp; cache bằng tính lại. Ba nhóm này chứng minh cơ chế đúng, không chứng minh mô hình học tốt hay chạy nhanh.
Đọc loss huấn luyện và loss kiểm định cùng nhau
Workload là 8 tài liệu huấn luyện và 2 câu kiểm định, batch 4 tài liệu, lr 0.01, seed 7; baseline là loss đoán đều ln 23 = 3.1355. Mỗi 50 bước hãy ghi loss batch, loss cả tập huấn luyện và loss kiểm định. Nếu hai đường loss cùng giảm thì mô hình đang học điều chuyển được sang dữ liệu chưa thấy; nếu loss huấn luyện giảm còn loss kiểm định tăng (như từ bước 50 ở lab) thì nghi học thuộc, nhưng cần thêm bằng chứng như kiểm định lớn hơn. Đừng viết loss giảm thành mô hình tốt.
Đếm byte bộ nhớ, không đo
Công thức 16 byte cho mỗi tham số, 4P + 12P/N, 2P + 14P/N và 16P/N chỉ tính trạng thái mô hình (trọng số, gradient, bộ tối ưu); chưa gồm activation, bộ đệm giao tiếp và phân mảnh bộ nhớ. Con số 14 GB của ZeRO-3 cho mô hình 7 tỉ tham số trên 8 GPU là phép tính tay. Khi làm trên phần cứng thật, hãy đo bộ nhớ đỉnh và thời gian mỗi bước cùng workload, vì ZeRO-3 thêm giao tiếp và recompute thêm một lượt forward (khoảng 33% theo mô hình đơn giản, thời gian thật phụ thuộc kernel và kích thước lớp).
Phân biệt khớp số học, loss thấp và chất lượng trả lời
Khớp số học (cache bằng tính lại, merge bằng không merge) cho biết cơ chế đúng. Loss thấp trên dữ liệu huấn luyện cho biết mô hình thuộc dữ liệu. Chất lượng trả lời cần bộ đánh giá riêng: ở lab, câu hỏi chưa gặp được trả lời bằng câu của một cặp khác, và DPO có thể làm xác suất câu được chọn giảm trong khi biên tăng. Hãy báo cả ba loại bằng chứng với đúng nhãn của chúng.

Mẹo đo lường: ghim số liệu trên workload thật, ghi cả p50/p95/p99 thay vì chỉ trung bình, và luôn ghi lại điều kiện chạy (batch, cache, đồng thời, phần cứng) trước và sau khi tối ưu. Số đo thiếu điều kiện chạy thì không so sánh được.

5Ví dụ

Vì sao bước AdamW đầu tiên gần bằng lr bất kể gradient lớn hay nhỏ?

Hãy dự đoán trước: tham số Q0 ở vị trí (0, 7) có gradient -0.017946 còn tham số khác có thể có gradient nhỏ hơn hàng nghìn lần, vậy bước đầu của chúng khác nhau bao nhiêu? Mở khối 6. Ở bước 1, m = 0.1 nhân g và v = 0.001 nhân g², sau hiệu chỉnh m̂ = g và v̂ = g², nên cập nhật bằng lr nhân g chia (căn của g² cộng ε) gần đúng bằng lr nhân dấu của g, tức -0.010000 ở tham số này. Toàn mô hình, mọi tham số có gradient khác 0 đổi 0.0098 đến 0.0101, còn SGD với cùng lr đổi từ 6e-09 đến 2.6e-03. Tự kiểm tra: AdamW luôn đổi mỗi tham số đúng lr ở mọi bước không? Đối chiếu: không, chỉ ở bước đầu m̂ và v̂ cùng đến từ một gradient; sau đó m là trung bình của nhiều gradient nên bước nhỏ lại khi gradient đổi dấu liên tục.

Prompt không được chấm thì loss của nó đi đâu?

Mở khối 9. Dự đoán: sau SFT chỉ chấm câu trả lời, loss trên token prompt giảm, đứng yên hay tăng? Kết quả đã chạy: từ 3.8980 lên 5.0391, trong khi loss trên token câu trả lời đi từ 4.1347 xuống 0.0036. Prompt không sinh gradient (hàng gradient ở vị trí bị che có tổng bằng 0), nên mô hình không có lý do giữ khả năng đoán các token của câu hỏi, và các thay đổi để trả lời đúng có thể làm nó tệ đi. Đó là điều mong muốn nếu mục tiêu là học cách đáp, nhưng cũng là lý do dữ liệu SFT cần đúng mẫu. Tự kiểm tra: đặt MASK_RESPONSE_ONLY = False thì loss trên token prompt sau SFT là bao nhiêu? Đối chiếu: 0.1851, vì lúc đó prompt cũng được chấm (xem bài tập 4).

Biên DPO tăng nhưng câu được chọn lại kém chắc chắn hơn

Mở khối 12 cho cặp đầu (hỏi "meo an gi", chosen "ca.", rejected "sua."). Dự đoán: sau DPO xác suất của cả câu chosen tăng hay giảm? Ban đầu policy bằng tham chiếu nên biên z = 0 và loss đúng bằng ln 2 = 0.6931. Sau 60 bước z trung bình là 8.317 và loss 0.0005, nhưng xác suất cả câu chosen đi từ 0.982 xuống 0.520, còn xác suất câu rejected đi từ 1.02e-03 xuống 2.70e-10. Biên chỉ đo hiệu giữa hai thay đổi, nên đẩy rejected xuống rất mạnh là đủ để biên tăng mạnh dù chosen không tăng. Tự kiểm tra: biên z dương lớn có đủ để kết luận câu trả lời của policy tốt hơn không? Đối chiếu: không; cần đánh giá câu trả lời trực tiếp và dữ liệu ưu tiên thật, còn ở đây chỉ có 6 cặp tự viết.

6Thuật ngữ

Giai đoạn của lab: Một khối [n · Tên] trong output, khớp một hoặc vài ô trên sơ đồ kiến trúc ở Tab 01 (ví dụ giai đoạn 3 khớp các ô 2A đến 2D).
Mặt nạ loss (loss mask): Danh sách đúng/sai cho từng vị trí, cho biết vị trí nào được tính vào loss; ở SFT chỉ các token của câu trả lời được chấm.
Sai phân hữu hạn (finite difference): Cách ước lượng gradient bằng cách dịch một tham số đi một lượng rất nhỏ rồi đo loss đổi bao nhiêu; lab dùng nó để kiểm tra gradient tính bằng công thức.
Checkpoint weights-only và checkpoint tiếp tục được: Bản weights-only chỉ có trọng số, cấu hình và bảng từ vựng, đủ để suy luận; bản tiếp tục được có thêm m, v của AdamW và số bước nên nặng gần gấp ba.
Biên DPO (margin): Hiệu giữa mức tăng log-xác suất của câu được chọn và mức tăng của câu bị loại, so với mô hình tham chiếu, nhân β; biên dương nghĩa là policy đã nghiêng về câu được chọn.
Đồng hồ và tốc độ: Lab này không đo thời gian hay tốc độ GPU; mọi số là loss, đếm số, đếm byte hoặc sai số tính thật trên mô hình đồ chơi.

Glossary (7)

Token ID: Số thứ tự của một token trong bảng từ vựng; mô hình dùng số này để tra vector nhúng, bản thân số chưa mang ý nghĩa.
Logits: Điểm số chưa chuẩn hóa của mọi token trong từ vựng ở một vị trí; softmax đổi chúng thành xác suất tổng bằng 1.
Nat: Đơn vị của loss khi dùng logarit tự nhiên; đoán đều trên 23 token cho loss ln 23 = 3.1355 nat.
Epoch: Một lượt đi hết toàn bộ dữ liệu huấn luyện; với 8 tài liệu và batch 4, một epoch là 2 bước.
Học thuộc (overfitting): Loss trên dữ liệu huấn luyện tiếp tục giảm trong khi loss trên dữ liệu chưa thấy đi lên; cần thêm bằng chứng khác trước khi kết luận chắc chắn.
KV cache: Bộ nhớ giữ key và value của mọi vị trí đã xử lý để mỗi bước decode chỉ tính vị trí mới.
Mô hình tham chiếu: Bản sao đóng băng của mô hình trước khi huấn luyện ưu tiên, dùng làm mốc đo xem policy đã dịch chuyển bao xa.

7Bài liên quan

Chuẩn bị dữ liệu để huấn luyện dự đoán token kế tiếp Tạo câu hỏi và đáp án cho mỗi vị trí trong chuỗi Bước training: forward từ token đến điểm số Query, key và value kết hợp context như thế nào Mạng truyền thẳng và cổng kích hoạt sau attention Chuẩn hóa vector và đường cộng residual trong transformer Hàm mất mát: model gán bao nhiêu xác suất cho đáp án đúng Tính gradient để biết tham số cần thay đổi ra sao Gradient đi qua một lớp transformer AdamW: cập nhật tham số từ gradient Vòng huấn luyện, đánh giá và lưu trạng thái tiếp tục Chia trạng thái huấn luyện để giảm bộ nhớ mỗi worker Tính lại biểu diễn trung gian để tiết kiệm memory training Pretraining và continued pretraining · học nền tảng rồi học tiếp Fine-tuning (SFT) · dạy tiếp model đã học bằng cặp yêu cầu – trả lời mẫu LoRA / QLoRA · học một cặp ma trận nhỏ thay vì sửa toàn bộ trọng số Knowledge distillation · model nhỏ học từ tín hiệu của model lớn Preference training · dạy model ưu tiên câu trả lời tốt hơn (reward model, PPO, DPO) Bước inference: tiếp nhận prompt và yêu cầu sinh Bước inference: chuyển prompt thành mã token Bước inference: xử lý prompt và chọn token đầu tiên Bước inference: vòng sinh và bổ sung context Bước inference: kết thúc và trả văn bản

8Nguồn

Decoupled Weight Decay Regularization (Loshchilov và Hutter, 2019), bài báo của AdamW ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (Rajbhandari và cộng sự, 2020) Training Deep Nets with Sublinear Memory Cost (Chen và cộng sự, 2016), activation checkpointing LoRA: Low-Rank Adaptation of Large Language Models (Hu và cộng sự, 2021) Distilling the Knowledge in a Neural Network (Hinton và cộng sự, 2015) Direct Preference Optimization (Rafailov và cộng sự, 2023)

9Bài lab

Huấn luyện một LLM thu nhỏ qua 13 giai đoạn của sơ đồ bằng NumPy

Thực hành NumPy · huấn luyện LLM thu nhỏ (13 giai đoạn)

Đọc phần hướng dẫn rồi bấm Chạy. Mỗi khối [n · Tên] khớp một hoặc vài ô trên sơ đồ kiến trúc ở Tab 01; hãy đọc dòng Quan sát, đối chiếu các con số, rồi mới đọc Kết luận và Không chứng minh.

Bạn sẽ học được gì?

Sau bài lab, bạn giải thích được bằng số cụ thể mười điều: (1) văn bản trở thành token ID và batch ra sao; (2) vì sao mỗi chuỗi tự sinh ra các bài toán dự đoán token kế tiếp; (3) forward đi qua embedding, attention, MLP và norm tới logits như thế nào; (4) cross-entropy loss nén bảng điểm thành một số ra sao; (5) backward cho gradient từng lớp và vì sao có thể kiểm chứng bằng sai phân hữu hạn; (6) AdamW đổi gradient thành bước cập nhật thế nào; (7) vòng lặp, đánh giá và checkpoint làm gì; (8) ZeRO và recompute tiết kiệm bộ nhớ mà không đổi toán học; (9) SFT, LoRA, distillation và DPO khác nhau ở phần nào của vòng huấn luyện; (10) checkpoint trở thành suy luận prefill rồi decode với KV cache ra sao.

Bạn không cần biết trước về học máy. Mỗi thuật ngữ được giải thích ở chỗ nó xuất hiện đầu tiên; bạn chỉ cần đọc được Python và biết một mảng NumPy có shape, ví dụ (10, 16) là 10 hàng, mỗi hàng 16 số.

Đầu vào và cấu hình thu nhỏ của bài lab

Dữ liệu do tác giả tự viết, không tải từ đâu. Kho huấn luyện nền gồm 11 dòng thô như "Meo an ca.", một dòng trống và hai dòng trùng; sau khi làm sạch còn 8 tài liệu. Tập kiểm định (validation: dữ liệu chỉ dùng để đo, không dùng để cập nhật) có 2 câu chưa gặp: "cho an ca." và "meo uong nuoc.". Dữ liệu SFT là 8 cặp hỏi đáp như "meo an gi" và "ca."; dữ liệu ưu tiên là 6 bộ (câu hỏi, câu được chọn, câu bị loại) như "meo an gi", "ca.", "sua.".

Mô hình là một bộ giải mã Transformer với 2 lớp, 1 head attention (head: một nhóm tính attention độc lập; ở đây chỉ có một nên head dùng toàn bộ chiều ẩn), chiều ẩn d = 16, mạng truyền thẳng SwiGLU (mạng mở rộng chiều rồi lọc bằng một cổng) rộng 32, RMSNorm (chuẩn hóa biên độ vector) và mã hóa vị trí cố định. Tokenizer tách theo ký tự nên từ vựng chỉ có 23 token: 19 ký tự cộng 4 token đặc biệt là pad (đệm), eos (kết thúc), user và assistant (đánh dấu lượt nói). Tổng cộng 5856 tham số, khởi tạo ngẫu nhiên với seed 7 nên chạy lại cho đúng các số cũ.

Các hằng số ở đầu file (STEPS, LR, ZERO_WORKERS, MASK_RESPONSE_ONLY, LORA_RANK, DISTILL_T, DPO_BETA và các hằng khác) là những núm bạn sẽ xoay ở phần bài tập. Chạy một lần mất vài giây trong trình duyệt.

Bảng ánh xạ giữa sơ đồ Tab 01 và 13 giai đoạn của lab

Giai đoạn 1 Dữ liệu → token IDs (ô 1A) = clean_corpus, encode_text và pack_lm; giai đoạn 2 Next-token targets (ô 1B) = cách pack_lm dịch nhãn và dựng mặt nạ; giai đoạn 3 Forward (ô 2A đến 2D) = block_forward và forward; giai đoạn 4 Cross-entropy loss (ô 3A) = cross_entropy; giai đoạn 5 Backward & gradient từng lớp (ô 3B, 3C) = block_backward, backward và check_gradients; giai đoạn 6 AdamW (ô 3D) = adamw_update; giai đoạn 7 Vòng lặp, eval, checkpoint (ô 3E, cùng ô 4C pretraining) = train_lm, save_checkpoint và load_checkpoint.

Giai đoạn 8 Bộ nhớ: ZeRO & recompute (ô 4A, 4B) = zero_bytes cùng hai phép mô phỏng chia dữ liệu và chia trạng thái AdamW, và nhánh backward tính lại từng lớp; giai đoạn 9 SFT (ô 1C, 4D) = chat_ids và pack_sft; giai đoạn 10 LoRA (ô 4E) = các cặp adapter A, B và hàm merged; giai đoạn 11 Distillation (ô 4F) = distill_loss_grad; giai đoạn 12 Preference (DPO) (ô 4G) = seq_logp_and_dlogits và dpo_grads; giai đoạn 13 Từ checkpoint sang suy luận (ô 5A đến 5C) = step_forward trên checkpoint nạp lại. Tên giai đoạn trong output khớp tên ô trên sơ đồ.

Cách thực hiện và đọc từng khối output

1. Trước khi bấm Chạy, hãy dự đoán: loss lúc chưa huấn luyện là bao nhiêu nếu mô hình đoán đều trên 23 token (gợi ý: ln 23)? Sau 300 bước, loss kiểm định sẽ thấp hơn hay cao hơn loss huấn luyện, và vì sao? Viết dự đoán ra giấy.

2. Bấm Chạy. Ở khối 1 và 2, đọc bảng từ vựng, các ID của tài liệu đầu và cặp input, nhãn: nhãn luôn là token đứng ngay sau input. Ở khối 3, tìm hàng attention chỉ có 4 phần tử khác 0, phép thử đổi token cuối cho sai lệch 0.0 và so shape (10, 16) của x với (10, 23) của logits. Ở khối 4, nhìn từng giá trị trừ ln p ở bốn vị trí đầu rồi so trung bình với ln 23.

3. Ở khối 5, đọc sai lệch giữa gradient giải tích và sai phân hữu hạn; ở khối 6, theo dõi một tham số cụ thể từ gradient g qua m, v tới bước cập nhật và so với SGD. Ở khối 7, đọc bảng loss theo bước: hai cột loss huấn luyện và loss kiểm định có đi cùng nhau không; rồi so lần chạy tiếp từ checkpoint đầy đủ (lệch 0.0) với lần chỉ nạp trọng số.

4. Ở khối 8, tự tính trước byte mỗi worker rồi đối chiếu ba dòng ZeRO và các sai lệch bằng 0. Ở khối 9, nhìn mặt nạ loss của cặp đầu và hai con số loss trên token câu trả lời với token prompt. Ở khối 10 đến 12, so số tham số học được, số câu đúng và các sai lệch kiểm tra; ở khối 12, đọc cả xác suất của câu chosen lẫn câu rejected chứ không chỉ biên.

5. Ở khối 13, đối chiếu ID của prompt, shape của KV cache, số vị trí có cache so với tính lại và dòng nạp nhầm checkpoint. Cuối cùng, đối chiếu dự đoán ban đầu với kết quả rồi làm các bài tập: sửa một hằng số ở đầu file, dự đoán, chạy lại và giải thích.

Kết quả này chứng tỏ điều gì và giới hạn ở đâu?

Các phép kiểm tra trong lab xác nhận, cho đúng mô hình và dữ liệu này: gradient giải tích khớp sai phân hữu hạn trong 1e-6 ở mọi tensor, cả với loss KL của distillation và loss DPO; gradient cả batch bằng tổng gradient từng worker; cập nhật AdamW chia theo ZeRO-1 giống hệt cập nhật không chia; gradient tính lại từng lớp giống hệt gradient lưu đầy đủ; mô hình LoRA gộp và chưa gộp cho cùng logits; trọng số gốc không đổi khi huấn luyện LoRA; logits có KV cache khớp logits tính lại cả chuỗi (sai lệch cỡ 1e-15, chặn trên 1e-9).

Lab không chứng minh: chất lượng ngôn ngữ (dữ liệu chỉ vài chục ký tự, câu trả lời được học thuộc); khả năng trả lời câu hỏi chưa gặp (xem câu meo bay o dau ở khối 9); tốc độ hay bộ nhớ đỉnh trên GPU thật (ZeRO chỉ là phép chia mảng và đếm byte trong một tiến trình, recompute chỉ so gradient); LoRA hay distillation tốt hơn tinh chỉnh toàn bộ ở mô hình thật (kết quả chỉ trên 8 cặp và một seed); và RLHF với dữ liệu ưu tiên của con người (DPO ở đây chỉ có 6 cặp tự viết).

Kết quả mong đợiOutput gồm phần mở đầu, 13 khối đánh số 1 đến 13 và dòng TỔNG KẾT. Mốc kiểm tra nhanh: từ vựng 23 token và 5856 tham số; loss ban đầu 3.2481 so với ln 23 = 3.1355; sau 300 bước loss huấn luyện 0.1014 và loss kiểm định 1.0469 (thấp nhất 0.8048 ở bước 50); bản tiếp tục được gấp 2.96 lần bản weights-only; ZeRO-3 với 2 worker còn 46848 byte mỗi worker; SFT đưa số câu đúng từ 0 lên 8 trên 8 còn LoRA 8 trên 8 với 18.6% tham số học được; distillation cho KL 0.0068 so với 0.1955; DPO nâng biên lên 8.317; cache cho 17 vị trí so với 75 và đầu ra sua.; cuối cùng là dòng PASS.

Tải training-end-to-end.py ↓

training-end-to-end.pyPYTHON

"""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()

Python outputSẵn sàng.

Bài tập 1 · Tăng ZERO_WORKERS từ 2 lên 4

Dựa vào khối 8: với 2 worker, mỗi worker cần 93696 byte (data parallel), 58560 (ZeRO-1), 52704 (ZeRO-2) và 46848 (ZeRO-3). Hãy dự đoán bốn con số này khi ZERO_WORKERS = 4 (gợi ý: công thức là 16P, 4P + 12P/N, 2P + 14P/N và 16P/N với P = 5856). Sau đó sửa dòng ZERO_WORKERS = 4, chạy lại và cho biết cột nào không đổi, cột nào không thể nhỏ hơn một mức nào đó và vì sao.

Đối chiếu sau khi tự trả lời

Với 4 worker: data parallel 93696 (không đổi vì mỗi worker vẫn giữ đủ mọi thứ), ZeRO-1 40992, ZeRO-2 32208 và ZeRO-3 23424. ZeRO-3 giảm đúng một nửa vì chia cả ba thành phần cho N. ZeRO-1 không thể nhỏ hơn 4P = 23424 và ZeRO-2 không thể nhỏ hơn 2P = 11712 vì phần trọng số và (với ZeRO-1) gradient vẫn được giữ đủ trên mỗi worker, dù N lớn bao nhiêu. Gradient của các worker cộng lại vẫn khớp gradient cả batch (lệch cỡ 1e-17) và cập nhật chia theo ZeRO-1 vẫn lệch 0.0: chia nhỏ trạng thái không đổi toán học. Đây là phép đếm byte, chưa phải đo bộ nhớ hay thời gian giao tiếp.

Bài tập 2 · Giảm STEPS từ 300 xuống 50

Dự đoán: nếu chỉ huấn luyện 50 bước thay vì 300, loss huấn luyện cuối cùng cao hơn hay thấp hơn 0.1014, và loss kiểm định cuối cùng cao hơn hay thấp hơn 1.0469? Sửa STEPS = 50, chạy lại và giải thích hiện tượng bằng bảng loss ở khối 7.

Đối chiếu sau khi tự trả lời

Loss huấn luyện cao hơn (0.2657, vì mô hình mới học ít bước) nhưng loss kiểm định thấp hơn (0.8048 so với 1.0469). Bảng ở lần chạy mặc định đã cho thấy điều đó: loss kiểm định thấp nhất ở bước 50 rồi tăng dần trong khi loss huấn luyện vẫn giảm, dấu hiệu mô hình bắt đầu học thuộc 8 tài liệu. Kết luận chỉ mang tính minh họa vì tập kiểm định chỉ có 2 câu; trên dữ liệu thật người ta dùng tập kiểm định đủ lớn để chọn bước dừng hoặc checkpoint.

Bài tập 3 · Tăng LR từ 0.01 lên 0.3

Dự đoán: learning rate lớn gấp 30 lần có làm huấn luyện nhanh hơn không? Sửa LR = 0.3 (hằng số này dùng cho cả huấn luyện nền và SFT), chạy lại và so loss huấn luyện sau 300 bước, câu mô hình sinh ra từ prompt meo an và số câu đúng sau SFT ở khối 9.

Đối chiếu sau khi tự trả lời

Không. Loss huấn luyện sau 300 bước chỉ xuống 2.0025 (so với 0.1014), prompt meo an cho ra chuỗi lặp " n n n n n n" thay vì " ca.", và SFT cho 0 trên 8 câu đúng (loss SFT chỉ xuống từ 2.5787 còn 1.7642). Bước quá dài làm mỗi lần cập nhật nhảy qua vùng loss thấp thay vì đi vào nó; AdamW chuẩn hóa độ dài bước theo lr nên lr là núm quyết định. Lưu ý các phép kiểm tra cơ chế (gradient khớp sai phân hữu hạn, ZeRO, merge, cache) vẫn đạt vì chúng không phụ thuộc vào việc mô hình học tốt hay không; dòng PASS không nói mô hình đã học tốt.

Bài tập 4 · Tắt MASK_RESPONSE_ONLY

Ở khối 9, loss trên token prompt đi từ 3.8980 lên 5.0391 sau SFT (token prompt không được chấm). Dự đoán: nếu đặt MASK_RESPONSE_ONLY = False (chấm cả prompt) thì loss trên token prompt sau SFT sẽ thế nào, loss trên token câu trả lời có tốt hơn không, và số câu đúng có đổi không? Chạy lại để kiểm tra.

Đối chiếu sau khi tự trả lời

Loss trên token prompt giảm mạnh xuống 0.1851 vì mô hình học luôn việc đoán các token của câu hỏi. Loss trên token câu trả lời không tốt hơn: 0.0055 so với 0.0036. Số câu đúng vẫn là 8 trên 8, nên trên dữ liệu nhỏ này khác biệt chủ yếu nằm ở nơi công sức học đi vào. Mặt nạ response-only dùng toàn bộ tín hiệu học cho việc trả lời; chấm cả prompt dạy mô hình sinh ra câu hỏi, điều thường không mong muốn khi huấn luyện một trợ lý. Kết quả này không chứng minh mặt nạ nào cho chất lượng hội thoại tốt hơn ở mô hình thật.

Bài tập 5 · Tăng LORA_RANK từ 2 lên 8

Dự đoán: số tham số học được và tỉ lệ so với 5856 sẽ là bao nhiêu khi hạng r tăng 4 lần (gợi ý: mỗi ma trận d_vào × d_ra thêm r × (d_vào + d_ra) tham số, tổng cho 14 ma trận)? Số câu đúng và loss sau LoRA có tốt hơn nhiều không? Sửa LORA_RANK = 8 rồi chạy lại.

Đối chiếu sau khi tự trả lời

Số tham số học được tăng gấp 4 lần, từ 1088 (18.6%) lên 4352 (74.3% số tham số gốc). Số câu đúng vẫn là 8 trên 8 và loss sau LoRA chỉ nhích từ 0.0120 xuống 0.0096. Với chiều ẩn nhỏ như d = 16, hạng 8 đã là một nửa chiều nên hạng thấp không còn thấp: chi phí tăng gấp 4 mà lợi ích gần như không đổi. Lợi thế của LoRA nằm ở chỗ hạng r nhỏ hơn rất nhiều so với d, điều mô hình thật (d hàng nghìn) đạt được còn mô hình đồ chơi này thì không; đừng suy ra cách chọn r cho mô hình thật từ kết quả trên.

Bài tập 6 · Giảm DPO_BETA từ 0.5 xuống 0.1

β nhân vào biên z nên β nhỏ làm loss giảm chậm hơn. Dự đoán: sau 60 bước biên z trung bình sẽ lớn hơn hay nhỏ hơn 8.317, và policy sẽ dịch xa hơn hay gần hơn mô hình tham chiếu (nhìn xác suất của câu chosen và rejected của cặp đầu)? Sửa DPO_BETA = 0.1, chạy lại rồi giải thích.

Đối chiếu sau khi tự trả lời

Biên z trung bình nhỏ hơn (4.608) và loss còn cao hơn (0.0130 so với 0.0005), nhưng policy dịch xa hơn: xác suất cả câu rejected của cặp đầu xuống tới 2.13e-18 (so với 2.70e-10) và xác suất câu chosen là 0.653 (so với 0.520). Vì z = β nhân với phần chênh log-xác suất, β nhỏ cần chênh lệch lớn hơn mới đạt cùng z, tức ràng buộc với mô hình tham chiếu lỏng hơn; β lớn giữ policy gần tham chiếu hơn. Ở cả hai lần, xác suất câu chosen vẫn giảm so với tham chiếu (0.982): biên dương không có nghĩa câu được chọn trở nên chắc chắn hơn.

Bài tập 7 · Giảm DISTILL_T từ 2.0 xuống 1.0

Ở khối 11, nhiệt độ T = 2 làm phân phối của teacher phẳng hơn. Dự đoán: nếu T = 1 (không làm phẳng) thì student học từ soft target có KL tới teacher thấp hơn hay cao hơn 0.0068 sau cùng số bước, và có còn trả lời đúng 8 trên 8 không? Sửa DISTILL_T = 1.0 rồi chạy lại.

Đối chiếu sau khi tự trả lời

KL cao hơn: 0.0166 so với 0.0068, còn số câu đúng vẫn là 8 trên 8, và student chỉ học nhãn cứng vẫn là 5 trên 8 với KL 0.1955. Ở T = 1 gần như toàn bộ xác suất của teacher nằm ở một token (0.996) nên soft target gần giống nhãn cứng; T = 2 làm lộ thứ tự ưu tiên giữa các token còn lại nên student bắt chước cả phân phối trong cùng số bước. Chênh lệch đo trên một seed và 8 cặp, nên đừng coi T = 2 là giá trị tốt nhất nói chung.