Bên trong một Transformer block: FFN, residual stream và normalization

Bài thứ bảy dừng lại ở mục 09, với một sơ đồ Transformer block gồm bốn phần: attention, MLP/FFN, residual và normalization. Bài đó chỉ đi sâu vào phần đầu tiên.

Có một chi tiết dễ bỏ qua: attention không phải phần nặng nhất của block. Trong Llama 3.1 8B, mỗi layer có khoảng 218 triệu tham số, và chỉ chưa tới 20% trong số đó thuộc về attention. Phần còn lại gần như toàn bộ nằm ở FFN, một khối mà nhiều bài giới thiệu Transformer chỉ nhắc qua trong một dòng.

Bài này lấp nửa còn lại đó. Có bốn câu hỏi cần trả lời:

1. Thông tin đi từ layer này sang layer khác bằng đường nào? Đó là residual stream.
2. FFN làm gì với từng token, và vì sao người ta gọi nó là một bộ nhớ key–value?
3. Vì sao hàm kích hoạt đổi từ ReLU sang GELU rồi SwiGLU, và vì sao vì thế mà d_ff của Llama 3 là 14.336 chứ không phải 4 × 4096?
4. Normalization đặt ở đâu, và vì sao chuyện “đặt ở đâu” lại quan trọng hơn “dùng loại nào”?

Cuối bài, ta ghép tất cả lại thành một phép đếm cụ thể: số tham số và FLOPs của một block Llama 3.1 8B, tính tay và khớp với con số Meta công bố.

Bài 1 đã có một mục tổng quan về block. Bài này không nhắc lại phần đó mà đi sâu hơn, và sửa lại vài con số quen thuộc nhưng đã lỗi thời.

01Nửa còn lại của block: attention trộn, FFN xử lý

Một decoder block trong LLM hiện đại có hai sub-layer, và mỗi sub-layer làm một việc rất khác nhau:

Cách nhìn hữu ích nhất: attention là bước giao tiếp, FFN là bước tính toán. Attention cho mỗi token lấy thông tin từ các token khác. Sau đó FFN xử lý riêng từng vị trí, như thể mỗi token được đưa qua cùng một hàm biến đổi lớn, độc lập với nhau.

Vì FFN không nhìn sang token khác, mọi thứ nó cần biết về ngữ cảnh phải được attention đưa vào trước đó. Ngược lại, attention chỉ trộn và di chuyển thông tin, không có một phép biến đổi phi tuyến mạnh nào cho từng token. Thiếu bên nào thì block cũng không đủ.

Hai sub-layer này không nối thẳng vào nhau. Chúng cùng nối vào một thứ thứ ba: residual stream.

02Residual stream: bus chung mà mọi layer đọc và ghi

Bài 1 đã giải thích residual connection ở mức cơ bản: output = x + f(x), và vì sao thiếu nó thì gradient tiêu biến khi model sâu. Mục này nhìn cùng phép cộng đó từ một góc khác, góc mà Elhage et al. 2021 dùng trong A Mathematical Framework for Transformer Circuits.

Viết lại toàn bộ forward pass của một model L layer theo kiểu Pre-LN (kiểu mà Llama dùng):

x_0 = Embedding(token)
x_{l+½} = x_l     + Attn_l( Norm(x_l) )
x_{l+1} = x_{l+½} + FFN_l ( Norm(x_{l+½}) )
logits  = lm_head( Norm(x_L) )

Không có dòng nào ghi đè x. Mỗi sub-layer chỉ cộng thêm một vector vào. Mở hết các phép cộng ra:

x_L = x_0 + Attn_1(…) + FFN_1(…) + Attn_2(…) + … + FFN_L(…)

Nói cách khác, hidden state ở cuối model là tổng của embedding ban đầu và tất cả đóng góp của mọi sub-layer. Elhage et al. gọi vector đang được cộng dồn này là residual stream, và mô tả nó như một kênh giao tiếp tuyến tính:

• mỗi sub-layer đọc từ stream bằng một phép chiếu tuyến tính (ma trận Q/K/V của attention, ma trận vào của FFN);
• rồi ghi vào stream bằng một phép chiếu khác (ma trận output) cộng với phép cộng.

