Hugging Face tối ưu huấn luyện GRPO: Dùng LoRA và Storage Bucket tăng tốc 3.9 lần, loại bỏ NCCL

Hỗ trợ LoRA gần đây đã được đưa vào `AsyncGRPOTrainer` của TRL thông qua PR #7017 và được phát hành cùng với TRL v1.14. Trình huấn luyện bất đồng bộ (asynchronous trainer) hiện có thể huấn luyện một adapter thay vì toàn bộ mô hình, và nó chỉ đồng bộ hóa LoRA adapter với vLLM. Bài viết này đề cập đến một dự án thực tế được xây dựng dựa trên cơ chế này, nơi quá trình huấn luyện và suy luận (inference) không còn dùng chung một máy.

Huấn luyện LoRA đặc biệt phù hợp với RL, như đã được chứng minh trong bài blog "LoRA Without Regret" của Thinking Machines. Họ cho thấy LoRA có thể đạt hiệu suất tương đương với tinh chỉnh toàn bộ (full fine-tuning) cho RL theo chính sách gradient (policy-gradient RL), ngay cả với rank 1. Điều này xuất phát từ thực tế là hàm lợi thế (advantage function) chỉ cung cấp khoảng ~O(1) bit thông tin mỗi tập (episode), vì vậy không có quá nhiều thứ để học từ mỗi bước nếu xét theo tổng số bit thông tin. Một adapter rank-1 có đủ khả năng để hấp thụ lượng thông tin đó.

Việc huấn luyện LoRA cũng mang lại hệ quả về mặt hệ thống. Một adapter rank-1 cho mô hình 1.5B chỉ nặng vài megabyte, trong khi toàn bộ mô hình nặng khoảng 3 GB. Thay vì gửi toàn bộ chính sách (policy) đến các worker suy luận sau mỗi lần cập nhật, chúng ta chỉ cần gửi adapter. vLLM cũng có thể giữ nhiều adapter được tải cùng lúc. Các lượt rollout cũ kết thúc với chính sách mà chúng đã bắt đầu, trong khi các lượt rollout mới sử dụng chính sách mới nhất.

`AsyncGRPOTrainer` của TRL đã tách biệt quá trình huấn luyện và tạo (generation). Trình huấn luyện và vLLM có thể chạy trên các máy khác nhau và với tốc độ riêng của chúng. Điều này rất dễ thực hiện trong môi trường đơn nút hoặc cụm (cluster) nơi cả hai tiến trình chia sẻ hệ thống tệp hoặc có thể tạo thành một nhóm NCCL.

Điều chúng tôi muốn là chạy cùng một thiết lập với Hugging Face Jobs. Về cơ bản, một HF Job là một container chạy trên một máy ảo (VM). Điều này có nghĩa là một Job không thể tạo ra nhiều nút (ít nhất là vào thời điểm hiện tại) để chứa một trình huấn luyện và một đội ngũ máy chủ vLLM (chúng tôi bị giới hạn tối đa 8xH200 mỗi nút). `AsyncGRPOTrainer` được xây dựng chính xác cho quy mô đó, vì vậy câu hỏi đặt ra là: chúng ta có thể tiến xa đến đâu nếu loại bỏ yêu cầu trình huấn luyện và các máy chủ suy luận phải dùng chung một nút?

Với việc đồng bộ hóa toàn bộ trọng số, câu trả lời sẽ là "không xa". Mỗi lần cập nhật sẽ phải di chuyển hàng gigabyte dữ liệu giữa các máy, đó là lý do NCCL tồn tại trong một cụm dày đặc, nhưng các Job không thể giao tiếp giữa các nút. Không có đĩa cục bộ chia sẻ và rõ ràng là không có localhost chia sẻ. Với LoRA, một lần đồng bộ chỉ tốn vài megabyte. Đối với phần hệ thống tệp, HF Jobs cung cấp các volume được hỗ trợ bởi Storage Buckets! Các bucket này sau đó có thể được mount như một hệ thống tệp FUSE trong mỗi Job và đủ để hoạt động như một hệ thống tệp chia sẻ giữa các nút. Không cần đường dẫn mạng giữa các Job.

