Bài Lab tổng hợp — LLM Training (Text)

Giới thiệu bài lab

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

Sơ đồ ở Tab 01 cho biết các ô huấn luyện nối với nhau thế nào; bài lab này cho bạn thấy mỗi ô tính ra con số gì. Một chương trình Python dùng NumPy (thư viện tính trên mảng số) chạy ngay trong trình duyệt và huấn luyện một mô hình ngôn ngữ rất nhỏ, gồm 5856 tham số (mỗi tham số là một con số mô hình học được) và từ vựng 23 token, qua 13 giai đoạn khớp với các ô 1A đến 5C.

Bạn cần bài lab vì đọc về loss, gradient hay LoRA chưa đủ để biết mô tả có đúng. Khi tự thấy loss lúc chưa huấn luyện là 3.2481 và sau 300 bước huấn luyện xuống 0.1014, bạn mới nắm được loss giảm nghĩa là gì và giảm bao nhiêu.

Bài lab gồm

Một bảng ánh xạ 13 hàng: mỗi giai đoạn của lab khớp ô nào ở Tab 01, bài học nào, lab dựng gì và kiểm chứng gì.
Mục Hướng dẫn đọc output gồm năm phần gấp lại: bạn sẽ học gì, đầu vào và cấu hình, bảng ánh xạ bằng lời, cách đọc từng khối output, và kết quả chứng tỏ gì cùng giới hạn của nó.
Trình chạy trực tiếp cho file training-end-to-end.py: bạn sửa được mã, chạy lại, và nhận output gồm phần mở đầu, 13 khối [n · Tên] rồi dòng TỔNG KẾT.
Dữ liệu tự viết vài chục ký tự, một mô hình 2 lớp với chiều ẩn d = 16 (mỗi vector có 16 số) và seed 7 cố định, nên chạy lại cho đúng các số cũ.
Các hằng số ở đầu file để thử, như STEPS, LR, ZERO_WORKERS, MASK_RESPONSE_ONLY, LORA_RANK, DISTILL_T, DPO_BETA; bảy bài tập có đáp án nằm ở trang bài đọc Lab đầy đủ.

Làm xong, bạn có thể

Đọc từng khối output và nói con số nào chứng tỏ điều gì, ví dụ loss ban đầu 3.2481 gần ln 23 = 3.1355 (logarit tự nhiên của 23, mức loss khi đoán đều trên 23 token) vì mô hình chưa học.
Giải thích vì sao gradient tính bằng công thức có thể được kiểm tra bằng sai phân hữu hạn: dịch một tham số một chút rồi đo loss đổi bao nhiêu.
Dự đoán, chạy rồi giải thích kết quả khi đổi một hằng số, ví dụ giảm STEPS từ 300 xuống 50 hay tăng LR từ 0.01 lên 0.3.
Phân biệt SFT, LoRA, distillation và DPO bằng số cụ thể: số câu trả lời đúng, tỉ lệ tham số học được (LoRA là 18.6%) và độ lệch so với mô hình lớn.
Nêu điều lab không chứng minh: chất lượng ngôn ngữ, tốc độ hay bộ nhớ trên GPU thật, và việc LoRA hay distillation tốt hơn ở mô hình thật.

Phạm vi: lab huấn luyện một mô hình văn bản thu nhỏ trên dữ liệu tự viết, không dùng ảnh hay âm thanh, và dòng PASS chỉ nói các phép kiểm tra cơ chế đạt chứ không nói mô hình trả lời giỏi hay chạy nhanh. Muốn xem kiến trúc đa phương thức, đọc bài Text-only / multimodal của LLM Foundations.

Cách dùng bài lab

Đọc bảng ánh xạ bên dưới: mỗi hàng là một giai đoạn, cho biết ô tương ứng ở Tab 01, bài học (bấm để mở), lab dựng gì và kiểm chứng gì. Giai đoạn 3 gộp các ô 2A đến 2D, giai đoạn 7 gộp 3E và 4C.
Mở mục Hướng dẫn đọc output, nhất là Đầu vào và cấu hình thu nhỏ của bài lab, để biết dữ liệu, mô hình và các hằng số là gì trước khi chạy.
Dự đoán trước khi bấm Chạy: 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, và sau 300 bước loss kiểm định cao hơn hay thấp hơn loss huấn luyện. Viết dự đoán ra giấy.
Bấm Chạy (mất vài giây). Đọc từng khối theo bảng: dòng Quan sát trước, đối chiếu các con số, rồi mới đọc Kết luận và Không chứng minh. Mốc kiểm tra nhanh: từ vựng 23 token, 5856 tham số, loss ban đầu 3.2481, và cuối cùng là dòng PASS.
Sửa một hằng số ở đầu file, dự đoán kết quả, chạy lại và giải thích; sau đó đối chiếu với đáp án của bài tập ở trang bài đọc Lab đầy đủ, nút ở thanh phía trên bảng.

Bạn chạy lại bao nhiêu lần cũng được: seed 7 cố định nên cùng một cấu hình cho đúng các số cũ, nên mọi thay đổi trong output đến từ chỗ bạn vừa sửa.

BẢNG ÁNH XẠ · MỖI Ô KIẾN TRÚC Ở TAB 01 ĐƯỢC DỰNG VÀ KIỂM CHỨNG TRONG LAB NHƯ THẾ NÀO?

