12Tham số, FLOP và bộ nhớ của Transformer#

Khi làm việc với mô hình lớn, nhiều câu hỏi quan trọng có thể trả lời trên giấy trong vài phút, trước khi chạy bất cứ thứ gì: mô hình có vừa bộ nhớ GPU không, huấn luyện mất bao lâu, tăng gấp đôi độ dài ngữ cảnh thì chi phí tăng bao nhiêu. Các công thức để trả lời đều dựng được từ những gì đã học ở Chương 9 và 10. Ta sẽ dựng chúng rồi kiểm chứng với các con số thật của GPT-2 và Llama 2; như sẽ thấy, công thức đếm tham số khớp tới từng tham số.

12.1Mục đích của việc ước lượng#

Câu hỏi thường gặp nhất là mô hình có vừa hay không. Số tham số nhân số byte mỗi tham số cho dung lượng trọng số, và cần cộng thêm KV cache khi suy luận, hoặc gradient, trạng thái của bộ tối ưu và giá trị kích hoạt khi huấn luyện. Câu hỏi thứ hai là tốn bao nhiêu thời gian: số phép tính chia cho tốc độ tính thực tế của phần cứng cho thời gian huấn luyện hoặc độ trễ. Câu hỏi thứ ba là khi thay đổi một cấu hình, như tăng độ dài ngữ cảnh, dùng GQA hay lượng tử hoá trọng số xuống 4 bit, thì chi phí thay đổi thế nào. Cả ba đều bắt đầu từ việc đếm tham số.

12.2Số tham số của một khối Transformer#

Với $d$ là số chiều của mô hình và $d_{\text{ff}}$ là số chiều ẩn của FFN, bỏ qua hệ số chặn và tham số của các lớp chuẩn hoá, bảng dưới đếm tham số của một khối.

Với $L$ khối, số tham số không tính embedding là

$$N \approx 12\,L\,d^2.$$

Đây là công thức Kaplan và cộng sự (2020) dùng trong nghiên cứu về quy luật co giãn. Phần embedding thêm $Vd$ tham số cho bảng embedding từ vựng, cộng $T_{\max} d$ nếu dùng embedding vị trí học được. Nếu lớp chiếu ra từ vựng không dùng chung trọng số với bảng embedding, cộng thêm $Vd$ nữa.

12.3Kiểm chứng trên GPT-2 và Llama 2#

Thí nghiệm trong code/models/experiments.py tính đầy đủ số tham số, kể cả hệ số chặn và tham số của LayerNorm, rồi so với số tham số đếm trực tiếp trên các mô hình GPT-2 mà OpenAI đã phát hành.

Công thức khớp chính xác tới từng tham số với cả ba kích thước. Các con số trong bài báo GPT-2 lại khác: 117 triệu, 345 triệu và 762 triệu. Kho mã của OpenAI về sau ghi chú rằng các con số ban đầu sai do một lỗi, và tên gọi 124M, 355M, 774M dùng hiện nay đến từ cách đếm trực tiếp như ở bảng trên. Phép đếm tay vì vậy còn có ích để kiểm tra chính các con số được công bố.

Chia nhỏ GPT-2 small cho thấy tham số nằm ở đâu.

Bảng chia nhỏ cho thấy hai điều. Embedding chiếm 31% số tham số của GPT-2 small: với mô hình nhỏ và từ vựng lớn, phần embedding chiếm tỉ trọng lớn, nên các nghiên cứu về quy luật co giãn tách riêng số tham số không tính embedding. Còn trong các khối Transformer, FFN chiếm 67% số tham số, đúng tỉ lệ $8d^2/12d^2$ ở Mục 12.2.

Cùng cách đếm áp dụng được cho kiến trúc kiểu Llama, chỉ cần sửa bốn chỗ: FFN dạng SwiGLU có ba ma trận, nên mỗi khối có $4d^2 + 3\,d\,d_{\text{ff}}$ tham số; RMSNorm chỉ có $\gamma$; không có hệ số chặn và không có embedding vị trí, vì vị trí được mã hoá bằng RoPE; và lớp chiếu ra không dùng chung trọng số với embedding. Kết quả cho Llama 2 7B là 6 738 415 616 và cho Llama 2 13B là 13 015 864 320, khớp chính xác với số tham số của trọng số đã phát hành. Bài tập 1 yêu cầu tự tính con số của bản 7B.