Thiết lập cuối cùng khá nhỏ gọn:

Kiến trúc: tận dụng Hugging Face Jobs và Storage Buckets 🪣

Đường dẫn đồng bộ chỉ-adapter mới trong `AsyncGRPOTrainer` hoạt động như sau. Trình huấn luyện không gửi tensor cho vLLM. Sau mỗi vài bước tối ưu hóa, nó lưu adapter vào `<output_dir>/.vllm_lora/trl-policy-v{N}`, xuất bản thư mục bằng lệnh đổi tên nguyên tử (atomic rename), sau đó gửi đường dẫn của nó đến endpoint `/v1/load_lora_adapter` của vLLM. vLLM tải các tệp từ đĩa, vì vậy rollout worker sau đó có thể yêu cầu `model="trl-policy-v{N}"`.

Đây là cách tải adapter trong thời gian chạy (runtime) hoạt động trong vLLM. Endpoint nhận một đường dẫn, không phải tensor, vì vậy trình huấn luyện và máy chủ được kỳ vọng sẽ chia sẻ một hệ thống tệp. Trên cụm Slurm, đó là hệ thống tệp mạng (network filesystem). Trên Jobs, chúng ta có kết quả tương tự bằng cách mount một Storage Bucket làm volume tại cùng một đường dẫn trong mỗi Job, như đã đề cập trước đó. Bên dưới, nó sử dụng `hf-mount`, hiển thị bucket dưới dạng hệ thống tệp POSIX bên trong container:

Không có gì trong TRL hoặc vLLM cần thay đổi cho việc này. Trình huấn luyện ghi vào `/lora/<run>/.vllm_lora/` và các máy chủ đọc từ cùng một đường dẫn. Đường dẫn được gửi trong yêu cầu POST đã hợp lệ bên trong mỗi container.

Lưu ý rằng chúng tôi cũng lưu trữ các checkpoint và adapter cuối cùng trong bucket. Các HF Jobs là tạm thời, nhưng một trình huấn luyện bị ngắt quãng có thể tiếp tục huấn luyện, vì adapter cuối cùng luôn được lưu vào bucket và không bao giờ bị mất khi Job dừng lại.

Các bản sao (replicas) vLLM

Mỗi bản sao sử dụng một GPU và image `vllm/vllm-openai` tiêu chuẩn. Chúng ta chỉ cần kích hoạt tải LoRA trong thời gian chạy và dự trữ đủ các slot adapter.

Số lượng slot adapter tuân theo `max_staleness`. Trong `AsyncGRPOTrainer`, mỗi lần đồng bộ trọng số sẽ tăng phiên bản chính sách lên một, và `max_staleness` là số lượng phiên bản mà một mẫu rollout có thể bị tụt hậu so với chính sách hiện tại trước khi trình huấn luyện loại bỏ nó. Với `max_staleness=4`, một mẫu được tạo dưới `trl-policy-v3` vẫn được sử dụng để huấn luyện trong khi trình huấn luyện đang ở v7. Một lượt rollout bắt đầu dưới v3 cũng phải có khả năng kết thúc dưới v3. Vì vậy, tại bất kỳ thời điểm nào, vLLM phải phục vụ chính sách hiện tại cộng với bốn phiên bản trước đó. Đó là lý do tại sao trình huấn luyện giữ `max_staleness + 1` phiên bản adapter đã đăng ký và giải phóng bất kỳ phiên bản cũ nào. Mỗi lần đồng bộ sẽ tải phiên bản mới trước khi giải phóng phiên bản cũ nhất, cần thêm một slot trong quá trình hoán đổi. Điều đó mang lại `--max-loras 6`. Với chỉ năm slot, vLLM sẽ âm thầm loại bỏ một chính sách vẫn còn các lượt rollout đang chạy tại mỗi lần đồng bộ.

Chúng tôi ghim vLLM ở phiên bản v0.27.1. vLLM thay đổi rất nhanh, và các cờ (flags) ở trên cùng các endpoint LoRA thời gian chạy là những gì phiên bản đó cung cấp, vì vậy hãy coi phiên bản này là một phần của công thức.

