Tokens of Curiosity

Hoang's exploring & learning place 📚

Tối Ưu Bộ Nhớ Khi Huấn Luyện LLMs – Giới Hạn Thực tế, PEFT, LoRA

CategorIes:

By

Hoang Tran

·

4–6 phút
  1. Giới thiệu
  2. Giới hạn thực tế khi huấn luyện LLMs
    1. Tại sao lại tốn nhiều memory?
    2. Giải pháp tạm thời
  3. Parameter-Efficient Fine-Tuning (PEFT) và LoRA
    1. Cách hoạt động của LoRA – “A bird’s-eye view”
    2. Lợi ích của LoRA
  4. Các Hyper-parameters quan trọng khi sử dụng LoRA
  5. Kết luận
  6. 📚 Tham khảo

Giới thiệu

Huấn luyện các mô hình ngôn ngữ lớn (LLMs) là một bài toán tốn kém về tài nguyên, một trong số đó là bộ nhớ GPU (VRAM). Ngay cả với những GPU mạnh nhất, việc fine-tune các mô hình này vẫn gặp nhiều rào cản thực tế. Trong bài viết này và một vài bài tiếp theo, chúng ta sẽ cùng khám phá các kỹ thuật tối ưu bộ nhớ để giúp quá trình huấn luyện hiệu quả hơn mà không làm giảm chất lượng model.

Giới hạn thực tế khi huấn luyện LLMs

Ở bài trước, chúng ta đã huấn luyện thành công một mô hình LLM đầu tiên trên một tập dữ liệu nhỏ với axolotl. Tuy nhiên, đây chỉ là bước khởi đầu. Khi mở rộng mô hình và tập dữ liệu, chúng ta sẽ nhanh chóng đạt đến giới hạn về memory.

Bây giờ, ta hãy thử chạy training Llama 3.1 Instruct phiên bản 8 tỷ parameters – một con số khá lớn, dù thực tế nó vẫn là một trong những phiên bản nhỏ nhất của Llama. Khi chạy training trên một GPU tiêu dùng như RTX 4090 với 24G VRAM, ngay lập tức chúng ta sẽ gặp lỗi out-of-memory (OOM).

Tại sao lại tốn nhiều memory?

Khi fine-tune một LLM, memory sử dụng sẽ phụ thuộc vào các yếu tố sau:

  1. Trọng số mô hình: Nếu mô hình có N tỷ tham số và sử dụng 16-bit precision, chỉ riêng trọng số đã tốn 2 × N GB VRAM.
  2. Gradient: Trong quá trình tối ưu, gradient của từng tham số cũng cần lưu trữ, chiếm thêm 2 × N GB.
  3. Trạng thái của optimizer: Ví dụ với Adam, mỗi tham số cần thêm 2 giá trị – momentum và variance để update tham số, làm tăng bộ nhớ lên thành 4 × N GB.
  4. Dữ liệu huấn luyện: Một batch dữ liệu được nạp vào bộ nhớ, kích thước của nó phụ thuộc vào batch size, độ dài chuỗi (sequence length) và kích thước embedding.

Tổng cộng, bộ nhớ cần thiết cho toàn bộ quá trình huấn luyện có thể gấp 4 – 6 lần so với giá trị trọng số mô hình.

Giải pháp tạm thời

Giảm batch size về mức 1 hoặc rút ngắn sequence length có thể giúp giảm yêu cầu bộ nhớ, nhưng điều này cũng làm giảm chất lượng quá trình training. Một giải pháp khác là scale vertically, hay sử dụng GPU mạnh hơn hoặc dùng nhiều GPU, nhưng đây là những phương án đắt đỏ và không dễ tiếp cận.

Vậy, có cách nào tối ưu hơn để fine-tune LLM lớn hơn với tài nguyên sẵn có?

Parameter-Efficient Fine-Tuning (PEFT) và LoRA

Thay vì fine-tune toàn bộ parameters của mô hình, chúng ta có thể điều chỉnh chỉ một phần nhỏ các tham số quan trọng nhất. Đây chính là mục tiêu của PEFT (Parameter-Efficient Fine-Tuning) – một nhóm các kỹ thuật giúp giảm đáng kể lượng tài nguyên cần thiết để fine-tune LLMs.

Một trong những phương pháp PEFT phổ biến nhất hiện nay là LoRA (Low-Rank Adaptation).

Cách hoạt động của LoRA – “A bird’s-eye view”

Thay vì cập nhật toàn bộ tham số mô hình trong quá trình fine-tuning, LoRA chỉ thêm một tập nhỏ các trainable matrices vào từng lớp của mô hình. Những ma trận này học được các điều chỉnh cần thiết để mô hình thích nghi với nhiệm vụ mới, mà không làm thay đổi toàn bộ parameters của model gốc.

Sau khi các ma trận nhỏ này được cập nhật, chúng chiếu (project) kiến thức mới lên toàn bộ mô hình, giúp thay đổi hành vi của LLM mà không cần phải huấn luyện lại toàn bộ từ đầu.

Lợi ích của LoRA

  • Tiết kiệm bộ nhớ: Giảm số lượng tham số cần cập nhật, tiết kiệm VRAM đáng kể
  • Giữ nguyên kiến thức cũ: Do chỉ tinh chỉnh một phần nhỏ, mô hình ít bị “quên” những gì đã học trong giai đoạn pre-training.
  • Hiệu quả với dữ liệu nhỏ: Có thể fine-tune tốt ngay cả khi tập dữ liệu không quá nhiều
  • Dễ triển khai: Chỉ cần thêm một số ma trận nhỏ vào các lớp của mô hình hiện có.

Các Hyper-parameters quan trọng khi sử dụng LoRA

  1. LoRA Rank
    • Xác định kích thước của các ma trận được thêm vào mô hình, rank càng lớn thì số lượng trainable params càng lớn.
    • Giá trị nhỏ (ví dụ: 8) phù hợp với các tác vụ không quá phức tạp, khá giống với những gì model đã biết trước đó (qua training data,…)
    • Giá trị lớn hơn (16, 32, 64) phù hợp với những nhiệm vụ phức tạp, nhưng lại làm tăng mức sử dụng bộ nhớ.
  2. Alpha
  3. Dropout
    • Giống như dropout thông thường, giúp tránh overfitting khi huấn luyện.
  4. Target Modules
    • Xác định các phần của mô hình sẽ được tinh chỉnh bởi LoRA.
Platypus: Quick, Cheap, and Powerful Refinement of LLMs (Lee et al. 2023)

Kết luận

LoRA giúp chúng ta tận dụng sức mạnh của LLMs mà không cần đến hạ tầng phần cứng quá đắt đỏ. Với phương pháp này, ta có thể fine-tune mô hình nhanh hơn, tốn ít bộ nhớ hơn, nhưng vẫn đạt hiệu suất tốt trên nhiều tác vụ khác nhau.

Ở các bài viết tiếp theo, chúng ta sẽ đi sâu vào cách triển khai thực tế LoRA khi fine-tune LLM. Cảm ơn mọi người đã đọc bài và hẹn gặp bạn ở bài viết tiếp theo!

📚 Tham khảo

  1. LoRA: Low-Rank Adaptation of Large Language Models (Hu et al. 2021)
  2. Natural Language Processing with Transformers (Lewis Tunstall, Leandro von Werra, Thomas Wolf)

Bình luận về bài viết này