Giai đoạn	Ô	Bài học (bấm để mở)	Lab NumPy dựng gì (`training-end-to-end.py`)	Cách đọc số liệu & phép kiểm chứng
1	1A	Dữ liệu → token IDs ↗	Làm sạch 11 dòng thô còn 8 tài liệu, tokenizer theo ký tự (23 token), đóng gói batch 4 tài liệu.	Tài liệu đầu thành 11 ID kết thúc bằng eos; một batch có 63 vị trí; 2 bước là một epoch.
2	1B	Next-token targets ↗	Dịch nhãn đi một vị trí và dựng mặt nạ nhân quả cắt theo ranh giới từng tài liệu.	Tài liệu 11 token cho 10 cặp (input, nhãn); tài liệu sau không nhìn thấy tài liệu trước (False).
3	2A–2D	Forward: embedding, attention, MLP, norm → logits ↗ Đọc thêm: Attention, MLP, Norm	Embedding, attention một head có mặt nạ nhân quả, SwiGLU, RMSNorm, residual và đầu ra logits; 2 lớp, d = 16.	Attention ở vị trí 3 chỉ có 4 phần tử khác 0, tổng 1; đổi token cuối không đổi logits phía trước (0.0); logits shape (10, 23).
4	3A	Cross-entropy loss ↗	Softmax, trừ ln p của token đúng, trung bình trên các vị trí được chấm.	Loss ban đầu 3.2481 gần ln 23 = 3.1355 vì mô hình chưa học.
5	3B–3C	Backward & gradient từng lớp ↗ Đọc thêm: Gradient từng lớp	Gradient giải tích cho 16 tensor, gom theo lớp, kèm kiểm tra bằng sai phân hữu hạn.	Gradient theo logits tổng mỗi hàng bằng 0; khớp sai phân hữu hạn (sai lệch cỡ 5e-10); lớp 0 là 0.9951, lớp 1 là 0.6791.
6	3D	AdamW ↗	AdamW viết tay với m, v, hiệu chỉnh ở các bước đầu và suy giảm trọng số tách riêng.	Bước đầu mỗi tham số đổi 0.0098 đến 0.0101 khi lr = 0.01; tính tay khớp mã; loss một bước 3.2746 xuống 2.8802.
7	3E, 4C	Vòng lặp, eval, checkpoint ↗ Đọc thêm: Pretraining	300 bước với batch 4 tài liệu, đánh giá mỗi 50 bước, lưu checkpoint weights-only và checkpoint tiếp tục được.	Loss huấn luyện 0.1014; loss kiểm định thấp nhất 0.8048 rồi 1.0469; checkpoint đầy đủ gấp 2.96 lần bản weights-only; chạy tiếp từ bản đầy đủ lệch 0.0.
8	4A–4B	Bộ nhớ: ZeRO & recompute ↗ Đọc thêm: Recompute	Đếm byte (16 mỗi tham số) cho data parallel và ZeRO-1/2/3; chia dữ liệu và trạng thái AdamW cho 2 worker; tính lại activation từng lớp khi backward.	ZeRO-3 còn 46848 byte mỗi worker; cập nhật chia và gradient tính lại lệch 0.0; 7 tỉ tham số trên 8 GPU từ 112 xuống 14 GB (tính tay, không đo).
9	1C, 4D	SFT: chat template + response mask ↗	Mẫu hội thoại có token user và assistant, mặt nạ chỉ chấm câu trả lời; 150 bước từ checkpoint nền.	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 3.8980 lên 5.0391; câu chưa gặp vẫn được trả lời bằng câu của cặp khác.
10	4E	LoRA ↗	Cặp adapter A, B hạng 2 trên 14 ma trận, B khởi tạo bằng 0, có thể gộp vào trọng số gốc.	1088 tham số học được (18.6%); trọng số gốc lệch 0.0; gộp và chưa gộp cho cùng logits (cỡ 1e-14); 8 trên 8 câu đúng.
11	4F	Distillation ↗ Đọc thêm: Response distillation	Student (d = 8, 1 lớp) học loss T² nhân KL với soft target của teacher ở T = 2.	KL 0.0068 và 8 trên 8 câu đúng, so với 0.1955 và 5 trên 8 khi chỉ học nhãn cứng; gradient KL khớp sai phân hữu hạn.
12	4G	Preference (DPO) ↗	Loss DPO với mô hình tham chiếu đóng băng trên 6 cặp (chosen, rejected), β = 0.5.	Khi policy bằng tham chiếu loss đúng ln 2 = 0.6931; sau 60 bước biên 8.317 nhưng xác suất câu chosen của cặp đầu lại giảm từ 0.982 xuống 0.520.
13	5A–5C	Từ checkpoint sang suy luận ↗ Đọc thêm: Tách token, Prefill, Decode, Output	Nạp checkpoint weights-only, đổi prompt thành token ID, prefill một lượt rồi decode với KV cache và giải mã văn bản.	Logits có cache khớp tính lại (cỡ 1e-15); 17 vị trí xử lý so với 75; đầu ra sua.; nạp nhầm checkpoint cùng cấu hình không báo lỗi.

Hướng dẫn đọc output

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).

Trình chạy Lab trực tiếp · `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.

Nguồn công bố tham chiếu

Attention Is All You Need (Vaswani và cộng sự, 2017) ↗ Kiến trúc Transformer: attention, mạng truyền thẳng, residual và norm Decoupled Weight Decay Regularization (AdamW, 2019) ↗ Bộ tối ưu AdamW và suy giảm trọng số tách riêng ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (2020) ↗ Chia trọng số, gradient và trạng thái bộ tối ưu cho nhiều GPU LoRA: Low-Rank Adaptation of Large Language Models (2021) ↗ Tinh chỉnh bằng cặp ma trận hạng thấp, đóng băng trọng số gốc Direct Preference Optimization (2023) ↗ Học từ cặp câu được chọn và câu bị loại, không cần mô hình thưởng riêng

Danh sách bài học (25)