Llama bắt đầu từ con số 0: cách hiện thực hóa bài báo mà không phải vừa làm vừa khóc

 (blog.briankitano.com)
2 điểm bởi GN⁺ 2023-08-10 | 1 bình luận | Chia sẻ qua WhatsApp
  • Brian Kitano tự xây dựng một phiên bản Llama thu nhỏ với TinyShakespeare, và rút ra rằng để hiện thực hóa bài báo một cách an toàn thì nên bắt đầu từ mô hình nhỏ, thay từng thành phần một và huấn luyện/đánh giá sau mỗi lần thay đổi
  • Trước hết, ông chuẩn bị các hàm hỗ trợ kiểm chứng như chia dữ liệu, tạo batch, đánh giá loss và hàm sinh văn bản; sau đó xác nhận một mô hình đơn giản có thể compile và huấn luyện rồi mới thêm các thành phần của Llama
  • Khi lần lượt gắn RMSNorm, RoPE, SwiGLU, ông kiểm tra tensor shape, tính chất công thức và attention map để xác nhận từng layer hoạt động đúng như kỳ vọng
  • Khi bỏ causal mask trong attention dùng RoPE, validation loss giảm xuống tới 0.16 nhưng chất lượng sinh văn bản lại tệ đi; nguyên nhân là rò rỉ thông tin do mô hình nhìn thấy token tương lai
  • Phiên bản Llama thu nhỏ cuối cùng có 4 block, khoảng 2.37 triệu tham số, hạ validation loss xuống khoảng 1.0; đồng thời cũng cần kiểm tra cả luồng gradient và lịch learning rate

Bắt đầu nhỏ và lặp lại để tích lũy độ chắc chắn

  • Cốt lõi khi hiện thực hóa bài báo là bắt đầu từ mô hình nhỏ, thay từng thành phần một, và mỗi lần thay đều lặp lại quá trình huấn luyện và đánh giá
  • Trước tiên cần chuẩn bị các hàm hỗ trợ để kiểm tra mô hình một cách định lượng
    • chia dữ liệu
    • vòng lặp huấn luyện
    • trực quan hóa loss
    • đánh giá validation loss
  • Thay vì chuyển toàn bộ các thành phần trong bài báo sang cùng lúc, ông còn chuẩn bị một hàm đánh giá định tính để xem trực tiếp kết quả sinh văn bản từ một mô hình đơn giản, nhanh và đã có kinh nghiệm hiện thực hóa
  • Với các layer tensor, ông kiểm tra bằng .shape, assert, plt.imshow; thay vì tối ưu hóa phép nhân ma trận ngay từ đầu, ông tự tính tay kết quả kỳ vọng để đối chiếu rồi mới tối ưu bằng các hàm của torch
  • Cần kiểm thử bằng cách thay đổi batch size, độ dài chuỗi và số chiều embedding; mã chỉ đúng với một kích thước có thể sẽ hỏng ở thời điểm suy luận

Dataset và thiết lập cơ bản

  • Đối tượng hiện thực hóa là một phiên bản thu nhỏ rất mạnh của Llama từ Meta AI, và dữ liệu huấn luyện là TinyShakespeare
  • Llama gốc được huấn luyện với 1.4T token, nhưng ở đây sử dụng TinyShakespeare với quy mô khoảng 1.11 triệu ký tự
  • Llama gốc dùng tokenizer byte-pair encoding của SentencePiece, nhưng bản hiện thực này dùng tokenizer theo ký tự đơn giản
    • vocabulary size là 65
    • do dataset nhỏ nên không tối ưu riêng cách lưu dữ liệu trong bộ nhớ
  • Từ điển MASTER_CONFIG dùng để quản lý các cấu hình mô hình như vocab_size, batch_size, context_window, d_model
    • mục tiêu là giảm hằng số và magic number, đồng thời giúp mã dễ đọc hơn
  • Hàm get_batches chia dữ liệu thành train 80%, val 10%, test 10%, rồi tạo đầu vào x và nhãn y lệch sau một ký tự từ các điểm bắt đầu ngẫu nhiên