Hình 1 ở đầu bài vẽ đúng điều này: stream là một trục đi thẳng qua mọi block, còn attention và FFN là các trạm rẽ ra đọc rồi nhập lại.

Vì sao chiều của stream không bao giờ đổi

Vì mọi sub-layer cộng vào cùng một vector, mọi output phải có cùng số chiều d_model. Với Llama 3.1 8B, d_model = 4096. Bên trong mỗi sub-layer, chiều có thể thay đổi:

Mọi khối đều phải chiếu về 4.096 trước khi ghi. Đó là “chuẩn giao tiếp” của bus.

Bus chung có nghĩa là gì trong thực tế

Ý “bus chung” có ba hệ quả cụ thể:

1. Layer sau đọc được những gì mọi layer trước đã ghi, không chỉ layer ngay trước nó. Một feature mà FFN ở layer 3 ghi vào vẫn còn trong stream khi tới layer 20, trừ khi có layer nào đó chủ động ghi một vector ngược dấu để xoá nó.
2. Các layer có thể ghi vào những subspace khác nhau. Stream có 4.096 chiều. Elhage et al. lập luận rằng các sub-layer có thể dùng những hướng khác nhau trong không gian đó, nên nhiều loại thông tin cùng tồn tại trong một vector mà không đè lên nhau.
3. Gradient có một đường đi thẳng. Vì x_L là tổng có chứa x_l, đạo hàm của x_L theo x_l luôn có một số hạng đơn vị. Đây chính là lý do kỹ thuật khiến residual giúp train model sâu, nhìn từ phía stream.

Từ góc nhìn này, câu hỏi “FFN làm gì” trở thành: FFN đọc gì từ stream, và ghi gì vào lại?

03FFN như một bộ nhớ key–value

Viết FFN cổ điển (kiểu GPT-2, hai ma trận, bỏ qua bias) cho một token x:

FFN(x) = f(x · W_in) · W_out

x      : 1 × d_model
W_in   : d_model × d_ff
W_out  : d_ff × d_model
f      : hàm kích hoạt (ReLU, GELU, …)

Giờ đặt nó cạnh công thức attention ở bài 7:

Attention(q, K, V) = softmax( q · Kᵀ ) · V
FFN(x)             =        f( x · K  ) · V     với K = W_in, V = W_out

Hai công thức có cùng một dạng. Geva et al. 2021 dựa vào đúng sự giống nhau này để đặt tên cho bài báo của họ: Transformer Feed-Forward Layers Are Key-Value Memories.

Đọc FFN theo cách đó:

• Mỗi cột của W_in là một key. Có d_ff key. Với GPT-2 small là 3.072 key mỗi layer, với Llama 3.1 8B là 14.336.
• x · W_in cho ra d_ff con số, mỗi con số đo xem token hiện tại khớp với key tương ứng đến đâu.
• Hàm kích hoạt f giữ lại các key khớp mạnh và dập các key không khớp. Với ReLU, phần âm bị cắt hẳn về 0.
• Mỗi hàng của W_out là một value. Output là tổng các value, mỗi value được nhân với hệ số khớp của key tương ứng.

Hai điểm khác attention đáng chú ý. Thứ nhất, key và value của FFN không phụ thuộc vào input: chúng là tham số, giống nhau cho mọi câu. Thứ hai, FFN dùng một hàm kích hoạt phi tuyến chứ không dùng softmax, nên các hệ số không bị ép tổng bằng 1. Nhiều key có thể cùng bật mạnh, hoặc không key nào bật.

Key bắt cái gì, value đẩy cái gì

Geva et al. 2021 không dừng ở phép so sánh hình thức. Họ lấy từng key trong một model ngôn ngữ 16 layer, tìm những câu trong tập train làm key đó bật mạnh nhất, rồi nhờ người đọc xem các câu đó có gì chung. Kết quả:

• Key bắt các mẫu (pattern) đọc được bằng mắt người. Ở layer thấp, mẫu thường nông: những câu kết thúc bằng cùng một từ hay cùng một n-gram. Ở layer cao, mẫu mang tính ngữ nghĩa hơn: những câu nói về cùng một chủ đề, chẳng hạn chương trình truyền hình.
• Value đẩy xác suất cho một nhóm token đầu ra. Chiếu một value qua ma trận output embedding sẽ ra một phân phối trên vocabulary. Ở các layer trên, phân phối đó thường khớp với token xuất hiện tiếp theo sau mẫu mà key tương ứng bắt.
• Output của FFN là hợp của nhiều memory. Một token thường bật hàng trăm key cùng lúc, và output là tổng của các value đó chứ không phải của một value đơn lẻ.

