Liger GRPO gặp TRL

Bài viết của khách

May 28, 2025
9 min read

🐯 Liger GRPO gặp TRL

TL;DR (Quá dài; Không đọc)

Liger tăng cường sức mạnh cho Group Relative Policy Optimization GRPO Trainer của TRL bằng cách giảm mức sử dụng bộ nhớ đi 40% mà không làm giảm chất lượng mô hình. Chúng tôi cũng đã thêm hỗ trợ cho FSDP và PEFT, giúp việc mở rộng GRPO trên nhiều GPU trở nên dễ dàng hơn bao giờ hết.

Động lực

Tinh chỉnh các mô hình ngôn ngữ bằng cách sử dụng học tăng cường (RL) là một bước quan trọng trong vòng đời đào tạo của mô hình để hướng các mô hình đến các hành vi mong muốn, phức tạp hơn so với những gì có thể đạt được thông qua tinh chỉnh có giám sát thông thường. RL theo truyền thống được áp dụng để tối ưu hóa các mô hình ngôn ngữ lớn (LLM) bằng cách sử dụng thuật toán Proximal Policy Optimization (PPO). Cách tiếp cận này, thường liên quan đến Reinforcement Learning from Human Feedback (RLHF), sử dụng một mô hình phần thưởng được đào tạo riêng để hướng dẫn việc tinh chỉnh mô hình chính.

Tuy nhiên, RLHF với PPO là một phương pháp rất tốn tài nguyên - PPO yêu cầu tải nhiều mô hình vào bộ nhớ (mô hình chính sách, giá trị, phần thưởng và tham chiếu), đồng thời yêu cầu một số lần lặp lại việc tinh chỉnh các mô hình phần thưởng và cơ sở để đạt được kết quả mong muốn. Sự thành công của RLHF cũng phụ thuộc vào khả năng của mô hình phần thưởng để phân biệt hiệu quả giữa hành vi mong muốn và không mong muốn từ mô hình của chúng ta.

Group Relative Policy Optimization (GRPO) đã chứng kiến sự phổ biến đáng kể gần đây cùng với mô hình R1 của DeepSeek. GRPO tránh các mô hình phần thưởng và giá trị được đào tạo trước được sử dụng trong RLHF và thay vào đó dựa vào các hàm phần thưởng có thể kiểm chứng có thể kiểm tra tính đúng đắn của đầu ra của mô hình một cách khép kín mà không cần mô hình phần thưởng bên ngoài. Điều này đã dẫn đến những cải tiến lớn khi sử dụng GRPO thay vì PPO để tinh chỉnh trên các miền có thể dễ dàng xác minh, chẳng hạn như dạy mô hình suy luận và thực hiện tốt các tác vụ toán học và mã hóa.

Sơ đồ sau đây cho thấy quy trình đào tạo GRPO so với PPO (tham khảo: Hình 4 của DeepSeekMath: Vượt qua giới hạn của suy luận toán học trong các mô hình ngôn ngữ mở):

Điều đó nói lên rằng, đào tạo RL vẫn ngốn rất nhiều bộ nhớ GPU, vì vậy vẫn còn rất nhiều chỗ cho các tối ưu hóa ở đây. Trong bài đăng trên blog này, chúng ta sẽ nói về một tối ưu hóa mà chúng tôi gần đây đã thêm vào TRL giúp giảm mức sử dụng bộ nhớ đỉnh lên 40% trong quá trình Đào tạo GRPO và chúng ta cũng sẽ đi sâu vào cách mở rộng GRPO sang nhiều GPU và nút mà không làm mất hiệu suất hoặc tính đúng đắn.

Cách Liger Kernel cắt giảm bộ nhớ cho GRPO

Chúng tôi đã mở rộng phương pháp Liger Chunked Loss cho GRPO Loss, cho phép chúng ta tránh phải lưu trữ đầy đủ các logits trong bộ nhớ cho mỗi bước đào tạo. Việc tính toán logits, liên quan đến đầu ra của mô hình, là một yếu tố đóng góp đáng kể vào mức sử dụng bộ nhớ đỉnh, đặc biệt khi xử lý từ vựng lớn, độ dài chuỗi dài hoặc kích thước lô lớn. Chúng tôi giải quyết vấn đề này bằng cách chia nhỏ đầu vào cho lm_head trên toàn lô và chạy chuyển tiếp từng khối một.