Có một thiết kế khả thi khác là trình huấn luyện chỉ giữ adapter mới nhất và luôn xuất bản nó dưới cùng một tên. Chúng tôi không đi theo hướng đó, vì vLLM sử dụng tên adapter làm khóa cho bộ nhớ đệm tiền tố (prefix cache). Với một tên duy nhất, các khối KV được tính toán dưới các trọng số trước đó vẫn sẽ khớp sau khi hoán đổi, vì vậy quá trình tiền điền (prefill) sẽ không được thực hiện lại và một lượt rollout có thể lấy tiền tố từ một phiên bản chính sách này và giải mã (decode) từ phiên bản tiếp theo. Trình huấn luyện sẽ không có cách nào để biết, và nó sẽ xuất hiện dưới dạng tỷ lệ bị lệch khỏi 1. Tên có phiên bản làm cho điều này không thể xảy ra: một tên luôn đại diện cho một bộ trọng số, và một tiền tố đã lưu trong bộ nhớ đệm không bao giờ có thể khớp với một phiên bản mới hơn.

Lựa chọn tập dữ liệu: tập Sanity

Chúng tôi đã chọn `sail/Sanity-Test-R1D-1.5B`, tập dữ liệu từ bài báo "Defeating the Training-Inference Mismatch via FP16" (Qi et al., 2025). Mã tái hiện nằm trong `sail-sg/Precision-RL`.

Các tác giả đã tạo 40 câu trả lời cho mỗi bài toán MATH với `DeepSeek-R1-Distill-Qwen-1.5B`. Họ giữ lại các bài toán có tỷ lệ thành công từ 20% đến 80%, thu được 1.460 câu hỏi. Tập dữ liệu này thực sự tốt cho việc xác thực RL vì các câu hỏi không quá dễ cũng không quá vô vọng đối với mô hình đó, nghĩa là mô hình có thể nhận được tín hiệu sớm tốt để huấn luyện và cải thiện.

Điều này rất tuyệt vời như một bài kiểm tra end-to-end mạnh mẽ: nếu một bản sao vLLM âm thầm phục vụ mô hình cơ sở dưới một tên adapter, chúng ta muốn thấy điều đó trên biểu đồ trong vòng vài chục bước. Ngoài ra, tập dữ liệu này đủ nhỏ để chạy hết trong vòng chưa đầy hai giờ.

Chúng tôi cũng lấy các siêu tham số từ các script LoRA của bài báo trong `oat/scripts/lora`: `Qwen/Qwen2.5-Math-1.5B`, LoRA rank 1 với alpha 2, tốc độ học 4e-5, 8 mẫu mỗi prompt, 128 câu trả lời mỗi bước, tối đa 3.000 token được tạo và ngữ cảnh 4.096 token.

Trình huấn luyện

Trình huấn luyện sử dụng cùng image `vllm/vllm-openai:v0.27.1` với TRL được cài đặt bên trên. Chúng tôi đã chạy nhánh PR vào thời điểm đó; cùng mã đó hiện đã có trong TRL v1.14. Script huấn luyện là một script `AsyncGRPOTrainer` bình thường. Các giá trị duy nhất dành riêng cho Job là thư mục đầu ra và URL máy chủ.

Trong quá trình khởi tạo, TRL gọi `/server_info`. Nếu tìm thấy `lora_config`, nó sẽ sử dụng đồng bộ chỉ-adapter. Các cấu hình mà vLLM không thể phục vụ trực tiếp, chẳng hạn như DoRA, `modules_to_save`, hoặc rank cao hơn `--max-lora-rank`, sẽ quay lại đồng bộ trọng số đã hợp nhất (merged-weight sync) kèm theo cảnh báo. Nhật ký sẽ chứa dòng "Adapter-only vLLM sync enabled".

Proxy

Bây giờ đến phần thú vị. Chúng ta cần một proxy giữa trình huấn luyện và các vLLM Jobs vì hai lý do:

Các cổng Job được công khai yêu cầu header `Authorization: Bearer <HF token>` trên mỗi yêu cầu. Proxy là nơi header đó được thêm vào, vì vậy TRL không cần biết về nó.