Một bài tiếp theo của cùng nhóm (Geva et al. 2022) đẩy ý này xa hơn: mỗi value có thể đọc như một lần “cập nhật” lên phân phối token trong không gian vocabulary, và nhiều value tương ứng với những khái niệm con người hiểu được.

Giới hạn của cách diễn giải này

Từ đây rất dễ trượt sang câu “FFN là nơi lưu tri thức của model”. Câu đó đi xa hơn bằng chứng:

• Geva et al. 2021 là quan sát trên một tập key được chọn và chấm bằng tay. Bài báo không chứng minh rằng phần lớn tri thức nằm ở FFN.
• Các phương pháp như Knowledge Neurons (Dai et al. 2022) và ROME (Meng et al. 2022) tìm thấy những vị trí trong FFN mà khi sửa thì model “nhớ” một fact khác. Đó là bằng chứng gợi ý rằng FFN đóng vai trò lớn trong việc gợi lại fact.
• Nhưng Hase et al. 2023 cho thấy chỗ mà phương pháp định vị chỉ ra chưa chắc là chỗ sửa hiệu quả nhất. Sửa ở một layer khác đôi khi tốt bằng hoặc tốt hơn. Vậy nên “định vị được” và “tri thức nằm ở đó” là hai mệnh đề khác nhau.

Câu an toàn hơn là: bằng chứng gợi ý rằng FFN hoạt động giống một bộ nhớ key–value và đóng vai trò lớn trong việc gợi lại kiến thức, còn attention lo việc đưa đúng ngữ cảnh tới trước. Ý “FFN là kho kiến thức, vậy chia kho ra thành nhiều expert” chính là động lực của Mixture of Experts. Bài 1 đã có một mục ngắn về MoE, còn phần đi sâu để dành cho một bài sau trong series.

04Hàm kích hoạt: từ ReLU đến GELU đến SwiGLU

Hàm f trong f(x · W_in) là thứ duy nhất làm FFN phi tuyến. Không có nó, x · W_in · W_out chỉ là một phép nhân ma trận, và xếp chồng bao nhiêu layer cũng vẫn chỉ là một phép biến đổi tuyến tính. Ba thế hệ hàm kích hoạt khác nhau ở cách chúng quyết định giữ bao nhiêu từ mỗi key.

ReLU: cắt cứng

Transformer gốc năm 2017 dùng ReLU(z) = max(0, z). Key nào khớp âm thì bị bỏ hẳn, key nào khớp dương thì giữ nguyên. Đơn giản và rẻ, nhưng có hai nhược điểm: đạo hàm bằng 0 trên toàn bộ nửa âm, và có một góc gãy ở 0.

GELU: nhân với xác suất giữ lại

Hendrycks và Gimpel đề xuất GELU với một trực giác xác suất: thay vì quyết định cứng “giữ hay bỏ”, hãy nhân input với xác suất giữ lại nó, mà xác suất này tăng theo độ lớn của input:

GELU(z) = z · Φ(z)        Φ là hàm phân phối tích luỹ của chuẩn N(0, 1)

Vài giá trị để thấy hình dạng:

Với z lớn, GELU gần như bằng ReLU. Quanh 0 nó mượt, và cho một chút giá trị âm nhỏ thay vì cắt hẳn. Paper cũng đưa ra một xấp xỉ bằng tanh để tính nhanh hơn:

GELU(z) ≈ 0,5 · z · (1 + tanh( √(2/π) · (z + 0,044715 · z³) ))

BERT và GPT-2 dùng GELU. Trong config.json của GPT-2, dòng "activation_function": "gelu_new" chính là bản xấp xỉ tanh này.

SwiGLU: thêm một cái cổng

Năm 2020, Shazeer thử một họ biến thể khác trong GLU Variants Improve Transformer. Ý tưởng lấy từ Gated Linear Unit (Dauphin et al. 2016): thay vì một nhánh, cho input đi qua hai phép chiếu song song, rồi dùng một nhánh làm “cổng” cho nhánh kia:

FFN_SwiGLU(x) = ( Swish(x · W_gate) ⊙ (x · W_up) ) · W_down

Swish(z) = z · σ(z)        σ là sigmoid
⊙        = nhân từng phần tử

Swish với hệ số β = 1 trùng với hàm SiLU (Elfwing et al. 2017; Ramachandran et al. 2017). Vì thế trong config.json của Llama và Qwen bạn sẽ thấy "hidden_act": "silu", dù kiến trúc được gọi là SwiGLU. Biến thể dùng GELU làm cổng thì gọi là GEGLU. Gemma 2 dùng biến thể này, nên config của nó ghi gelu_pytorch_tanh.

Đọc theo ngôn ngữ key–value của mục trước: nhánh x · W_up vẫn là “độ khớp với key”, còn nhánh gate quyết định chiều nào được phép đi qua và đi qua bao nhiêu. Vì cổng cũng phụ thuộc vào x, mỗi token có một mặt nạ mềm riêng trên d_ff chiều.

Shazeer đo perplexity trên cùng một ngân sách tham số và compute: GEGLU và SwiGLU tốt hơn ReLU và GELU. Bài báo kết thúc bằng một câu thường được trích:

“We offer no explanation as to why these architectures seem to work; we attribute their success, as all else, to divine benevolence.”

Nói cách khác, tác giả không đưa ra lý do lý thuyết cho việc SwiGLU tốt hơn. Đây là một lựa chọn thực nghiệm, sau đó được LLaMA, Llama 2, Llama 3, Qwen, Mistral và nhiều model khác dùng theo.

Quy tắc 2/3: vì sao d_ff của Llama 3 là 14.336

SwiGLU có ba ma trận thay vì hai. Nếu giữ d_ff = 4 · d_model như FFN cổ điển, số tham số FFN sẽ tăng 1,5 lần. Để so sánh công bằng, Shazeer giảm d_ff xuống còn 2/3, để 3 · d · (2/3 · 4d) = 8d², đúng bằng FFN hai ma trận với 4d. Paper LLaMA ghi rõ: dùng SwiGLU với chiều 2/3 · 4d thay cho 4d.

Code tham chiếu của Meta (llama/model.py) tính d_ff theo ba bước: lấy 2/3 · 4d, nhân thêm một hệ số tuỳ chọn ffn_dim_multiplier, rồi làm tròn lên bội số của multiple_of:

Với Llama 3: 10.922 × 1,3 ≈ 14.198, làm tròn lên bội số của 1.024 thành 14.336. Qwen2.5-7B không theo công thức của Meta. Con số 18.944 lấy thẳng từ intermediate_size trong config.

Bài học ở đây: “d_ff = 4d” chỉ đúng với FFN hai ma trận kiểu GPT-2. Với các LLM dùng SwiGLU, tỉ lệ có thể là 2,7×, 3,5× hay 5,3×, tuỳ nhóm phát triển muốn dồn bao nhiêu tham số vào FFN.

05Normalization: LayerNorm, RMSNorm và chỗ đặt norm

Mỗi sub-layer cộng thêm một vector vào stream. Sau vài chục lần cộng, độ lớn của stream có thể trôi rất xa so với lúc đầu, và mỗi chiều có thể ở một thang khác nhau. Normalization đưa vector về một thang ổn định trước khi đưa vào một phép tính nhạy với độ lớn, chẳng hạn tích vô hướng trong attention.

Bài 1 đã giới thiệu ngắn hai loại norm. Ở đây ta viết lại công thức cho đủ, có cả ε, rồi tập trung vào câu hỏi quan trọng hơn: đặt norm ở đâu.

LayerNorm và RMSNorm

Với một vector x có d chiều:

LayerNorm(x) = (x − μ) / √(σ² + ε) · γ + β
  μ  = trung bình các chiều của x
  σ² = phương sai các chiều của x

RMSNorm(x)   = x / √( mean(x²) + ε ) · γ

γ và β là tham số học được, mỗi cái có d phần tử. ε là một hằng số nhỏ để tránh chia cho 0: Llama 3.1 dùng 1e-5, Qwen2.5 dùng 1e-6.

