Mô hình Llama 3 được triển khai bằng NumPy thuần

• Hiểu chính xác cấu trúc thông qua một bản triển khai thực sự chạy được của mô hình Llama 3

Tổng quan

• Mô hình Llama 3 do Meta công bố đang thu hút nhiều sự chú ý.
• Mô hình này sở hữu quy mô và hiệu năng vượt trội với 24K GPUs, 15T dữ liệu huấn luyện, 10M dữ liệu chỉ dẫn và 1.3M giờ GPU.
• Cấu trúc mô hình nhìn chung không thay đổi nhiều. Llama 3 sử dụng GQA, nhưng điều này đã được triển khai trước đó trong Llama 2 70B.
• Bài viết triển khai mô hình chỉ bằng NumPy để có thể hiểu trực quan cấu trúc của mô hình.
• Tác giả sử dụng mô hình stories15M do Andrej Karpathy huấn luyện theo cấu trúc Llama 2 và chuyển đổi nó sang định dạng nén NumPy.

Cấu trúc

• Cấu trúc mô hình Llama 3 giống với 42dot LLM.
• Các tham số mô hình:
  • dim: 288
  • n_layers: 6
  • n_heads: 6
  • vocab_size: 32000
  • max_seq_len: 256
  • max_new_tokens: 50

RoPE #1

• Cos và sin được tính trước cho embedding RoPE.
• Các giá trị này được dùng cho Q và K.
• Kết quả tính toán được nhân bằng np.outer, rồi tính ra cos và sin.

RMSNorm

• RMSNorm chuẩn hóa giá trị kích hoạt bằng Root Mean Square thay vì dùng thống kê Mini Batch hoặc Layer truyền thống.
• Cách này cung cấp khả năng scale activation nhất quán.

QKV

• Khác với GPT tính QKV bằng cách nhân ma trận với một trọng số duy nhất rồi tách ra, Llama có trọng số riêng cho từng Q, K và V.
• Sau đó mỗi giá trị được tái cấu trúc để phục vụ Multi-Head Attention.

RoPE #2

• RoPE có cả đặc tính mã hóa vị trí tuyệt đối lẫn tương đối.
• Nó chỉ được áp dụng cho Q và K; đầu vào được chia ra, nhân với cos và sin, rồi cộng trừ kết quả để tái cấu trúc lại.

KV cache

• Các mô hình sinh kiểu GPT dùng Masked Attention nên có thể áp dụng KV cache.
• Vì kết quả trước đó luôn giống nhau, K và V được cache lại, còn Q chỉ tính cho giá trị cuối cùng.

GQA(Grouped-Query Attention)

• GQA là kỹ thuật được giới thiệu từ Llama 2, giúp tiết kiệm bộ nhớ và cải thiện hiệu năng.
• Trong Llama 3, GQA được áp dụng cho tất cả các mô hình từ 8B trở lên.

Scaled Dot-Product Attention

• Mỗi attention được tính bằng Multi-Head Attention.
• Kết quả thu được qua softmax và matmul.

Feed Forward

• Feed Forward của mô hình Llama dùng 3 lớp tuyến tính và không có bias.
• Nó tạo ra giá trị swish, nhân với x_V, rồi scale xuống lại.

SwiGLU

• SwiGLU cải thiện hiệu năng mô hình nhờ sự kết hợp đặc biệt của nhiều lớp feed-forward.

Linear

• Đầu ra cuối cùng chỉ tính matmul cho logit cuối để tăng tốc độ.

Sinh văn bản

• Các token được sinh lần lượt từng cái một bằng logit đã trích xuất.
• Quá trình được chia thành Prefill Phase và Decode Phase.
• Trong Prefill Phase, toàn bộ đầu vào được đưa vào; trong Decode Phase, chỉ truyền token ID cuối cùng để lấy kết quả.

Ví dụ

• Có thể chạy như sau:

  $ python llama3.py "I have a dream"  
  

GitHub

• Có thể xem toàn bộ mã nguồn tại likejazz/llama3.np.

Tài liệu tham khảo

1. Exploring and Building the Llama 3 Architecture
2. Rotation Matrix
3. Mastering LLM Techniques: Inference Optimization
4. arXiv:2305.13245

Ý kiến của GN⁺

• Cấu trúc và hiệu năng của mô hình Llama 3: Mô hình Llama 3 vẫn giữ cấu trúc của Llama 2 hiện có nhưng cải thiện hiệu năng đáng kể. Đây là kết quả của việc đồng thời cân nhắc khả năng mở rộng và hiệu quả của mô hình.
• Lý do triển khai bằng NumPy: Việc triển khai mô hình bằng NumPy giúp hiểu trực quan hơn về cấu trúc và cách vận hành của mô hình. Điều này rất hữu ích cho người học và nhà nghiên cứu.
• Việc áp dụng GQA: GQA là công nghệ vừa tiết kiệm bộ nhớ vừa nâng cao hiệu năng, và việc áp dụng nó cho tất cả mô hình trong Llama 3 đã tối đa hóa hiệu quả của mô hình.
• Tầm quan trọng của KV cache: KV cache đóng vai trò quan trọng trong các mô hình sinh kiểu GPT, giúp tăng mạnh hiệu quả tính toán của mô hình.
• Trường hợp sử dụng thực tế: Có thể chạy thử mô hình thực tế thông qua mã ví dụ, đây là cơ hội tốt để trực tiếp kiểm chứng hiệu năng của mô hình.

[xguru]

Bài đăng trên Hacker News là bản tiếng Anh, nhưng tôi đã đổi sang liên kết do tác giả gốc Likejazz viết bằng tiếng Hàn.