Xác nhận compile và huấn luyện bằng mô hình cơ bản

  • Mô hình đầu tiên là SimpleBrokenModel, gồm embedding và một mạng feed-forward đơn giản
    • nn.Embedding
    • Linear
    • ReLU
    • Linear
  • Trong việc hiện thực hóa bài báo, nói mô hình “chạy được” nghĩa là phải thỏa cả hai điều kiện
    • compile: tensor shape khớp giữa các layer
    • huấn luyện: loss thực sự giảm
  • Hàm evaluate_loss lấy mẫu 10 batch từ train và val split để tính loss trung bình
  • Sau 1000 epoch huấn luyện, SimpleBrokenModel có validation loss ở mức 3.94, hầu như không giảm nhiều so với cross-entropy ban đầu là 4.17
  • Nguyên nhân là ông đã đưa giá trị đã qua softmax vào F.cross_entropy
    • F.cross_entropy của PyTorch nhận trực tiếp logits chưa chuẩn hóa
    • SimpleModel sau khi bỏ softmax đã hạ validation loss xuống khoảng 2.51
  • Sau đó ông thêm hàm generate để kiểm tra trực tiếp các ký tự mô hình tạo ra; mô hình cơ bản chưa hoàn hảo nhưng đã ở trạng thái validation loss giảm đúng hướng

Thành phần Llama 1: RMSNorm

  • So với Transformer gốc, Llama dùng ba thay đổi kiến trúc chính
    • RMSNorm pre-normalization
    • Rotary embeddings
    • hàm kích hoạt SwiGLU
  • Transformer gốc dùng BatchNormalization, còn Llama dùng RMSNorm, tức là scale theo variance mà không centering vector
  • Trong khi Transformer gốc áp dụng normalization theo kiểu post-normalization ở đầu ra của attention layer, Llama dùng pre-normalization, tức là áp dụng trước trên đầu vào
  • RMSNorm được hiện thực với giả định input shape là (batch, seq_len, d_model)
  • Kết quả RMSNorm được kiểm thử bằng tính chất norm của layer bằng căn bậc hai của số phần tử trong layer
    • assert
    • row-wise comparison
    • torch.allclose
  • SimpleModel_RMS, tức mô hình cơ bản có thêm RMSNorm, giảm validation loss nhẹ xuống khoảng 2.5015

Thành phần Llama 2: RoPE và causal mask

  • RoPE là phương pháp positional encoding cho Transformer, biểu diễn vị trí token bằng phép quay embedding
  • get_rotary_matrix tạo ma trận quay theo từng vị trí cho context window và embedding dimension
  • Việc hiện thực RoPE được kiểm thử bằng tính chất sau
    • tích vô hướng của hai vector đã quay tại vị trí m, n phải khớp với phép quay theo vị trí tương đối n-m
  • RoPEAttentionHead tạo w_q, w_k, w_v, áp dụng phép quay RoPE cho query và key, rồi dùng F.scaled_dot_product_attention
  • Cần chú ý tới khác biệt tensor shape giữa thời điểm huấn luyện và suy luận
    • khi huấn luyện thường là shape khớp cấu hình như (config['batch_size'], config['context_window'], config['d_model'])
    • khi suy luận có thể xử lý một ví dụ đơn lẻ như (1, 1, config['d_model'])
    • trong forward, cần index theo shape lấy từ input chứ không phải theo giá trị cấu hình của mô hình
  • Mô hình thêm multi-head attention dùng RoPE nhưng không có causal mask đã làm validation loss giảm mạnh xuống 0.1623, nhưng kết quả sinh văn bản lại kém như OOOO..., IIII...
  • Khi kiểm tra attention map, mọi vị trí đều tham chiếu tới mọi vị trí khác, gây ra rò rỉ thông tin vì mô hình nhìn được token tương lai trong tác vụ dự đoán token tiếp theo
  • Khi đổi sang RoPEMaskedAttentionHead với is_causal=True trong F.scaled_dot_product_attention, attention ở phần tam giác trên tương ứng với tương lai gần như về 0
  • Sau khi áp dụng causal mask, validation loss là 2.0815, và khi huấn luyện lâu hơn thì giảm xuống 1.8985