LayerNorm (Ba et al. 2016) làm hai việc: re-centering (trừ trung bình) và re-scaling (chia cho độ lệch chuẩn). Zhang và Sennrich 2019 đặt câu hỏi: liệu cả hai có cần thiết không? RMSNorm của họ bỏ bước trừ trung bình và bỏ luôn β, chỉ chia cho căn bậc hai của trung bình bình phương (root mean square). Paper báo cáo thời gian chạy giảm từ 7% đến 64% tuỳ model, với chất lượng tương đương LayerNorm.

Một ví dụ nhỏ với x = [2, 4, 6], bỏ qua γ, β và ε:

LayerNorm dời vector về tâm 0, nên thông tin “mọi chiều đều dương” biến mất. RMSNorm chỉ co giãn, nên hướng của vector được giữ nguyên. Thực nghiệm cho thấy bước re-centering không cần thiết cho chất lượng, nên bỏ nó đi là một khoản tiết kiệm gần như miễn phí. Llama, Qwen, Mistral và Gemma đều dùng RMSNorm.

Post-LN và Pre-LN

Transformer gốc năm 2017 đặt norm sau phép cộng residual. Cách này gọi là Post-LN:

Post-LN:  x_{l+1} = Norm( x_l + Sub(x_l) )
Pre-LN:   x_{l+1} = x_l + Sub( Norm(x_l) )

Khác biệt nằm ở chỗ norm có nằm trên đường residual hay không:

• Với Post-LN, norm nằm ngay trên stream. Mọi tín hiệu (và mọi gradient) đi từ layer cuối về layer đầu đều phải đi qua L lần norm. Stream không còn là một tổng đơn giản nữa.
• Với Pre-LN, norm chỉ nằm ở nhánh đi vào sub-layer. Stream vẫn là tổng thuần tuý như ở mục 02, và đường gradient thẳng vẫn còn nguyên.

Xiong et al. 2020 phân tích hai cách này lúc khởi tạo. Với Post-LN, gradient của các tham số gần lớp output rất lớn. Nếu dùng ngay learning rate lớn, training dễ phân kỳ. Đó là lý do Post-LN cần learning rate warmup: bắt đầu với learning rate rất nhỏ rồi tăng dần. Với Pre-LN, gradient “well-behaved” ngay từ đầu. Paper cho thấy có thể bỏ hẳn warmup và train nhanh hơn.

GPT-2 là một trong những model lớn đầu tiên chuyển sang Pre-LN. Paper GPT-2 ghi rằng LayerNorm được chuyển lên input của mỗi sub-block, và có thêm một LayerNorm sau block cuối cùng. Chi tiết “thêm một norm ở cuối” không phải ngẫu nhiên. Với Pre-LN, stream không bao giờ được chuẩn hoá trên đường chính, nên trước khi vào lm_head phải có một final norm. Trong code Hugging Face của Llama, đó là module model.norm.

Vị trí đặt norm vẫn chưa ngã ngũ

Từ đây dễ kết luận “mọi LLM hiện đại đều Pre-LN, hết chuyện”. Thực tế phức tạp hơn:

• Post-LN có thể cho kết quả tốt hơn, nếu train được. Liu et al. 2020 quan sát rằng khi Post-LN train ổn định, nó thường cho kết quả cuối tốt hơn Pre-LN. Họ đề xuất một cách khởi tạo (Admin) để giữ ưu điểm đó mà không bị phân kỳ.
• Post-LN có thể sâu tới 1.000 layer. DeepNet (Wang et al. 2022) dùng DeepNorm, một biến thể Post-LN nhân hệ số α > 1 lên đường shortcut x trước khi cộng, kèm một cách khởi tạo riêng, để train Transformer 1.000 layer.
• Gemma 2 chuẩn hoá cả hai phía. Paper ghi: “use RMSNorm to normalize the input and output of each transformer sub-layer”. Công thức thành x + Norm(Sub(Norm(x))). Hàng dưới của Hình 4 vẽ đúng kiểu này. Gemma 3 giữ cách làm đó, dùng RMSNorm cả trước lẫn sau sub-layer (Gemma Team 2025).
• OLMo 2 chuẩn hoá output của attention và FFN thay vì input. Peri-LN (Kim et al. 2025) phân tích một cách hệ thống kiểu đặt norm ở cả hai phía sub-layer và gọi nó là Peri-LN.