Chúng tôi muốn nhiều hơn một GPU thực hiện tạo. Trên một máy chủ vLLM đơn lẻ, cách thông thường để đạt được điều đó là `--data-parallel-size > 1`, nhưng TRL từ chối đồng bộ chỉ-adapter ở chế độ đó, vì một lý do chính đáng: một lệnh gọi đến `/v1/load_lora_adapter` chỉ đến được DP rank trả lời nó, vì vậy các rank khác sẽ tiếp tục phục vụ mô hình cơ sở dưới tên chính sách mới. Trên Jobs, câu hỏi đó thậm chí không nảy sinh, vì mỗi bản sao là một máy riêng biệt. Vì vậy, tính song song dữ liệu phải nằm ở cấp độ cao hơn, trong một thứ gì đó phân phối việc tải adapter đến mọi bản sao.

Do đó, chúng tôi chạy một proxy nhỏ tại `127.0.0.1:8000` trên Job trình huấn luyện và trỏ TRL đến nó như thể đó là một máy chủ vLLM duy nhất. Ngoài việc thêm header, proxy thực hiện hai việc về mặt chức năng:

Định tuyến các lượt rollout theo tiền tố KV

Một lời nhắc nhanh về lý do tại sao điều này quan trọng. Việc tạo một câu trả lời có hai giai đoạn với cấu hình khối lượng công việc rất khác nhau:

Các khóa và giá trị đó là bộ nhớ đệm KV (KV cache). Vì attention là nhân quả (causal), KV của một token chỉ phụ thuộc vào các token trước nó, không phải những gì đến sau. Do đó, hai yêu cầu chia sẻ cùng một tiền tố sẽ chia sẻ KV của tiền tố đó, và một bản sao đã có nó trong bộ nhớ đệm có thể bỏ qua hoàn toàn phần tiền điền (prefill). Toàn bộ trò chơi bây giờ là tìm bản sao đó, để một yêu cầu có thể hưởng lợi từ việc rơi vào bản sao đã thấy tiền tố của nó.

vLLM lưu trữ bộ nhớ đệm KV tiền tố của nó trong các khối 16 token. Do GRPO, rollout worker gửi G yêu cầu với cùng một prompt (trong trường hợp của chúng tôi G=8). Nếu tất cả chúng đến cùng một bản sao, yêu cầu đầu tiên sẽ tính toán tiền điền và bảy yêu cầu tiếp theo sẽ sử dụng lại nó. Với định tuyến round-robin, một nửa sẽ đi đến một bản sao không có tiền tố trong bộ nhớ đệm, và bốn yêu cầu đó sẽ phải thực hiện lại công việc tiền điền và lãng phí tài nguyên tính toán GPU quý giá.

Công việc của bộ định tuyến (router) của chúng tôi là theo dõi bản sao nào đã thấy mã băm (hash) khối nào. Một chi tiết quan trọng là các mã băm được liên kết với nhau, vì vậy mã băm của khối 3 đại diện cho khối 1, 2 và 3, không chỉ khối 3. Điều này phản ánh attention nhân quả: KV của khối 3 chỉ hợp lệ nếu khối 1 và 2 cũng giống nhau. Chúng tôi cũng gieo chuỗi với tên adapter vì bộ nhớ đệm KV cũng phụ thuộc vào adapter đã tạo ra nó: một tiền tố được lưu cho chính sách v3 là vô dụng đối với chính sách v4!

Video hướng dẫn toàn bộ quá trình ra quyết định để chọn một bản sao. Các bước dưới đây đi qua một ví dụ yêu cầu hoàn thành 135 token thực tế (từ các bài toán trong tập dữ liệu Sanity):

1. Chia prompt thành các khối. Bộ định tuyến nhận các id token và cắt chúng thành các khối 16 token, giống như vLLM. Nó chỉ băm các khối hoàn chỉnh, vì vậy 7 token cuối cùng bị bỏ qua ở đây.

Bài gốc còn tiếp — xem tiếp tại bài gốc ↗