Thành phần Llama 3: SwiGLU và chồng block

  • Llama thay phi tuyến ReLU bằng hàm kích hoạt SwiGLU
  • SwiGLU được hiện thực như một Swish-gated linear unit, dùng hai phép biến đổi tuyến tính và tham số beta có thể học được
  • RopeModel với SwiGLU trong phần feed-forward có 592,706 tham số và validation loss ở mức 1.8963
  • Sau đó ông tạo LlamaBlock để gom cấu trúc sau thành một block
    • RMSNorm pre-normalization
    • masked RoPE multi-head attention
    • residual connection
    • RMSNorm pre-normalization
    • feed-forward SwiGLU
    • residual connection
  • Mô hình Llama cuối cùng đặt n_layers=4 và chồng 4 LlamaBlock bằng nn.Sequential dựa trên OrderedDict
  • Tổng số tham số của mô hình cuối cùng là 2,370,246, và kết quả huấn luyện như sau
    • sau lần huấn luyện đầu tiên với 4 layer, validation loss là 1.5532
    • sau khi huấn luyện thêm tới 10,000 epoch, validation loss là 1.1479
    • sau khi huấn luyện thêm nữa, validation loss là 0.9997
    • loss của một batch trong test split là 1.2358

Kết quả sinh văn bản và các điểm cần debug

  • Mô hình cuối cùng có thể tạo ra tên riêng, xuống dòng và các mảnh từ giống phong cách Shakespeare, nhưng chất lượng câu thực tế vẫn còn hạn chế
  • Có thể trực quan hóa cross-entropy loss dưới góc độ lựa chọn token
    • loss ban đầu 4.17 gần tương đương chọn ngẫu nhiên trong vocabulary size 65
    • loss 1.08 có thể hiểu như chọn ngẫu nhiên trong khoảng 2.9 token
  • Luồng gradient được kiểm tra bằng hàm show_grads
    • hàm này tính tỷ lệ gradient có giá trị tuyệt đối nhỏ tại từng tham số
    • nếu phần lớn gradient của các tham số không gần 0 thì luồng gradient đang ở trạng thái ổn
  • Llama gốc dùng lịch learning Cosine Annealing, nhưng trong bản hiện thực này kết quả thực nghiệm lại kém hơn
  • Trong thí nghiệm với Cosine Annealing, ngay cả ở tolerance rất thấp, attention bias gần như không nhận được tín hiệu; nguyên nhân chưa rõ, nên trong hiện thực thực tế, bắt đầu đơn giản vẫn là cách an toàn hơn

Bài viết liên quan