Nhưng nếu bạn chỉ triển khai nó một cách đơn giản, bạn sẽ thực sự không thể giảm bộ nhớ đỉnh vì bạn vẫn cần giữ tất cả các logits trong bộ nhớ GPU cho chuyển ngược. Để giải quyết vấn đề đó, chúng tôi tính toán gradient cho mỗi khối tổn thất (liên quan đến khối input và trọng số lm_head) trong quá trình chuyển tiếp và sau đó tích lũy chúng khi chúng ta đi qua từng khối.

Đây là hình ảnh trực quan của tối ưu hóa (tham khảo: Byron Hsu):

Tích hợp Plug-and-Play với TRL

Chúng tôi gần đây đã tích hợp Liger GRPO với TRL trong PR #3184, vì vậy giờ đây bạn có thể sử dụng tổn thất Liger GRPO chỉ bằng cách đặt use_liger_loss thành True trong GRPOConfig của bạn và tận hưởng những tiết kiệm bộ nhớ!

Xin lưu ý: các tính năng này chưa có trong bản phát hành TRL mới nhất, vì vậy bạn sẽ cần cài đặt TRL từ nguồn ngay bây giờ:

pip install "trl[liger] @ git+https://github.com/huggingface/trl.git"

và sau đó bạn có thể sử dụng nó như thế này:

from trl import GRPOConfig, GRPOTrainer
from datasets import load_dataset


train_dataset = load_dataset("trl-lib/tldr", split="train")
training_args = GRPOConfig(output_dir="Qwen3-0.6B-GRPO", use_liger_loss=True)

def reward_len(completions, **kwargs):
    return [-abs(20 - len(completion)) for completion in completions]

trainer = GRPOTrainer(
    model="Qwen/Qwen3-0.6B-Instruct",
    reward_funcs=reward_len,
    args=training_args,
    train_dataset=train_dataset,
)
trainer.train()

Điểm chuẩn

Chúng tôi đã chạy một loạt các thí nghiệm GRPO có và không có Liger GRPO Loss để xem mọi thứ so sánh như thế nào. Đối với mô hình chính sách, chúng tôi đã sử dụng Qwen3-0.6B và chơi với các kích thước lô khác nhau. Tất cả các thí nghiệm đều được chạy trên tập dữ liệu gsm8k bằng các hàm phần thưởng của nó.

Đây là các biểu đồ về mức sử dụng bộ nhớ đỉnh so với kích thước lô cho cả đào tạo FP32 và BF16. Như dự kiến, mức tiết kiệm bộ nhớ sẽ tốt hơn với các kích thước lô lớn hơn vì chúng ta chia theo chiều lô. Vì vậy, khi kích thước lô tăng lên, tổn thất chia nhỏ Liger sẽ sử dụng ít bộ nhớ hơn nhiều, ít hơn tới 40%, so với phiên bản thông thường (không phải liger).

Lưu ý nhanh: Hiện tại, chúng tôi chỉ hỗ trợ FP32, nhưng chúng tôi đang nỗ lực mã nguồn mở hỗ trợ BF16 cho Liger GRPO trong TRL. Các kết quả BF16 được hiển thị ở đây là từ các bản vá lỗi nội bộ mà chúng tôi đã thử nghiệm.

Chúng tôi cũng chỉ ra rằng Liger Loss thực sự chính xác. Như được thấy trong biểu đồ, phần thưởng qua các bước đào tạo gần như giống với những gì bạn thấy khi sử dụng triển khai TRL tiêu chuẩn.

Mở rộng hơn nữa với FSDP và PEFT

Chúng tôi cũng đã thêm hỗ trợ FSDP và PEFT vào Liger GRPO Loss trong PR #3260 và PR #3355, tương ứng, cho phép người dùng dễ dàng mở rộng các thí nghiệm của họ trên nhiều GPU hoặc nút. Các kỹ thuật PEFT như LoRA và QLoRA làm giảm số lượng tham số có thể đào tạo bằng cách chỉ điều chỉnh trọng số của các trọng số bộ điều hợp nhỏ hơn trên mô hình ban đầu, giảm đáng kể áp lực bộ nhớ vì gradient, kích hoạt và trạng thái tối ưu hóa cho toàn bộ mô hình không cần phải được giữ trong bộ nhớ. Ngoài ra, việc sử dụng PEFT trong GRPO cho phép người ta bỏ qua việc tải một mô hình tham chiếu riêng biệt trong quá trình đào tạo, vì chúng ta có thể có được mô hình ban đầu, chưa được sửa đổi trong quá trình đào tạo bằng cách chỉ cần tắt các bộ điều hợp LoRA.