12.4Số FLOP#

Phép nhân ma trận $(m \times k)$ với $(k \times n)$ tốn $2mkn$ FLOP, vì mỗi phần tử kết quả cần $k$ phép nhân và $k$ phép cộng. Với mỗi token, mỗi tham số của một lớp tuyến tính tham gia đúng một phép nhân và một phép cộng ở lượt xuôi, nên

$$C_{\text{xuôi}} \approx 2N \quad \text{FLOP mỗi token}.$$

Lượt ngược tốn khoảng gấp đôi lượt xuôi: một lần cho gradient theo đầu vào của mỗi lớp, để truyền tiếp về phía trước, và một lần cho gradient theo trọng số. Vì vậy huấn luyện tốn khoảng $6N$ FLOP mỗi token, và toàn bộ quá trình huấn luyện trên $D$ token tốn

$$C \approx 6\,N\,D \quad \text{FLOP}.$$

Với GPT-2 small, $N = 84\,934\,656$, nên suy luận tốn khoảng 170 triệu FLOP mỗi token và huấn luyện khoảng 510 triệu FLOP mỗi token. Đó là con số theo quy ước của Kaplan và cộng sự, chỉ tính các tham số không thuộc embedding. Bảng embedding đầu vào đúng là không tốn phép tính, vì chỉ là phép tra bảng, nhưng lớp chiếu ra từ vựng là một phép nhân ma trận thật. Với GPT-2 small, lớp này tốn thêm $2Vd \approx 77$ triệu FLOP mỗi token ở lượt xuôi, gần một nửa con số 170 triệu, dù ma trận của nó dùng chung với bảng embedding.

Ví dụ 12.1 (Ngân sách huấn luyện Llama 2 7B)

Llama 2 7B có $N = 6\,476\,005\,376$ tham số không tính embedding và được huấn luyện trên $D = 2 \times 10^{12}$ token, nên $C \approx 6ND = 7{,}77 \times 10^{22}$ FLOP. GPU A100 đạt tối đa 312 TFLOPS với số bf16, nhưng trong thực tế chỉ dùng được một phần; tỉ lệ này gọi là mức sử dụng phần cứng (model FLOPs utilization, MFU), thường 30% tới 50% với mô hình lớn. Với MFU 30%, 37,6% và 40%, ước lượng lần lượt là 230 627, 184 011 và 172 970 giờ-GPU. Meta công bố 184 320 giờ-GPU (Touvron và cộng sự, 2023), ứng với MFU khoảng 37,5%, nằm trong khoảng thông thường. Công thức $6ND$ vì vậy cho đúng bậc độ lớn của chi phí huấn luyện; con số chính xác phụ thuộc hiệu suất của hệ thống.

Con số 37,5% còn phụ thuộc vào cách đếm. Nếu lấy $N$ là tổng số tham số, 6 738 415 616, MFU thành khoảng 39%. Nếu cộng thêm phép tính của attention theo quy ước của PaLM (Chowdhery và cộng sự, 2022), tức $12LdT$ FLOP mỗi token khi huấn luyện với $T = 4\,096$, MFU thành khoảng 45%. Khi so MFU giữa các báo cáo, cần xem họ đếm những phép tính nào.

Cùng phép tính cho phép trả lời nhanh những câu hỏi mà bài giảng CS336 dùng làm bài tập nhẩm. Huấn luyện một mô hình 70 tỉ tham số trên 15 nghìn tỉ token bằng 1 024 GPU H100 mất bao lâu? Tổng là ${6 \cdot 70 \times 10^9 \cdot 15 \times 10^{12} = 6{,}3 \times 10^{24}}$ FLOP. Một H100 đạt khoảng 989 TFLOPS với bf16 khi không dùng tính thưa; với MFU 50%, 1 024 GPU làm được khoảng $4{,}4 \times 10^{22}$ FLOP mỗi ngày, nên cần khoảng 144 ngày. Karpathy làm phép tính tương tự cho GPT-2 small trong nanoGPT: 124 triệu tham số, 300 tỉ token và 8 GPU A100 với MFU 30% cho khoảng 3,5 ngày, còn lần huấn luyện thật hội tụ sau khoảng 4 ngày.

