Huấn luyện mô hình ngôn ngữ tác tử với FlashREINFORCE

FlashREINFORCE là gì?

FlashREINFORCE là một framework học tăng cường không dùng critic để huấn luyện bất đồng bộ các mô hình ngôn ngữ có khả năng thực hiện tác vụ. Framework này được thiết kế cho các tác vụ có chuỗi thời gian dài, trong đó các trajectory hoàn tất vào những thời điểm không đều do độ dài suy luận thay đổi, lời gọi công cụ hoặc tương tác với môi trường.

Thay vì tạo nhiều rollout cùng cấp cho mỗi prompt, FlashREINFORCE sử dụng một rollout cho mỗi prompt. Vì vậy, một batch chứa B trajectory có thể bao phủ B prompt khác nhau, đồng thời các trajectory đã hoàn tất có thể được gửi đến learner mà không cần chờ nhiều phản hồi cho cùng một prompt.

Repository cung cấp loss tham chiếu PyTorch độc lập, ví dụ cập nhật optimizer trên CPU, các bài kiểm thử, launcher Molt cho R1 và Qwen2.5-Math, cùng các thiết lập thử nghiệm cho suy luận, công cụ Python, mô hình mixture-of-experts và ALFWorld. Các launcher huấn luyện bất đồng bộ đầy đủ sử dụng Molt.

Học tăng cường nên thực hiện REINFORCE.

Tại sao sử dụng một rollout?

Các phương pháp tương đối theo nhóm phụ thuộc vào nhiều rollout từ cùng một prompt. Trong các workload tác vụ bất đồng bộ, yêu cầu này có thể tạo ra các rào cản đồng bộ hóa và làm giảm số lượng prompt khác nhau được bao phủ với một ngân sách rollout cố định.

Học với một rollout loại bỏ yêu cầu về các rollout cùng cấp, nhưng tạo ra ba vấn đề về độ ổn định:

• Không có nhóm cùng cấp để xây dựng baseline ở cấp prompt.
• Các trajectory được thu thập bất đồng bộ có thể được tạo ra bởi một behavior policy cũ hơn.
• Các trajectory thất bại dài có thể chi phối một lần cập nhật nếu mọi token nhận cùng trọng số ở cấp batch.

FlashREINFORCE giải quyết các vấn đề này bằng One-Batch REINFORCE, Sequence Trust Region và Sample-Mean Optimization.

Các thành phần cốt lõi

1. One-Batch REINFORCE

Với mỗi batch mới, learner chuẩn hóa các reward vô hướng trên các prompt độc lập. Nếu R_i là reward của trajectory i, advantage của nó là reward của trajectory trừ đi trung bình batch:

mean_reward = sum(R_j for j in batch) / B
A_i = R_i - mean_reward

Các trajectory cao hơn mức trung bình nhận advantage dương, trong khi các trajectory thấp hơn mức trung bình nhận advantage âm. Cách này cung cấp phản hồi có dấu mà không cần huấn luyện critic. Learner thực hiện một lần cập nhật optimizer trên batch rồi loại bỏ batch đó.

2. Sequence Trust Region

Mỗi trajectory phải lưu xác suất của behavior policy đã thực sự tạo ra từng token được lấy mẫu. Learner tính lại xác suất của current policy và tạo ra tỷ lệ importance ở cấp token:

rho_i_t = exp(current_log_probability_i_t
              - behavior_log_probability_i_t)

Tỷ lệ này hiệu chỉnh phân phối hành động được lấy mẫu tại một history đã lưu. Nó không hiệu chỉnh thực tế rằng chính history đó được tạo ra từ một behavior policy cũ, vì vậy FlashREINFORCE cũng đánh giá độ lệch tích lũy trên toàn bộ trajectory.

Với mỗi action được lấy mẫu, framework tính một giá trị gần đúng KL Bernoulli từ xác suất behavior p và xác suất hiện tại q:

d_i_t = p * log(p / q)
        + (1 - p) * log((1 - p) / (1 - q))

Các giá trị ở cấp token được lấy trung bình trên trajectory. Một trajectory chỉ được chấp nhận khi giá trị gần đúng trung bình không lớn hơn ngưỡng tin cậy đã cấu hình:

mean_drift_i = sum(d_i_t for t in trajectory_i) / T_i
mask_i = 1 if mean_drift_i <= delta else 0

Trajectory bị từ chối không đóng góp loss ở cấp token. Do đó, quyết định tin cậy được áp dụng cho toàn bộ history đã lưu thay vì các token riêng lẻ.

3. Sample-Mean Optimization

Trung bình batch ở cấp token tự động khiến các trajectory dài có ảnh hưởng lớn hơn. Điều này có thể gây hại khi một rollout thất bại dài chứa các bước trung gian hữu ích nhưng kết thúc với reward âm.

FlashREINFORCE trước tiên lấy trung bình các đóng góp ở cấp token trong từng trajectory, sau đó lấy trung bình trên các trajectory:

trajectory_term_i = (
    mask_i * advantage_i / trajectory_length_i
    * sum(importance_ratio_i_t for each token t)
)
objective = sum(trajectory_term_i for i in batch) / B