Ở đây, chúng tôi hiển thị một biểu đồ đào tạo GRPO nhiều GPU bằng FSDP và PEFT, trong đó chúng tôi so sánh kích thước lô đào tạo tối đa có thể có và không có Liger Loss trên các kích thước mô hình Qwen3 khác nhau. Chúng tôi nhận thấy rằng với Liger, chúng tôi có thể tăng kích thước lô lên khoảng 1,5 đến 1,8 lần!

Mở rộng hơn nữa với vLLM

Để tăng tốc độ tạo văn bản trong quá trình đào tạo, Liger Loss có thể được kết hợp hiệu quả với máy chủ vLLM tích hợp của TRL. Điều này tăng tốc đáng kể việc thu thập dữ liệu triển khai với chi phí tối thiểu và mang lại trải nghiệm tích hợp liền mạch.

Đây là cách thiết lập nó:

Khởi động Máy chủ vLLM: Đầu tiên, khởi chạy máy chủ vLLM. Máy chủ này sẽ xử lý các yêu cầu tạo từ tập lệnh đào tạo của bạn. Mở một thiết bị đầu cuối và chạy:

CUDA_VISIBLE_DEVICES=1 trl vllm-serve --model "Qwen/Qwen3-0.6B"

Lưu ý: Chúng tôi gán CUDA_VISIBLE_DEVICES=1 để chạy máy chủ vLLM trên một GPU cụ thể (GPU 1 trong trường hợp này), để các GPU khác rảnh cho đào tạo. 2. Định cấu hình và Chạy Tập lệnh Đào tạo của bạn: Tiếp theo, sửa đổi tập lệnh đào tạo của bạn để sử dụng máy chủ vLLM. Thay đổi quan trọng là đặt use_vllm=True trong GRPOConfig của bạn.

from trl import GRPOConfig, GRPOTrainer
from datasets import load_dataset


def reward_len(completions, **kwargs):
    return [-abs(20 - len(completion)) for completion in completions]

dataset = load_dataset("trl-lib/tldr", split="train[:1%]")
training_args = GRPOConfig(
    output_dir="Qwen3-0.6B-GRPO", 
    use_liger_loss=True, 
    use_vllm=True, # Bật tích hợp vLLM
    logging_steps=10
)
trainer = GRPOTrainer(
    model="Qwen/Qwen3-0.6B", # Đảm bảo điều này khớp với mô hình được phục vụ bởi vLLM
    reward_funcs=reward_len,
    args=training_args,
    train_dataset=dataset,
)
trainer.train()

Khởi chạy Đào tạo: Cuối cùng, chạy tập lệnh đào tạo của bạn bằng accelerate launch (hoặc python nếu không sử dụng Accelerate cho đào tạo đa GPU/phân tán). Đảm bảo nhắm mục tiêu một GPU khác cho đào tạo nếu máy chủ vLLM của bạn đang chiếm một GPU.

CUDA_VISIBLE_DEVICES=0 accelerate launch train.py

(Giả sử tập lệnh của bạn có tên train.py và bạn muốn chạy đào tạo trên GPU 0).

Bằng cách làm theo các bước này, bạn có thể tận dụng vLLM để có thời gian phản hồi thế hệ nhanh hơn trong quá trình đào tạo GRPO của bạn với Liger Loss.

Kết luận

Với việc tích hợp Liger-GRPO vào TRL, cùng với hỗ trợ FSDP và PEFT, việc tinh chỉnh các mô hình ngôn ngữ bằng GRPO giờ đây hiệu quả về bộ nhớ và có khả năng mở rộng hơn bao giờ hết. Chúng tôi khuyến khích cộng đồng dùng thử các tính năng mới này và chia sẻ phản hồi của họ để giúp chúng tôi cải thiện hơn nữa quá trình đào tạo RL cho LLM.

AI Today - SkyAI