Công thức $6ND$ dẫn tới một câu hỏi thiết kế: với một ngân sách tính toán $C$ cố định, nên chia $C$ giữa kích thước mô hình $N$ và lượng dữ liệu $D$ thế nào? Hoffmann và cộng sự (2022) huấn luyện hơn 400 mô hình và kết luận rằng $N$ và $D$ nên tăng cùng tỉ lệ khi $C$ tăng, với khoảng 20 token dữ liệu cho mỗi tham số; mô hình Chinchilla 70B huấn luyện trên 1,4 nghìn tỉ token theo nguyên tắc này tốt hơn các mô hình lớn hơn nhiều nhưng ít dữ liệu hơn. Kết quả này được gọi là quy luật co giãn tối ưu về tính toán. Trong thực tế, nhiều mô hình được huấn luyện trên lượng dữ liệu vượt xa tỉ lệ 20 token mỗi tham số, vì mô hình nhỏ hơn rẻ hơn khi phục vụ; Llama 2 7B dùng khoảng 300 token mỗi tham số, còn Llama 3 8B dùng hơn 15 nghìn tỉ token, gần 1 900 token mỗi tham số (Llama Team, 2024).

Khi đọc các quy luật co giãn cần để ý $N$ được đếm thế nào. Kaplan và cộng sự dùng số tham số không tính embedding, còn Hoffmann và cộng sự dùng tổng số tham số; với GPT-2 small, hai con số là 85 triệu và 124 triệu. Hai nhóm đi tới hai khuyến nghị khác nhau: Kaplan và cộng sự khuyên tăng $N$ nhanh hơn $D$ nhiều khi ngân sách tăng, còn Chinchilla khuyên tăng như nhau. Pearce và Song (2024) cho thấy phần lớn chênh lệch đó đến từ chính khác biệt trong cách đếm $N$, cộng với việc Kaplan và cộng sự làm thí nghiệm ở quy mô nhỏ, nơi embedding chiếm tỉ trọng lớn.

Công thức $2N$ chỉ tính các phép nhân với ma trận trọng số, chưa tính hai phép nhân phụ thuộc độ dài ngữ cảnh trong attention: điểm số $QK^\top$ và tích với $V$. Với mỗi token và mỗi lớp, mỗi phép tốn khoảng $2Td$ FLOP khi tính trên cả ngữ cảnh $T$ token, tổng $4LTd$ cho cả mô hình. So với $2N = 24Ld^2$, ta có

$$\frac{\text{phần attention}}{\text{phần còn lại}} \approx \frac{4LTd}{24Ld^2} = \frac{T}{6d}.$$

Với GPT-2 small ($d = 768$, $T = 1\,024$), tỉ lệ là 0,22, tức attention chiếm phần nhỏ. Với $d = 4\,096$, tỉ lệ là 0,17 ở $T = 4\,096$, nhưng 1,33 ở $T = 32\,768$ và 5,33 ở $T = 131\,072$, khi attention trở thành phần chính. Như vậy chi phí tăng theo bình phương độ dài chỉ thực sự chi phối khi $T$ vượt khoảng $6d$. Nếu cài đặt bỏ qua phần bị mặt nạ nhân quả che, như FlashAttention, chi phí attention giảm khoảng một nửa và ngưỡng thành $T \approx 12d$, là con số mà Kaplan và cộng sự (2020) dùng.

12.5Bộ nhớ#

Bộ nhớ gồm nhiều khoản, mỗi khoản phụ thuộc vào những đại lượng khác nhau, và nhầm lẫn giữa chúng là nguyên nhân của nhiều sự cố hết bộ nhớ. Bảng dưới tách riêng từng khoản.

Khi huấn luyện với độ chính xác hỗn hợp và Adam, mỗi tham số cần khoảng 16 byte (Rajbhandari và cộng sự, 2020): 2 byte cho trọng số 16 bit, 2 byte cho gradient 16 bit, 4 byte cho bản sao trọng số FP32 mà bộ tối ưu cập nhật, và 8 byte cho hai trạng thái $m$, $v$ của Adam ở FP32. Với mô hình 7 tỉ tham số, con số là 112 GB, chưa tính giá trị kích hoạt, nên không vừa một GPU 80 GB. Các kỹ thuật như ZeRO (Rajbhandari và cộng sự, 2020) chia gradient và trạng thái bộ tối ưu cho nhiều GPU để giải quyết vấn đề này. Huấn luyện hoàn toàn bằng FP32 cũng tốn đúng 16 byte mỗi tham số, gồm 4 cho trọng số, 4 cho gradient và 8 cho Adam, nên độ chính xác hỗn hợp không giảm các khoản này; lợi ích của nó nằm ở tốc độ tính và ở bộ nhớ cho giá trị kích hoạt, vốn được lưu ở 16 bit. Phép tính theo chiều ngược lại cũng hữu ích, và CS336 dùng nó làm bài tập nhẩm thứ hai: 8 GPU 80 GB, với 16 byte mỗi tham số và mọi khoản được chia đều cho các GPU, chứa được khoảng ${8 \cdot 80 \times 10^9 / 16 = 4 \times 10^{10}}$ tham số, tức 40 tỉ, trước khi tính giá trị kích hoạt.