1 bình luận

 
GN⁺ 2023-08-10
Ý kiến trên Hacker News

  • Có vẻ có lỗi trong phần triển khai SwiGLU: trong bài báo tham khảo, beta của feed-forward network là một hằng số chứ không phải giá trị có thể học, và được đặt là FFnSwiGLU = Swish1...
    Dựa theo công thức 6 trong https://arxiv.org/pdf/2002.05202.pdf
    Trong bản triển khai llama chính thức, hằng số beta cũng đã bị loại bỏ: https://github.com/facebookresearch/llama/blob/main/llama/mo...
    Nhìn các dòng "feedforward.1.beta', 0.0" trong log của blog thì beta đã suy biến về 0 trong quá trình huấn luyện, trong khi lẽ ra phải là hằng số 1

    • Điều này cho thấy việc triển khai chính xác mạng nơ-ron Transformer khó đến mức nào. Có thể mắc lỗi ở nhiều bước, và thường chỉ biểu hiện ở mức “hiệu năng hơi kém hơn bản gốc”, nên khó nhận ra chắc chắn
      Mạng cũng thường thích nghi với các thay đổi dù có chủ ý hay không, và sau khi huấn luyện, nhiều biến thể kiến trúc cũng hoạt động tương tự nhau, nên đôi khi không rõ có nhất thiết phải khớp với bản gốc hay không
      Một cách để tìm các lỗi kiểu này là khớp chính xác đầu ra với bản triển khai tham chiếu. Ngay cả với trọng số ngẫu nhiên như các mô hình tiny-random của HuggingFace, đầu ra vẫn phải giống hệt; nếu khác thì đó là tín hiệu có lỗi
      Tuy nhiên cách này chủ yếu hiệu quả với lỗi phát sinh khi suy luận; các vấn đề chỉ xảy ra trong xử lý dữ liệu, optimizer hoặc khi huấn luyện thì khó bắt hơn
    • Tôi cho rằng trong Transformer, giá trị bias nhìn chung không khớp tốt lắm
      Cá nhân tôi nghĩ là do tính tự hồi quy và đặc tính giống ODE, nhưng cũng chưa đủ chắc chắn

  • Công việc rất tuyệt, nhưng SimpleBrokenModel và SimpleModel ban đầu có khá nhiều phép tính lãng phí. Thứ tự là embedding 65 -> 128, linear 128 -> 128, ReLU, linear 128 -> 65; giữa hai tầng đầu không có phi tuyến tính và cả hai đều tuyến tính, nên tầng tuyến tính thứ hai thực chất vô dụng
    Mô hình này rốt cuộc tương đương với một MLP một tầng ẩn cổ điển, và tính theo FLOPS thì 128*128=16k phép tính bị lãng phí trên tổng 128*128+65*128=24k

    • Có vẻ không chỉ mình tôi vẫn đang học về phi tuyến tính. Tôi thắc mắc cách sửa tốt nhất ở đây là chèn ReLU hoặc SwiGLU giữa embedding và tầng tuyến tính đầu tiên, hay đơn giản là xóa tầng tuyến tính đó
      Tầng embedding là cấu trúc đặc biệt để chuyển chỉ số token thành vector embedding, nên có lẽ không thể bỏ nó

  • Nhìn chung, bài này minh họa tốt các nguyên tắc cơ bản. Tôi đặc biệt thích câu “hãy dùng .shape một cách sùng đạo. assert và plt.imshow là bạn của bạn”, và nên luôn assert các điều kiện trước/sau của shape
    Tôi cũng thắc mắc liệu bear hay typeguard có hỗ trợ các kiểm tra này bằng decorator không
    Tuy vậy, đoạn “hãy chọn một mô hình nhỏ, đơn giản, nhanh và tạo helper để đánh giá định tính” có lẽ ý là đánh giá định lượng. Như vậy mới có baseline bằng số để so sánh với các kỹ thuật cao cấp hơn
    Lời khuyên triển khai từng thành phần của bài báo cũng nên chính xác hơn. Bài báo thường thử nhiều thay đổi cùng lúc rồi dùng ablation study để chỉ ra đóng góp của từng yếu tố, nên tốt hơn là bắt đầu từ thay đổi kiến trúc cốt lõi, rồi theo thứ tự ảnh hưởng lớn trong ablation study, giữ đúng quan hệ phụ thuộc và đánh giá sau mỗi thay đổi nguyên tử

    • Thay vì bear hay typeguard, nhờ https://peps.python.org/pep-0646/ có thể đẩy một phần việc này trực tiếp vào chú thích kiểu của Python
      Ví dụ có thể biểu diễn shape theo từng trục trong kiểu như ndarray[float, Dim1, *Shape], và overload shape trả về tùy theo giá trị axis
    • Tôi không rõ PyTorch thế nào, nhưng lần cuối tôi kiểm tra thì chưa có; còn Jax hỗ trợ kiểm tra runtime cơ bản cho shape ma trận thông qua bear / typeguard
      Dù vậy, Python có vẻ khó mà tốt bằng Julia. Hệ thống kiểu của Julia có thể đảm bảo kích thước ma trận khớp nhau dễ dàng hơn nhiều

  • Tôi tò mò nguyên tắc dùng SwiGLU thay cho ReLU là gì. Không biết các tác giả chỉ thử hết các hàm phi tuyến khả dĩ, hay có lý do sâu xa hơn

    • Cũng như nhiều nghiên cứu khác, nếu không có lời giải thích rõ ràng được hậu thuẫn bởi nghiên cứu nghiêm ngặt, có lẽ họ đã ngẫu nhiên tìm kiếm kiểu leo đồi với các thay đổi một dòng trông có vẻ hay, rồi dừng lại khi đến lúc viết bài và làm ablation study

  • bearblog đang bị DDoS nên để lại repository: https://github.com/bkitano/llama-from-scratch

  • Từ góc độ người đang học AI, tôi thử tóm tắt ngắn gọn các thuật ngữ xuất hiện trong bài. Token là một định danh số nguyên biểu thị một mảnh văn bản; trong LLM, các mảnh ký tự thường dùng được gom lại để sử dụng trong phạm vi kích thước từ vựng giới hạn
    Hàm mất mát là giá trị đo chênh lệch giữa dự đoán và đáp án đúng, càng thấp càng tốt. PyTorch là thư viện để xử lý tensor và mạng nơ-ron, còn tensor là mảng số đa chiều bao gồm vô hướng, vector và ma trận
    Mạng nơ-ron là cấu trúc liên kết các neuron có trọng số và bias, còn lớp tuyến tính là cấu trúc đơn giản trong đó mọi đầu vào và đầu ra đều được kết nối. ReLU là một hàm kích hoạt như Math.max(0, x); nếu chỉ xếp chồng các lớp tuyến tính thì cuối cùng cũng tương đương một hàm tuyến tính duy nhất, nên nó thêm tính phi tuyến để tăng khả năng học
    Gradient là lượng thay đổi số được tính trong quá trình huấn luyện để làm mô hình chính xác hơn, còn chuẩn hóa batch là phương pháp điều chỉnh các con số đang truyền qua để hỗ trợ việc học. Mã hóa vị trí cho biết vị trí tương đối của các token bằng vector
    Toán tử @ trong Python là alias của __matmul__, dùng cho phép nhân ma trận. Epoch là việc huấn luyện một lượt trên toàn bộ dataset, còn batch là số lượng dữ liệu đưa vào cùng lúc trước khi cập nhật tham số
    Attention là cốt lõi giúp LLM hoạt động: nó xử lý song song các token đầu vào để tạo tensor trung gian rồi dùng chúng để sinh các token đầu ra

    • Với người ngoài lĩnh vực, có thể không biết “Karpathy” nghĩa là gì. Nếu giới thiệu Andrej Karpathy kèm ngữ cảnh như “nhà truyền thông khoa học kiêm nhà nghiên cứu”, ý rằng nên tham khảo bài viết hoặc video của ông sẽ rõ ràng hơn
    • Với người mới bắt đầu, sẽ chính xác hơn nếu xem token không chỉ đơn giản là định danh số nguyên của một mảnh văn bản, mà gần với mảnh từ đủ phổ biến để tự nó hữu ích
      Ví dụ writ, phần chung trong writing, written, writer, có thể trở thành một token, và writer có thể được token hóa thành writ và er
      Embedding là bước chuyển những token như vậy thành biểu diễn số riêng biệt
    • Hợp thành các hàm tuyến tính thì kết quả vẫn là một hàm tuyến tính. Vì vậy nếu tất cả đều tuyến tính, dù xếp nhiều lớp thì các lớp ngoài một lớp sẽ thành lãng phí; để tránh điều đó cần có tính phi tuyến
    • Ngoài loạt video của Karpathy và repo đi kèm, tôi cũng tò mò liệu trong hành trình học còn có tài liệu hay sách nào đặc biệt hữu ích nữa không
    • Tôi muốn biết chính xác chuẩn hóa batch làm gì và nó giúp ích như thế nào

  • Nếu có bản triển khai và checkpoint hiện có của mô hình, cách hiệu quả nhất để kiểm tra triển khai của mình có đúng không là tải checkpoint đó và so sánh giá trị đầu ra
    Nếu đầu ra không khớp, thường là do triển khai sai ở chi tiết nào đó; bạn có thể lần theo từng lớp một cách có hệ thống để tìm khác biệt thực tế. Trong quá trình đó cũng có thể phát hiện điểm kỳ lạ trong bản triển khai hiện có
    Đây là chuyện về chính mô hình, còn huấn luyện là một trục riêng. Dù vậy, nếu đã chỉnh hyperparameter tương đối giống, khi triển khai mô hình chính xác thì nhìn chung sẽ ổn

  • Cả cách đọc paper lẫn nội dung paper này đều hay, và tôi cũng khuyên xem series Makemore của Karpathy

  • Các lời khuyên tóm tắt đều rất hay, và tôi nghĩ lời khuyên assert shape của tensor áp dụng được cho bất kỳ thư viện đại số tuyến tính phổ biến nào. Khi viết mã đại số tuyến tính phức tạp, đi từng bước nhỏ và lập trình phòng thủ là cực kỳ quan trọng
    Việc lập trình đại số tuyến tính trong các ngôn ngữ chủ đạo thật kinh khủng vì không có kiểm tra shape tại compile time. Shape của tensor nên là một phần của kiểu, và nếu cố nhân 3x4 với 3x4 mà không chuyển vị thì đáng lẽ không nên biên dịch được ngay từ đầu
    Chạy một phép tính dài rồi mới thất bại ở một phép toán không khớp chiều thật sự là tệ nhất
    Tôi cũng nghĩ thiết bị của tensor PyTorch nên được kiểu hóa tĩnh. Hiện tại, nếu cố nhân tensor trong bộ nhớ CPU với tensor trong bộ nhớ GPU thì sẽ gặp lỗi runtime