Do đó, mỗi trajectory nhận trọng số 1/B bất kể độ dài phản hồi. Loss thực tế sử dụng các importance ratio đã tách khỏi đồ thị tính toán. Phương pháp cơ sở không sử dụng critic được học, cắt ratio hoặc lượt chạy forward của reference model.

Cài đặt và thiết lập cục bộ

Clone hoặc tải repository FlashREINFORCE, đi vào thư mục gốc rồi cài đặt ở chế độ editable cùng các dependency dành cho kiểm thử:

pip install -e '.[test]'

Tiếp theo, chạy ví dụ CPU đi kèm:

python examples/toy_update.py

Ví dụ này kiểm tra rằng objective tham chiếu có thể tạo ra một lần cập nhật optimizer trên CPU. Đây là ví dụ xác minh chức năng, không phải bản tái hiện kết quả benchmark của bài báo.

Cuối cùng, chạy bộ kiểm thử:

python -m pytest -q

Ví dụ và bộ kiểm thử chạy thành công xác nhận rằng implementation tham chiếu cục bộ cùng các dependency của nó đã sẵn sàng.

Khám phá implementation tham chiếu

Objective tham chiếu trung tâm nằm trong flashreinforce/loss.py. Nó bao quát các thao tác được mô tả trong bài báo và README:

• Chuẩn hóa reward trên toàn batch.
• Importance sampling ở cấp token.
• Chấp nhận theo vùng tin cậy ở cấp sequence.
• Tối ưu hóa theo sample mean.
• Lọc token thất bại tùy chọn dựa trên entropy.

Bắt đầu với examples/toy_update.py để xem loss tham chiếu tham gia như thế nào vào một lần cập nhật optimizer. Sau đó kiểm tra flashreinforce/loss.py để liên hệ các tensor và mask trong ví dụ với từng giai đoạn riêng lẻ của objective.

Ví dụ CPU là cách sử dụng cơ bản an toàn nhất vì không yêu cầu toàn bộ hạ tầng triển khai phân tán.

Tìm hiểu về bản cập nhật bất đồng bộ một lượt

Một bản cập nhật FlashREINFORCE trong môi trường production trải qua sáu giai đoạn:

1. Thu thập: Các worker rollout độc lập gửi một quỹ đạo đã hoàn tất cho mỗi prompt. Mỗi quỹ đạo bao gồm phần thưởng và xác suất của chính sách hành vi thực tế được dùng để lấy mẫu các token.
2. Gộp batch: Learner lấy B quỹ đạo đã hoàn tất tiếp theo và chuẩn hóa phần thưởng bằng cách lấy trọng tâm để thu được các lợi thế có dấu.
3. Hiệu chỉnh: Learner tính lại log-probability của chính sách hiện tại và tạo các tỷ lệ importance theo token.
4. Kiểm soát: Learner tính proxy KL trung bình của các hành động được lấy mẫu cho từng quỹ đạo và chỉ chấp nhận toàn bộ quỹ đạo nếu giá trị này nằm trong ngưỡng đã cấu hình.
5. Tối ưu: Learner lấy trung bình trong từng quỹ đạo được chấp nhận, lấy trung bình trên toàn batch và thực hiện một bước cập nhật optimizer.
6. Loại bỏ: Learner không phát lại batch đã thu thập để thực hiện thêm các bản cập nhật.

Bước loại bỏ cuối cùng là một phần quan trọng của phương pháp. FlashREINFORCE sử dụng một bản cập nhật trên toàn batch cho mỗi batch vừa thu thập, thay vì chạy nhiều bản cập nhật minibatch tuần tự từ cùng một snapshot của chính sách hành vi.

Xem trước cấu hình huấn luyện Molt

Repository bao gồm scripts/train_molt.py để khởi chạy quá trình huấn luyện thông qua triển khai bất đồng bộ của Molt. Trước khi phân bổ tài nguyên cho một lần chạy huấn luyện đầy đủ, hãy sử dụng chế độ dry-run để xem trước các cờ huấn luyện:

python scripts/train_molt.py --dry-run

Tham khảo docs/training.md để biết cách thiết lập GPU, chuẩn bị dữ liệu, các lệnh huấn luyện, cách ánh xạ từ bài báo sang mã nguồn và những hạn chế tái hiện đã biết. Hướng dẫn examples/README.md bao gồm các preset tác vụ, các phép ablation và trình khởi chạy native trainer kèm kiểm tra khả năng.

Các trình khởi chạy R1 và Qwen2.5-Math đã được ghim phiên bản trong repository. Các thiết lập sử dụng công cụ và ALFWorld cũng yêu cầu agent và môi trường tương thích, vì vậy hãy xem xét các yêu cầu runtime của chúng trước khi khởi chạy những thí nghiệm đó.

Dữ liệu mỗi quỹ đạo phải bảo toàn