Giá trị kích hoạt có thể còn lớn hơn thế. Korthikanti và cộng sự (2022) ước lượng mỗi khối Transformer huấn luyện với số 16 bit cần khoảng $sbd\,(34 + 5as/d)$ byte cho giá trị kích hoạt, với $s$ là độ dài chuỗi, $b$ là kích thước lô và $a$ là số đầu attention. Với cấu hình của Llama 2 7B ($d = 4\,096$, $a = 32$, 32 khối), một chuỗi 4 096 token cần khoảng 97 GiB, trong đó 80 GiB là số hạng $5as/d$ đến từ các ma trận attention $T \times T$. FlashAttention (Mục 10.7) không lưu các ma trận này, đưa con số xuống khoảng 17 GiB, và gradient checkpointing (Mục 5.3) giảm tiếp bằng cách tính lại.

Khi suy luận, trọng số và KV cache là hai khoản chính. Với Llama 2 7B dùng MHA (32 lớp, 32 đầu, $d_{\text{head}} = 128$) ở FP16, KV cache tốn $2 \cdot 32 \cdot 32 \cdot 128 \cdot 2 = 524\,288$ byte, tức 0,5 MiB mỗi token; một chuỗi 4 096 token cần 2 GiB, bằng khoảng 16% dung lượng trọng số 13,5 GB. Với nhiều chuỗi và ngữ cảnh dài, KV cache có thể vượt dung lượng trọng số. Mục 3.3 của Ứng dụng LLM dùng các con số này để ước lượng bộ nhớ khi tự triển khai mô hình.

Giảm số byte mỗi số là cách tác động trực tiếp lên mọi dòng của bảng trên, và là chủ đề của giáo trình Quantization.

12.6Tóm tắt#

Các công thức của chương được gom trong bảng dưới.

Công thức đếm tham số khớp tới từng tham số với GPT-2 và Llama 2, còn công thức $6ND$ cho đúng bậc độ lớn của 184 320 giờ-GPU mà Meta công bố cho Llama 2 7B khi mức sử dụng phần cứng khoảng 37,5%. Khi dùng $6ND$ cần nói rõ $N$ và các phép tính được đếm gồm những gì: cùng số giờ-GPU đó ứng với MFU 37,5%, 39% hay 45% tuỳ cách đếm, và hai quy luật co giãn của Kaplan và của Chinchilla khác nhau phần lớn vì đếm $N$ khác nhau. Attention chỉ chi phối chi phí tính toán khi độ dài ngữ cảnh vượt khoảng $6d$. Về bộ nhớ, huấn luyện tốn khoảng 16 byte mỗi tham số cộng giá trị kích hoạt, còn suy luận tốn trọng số cộng KV cache, và mỗi khoản phụ thuộc vào những đại lượng khác nhau.

Chương này khép lại phần kiến trúc của giáo trình. Đi từ các họ mô hình qua cây, mạng nhiều lớp, mạng tích chập, mạng hồi quy tới Transformer, ta đã thấy cùng một câu hỏi lặp lại dưới nhiều dạng: giả định nào của mô hình khớp với cấu trúc của dữ liệu, và làm sao để gradient đi được qua một mạng sâu. Hai giáo trình tiếp theo dùng trực tiếp các kết quả ở đây: Biểu diễn & Căn chỉnh xét những gì một mô hình như vậy học được và cách căn chỉnh nó theo ý người dùng, còn Quantization giảm số byte mỗi số trong mọi dòng của bảng bộ nhớ ở Mục 12.5. Chương 13 gồm các bài tập để luyện các phép tính của giáo trình, và Chương 14 gom các câu hỏi phỏng vấn thường gặp.