Vậy câu chính xác hơn là: đa số LLM mở hiện nay dùng Pre-LN với RMSNorm, nhưng vị trí đặt norm vẫn là một chủ đề nghiên cứu đang thay đổi. Điểm chung của các cách làm mới là đều giữ đường residual thẳng, không đặt norm lên trên nó.

06Đếm tham số và FLOPs của một block

Đến đây ta có đủ mảnh để đếm. Mục tiêu là tự tính được số tham số của một model thật từ config.json, khớp với con số công bố, và hiểu vì sao các công thức quen thuộc cho ra kết quả khác.

Công thức của Kaplan và các giả định của nó

Kaplan et al. 2020 (Scaling Laws for Neural Language Models, Bảng 1) đưa ra một ước lượng nổi tiếng cho số tham số không tính embedding:

N ≈ 2 · d_model · n_layer · (2 · d_attn + d_ff)
  ≈ 12 · n_layer · d_model²        khi d_attn = d_model và d_ff = 4 · d_model

Số 12 đến từ: attention có bốn ma trận d × d (Q, K, V, O) là 4d²; FFN có hai ma trận d × 4d là 8d². Cộng lại là 12d² mỗi layer.

Công thức này dựa trên ba giả định: MHA (K, V có cùng số head với Q), FFN hai ma trận với d_ff = 4d, và bỏ qua bias, norm, embedding. LLM hiện đại phá cả hai giả định đầu.

GPT-2 small: nguồn gốc của con số “2/3”

Với GPT-2 small (d = 768), đúng kiểu mà Kaplan giả định:

Đây là nguồn gốc của câu “FFN chiếm khoảng 2/3 tham số”. Câu đó đúng với kiến trúc kiểu GPT-2, và chỉ đúng với kiến trúc đó.

Llama 3.1 8B: đếm từng ma trận

Cấu hình từ config.json: hidden_size = 4096, intermediate_size = 14336, num_hidden_layers = 32, num_attention_heads = 32, num_key_value_heads = 8, vocab_size = 128256, tie_word_embeddings = false, không có bias. Mỗi head có 4096 / 32 = 128 chiều.

FFN chiếm 80,8% của block, không phải 66,7%. Có hai lý do cùng đẩy theo một hướng:

• GQA làm attention nhẹ đi. K và V chỉ có 8 head thay vì 32, nên attention còn 2,5d² thay vì 4d².
• SwiGLU với 3,5× làm FFN nặng lên. Ba ma trận d × 3,5d là 10,5d² thay vì 8d².

Tổng mỗi block là 13d², so với 12d² của Kaplan, với tỉ lệ bên trong khác hẳn. Cộng lên toàn model:

Con số này khớp với “8B” mà Meta công bố. FFN chiếm khoảng 70% toàn model, phần lớn của 30% còn lại là embedding và lm_head. Công thức 12 · n_layer · d² sẽ cho 6,44B phần không tính embedding, thấp hơn khoảng 8% so với con số thật 6,98B.

Qwen2.5-7B để đối chiếu

Qwen2.5-7B (d = 3584, d_ff = 18944, 28 layer, 28 head Q / 4 head KV, có bias ở Q, K, V) đẩy tỉ lệ còn xa hơn: FFN chiếm 87,4% của mỗi block. Tổng tính ra là 7.615.616.512, phần không tính embedding là 6,53B. Model card của Qwen ghi 7,61B và 6,53B, khớp với phép đếm. Tỉ lệ d_ff / d_model ≈ 5,3× cùng với chỉ 4 head KV là lý do FFN lấn át mạnh hơn cả Llama.

Tự đếm bằng code

Đoạn code dưới đây chỉ dùng Python chuẩn (đã chạy với Python 3.12, không cần thư viện ngoài). Cấu hình được chép từ config.json của từng model:

def count_params(cfg, swiglu=True, qkv_bias=False):
    d, ff, L = cfg["hidden_size"], cfg["intermediate_size"], cfg["num_hidden_layers"]
    head_dim = d // cfg["num_attention_heads"]
    kv_dim = cfg["num_key_value_heads"] * head_dim

    attn = 2 * d * d + 2 * d * kv_dim                 # W_Q, W_O, W_K, W_V
    if qkv_bias:
        attn += d + 2 * kv_dim                        # bias của Q, K, V
    ffn = (3 if swiglu else 2) * d * ff               # gate, up, down
    block = attn + ffn + 2 * d                        # + hai RMSNorm

    embed = cfg["vocab_size"] * d
    head = 0 if cfg["tie_word_embeddings"] else embed
    total = L * block + d + embed + head              # + final norm
    return block, ffn / block, total, total - embed - head

llama31_8b = dict(hidden_size=4096, intermediate_size=14336, num_hidden_layers=32,
                  num_attention_heads=32, num_key_value_heads=8,
                  vocab_size=128256, tie_word_embeddings=False)
qwen25_7b = dict(hidden_size=3584, intermediate_size=18944, num_hidden_layers=28,
                 num_attention_heads=28, num_key_value_heads=4,
                 vocab_size=152064, tie_word_embeddings=False)

for name, cfg, bias in [("Llama 3.1 8B", llama31_8b, False), ("Qwen2.5-7B", qwen25_7b, True)]:
    block, ffn_share, total, non_emb = count_params(cfg, qkv_bias=bias)
    print(f"{name}: block={block:,}  FFN={ffn_share:.1%}  "
          f"total={total / 1e9:.2f}B  non-embedding={non_emb / 1e9:.2f}B")

Output:

Llama 3.1 8B: block=218,112,000  FFN=80.8%  total=8.03B  non-embedding=6.98B
Qwen2.5-7B: block=233,057,792  FFN=87.4%  total=7.62B  non-embedding=6.53B

Với một model khác, chỉ cần thay dict bằng các trường tương ứng trong config.json của nó. Đoạn code chỉ đúng cho kiến trúc dense kiểu Llama: không áp dụng nguyên xi cho MoE, cho model có bias ở mọi lớp, hay cho model chia sẻ trọng số giữa các layer.

FLOPs: 2N khi forward, 6N khi train

Một phép nhân ma trận (1 × a) · (a × b) cần a · b phép nhân và a · b phép cộng, tức khoảng 2 FLOPs cho mỗi tham số. Từ đó Kaplan et al. ước lượng chi phí forward cho mỗi token:

C_forward ≈ 2N + 2 · n_layer · n_ctx · d_attn

Số hạng đầu là các phép nhân với ma trận trọng số. Số hạng sau là phần attention phụ thuộc vào độ dài context: Q · Kᵀ và nhân với V, đúng phần chi phí mà bài 7, mục 10 mô tả. Khi train, backward pass tốn khoảng gấp đôi forward, nên tổng chi phí xấp xỉ 6N FLOPs mỗi token.

Áp dụng cho Llama 3.1 8B, với N là số tham số không tính embedding (6,98B) như cách Kaplan định nghĩa:

Bảng này cho thấy ba điều:

1. Ở context vừa phải, FLOPs gần như hoàn toàn đến từ các ma trận, và phần lớn là FFN. Ở 8K token, phần phụ thuộc context chỉ khoảng 15% so với 2N.
2. Ở context rất dài, thứ tự đảo ngược. Ở 128K token, phần attention theo context (34,4 GFLOPs) lớn hơn cả toàn bộ phép nhân ma trận. Đây chính là chi phí bậc hai mà bài 7 nói tới.
3. GQA không làm giảm số hạng n_ctx này. Query vẫn có 32 head × 128 = 4.096 chiều, nên Q · Kᵀ vẫn tốn ngần ấy phép tính. GQA giảm tham số của W_K, W_V và giảm KV cache, không giảm FLOPs của bản thân phép attention.

Con số 6N cũng là thứ mà các paper scaling, như Chinchilla, dùng để quy đổi giữa số tham số, số token và ngân sách compute. Khi đọc một con số “6N”, hãy kiểm tra N có tính embedding hay không: với Llama 3.1 8B, chênh lệch là 42 so với 48 GFLOPs mỗi token.

07Những hiểu lầm thường gặp

Hai nửa của một block, và một bus nối chúng

Một Transformer block có hai nửa. Attention, đề tài của bài 7, là bước giao tiếp: mỗi token lấy thông tin từ các token khác. Bài này là nửa còn lại, và nhìn kỹ thì đó mới là nửa chứa phần lớn tham số và compute.