Bản cập nhật phụ thuộc vào thông tin được ghi lại trong quá trình sinh. Tối thiểu, luồng huấn luyện về mặt khái niệm cần chuỗi token đã hoàn tất, phần thưởng vô hướng của chuỗi và xác suất hành vi tương ứng với mọi hành động được lấy mẫu. Sau đó, learner lấy các xác suất tương ứng của chính sách hiện tại bằng cách đánh giá các hành động và lịch sử đã lưu.

Không thay thế các xác suất lấy mẫu đã ghi bằng những xác suất được tái dựng về sau. Việc đánh giá mô hình ở thời điểm sau đó không đảm bảo tái tạo đúng xác suất mà engine suy luận thực sự đã sử dụng trong quá trình lấy mẫu.

Mẹo ổn định nâng cao

Bảo toàn miền hỗ trợ của quá trình lấy mẫu

Hiệu chỉnh importance chính xác giả định rằng chính sách đích không gán khối lượng xác suất cho các giá trị nằm ngoài miền hỗ trợ của chính sách hành vi. Việc cắt ngưỡng top-k hoặc top-p mạnh có thể vi phạm giả định này. Vì vậy, nên lựa chọn các thiết lập lấy mẫu có cân nhắc đến yêu cầu về miền hỗ trợ.

Không xem importance sampling theo token là hiệu chỉnh toàn bộ quỹ đạo

Các tỷ lệ importance theo token hiệu chỉnh các hành động được lấy mẫu dựa trên những lịch sử đã lưu. Chúng không hiệu chỉnh phân phối của các lịch sử đó. Hãy sử dụng Sequence Trust Region để sàng lọc các quỹ đạo có độ lệch tích lũy từ hành vi đến learner quá lớn.

Giữ cho các bản cập nhật luôn mới

Thực hiện một bản cập nhật optimizer cho mỗi batch đã thu thập rồi loại bỏ batch đó. Việc sử dụng lại một batch cũ cho nhiều bản cập nhật learner sẽ lệch khỏi phương pháp một lượt được tài liệu hóa và có thể làm tăng độ không khớp giữa các chính sách.

Bảo toàn trọng số ở cấp quỹ đạo

Lấy trung bình các đóng góp của token trong từng quỹ đạo trước khi lấy trung bình các quỹ đạo trên toàn batch. Việc dàn phẳng mọi token hợp lệ thành một giá trị trung bình duy nhất trên toàn batch sẽ đưa trở lại hệ số nhân theo độ dài phản hồi và khiến các quỹ đạo dài nhận trọng số lớn không tương xứng.

Diễn giải đúng mặt nạ tin cậy

Sequence Trust Region đưa ra một quyết định chấp nhận duy nhất cho toàn bộ quỹ đạo. Nếu một quỹ đạo vượt quá ngưỡng, mọi loss token của quỹ đạo đó sẽ bị loại khỏi bản cập nhật tương ứng.

Tách biệt việc xác minh và tái hiện

Ví dụ CPU chỉ xác minh một bản cập nhật optimizer. Để tái hiện các kết quả nghiên cứu, cần có GPU, dữ liệu, mô hình, agent, môi trường và các cấu hình runtime bất đồng bộ được nêu trong tài liệu huấn luyện và hướng dẫn examples.

Các lĩnh vực ứng dụng được báo cáo

Các thí nghiệm được báo cáo bao gồm suy luận toán học, sử dụng công cụ Python nhiều lượt, huấn luyện mixture-of-experts và ra quyết định tương tác trong ALFWorld. README báo cáo quá trình huấn luyện chuỗi suy luận dài ổn định qua 6.000 bản cập nhật cho DeepSeek-R1-Distill-Qwen-1.5B, kết quả suy luận toán học cho Qwen2.5-Math-1.5B, các thí nghiệm dùng công cụ Python với Qwen2.5-7B-Instruct và Qwen3-30B-A3B, cùng các thí nghiệm ALFWorld với Qwen2.5-7B-Instruct.

Những kết quả này cho thấy phạm vi dự kiến của framework, nhưng không nên kỳ vọng đạt được chúng từ ví dụ cập nhật CPU quy mô nhỏ.

Tài nguyên hữu ích

• Kho lưu trữ FlashREINFORCE
• Trang web dự án
• Mã nguồn NVIDIA NeMo Molt
• Hướng dẫn huấn luyện
• Cấu hình và ví dụ thử nghiệm

Kết luận

FlashREINFORCE kết hợp khả năng bao phủ prompt trong một lần rollout với một bước cập nhật gọn nhẹ không cần critic. Phần thưởng tập trung theo batch cung cấp các lợi thế có dấu, lấy mẫu quan trọng theo token xử lý các hành động lỗi thời, Sequence Trust Region kiểm soát độ lệch tích lũy của quỹ đạo, còn tối ưu hóa theo trung bình mẫu ngăn độ dài phản hồi quyết định trọng số quỹ đạo.

Hãy bắt đầu với bản cài đặt có thể chỉnh sửa, ví dụ cập nhật CPU và các bài kiểm thử. Khi chúng hoạt động, hãy xem trước trình khởi chạy Molt bằng --dry-run và làm theo hướng dẫn huấn luyện để chuẩn bị GPU và dữ liệu theo từng mô hình. Dự án được phát hành theo Apache License 2.0.