Những gì nên mang theo:

• residual stream là một vector d_model chiều đi thẳng qua cả model; mỗi sub-layer đọc bằng phép chiếu và ghi bằng phép cộng, không bao giờ ghi đè;
• FFN có cùng dạng công thức với attention, f(x · K) · V, nhưng key và value là tham số cố định, và có bằng chứng cho thấy key bắt mẫu văn bản còn value đẩy xác suất token;
• “FFN là kho tri thức” là một cách diễn giải có bằng chứng, chưa phải định lý;
• GELU làm mượt ReLU; SwiGLU thêm một nhánh cổng, cần ba ma trận, và vì thế d_ff bị giảm còn khoảng 2/3 · 4d trước khi mỗi nhóm điều chỉnh lại;
• RMSNorm bỏ bước trừ trung bình mà không mất chất lượng; Pre-LN giữ đường residual thẳng nên train ổn định, nhưng vị trí đặt norm vẫn đang được thử nghiệm;
• con số “FFN ≈ 2/3” và “d_ff = 4d” là của GPT-2; với Llama 3.1 8B, FFN chiếm 80,8% mỗi block và tổng 8,03B tham số tự đếm được từ config.json;
• FLOPs xấp xỉ 2N mỗi token khi forward và 6N khi train, cộng thêm phần attention tăng theo context, phần đó vượt 2N khi context đủ dài.

Sau block cuối cùng, stream đi qua final norm rồi vào lm_head để thành logits, một con số cho mỗi token trong vocabulary. Từ logits đến một câu trả lời cụ thể còn một bước nữa: chọn token nào. Bài 9 sẽ nói về bước đó, decoding và sampling: greedy, temperature, top-k, top-p, và vì sao cùng một model có thể trả lời rất khác nhau chỉ vì cách chọn token.

Tài liệu tham khảo

• Elhage et al. 2021 — A Mathematical Framework for Transformer Circuits
• He et al. 2015 — Deep Residual Learning for Image Recognition
• Geva et al. 2021 — Transformer Feed-Forward Layers Are Key-Value Memories
• Geva et al. 2022 — Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space
• Dai et al. 2022 — Knowledge Neurons in Pretrained Transformers
• Meng et al. 2022 — Locating and Editing Factual Associations in GPT (ROME)
• Hase et al. 2023 — Does Localization Inform Editing? Surprising Differences in Causality-Based Localization vs. Knowledge Editing in Language Models
• Hendrycks and Gimpel 2016 — Gaussian Error Linear Units (GELUs)
• Dauphin et al. 2016 — Language Modeling with Gated Convolutional Networks
• Ramachandran et al. 2017 — Searching for Activation Functions
• Elfwing et al. 2017 — Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning
• Shazeer 2020 — GLU Variants Improve Transformer
• Ba et al. 2016 — Layer Normalization
• Zhang and Sennrich 2019 — Root Mean Square Layer Normalization
• Xiong et al. 2020 — On Layer Normalization in the Transformer Architecture
• Liu et al. 2020 — Understanding the Difficulty of Training Transformers
• Wang et al. 2022 — DeepNet: Scaling Transformers to 1,000 Layers
• Kim et al. 2025 — Peri-LN: Revisiting Normalization Layer in the Transformer Architecture
• Radford et al. 2019 — Language Models are Unsupervised Multitask Learners (GPT-2)
• Touvron et al. 2023 — LLaMA: Open and Efficient Foundation Language Models
• Llama Team, Meta 2024 — The Llama 3 Herd of Models
• Meta — llama/model.py trong repo meta-llama/llama3
• Gemma Team 2024 — Gemma 2: Improving Open Language Models at a Practical Size
• Gemma Team 2025 — Gemma 3 Technical Report
• OLMo Team 2024 — 2 OLMo 2 Furious
• Kaplan et al. 2020 — Scaling Laws for Neural Language Models
• Hoffmann et al. 2022 — Training Compute-Optimal Large Language Models (Chinchilla)
• Hugging Face — config.json của Llama 3.1 8B (bản mirror unsloth, vì repo của Meta bị gate)
• Hugging Face — config.json của Qwen2.5-7B
• Hugging Face — config.json của GPT-2
• Hugging Face — config.json của Gemma 2 9B (bản mirror unsloth)