Trainer GRPO bất đồng bộ

Tổng quan

AsyncGRPOTrainer triển khai cùng thuật toán GRPO nhưng tách rời việc tạo rollout khỏi quá trình huấn luyện. Một worker nền liên tục stream các completion từ một máy chủ vLLM trong khi vòng lặp huấn luyện tiêu thụ chúng, nhờ đó việc tạo và cập nhật gradient diễn ra chồng chéo thay vì luân phiên.

Trainer này được đóng góp bởi Amine Dirhoussi và Quentin Gallouédec.

Bắt đầu nhanh

Ví dụ này minh họa cách huấn luyện một mô hình bằng phương pháp Asynchronous GRPO. Chúng ta huấn luyện mô hình Qwen3.5 2B với các prompt từ dataset DeepMath-103K. Bạn có thể xem dữ liệu trong dataset tại đây:

Dưới đây là script để huấn luyện mô hình. Trainer cần vLLM và kernels, và accuracy_reward cần math-verify, bạn có thể cài đặt bằng pip install trl[vllm,kernels,math_verify].

# train_async_grpo.py
from datasets import load_dataset
from trl import AsyncGRPOTrainer
from trl.rewards import accuracy_reward

if __name__ == "__main__":  # the rollout worker is a spawned process that re-imports this script
    dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

    trainer = AsyncGRPOTrainer(
        model="Qwen/Qwen3.5-2B",
        reward_funcs=accuracy_reward,
        train_dataset=dataset,
    )
    trainer.train()

Quá trình sinh (Generation) chạy trên một máy chủ vLLM, trên các GPU khác với máy huấn luyện (trainer). Khởi động nó trên GPU 0:

CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3.5-2B \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --api-server-count 8 \
    --max-model-len 4096

Sau đó, thực thi script trên GPU 1 bằng lệnh sau:

CUDA_VISIBLE_DEVICES=1 python train_async_grpo.py

Tương thích và ràng buộc

Một số tổ hợp tham số bị hạn chế có chủ đích trong triển khai AsyncGRPOTrainer hiện tại:

• Việc sinh (Generation) luôn được chạy trên một máy chủ vLLM riêng biệt, được khởi động bằng vllm serve. Không có chế độ colocate.
• Prompts from the dataset must be in conversational format. With an environment_factory, the dataset may have no prompt column, and reset() provides it.
• Hàm thưởng phải là các callable Python. Các mô hình thưởng được cung cấp dưới dạng model id hoặc PreTrainedModel không được hỗ trợ, xem Sử dụng hàm thưởng tùy chỉnh.
• Trainer cần transformers>=5.2.0.
• Huấn luyện phân mảnh (Sharded training) sử dụng FSDP2. DeepSpeed không được hỗ trợ, và cả song song ngữ cảnh lẫn song song chuỗi (cp_size > 1 hoặc sp_size > 1 trong parallelism_config) cũng không được hỗ trợ.
• use_liger_kernel=True không được hỗ trợ.
• Các bộ thích ứng PEFT nhắm vào lm_head không được hỗ trợ.
• Các mô hình có attention sinks không được hỗ trợ. Các lớp không phải attention duy nhất được hỗ trợ là các lớp gated DeltaNet của các mô hình hybrid (Qwen3.5, Qwen3.6, Qwen3-Next, …), cần transformers>=5.15.0.
• Hình ảnh không được hỗ trợ, xem Đào tạo Mô hình Ngôn ngữ Thị giác.

Training nạp attention kernel từ Hub thông qua thư viện kernels: kernels-community/aiter-flash-attn trên GPU AMD (ROCm), kernels-community/flash-attn4 trên GPU NVIDIA Blackwell và kernels-community/flash-attn3 trên các GPU NVIDIA khác. Các mô hình giới hạn attention logits (attn_logit_softcapping), chẳng hạn như Gemma 2, không được hỗ trợ trên GPU AMD hoặc GPU NVIDIA Blackwell.

Đi sâu hơn vào phương pháp GRPO bất đồng bộ

Trong GRPOTrainer, quá trình sinh và huấn luyện diễn ra tuần tự: sinh một lô dữ liệu, tính toán loss, cập nhật trọng số, lặp lại. Ngay cả trong colocate mode, một pha phải hoàn tất trước khi pha kia bắt đầu. Trong AsyncGRPOTrainer, một worker rollout sinh và chấm điểm các câu trả lời trong một tiến trình nền trong khi vòng lặp huấn luyện tính toán loss và cập nhật mô hình. Để hiểu cách hoạt động, quy trình có thể được chia thành ba bước chính: Sinh các câu trả lời, tính toán advantage và tính toán loss. Không có mô hình tham chiếu, do đó không có số hạng KL divergence (mặc định của GRPO, beta=0.0).

Tạo các câu trả lời hoàn chỉnh

Worker rollout duy trì tối đa max_concurrent_rollouts yêu cầu đang chạy tới máy chủ vLLM. Với mỗi prompt, nó tạo ra một nhóm $G$ bổ sung (num_generations) và, khi toàn bộ nhóm hoàn tất, chấm điểm và tính toán lợi thế của chúng. Các mẫu huấn luyện kết quả được đưa vào hàng đợi, nơi vòng lặp huấn luyện sẽ lấy ra. Quá trình tạo chỉ chờ khi nó đi trước huấn luyện max_staleness bước (xem Độ song song), và một bước huấn luyện chỉ chờ khi hàng đợi trống.

Mỗi weight_sync_steps bước huấn luyện, trainer gửi trọng số đã cập nhật của nó đến máy chủ vLLM qua NCCL, để các completion tiếp theo được tạo từ chính sách mới nhất. Do đó, một completion có thể đến từ một chính sách cũ hơn vài bước so với chính sách đang được huấn luyện: một mẫu bị bỏ nếu nó cũ hơn max_staleness lần cập nhật trọng số tại thời điểm trainer lấy nó.

Các chỉ số và phần còn lại của trang này sử dụng các thuật ngữ sau cho các giai đoạn mà một completion trải qua:

• a nhóm là một prompt và các $G$ rollouts của nó. Lợi thế được tính toán bên trong nhóm đó.
• Một rollout là một cuộc hội thoại hoàn chỉnh được tạo ra cho một prompt, bao gồm mọi lần gọi công cụ. Nó được cấu thành từ một hoặc nhiều turns, mỗi turn tương ứng với một yêu cầu vLLM.
• sample là thứ mà loss được tính trên đó: các token của một rollout, kèm theo advantage của nó. Một rollout đa lượt có thể tạo ra nhiều hơn một sample (xem Multi-turn rollouts).
• Một dòng là những gì một tiến trình chuyển tiếp trong một micro-batch: một vài mẫu được đóng gói cùng nhau (xem Đóng gói cây). Một micro-batch chứa một dòng cho mỗi tiến trình, và gradient_accumulation_steps micro-batches tạo thành một bước, tức là bước tối ưu hóa.

Số lượng token được chia thành ba loại: token được tạo ra là những token do mô hình tạo ra, token được chuyển tiếp là tất cả các token của một mẫu (bao gồm cả prompt và kết quả công cụ), và token được huấn luyện là những token mà loss được tính trên đó.

Tính toán lợi thế

Với mỗi $G$ bổ sung của một nhóm, phần thưởng là tổng các phần thưởng từ mỗi hàm phần thưởng. Lợi thế được chuẩn hóa trong nhóm, giống như trong GRPO: $$\hat{A}_{i,t} = \frac{r_i - \text{mean}(\mathbf{r})}{\text{std}(\mathbf{r})}$$

Không có tùy chọn scale_rewards: lợi thế luôn được nhân với độ lệch chuẩn của nhóm. Một kết quả mà mọi hàm phần thưởng đều trả về None sẽ bị loại khỏi trung bình và độ lệch chuẩn của nhóm và nhận lợi thế bằng không.

Một nhóm có phần thưởng đều nhau sẽ bị loại bỏ: với lợi thế mặc định, giá trị của nó bằng không ở mọi vị trí, nên huấn luyện trên nhóm đó sẽ tốn một lần truyền thuận và một lần truyền ngược mà không thay đổi gì. Tỷ lệ các nhóm như vậy được ghi log dưới dạng frac_zero_std_groups. Bước huấn luyện được lấp đầy từ các nhóm được giữ lại, do đó nhiều nhóm bị loại nghĩa là phải sinh thêm dữ liệu cho mỗi bước. Một nhóm bị loại vẫn được tính vào epoch, vốn chỉ là một lượt quét qua các prompt.

Tính toán tổn thất

Log-probabilities của server vLLM không bao giờ khớp hoàn toàn với trainer: các kernel khác nhau, độ chính xác khác nhau, và một policy chậm hơn vài bước (xem Xử lý sự không khớp giữa huấn luyện và suy luận). Trainer sử dụng log-probabilities của sampler làm policy cũ, do đó một tỷ lệ quan trọng duy nhất bao phủ cả sự không khớp lẫn độ cũ: $$r_{i,t}(\theta) = \frac{\pi_\theta(o_{i,t} \mid q, o_{i,< t})}{\pi_{\text{sampler}}(o_{i,t} \mid q, o_{i,< t})}.$$

Các token có tỷ lệ nằm ngoài [vllm_importance_sampling_clip_min, vllm_importance_sampling_clip_max] sẽ bị loại khỏi loss. Để biết thêm thông tin, xem Masked Importance Sampling.

Mặc định (score_centering=True), trainer sử dụng Score Centering. Chỉ riêng trọng số quan trọng không loại bỏ được sự lệch khớp: điểm $\nabla_\theta \log \pi_\theta$ của một token được lấy mẫu không có trung bình bằng không dưới bộ lấy mẫu $\mu$, nên mỗi lần cập nhật cũng kéo chính sách về phía bộ lấy mẫu. Score centering trừ đi kỳ vọng đó tại mỗi vị trí: $$\nabla_\theta \mathcal{L} = - \mathbb{E}_{o \sim \mu} \left[ \hat{A} \left( f(r_t) \nabla_\theta \log \pi_\theta(o_t) - \mathbb{E}_{v \sim \mu} \left[ f(r_v) \nabla_\theta \log \pi_\theta(v) \right] \right) \right], \quad r_v = \frac{\pi_\theta(v)}{\mu(v)},$$

trong đó $f(r) = r$ bên trong $[C_{\min}, C_{\max}]$ và $0$ bên ngoài. Kỳ vọng bên trong được tính trên num_sampler_logprobs các token có khả năng cao nhất tại mỗi vị trí, do worker rollout yêu cầu từ server.

Với score_centering=False, trainer sẽ sử dụng hàm mục tiêu surrogate bị cắt của GRPO thay thế, với tỷ lệ được che (masked ratio) tương tự: $$\mathcal{L}(\theta) = - \frac{1}{\sum_{i} |o_i|} \sum_{i} \sum_{t=1}^{|o_i|} \mathbb{1}\left[ C_{\min} \le r_{i,t}(\theta) \le C_{\max} \right] \min \left( r_{i,t}(\theta) \hat{A}_{i,t}, \, \text{clip}\left( r_{i,t}(\theta), 1 - \epsilon_{\text{low}}, 1 + \epsilon_{\text{high}} \right) \hat{A}_{i,t} \right),$$

trong đó $\epsilon_{\text{low}}$ là epsilon và $\epsilon_{\text{high}}$ là epsilon_high.

Trong cả hai trường hợp, loss được chuẩn hóa theo số lượng token được huấn luyện trong micro-batch, được đếm trên tất cả các process, và lấy trung bình trên gradient_accumulation_steps micro-batches của một bước.

Các chỉ số được ghi log

Trong quá trình huấn luyện, chúng ta ghi lại các chỉ số sau:

• num_tokens: Tổng số token đã được chuyển tiếp cho đến nay, bao gồm prompts và completions. Với tree packing, một prefix được chia sẻ bởi nhiều mẫu chỉ được tính một lần.
• completions/mean_length: Độ dài trung bình của các phần hoàn thành được tạo ra. Khi sử dụng công cụ, chỉ các token không phải của công cụ mới được tính.
• completions/min_length: Độ dài tối thiểu của các phần hoàn thành được tạo ra. Khi sử dụng công cụ, chỉ các token không phải của công cụ mới được tính.
• completions/max_length: Độ dài tối đa của các phần hoàn thành được tạo ra. Khi sử dụng công cụ, chỉ các token không phải của công cụ mới được tính.
• completions/mean_terminated_length: Độ dài trung bình của các chuỗi hoàn thành được tạo ra kết thúc bằng EOS. Khi sử dụng công cụ, chỉ các token không phải công cụ được tính.
• completions/min_terminated_length: Độ dài tối thiểu của các phần hoàn thành được tạo ra kết thúc bằng EOS. Khi sử dụng công cụ, chỉ các token không phải công cụ được tính.
• completions/max_terminated_length: Độ dài tối đa của các completion được tạo ra kết thúc bằng EOS. Khi sử dụng công cụ, chỉ các token không phải công cụ được tính.
• completions/clipped_ratio: Tỷ lệ các câu trả lời bị cắt (clipped), trong đó lượt cuối không kết thúc bằng token EOS.
• completions/duplicate_frac: Tỷ lệ các rollout lặp lại một rollout trước đó của nhóm chúng theo từng token. Xem Bão hòa và bản sao.
• tools/call_frequency: Số lần gọi công cụ trung bình cho mỗi lần hoàn thành, tính trên mọi lượt rollout, bao gồm cả các nhóm sau đó bị loại bỏ. Chỉ được ghi log khi trainer có công cụ.
• tools/failure_frequency: Tỷ lệ các lần gọi công cụ bị thất bại (công cụ không được tìm thấy hoặc gây ra ngoại lệ) trên tổng số lần gọi công cụ, bao gồm cả các nhóm sau đó bị loại bỏ. Giá trị là NaN trong một cửa sổ ghi log không có lần gọi công cụ nào, nơi GRPO ghi log 0.0. Chỉ được ghi log khi trainer có công cụ.
• rewards/{reward_func_name}/mean: Phần thưởng trung bình từ một hàm phần thưởng cụ thể, tính trên mọi completion mà nó đã chấm điểm, bao gồm cả các nhóm sau đó bị loại bỏ. Khi một môi trường sở hữu phần thưởng thông qua get_reward, {reward_func_name} là tên lớp của môi trường đó.
• rewards/{reward_func_name}/group_std: Độ lệch chuẩn của phần thưởng từ một hàm phần thưởng cụ thể, được tính trong từng nhóm và lấy trung bình trên các mẫu đã huấn luyện.
• reward: Phần thưởng trung bình tổng thể sau khi cộng phần thưởng từ các hàm, trên mọi hoàn thành được chấm điểm, bao gồm cả các nhóm sau đó bị loại bỏ. Đây là phần thưởng của chính sách, vì vậy đây là chỉ số cần theo dõi: reward_trained loại trừ các nhóm có phần thưởng đều bằng nhau.
• reward_group_std: Độ lệch chuẩn của phần thưởng tổng hợp trên các hàm, được tính trong từng nhóm và lấy trung bình trên các mẫu đã huấn luyện. Đây là giá trị mà lợi thế được chia cho, khác với reward_std của GRPO, vốn được tính trên toàn bộ lô (batch).
• reward_trained: Phần thưởng trung bình của các mẫu mà trình huấn luyện đã dùng để huấn luyện. Giá trị này loại trừ các nhóm có phần thưởng đều bằng nhau và các mẫu bị loại do vượt quá max_staleness.
• reward_zero_std_groups: Phần thưởng trung bình của các nhóm có phần thưởng đều bằng nhau, những nhóm này bị loại bỏ: mức phần thưởng mà chúng bị mắc kẹt. Xem Bão hòa và bản sao để đọc nó cùng với frac_zero_std_groups và reward.
• frac_zero_std_groups: Tỷ lệ các nhóm có độ lệch chuẩn phần thưởng bằng 0 (tất cả các câu trả lời đều đúng hoặc đều sai), hoặc có ít hơn hai câu hoàn thành có thể chấm điểm. Những nhóm này sẽ bị loại bỏ.
• sampling/sampling_logp_difference/mean: Chênh lệch tuyệt đối trung bình giữa các xác suất log do trình lấy mẫu (vLLM) trả về và các xác suất log được tính lại bởi mô hình huấn luyện, trên các token đã được huấn luyện.
• sampling/sampling_logp_difference/max: Chênh lệch tuyệt đối lớn nhất như vậy.
• sampling/importance_sampling_ratio/min: Tỷ lệ importance sampling nhỏ nhất $r_{i,t}(\theta)$ giữa mô hình và bộ lấy mẫu, trên các token đã huấn luyện, trước khi loại bỏ các token nằm ngoài [vllm_importance_sampling_clip_min, vllm_importance_sampling_clip_max] .
• sampling/importance_sampling_ratio/mean: Trung bình của tỷ lệ lấy mẫu quan trọng như vậy.
• sampling/importance_sampling_ratio/max: Tỷ lệ lấy mẫu quan trọng lớn nhất như vậy.
• sampling/masked_mean: Tỷ lệ các token được huấn luyện bị loại khỏi loss do tỷ lệ importance sampling của chúng nằm ngoài khoảng [vllm_importance_sampling_clip_min, vllm_importance_sampling_clip_max].
• sampling/sampler_mass: Khối lượng xác suất trung bình mà các token có khả năng cao nhất theo num_sampler_logprobs của bộ lấy mẫu bao phủ. Chỉ được ghi lại khi score_centering=True.
• entropy: Entropy trung bình của các dự đoán token trên các token đã được huấn luyện.
• aux_loss: Tổn phụ trợ cân bằng tải của mô hình Mixture-of-Experts, trước khi được nhân với router_aux_loss_coef và cộng vào tổn. Chỉ được ghi log khi mô hình là mô hình MoE và router_aux_loss_coef khác không.
• clip_ratio/region_mean: Tỷ lệ xác suất của các token mà mục tiêu GRPO bị cắt (clip) để giữ trong vùng tin cậy (trust region): $\text{clip}\left( r_{i,t}(\theta), 1 - \epsilon_\mathrm{low}, 1 + \epsilon_\mathrm{high} \right)$. Với score_centering=True thì loss không bị cắt, do đó các clip_ratio/* metrics đếm số token nằm ngoài vùng tin cậy.
• clip_ratio/low_mean: Tỷ lệ trung bình của xác suất token bị cắt ở cận dưới của vùng tin cậy: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/low_min: Tỷ lệ phần nhỏ nhất trên mỗi completion của các token bị cắt ở cận dưới của vùng tin cậy: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/high_mean: Tỷ lệ trung bình của xác suất token bị cắt ở cận trên của vùng tin cậy: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.
• clip_ratio/high_max: Tỷ lệ phần token bị cắt ở cận trên của vùng tin cậy cho mỗi lần hoàn thành lớn nhất: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.

Các chỉ số khác mô tả đường ống xử lý bất đồng bộ.

Các chỉ số được ghi log một lần mỗi cửa sổ ghi log: tất cả những gì được ghi lại kể từ lần log trước, bao phủ logging_steps bước tối ưu hóa gần nhất. Trong cửa sổ:

• Một giá trị trung bình được đo cho mỗi lượt rollout, mẫu hoặc cuộc gọi, chẳng hạn như reward, rollout/turns_mean và tools/latency_s, sẽ được tính trung bình trên tất cả các lượt rollout, mẫu hoặc cuộc gọi trong cửa sổ. Một tỷ lệ như rollout/generated_tokens_per_s là tổng số token chia cho tổng số giây.
• Ngược lại, một tên mà phần cuối chứa từ count sẽ được cộng dồn: sample/dropped_stale_count là số lượng mẫu bị loại kể từ lần ghi log cuối, không phải kể từ khi bắt đầu huấn luyện.
• Tên mà phần cuối chứa từ max hoặc min, chẳng hạn như sample/staleness_max hoặc completions/max_length, sẽ giữ lại giá trị lớn nhất hoặc nhỏ nhất.
• Bất kỳ chỉ số nào khác đều là giá trị trung bình của các giá trị được ghi lại trong cửa sổ.

Đơn vị:

• _s: giây.
• tokens_per_s: số token mỗi giây.
• _frac: tỷ lệ nằm trong khoảng từ 0 đến 1.
• _per_step: lượng trong một bước tối ưu hóa, được trung bình hóa trên các bước trong cửa sổ.

Chỉ số triển khai

• rollout/duration_s: Thời gian thực tế của một lượt rollout, tính từ khi gửi lệnh cho đến khi vòng lặp rollout thu thập kết quả.
• rollout/generated_tokens_per_s: Thông lượng sinh, tính bằng số token được sinh mỗi giây, kể từ lần báo cáo trước.
• rollout/concurrency: Số lượng rollout đang chạy đến máy chủ vLLM.
• rollout/concurrency_cap: Giới hạn hiện tại cho các rollout đang chạy, tuân theo tải của máy chủ. Xem Độ song song.
• rollout/turns_mean, rollout/turns_max: Số lượt trung bình và tối đa cho mỗi rollout.
• rollout/samples_per_rollout: Số lượng mẫu huấn luyện cho mỗi rollout. Lớn hơn 1.0, một số rollout đã được phân nhánh; nhỏ hơn 1.0, một số không tạo ra token nào được huấn luyện. Xem Multi-turn rollouts.
• rollout/fork_frac: Tỷ lệ các bước chuyển đổi lượt mà prompt của chúng không bắt đầu bằng các token được giữ lại, do đó lượt đó đã bắt đầu một mẫu mới.
• rollout/drift_tokens_mean, rollout/drift_tokens_max: Số lượng token được giữ lại mà một fork đã render lại, tính trung bình trên các lần chuyển đổi lượt (một lượt không fork được tính là 0). Mẫu mới chia sẻ các token trước điểm phân kỳ, do đó đây là phần mà mỗi fork thêm vào quá trình forward.
• rollout/score_queue_size: Số nhóm đã hoàn thành đang chờ được chấm điểm.
• rollout/queue_size: Số lượng mẫu đã được chấm điểm đang chờ trong hàng đợi rollout, được đo tại vị trí trainer lấy chúng. Xem Hàng đợi rollout.
• rollout/score_s: Thời gian để chấm điểm một nhóm.
• rollout/score_wait_s: Thời gian một nhóm đã hoàn thành chờ trước khi được chấm điểm.
• rollout/score_block_s: Thời gian sinh bị chặn do hàng đợi chấm điểm đã đầy.
• rollout/backpressure_s: Thời gian sinh chờ để bắt đầu một rollout vì nó đã đi trước huấn luyện max_staleness bước, tính trung bình trên các lần chờ trong cửa sổ ghi log. Xem Độ song song.
• rollout/vllm_retry_count: Số lượng yêu cầu vLLM đã được thử lại. Nếu không có chỉ số này, một máy chủ bị suy giảm hiệu năng sẽ trông giống như sự chậm trễ không rõ nguyên nhân.
• rollout/vllm_kv_cache_usage_frac: Mức sử dụng bộ đệm KV của vLLM (giá trị tối đa trên các engine), được lấy trung bình trên các lần lấy mẫu trong cửa sổ ghi log.
• rollout/vllm_running_requests: Số lượng yêu cầu mà các engine của server đang giải mã (tổng của num_requests_running của vLLM trên tất cả các engine). Chênh lệch so với rollout/concurrency là phần nằm trong tiến trình API server, đang xây dựng các phản hồi.
• rollout/vllm_queued_requests: Số lượng yêu cầu mà máy chủ đã chấp nhận nhưng chưa được lên lịch (tổng của num_requests_waiting của vLLM trên các engine).
• rollout/vllm_preemption_count: Số lượng yêu cầu mà máy chủ đã tạm dừng (tổng của num_preemptions_total của vLLM trên tất cả các engine). Việc đồng bộ trọng số sẽ đặt lại bộ nhớ đệm của máy chủ, dẫn đến việc tạm dừng mọi yêu cầu đang chạy, do đó chỉ số này không phải là tín hiệu cho thấy áp lực hệ thống.

Chỉ số công cụ

• tools/latency_s, tools/{tool_name}/latency_s: Thời gian thực hiện một lệnh gọi công cụ, tính trên tất cả các công cụ và theo từng công cụ.
• tools/{tool_name}/call_count, tools/{tool_name}/failure_count: Số lượng các cuộc gọi và các cuộc gọi thất bại cho mỗi công cụ.
• tools/unknown_name_count: Số lần gọi một công cụ không tồn tại. Khác với một công cụ đã chạy và gây ra lỗi, đây là một lỗi chính sách.
• tools/parallel_calls_mean: Số lượng trung bình các lần gọi công cụ trong một thông điệp của trợ lý có gọi công cụ.
• tools/loop_exhausted_frac: Tỷ lệ các rollout bị cắt ở max_tool_calling_iterations trong khi vẫn đang gọi công cụ, trên tổng số mọi rollout, bao gồm cả các rollout không có lần gọi nào và các nhóm bị loại bỏ sau đó. Chỉ được ghi log khi trainer có công cụ.

HarnessAsyncGRPOTrainer ghi log hầu hết các mục này dưới cùng tên gọi. Xem Training on harnesses để biết các khác biệt.

Chỉ số môi trường

• env/reset_s: Thời gian dành cho reset() của môi trường trước khi một rollout bắt đầu. Chỉ được ghi log khi environment_factory được cung cấp.
• env/reward_s: Thời gian dành cho get_reward() của môi trường sau khi một rollout kết thúc. Chỉ được ghi log khi môi trường định nghĩa get_reward.

Chỉ số mẫu

• sample/forwarded_tokens_mean, sample/forwarded_tokens_max: Số lượng trung bình và tối đa của các token được chuyển tiếp trên mỗi mẫu.
• sample/trained_tokens_mean: Số lượng token được huấn luyện trung bình trên mỗi mẫu.
• sample/time_in_queue_s: Thời gian một mẫu chờ trong hàng đợi rollout trước khi được huấn luyện.
• sample/staleness_mean, sample/staleness_max: Số lần cập nhật trọng số giữa chính sách tạo ra mẫu và chính sách đang huấn luyện trên mẫu đó. Các chỉ số sampling/importance_sampling_ratio/* và sampling/sampling_logp_difference/* kết hợp độ cũ (staleness) này với sự không khớp giữa huấn luyện và suy luận: một tỷ lệ lệch khỏi 1.0 hoặc một chênh lệch log-xác suất lệch khỏi 0.0 ở mức độ cũ bằng không cho thấy sự không khớp, chẳng hạn như sự khác biệt về dtype với máy chủ.
• sample/dropped_stale_count: Số lượng mẫu bị loại do vượt quá max_staleness.

Chỉ số theo lô

• batch/forwarded_tokens_per_step, batch/trained_tokens_per_step: Số lượng token được forward và train trong mỗi bước, trên tất cả các tiến trình.
• batch/samples_per_step, batch/groups_per_step: Số lượng mẫu và số lượng nhóm riêng biệt trong mỗi bước. Chúng khác nhau bởi số lượng mẫu trong mỗi nhóm.
• batch/microbatches_per_step: Số lượng micro-batch cho mỗi bước.
• batch/masked_token_frac: Tỷ lệ các token được chuyển tiếp nhưng không được huấn luyện.
• batch/samples_per_row, batch/row_tokens_mean, batch/row_tokens_max: Số lượng mẫu trên mỗi hàng, cùng với số lượng token được chuyển tiếp trung bình và tối đa trên mỗi hàng.
• batch/row_fill_frac: Tỷ lệ token_budget mà các hàng sử dụng.
• batch/row_imbalance: Tỷ lệ giữa chi phí attention của hàng đắt nhất so với chi phí trung bình. Tiến trình có hàng đắt nhất sẽ làm các tiến trình khác bị dừng lại tại bước all-reduce gradient; 1.0 là cân bằng hoàn hảo.
• batch/packing_ratio: Số lượng token trước khi đóng gói trên mỗi token được truyền đi: hệ số mà đóng gói làm thu nhỏ bước truyền đi, 1.0 khi không có gì được chia sẻ.
• batch/pad_frac: Tỷ lệ padding được thêm vào để các hàng có cùng độ dài. Nó chỉ tốn chi phí truyền tải và sẽ được loại bỏ trước khi thực hiện forward.
• batch/dropped_oversize_count: Số lượng mẫu bị loại do vượt quá token_budget.

Các chỉ số hiệu năng

• perf/step_s: Thời gian thực tế của một bước, bao gồm optimizer, đồng bộ trọng số và thời gian chờ mẫu.
• perf/fwd_bwd_s, perf/fwd_s: Thời gian dành cho các bước truyền thuận và truyền ngược của một bước, và chỉ riêng các bước truyền thuận.
• perf/optimizer_s: Thời gian dành cho bước tối ưu hóa.
• perf/rollout_wait_s: Thời gian huấn luyện phải chờ do hàng đợi rollout trống. Xem Hàng đợi rollout.
• perf/weight_sync_s: Thời gian của một lần đồng bộ trọng số, được chia thành perf/weight_sync_pause_s (chờ vLLM tạm dừng), perf/weight_sync_barrier_s (chờ các tiến trình khác) và perf/weight_sync_transfer_s (gửi các trọng số).
• perf/fwd_bwd_forwarded_tokens_per_s, perf/wall_clock_forwarded_tokens_per_s: Số lượng token được forward mỗi giây, tính trên các bước forward và backward hoặc trên toàn bộ bước.
• perf/wall_clock_trained_tokens_per_s: Số lượng token được huấn luyện mỗi giây, tính trên toàn bộ bước.
• perf/fwd_bwd_mfu, perf/wall_clock_mfu: Mức sử dụng FLOPs của mô hình trên các thiết bị huấn luyện, tính trên các bước truyền thuận và truyền ngược hoặc trên toàn bộ bước. Chúng loại trừ máy chủ vLLM. Chỉ được ghi log khi công suất tính toán đỉnh của thiết bị huấn luyện ở độ chính xác huấn luyện được biết đến, và khi ước tính FLOPs hỗ trợ kiến trúc của mô hình (chẳng hạn như Llama, Qwen2, Qwen3, Qwen3-MoE hoặc Qwen3-VL).

Tùy chỉnh

Cấu hình máy chủ vLLM

Trainer tạo ra bằng một máy chủ vLLM, được truy cập tại vllm_server_base_url (mặc định là http://localhost:8000):

VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --api-server-count 8 \
    --max-model-len 4096

• VLLM_SERVER_DEV_MODE=1 và --weight-transfer-config mở các endpoint mà trình huấn luyện dùng để gửi trọng số của nó.
• --logprobs-mode processed_logprobs khiến server trả về log-probabilities của phân phối mà nó đã lấy mẫu, sau khi áp dụng temperature và top-k/top-p, đây là giá trị mà importance ratio so sánh.
• --api-server-count 8 giữ lại các log-probabilities hàng đầu để các yêu cầu score centering có chi phí phục vụ thấp (xem cảnh báo bên dưới).
• --max-model-len là tổng độ dài chuỗi tối đa (prompt và completion) mà bạn dự kiến. Giá trị thấp hơn sẽ giảm mức sử dụng bộ nhớ GPU trên server, giải phóng thêm bộ nhớ cho KV cache và tăng thông lượng. Điểm khởi đầu tốt là độ dài prompt cộng với max_completion_length từ config của bạn. Khi có tools, mỗi lượt thêm tối đa max_completion_length tokens cùng với kết quả tool, vì vậy hãy đặt nó bằng cuộc trò chuyện dài nhất mà bạn dự kiến. Nếu prompt của một lượt cộng với max_completion_length dài hơn --max-model-len, server sẽ từ chối yêu cầu và quá trình chạy sẽ dừng lại với lỗi từ server. Với top_k hoặc top_p, server sẽ dừng lượt đó khi đạt đến --max-model-len, và chỉ từ chối một prompt vốn đã dài bằng mức đó. Đây cũng là giá trị mặc định của token_budget.

Hãy đảm bảo rằng server đang sử dụng các GPU khác với trainer, nếu không bạn có thể gặp lỗi NCCL. Bạn có thể chỉ định các GPU cần sử dụng bằng biến môi trường CUDA_VISIBLE_DEVICES.

Việc serialize các log-probability hàng đầu là công việc trên CPU trong API server của vLLM, và với cấu hình mặc định là một API server duy nhất, num_sampler_logprobs=8 làm giảm một nửa thông lượng sinh. Khởi động server với --api-server-count 8 trên vLLM 0.32 trở lên: đo trên một H100 với các mô hình từ 1.5B đến 7B, các log-probability hàng đầu (top-8) khi đó chỉ tốn dưới 3% thông lượng sinh. Trên vLLM 0.31, cùng cờ này vẫn để lại mức tổn thất 20 đến 40%, đã được sửa ở upstream bởi vllm-project/vllm#58181.

Sự đồng thời

Ba giới hạn quyết định số lượng rollout chạy đồng thời. Hai giới hạn đầu tiên tự động điều chỉnh; giới hạn thứ ba do bạn thiết lập.

Độ cũ (Staleness). max_staleness xác định số bước huấn luyện mà quá trình sinh có thể chạy trước trainer. Worker theo dõi điều này theo token, vì số lượng mẫu trong một bước thay đổi khi các completion dài hơn hoặc ngắn hơn. Sau mỗi bước, trainer đếm số token đã sử dụng và cho phép worker giữ tối đa max(max_staleness, 1) lần số token đó trong hàng đợi hoặc đang xử lý. Trước bước đầu tiên, nó giả định rằng mỗi hàng đều đầy. Worker chỉ bắt đầu một nhóm rollout mới khi vẫn dưới ngưỡng đó. Ngay từ đầu, nó chưa biết một rollout tạo ra bao nhiêu token, nên nó bắt đầu từng nhóm một cho đến khi nhóm đầu tiên hoàn thành. Một mẫu mà cuối cùng vẫn cũ hơn max_staleness bước sẽ bị loại bỏ.

Tải vLLM. Gửi quá nhiều yêu cầu cho vLLM cùng lúc sẽ làm nó chậm hơn: bộ nhớ đệm KV bị đầy và vLLM phải tạm dừng rồi làm lại một số yêu cầu. Vì vậy, worker cũng duy trì một giới hạn trực tiếp về số lượng rollout chạy đồng thời, được ghi log dưới dạng rollout/concurrency_cap, và kiểm tra /metrics của máy chủ mỗi 5 giây để điều chỉnh. Giới hạn bắt đầu bằng số yêu cầu vừa đủ trong bộ nhớ đệm KV ở độ dài tối đa (KV cache tokens / max_model_len), để máy chủ không bị quá tải ngay từ đầu. Trong khi mức sử dụng bộ nhớ đệm KV duy trì dưới 60% và không có yêu cầu nào đang chờ, giới hạn tăng lên 25%, tối đa một lần cho mỗi lượng rollout hoàn thành tương đương một cap. Khi mức sử dụng vượt quá 80%, giới hạn giảm xuống 30%: các yêu cầu đang chạy tiếp tục, và các yêu cầu mới sẽ chờ. Nếu máy chủ không có endpoint /metrics, giới hạn sẽ là max_concurrent_rollouts.

Tài nguyên. max_concurrent_rollouts (mặc định là 512) là giới hạn trên về số lượng rollout chạy đồng thời. Đây là giới hạn duy nhất do bạn tự thiết lập, vì trình huấn luyện không thể biết môi trường và công cụ của bạn có thể xử lý được bao nhiêu. Mỗi rollout đang chạy sẽ sử dụng một môi trường, có thể thực hiện các lời gọi công cụ đồng bộ trên một nhóm luồng có kích thước bằng giá trị này, và giữ một kết nối mở với máy chủ vLLM. Hãy đặt giá trị này theo khả năng xử lý của các thành phần đó, ví dụ như số phiên mà máy chủ môi trường của bạn chấp nhận:

from trl import AsyncGRPOConfig

training_args = AsyncGRPOConfig(
    ...,
    max_concurrent_rollouts=64,
)

Nếu giá trị quá thấp, quá trình sinh sẽ chậm lại. Một rollout giữ nguyên slot của nó trong khi chạy công cụ hoặc chờ giữa các lượt, không chỉ trong khi vLLM sinh, do đó các rollout gọi nhiều công cụ cần nhiều slot hơn để giữ cho vLLM bận. Nếu giá trị cao hơn nhiều so với mặc định, máy chủ vLLM có thể bắt đầu ngắt kết nối.

Các chỉ số được ghi lại cho thấy giới hạn nào đang kìm hãm quá trình tạo:

• rollout/backpressure_s lớn hơn 0: giới hạn độ cũ. Việc sinh đang đi trước huấn luyện, điều này là dự kiến khi suy luận nhanh hơn huấn luyện.
• rollout/concurrency tại rollout/concurrency_cap, thấp hơn max_concurrent_rollouts: tải của máy chủ vLLM.
• rollout/concurrency ở mức max_concurrent_rollouts, với rollout/vllm_kv_cache_usage_frac thấp và perf/rollout_wait_s cao: đây là giới hạn trên. Hãy tăng max_concurrent_rollouts nếu môi trường và công cụ của bạn cho phép.

Hàng đợi triển khai

Worker đẩy các mẫu đã chấm điểm vào hàng đợi và trainer lấy từ đó. Bốn chỉ số mô tả hàng đợi đó:

• rollout/queue_size: số lượng mẫu đang chờ.
• sample/time_in_queue_s: thời gian một mẫu đơn lẻ chờ trước khi được huấn luyện, phần thời gian của độ cũ (staleness) của nó.
• perf/rollout_wait_s: thời gian huấn luyện phải chờ do hàng đợi trống.
• rollout/backpressure_s: thời gian chờ của quá trình sinh do nó đi trước quá trình huấn luyện max_staleness bước.

Hai chỉ số cuối cùng là hình ảnh phản chiếu của nhau. Khi đọc cùng với kích thước hàng đợi, chúng cho bạn biết phía nào là nút thắt cổ chai:

• Hàng đợi gần như trống, perf/rollout_wait_s cao: giới hạn bởi quá trình sinh. Trình huấn luyện đang thiếu dữ liệu; hãy xem rollout/generated_tokens_per_s, rollout/concurrency và rollout/score_queue_size.
• hàng đợi chứa khoảng max_staleness bước token, rollout/backpressure_s cao: trainer-bound. Việc sinh bị giữ lại, và sample/staleness_mean nằm gần giới hạn của nó.
• cả hai đều cao: chính giới hạn max_staleness là quá chặt. Các rollout đang chạy sẽ lấp đầy giới hạn này trước khi mẫu của chúng đến hàng đợi, khiến trình huấn luyện phải chờ trong khi quá trình sinh bị giữ lại. Hãy tăng max_staleness.
• cả hai đều gần bằng không: cân bằng.

Bão hòa và các bản sao

Chính sách học từ sự khác biệt giữa các rollout của một nhóm: mỗi lợi thế (advantage) được tính so với phần thưởng trung bình của nhóm. Khi mọi rollout trong một nhóm nhận cùng một phần thưởng, tất cả đều có lợi thế bằng không, do đó nhóm đó bị loại bỏ (xem Các chỉ số được ghi log) và không dạy cho chính sách bất kỳ điều gì. Khi một lượt chạy bão hòa, ngày càng nhiều nhóm rơi vào tình trạng này, hoặc vì mọi rollout đều giải quyết được prompt, hoặc vì không rollout nào giải quyết được. Bốn chỉ số giúp phân biệt các trường hợp này, và mỗi chỉ số trả lời một câu hỏi riêng, nên bạn có thể đọc chúng trực tiếp từ biểu đồ:

Đọc ba điều đầu tiên cùng nhau, đối chiếu với phạm vi phần thưởng của bạn:

completions/duplicate_frac là tỷ lệ các rollout lặp lại một rollout trước đó của nhóm chúng, khớp từng token, trên tất cả các lượt của chúng. Các rollout giống hệt nhau nhận phần thưởng giống hệt nhau, do đó các bản sao cũng đẩy frac_zero_std_groups lên cao.

Các chỉ số này được ghi lại khi worker chấm điểm một nhóm, chứ không phải khi các mẫu của nhóm đó được huấn luyện, do đó chúng chạy hơi nhanh hơn quá trình huấn luyện và bao gồm cả những nhóm mà các mẫu sau đó bị loại do vượt quá max_staleness. Một nhóm mà không có rollout nào được chấm điểm, vì mọi hàm phần thưởng đều trả về None, sẽ bị loại và được tính vào frac_zero_std_groups, nhưng vì nhóm này không có phần thưởng nên reward_zero_std_groups sẽ không bao gồm nó.

Đóng gói cây

Các mẫu của một micro-batch được đóng gói vào một hàng cho mỗi tiến trình, không có padding. Một hàng là một prefix trie, không phải một phép nối: một token được chia sẻ bởi nhiều mẫu chỉ được truyền đi một lần, và attention vẫn hiển thị cho mỗi mẫu chính xác lịch sử của riêng nó. Các rollout của một nhóm đều lặp lại prompt của nhóm đó, và mỗi lượt của một rollout nhiều lượt lặp lại cuộc trò chuyện cho đến thời điểm đó, do đó phép truyền đi (forward) được thu nhỏ theo tỷ lệ đóng gói (batch/packing_ratio). Đây chỉ là một thay đổi về bố cục: mỗi token được huấn luyện giữ nguyên target, advantage và old log-probability của riêng nó, và log-probability của nó khớp với việc truyền đi mẫu đó một cách độc lập.

Với hai lần chạy trên một prompt, trong đó p là prompt và a, b là hai phần hoàn thành:

1. Rollouts         r0:  p1 p2 p3 a1 a2                10 tokens
                    r1:  p1 p2 p3 b1 b2

2. Trie             S0 [p1 p2] ─┬─ S1 [p3 a1 a2]       p3 is not shared: a position predicts
                                └─ S2 [p3 b1 b2]       one token, and p3 predicts a1 or b1

3. Packed row       index         0  1 |  2  3  4 |  5  6  7
   (depth-first)    input_ids    p1 p2 | p3 a1 a2 | p3 b1 b2    8 tokens: everything but attention
                    position_ids  0  1 |  2  3  4 |  2  3  4    runs on 8 instead of 10
                    segment      S0    | S1       | S2

4. Attention        segment   queries    keys and values: the segment's path in the trie
   (one call)       S0        p1 p2      p1 p2              index 0 1
                    S1        p3 a1 a2   p1 p2 p3 a1 a2     index 0 1 2 3 4
                    S2        p3 b1 b2   p1 p2 p3 b1 b2     index 0 1 5 6 7

5. Loss             r0: a1 a2 predicted at index 2 3   each trained token is predicted at the index
                    r1: b1 b2 predicted at index 5 6   of the token before it in its rollout

Một vài hệ quả:

• Mọi tiền tố chung đều được chia sẻ: các rollout của một prompt, và giữa các prompt là một system prompt hoặc tool schema chung. Một sample sẽ được đưa vào hàng đã chứa nhóm của nó, nơi prompt của nó đã có sẵn. Khi hàng đó đầy, sample sẽ được đưa sang hàng khác nếu một process vẫn còn một hàng trống, hoặc nếu điều đó không làm tăng tải của hàng có tải cao nhất (hàng này quyết định thời gian của micro-batch); nếu không, micro-batch sẽ đóng lại.
• Không chia sẻ gì là đóng gói thuần túy: mỗi mẫu sau đó sẽ là một phân đoạn riêng, tức là đóng gói không đệm thông thường.
• Chỉ K/V bị nhân bản: một tiền tố chung được thu thập một lần cho mỗi phân đoạn con, tương đương với những gì huấn luyện không đóng gói sẽ giữ lại, và khá nhỏ khi sử dụng chú ý truy vấn nhóm.
• Các hybrid Gated DeltaNet được hỗ trợ (Qwen3.5, Qwen3.6, Qwen3-Next, …): phần convolution của chúng đọc các token tiền nhiệm của chính mỗi token, và quy tắc delta dạng chunk của chúng khởi tạo mỗi đoạn từ trạng thái cuối cùng của đoạn cha.

token_budget giới hạn số token của một hàng, trong đó một token dùng chung chỉ được tính một lần. Mặc định, giá trị này là max_model_len của server, nên không mẫu nào có thể vượt quá nó. Khi không có gì dùng chung, tree packing chạy với tốc độ xấp xỉ bằng packing không padding thông thường, và mức tăng tốc tăng xấp xỉ theo tỷ lệ packing (khoảng 2,8 đến 3 lần ở tỷ lệ 3,4 trên Qwen3-0.6B và 1.7B, trên H100). Tỷ lệ này bằng tiền tố dùng chung nhân với số mẫu trên mỗi hàng, nên batch/samples_per_row là chỉ báo chính: khi gần bằng 1, không có gì để dùng chung, và đòn bẩy là token_budget. batch/row_fill_frac quan trọng khi các mẫu dài: một mẫu 10k token lấp đầy ngân sách 32k kém (ba mẫu vừa, bốn mẫu thì không bao giờ vừa, nên planner thường chỉ nhận được hai và hàng chạy ở mức khoảng 77% đầy), trong khi các mẫu 1k token lấp đầy nó gần như hoàn hảo.

Đánh giá việc đóng gói cây (tree packing) dựa trên thông lượng, batch/trained_tokens_per_step / perf/fwd_bwd_s. Một micro-batch đã đóng gói chứa nhiều mẫu hơn, do đó perf/fwd_bwd_s tăng lên và perf/fwd_bwd_mfu, vốn được chia cho số token đã truyền, giảm xuống: cả hai đều không phải là sự thụt lùi. Nhiều mẫu hơn trên mỗi hàng cũng có nghĩa là nhiều rollout được lấy ra hơn ở mỗi bước, khiến một lần chạy bị giới hạn bởi quá trình sinh (generation-bound) trở nên bị giới hạn hơn nữa.

Các chỉ số theo batch được tổng hợp như sau: một bước giữ gradient_accumulation_steps × num_processes dòng, do đó batch/samples_per_step xấp xỉ bằng số lần đó nhân với batch/samples_per_row, và tương tự cho số token. Hai phía chỉ khớp nhau trong phạm vi biến động giữa các micro-batch của một bước, vì các chỉ số theo bước là tổng còn các chỉ số theo dòng là trung bình.

Bất kỳ yếu tố nào làm cho các hàng dài hơn hoặc ít hơn sẽ loại bỏ việc chia sẻ. Một cuộc hội thoại nhiều lượt không bao giờ phân nhánh là một mẫu dài duy nhất, và rollout/samples_per_rollout cho biết một cuộc hội thoại tạo ra bao nhiêu mẫu.

Mẫu top-k và top-p

Top-k và top-p chuẩn hóa lại phân phối lấy mẫu trên một tập ứng viên bị hạn chế thay vì toàn bộ bảng từ vựng. Log-probabilities của máy chủ phản ánh điều đó (--logprobs-mode processed_logprobs), do đó trainer phải chuẩn hóa trên cùng một tập, nếu không mọi tỷ lệ quan trọng đều bị thiên lệch: xác suất của một token trên toàn bộ bảng từ vựng thấp hơn so với trên tập ứng viên mà nó được lấy ra (DeepSeek-V3.2, mục 3.1). Với top_p < 1 hoặc top_k > 0, worker sẽ yêu cầu máy chủ cung cấp tập ứng viên của mỗi token được lấy mẫu và trainer sẽ chuẩn hóa lại log-probabilities của nó trên tập đó. Điều này cần thêm hai cờ máy chủ, và top_k phải được đặt, vì vLLM giới hạn kích thước của tập ứng viên bằng tham số này:

VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --api-server-count 8 \
    --max-model-len 4096 \
    --return-sampling-mask \
    --enable-scale-out

from trl import AsyncGRPOConfig

training_args = AsyncGRPOConfig(
    ...,
    top_p=0.95,
    top_k=512,
)

--return-sampling-mask khiến vLLM ghi lại tập ứng viên (vLLM 0.28 trở lên, trên trình chạy mô hình V2 của nó: VLLM_USE_V2_MODEL_RUNNER=1 nếu máy chủ từ chối cờ này), và --enable-scale-out phơi bày điểm cuối token-in-token-out để trả về tập đó. Worker kiểm tra cả hai trên máy chủ khi nó khởi động. Một phiên bản vLLM cũ hơn không thể trả về tập này, do đó trainer tiếp tục chuẩn hóa trên toàn bộ bảng mã và cảnh báo về tỷ lệ có thiên lệch.

Huấn luyện adapter với PEFT

Chúng tôi hỗ trợ tích hợp chặt chẽ với thư viện 🤗 PEFT, cho phép bạn huấn luyện các adapter và chia sẻ chúng trên Hub thay vì huấn luyện toàn bộ mô hình.

from datasets import load_dataset
from peft import LoraConfig
from trl import AsyncGRPOTrainer
from trl.rewards import accuracy_reward

if __name__ == "__main__":
    dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

    trainer = AsyncGRPOTrainer(
        model="Qwen/Qwen3.5-2B",
        reward_funcs=accuracy_reward,
        train_dataset=dataset,
        peft_config=LoraConfig(r=32, lora_alpha=64, target_modules="all-linear"),
    )
    trainer.train()

Mọi lần đồng bộ trọng số đều phải gửi chính sách đã cập nhật đến máy chủ. Mặc định, trainer sẽ gộp adapter vào mô hình gốc và gửi toàn bộ trọng số đã gộp, tương đương khoảng 15 GB cho một mô hình 7B ở mỗi lần đồng bộ. Khởi động máy chủ vLLM với --enable-lora và nó sẽ chỉ gửi adapter, khoảng 160 MB ở r=32:

VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 vllm serve <model_name> \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --api-server-count 8 \
    --max-model-len 4096 \
    --enable-lora \
    --max-lora-rank 32 \
    --max-loras 6

VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 mở endpoint mà trainer gửi mỗi adapter mới tới. Giữ lại --weight-transfer-config ngay cả khi dùng --enable-lora: trainer chỉ chọn chế độ đồng bộ khi nó khởi động, lúc đó server đã chạy sẵn, và đồng bộ gộp là phương án dự phòng.

--max-lora-rank phải là một trong các giá trị 1, 8, 16, 32, 64, 128, 256, 320, 512. Tham số này đặt rank cao nhất mà server có thể phục vụ, chứ không phải rank mà server sẽ phục vụ, do đó adapter r=4 hoạt động bình thường dưới 8. --max-loras phải ít nhất bằng max_staleness + 2: trainer giữ max_staleness + 1 phiên bản adapter được đăng ký để một rollout bắt đầu dưới một policy cũ hơn có thể hoàn tất dưới policy đó thay vì chuyển đổi policy giữa chừng quá trình sinh, và mỗi lần đồng bộ sẽ nạp phiên bản tiếp theo trước khi gỡ phiên bản cũ nhất. Trainer kiểm tra cả hai giá trị khi khởi động và cho bạn biết cần khởi động lại server với giá trị nào.

vLLM nạp adapter từ đĩa thay vì qua mạng, do đó trainer ghi mỗi phiên bản dưới <output_dir>/.vllm_lora/ và cung cấp đường dẫn cho server. Nếu server chạy trên một máy khác, thư mục đó phải nằm trên một hệ thống tệp mà cả hai đều có thể truy cập.

<output_dir>/.vllm_lora/ là bộ nhớ đệm phục vụ, không phải nơi lưu adapter đã huấn luyện. Thư mục này chỉ giữ các phiên bản mới nhất và mỗi phiên bản sẽ bị xóa ngay sau khi hết cửa sổ thời gian hợp lệ. Checkpoint mới là thứ bảo toàn adapter bạn đã huấn luyện, vì vậy đừng tắt tính năng này: với save_strategy="no", một phiên chạy kéo dài nhiều giờ sẽ kết thúc mà không để lại gì. output_dir cũng phải có thể đọc được bởi server và phải tồn tại lâu hơn thời gian chạy của job, nên không thể dùng vùng lưu tạm cục bộ trên node, dù cách đó sẽ tránh được việc đồng bộ hóa trên hệ thống tệp chia sẻ chậm.

Adapter sync không phải lúc nào cũng nhanh hơn. vLLM tạo ra token chậm hơn khoảng 1,33 đến 1,39 lần khi phục vụ một adapter LoRA so với khi phục vụ các trọng số đã hợp nhất, và quá trình tạo ra thường là yếu tố chi phối thời gian thực tế trong RL. Đổi lại, bạn có được thời gian tạm dừng ngắn hơn đáng kể ở mỗi lần đồng bộ và khoảng 1% băng thông đồng bộ. Nếu thông lượng tạo ra là nút thắt của bạn, đồng bộ hợp nhất thường vẫn là lựa chọn tốt hơn.

Sử dụng hàm phần thưởng tùy chỉnh

AsyncGRPOTrainer hỗ trợ sử dụng các hàm thưởng tùy chỉnh. Để đảm bảo tính tương thích, hàm thưởng của bạn phải đáp ứng các yêu cầu sau:

Các hàm reward có thể là các callable Python đồng bộ hoặc các coroutine async def bất đồng bộ. Khi bạn cung cấp nhiều hàm reward, chúng sẽ chạy song song: các hàm bất đồng bộ được await cùng nhau và các hàm đồng bộ chạy trong các thread, do đó độ trễ của chúng sẽ chồng lên nhau.

1. Các đối số đầu vào:

   • Hàm phải nhận các đối số sau dưới dạng đối số từ khóa:

     • prompts (chứa các prompts),
     • completions (chứa các phần hoàn thành đã tạo),
     • completion_ids (chứa các completion đã được token hóa),
     • Tất cả tên cột (ngoại trừ prompt) mà bộ dữ liệu có thể có. Ví dụ, nếu bộ dữ liệu chứa một cột có tên ground_truth, hàm sẽ được gọi với ground_truth như một đối số từ khóa.

     Cách dễ nhất để đáp ứng yêu cầu này là sử dụng **kwargs trong chữ ký hàm.

   • Bộ dữ liệu ở định dạng hội thoại, do đó prompts và completions là các danh sách từ điển tin nhắn.

2. Giá trị trả về: Hàm phải trả về một danh sách các số float. Mỗi số float đại diện cho phần thưởng tương ứng với một kết quả hoàn thành duy nhất.

Các hàm reward được thực thi trong rollout worker, một tiến trình riêng biệt, nên mọi thứ được truyền vào nó đều phải pickled. Mỗi hàm reward, công cụ và environment_factory (cùng với bất kỳ thứ gì chúng đóng gói) đều phải có thể picklable: hãy sử dụng một hàm ở cấp module, functools.partial, hoặc một thực thể lớp callable. Các lambda và closure sẽ gây ra TypeError khi gọi trainer.train(). Đây là điểm khác biệt so với GRPOTrainer, nơi các hàm reward được gọi trong tiến trình và closure hoạt động bình thường.

Quy trình rollout cũng được chạy với CUDA_VISIBLE_DEVICES="", do đó không thể sử dụng GPU. Một mô hình phần thưởng dựa trên GPU (ví dụ: bộ chấm điểm AutoModelForSequenceClassification) sẽ được nạp mà không có lỗi nhưng âm thầm chuyển sang chạy trên CPU. Hãy giữ các hàm phần thưởng ở mức nhẹ (các bộ xác minh như accuracy_reward, kiểm tra định dạng hoặc độ dài). Nếu bạn cần một mô hình phần thưởng trên GPU, hãy triển khai nó sau một engine suy luận riêng (vLLM, TGI, …) trên các GPU riêng biệt và để một hàm phần thưởng có thể picklable gọi nó qua HTTP.

Ví dụ 1: Phần thưởng cho các câu trả lời dài hơn

Dưới đây là một ví dụ về hàm phần thưởng (reward function) thưởng cho các câu trả lời dài hơn:

def reward_func(completion_ids, **kwargs):
    """Reward function that assigns higher scores to longer completions (in terms of token count)."""
    return [float(len(ids)) for ids in completion_ids]

Bạn có thể kiểm tra nó như sau:

>>> prompts = [[{"role": "user", "content": "The sky is"}], [{"role": "user", "content": "The sun is"}]]  # not used in the reward function, but the trainer will pass it
>>> completions = [[{"role": "assistant", "content": " blue."}], [{"role": "assistant", "content": " in the sky."}]]  # not used in the reward function, but the trainer will pass it
>>> completion_ids = [[6303, 13], [304, 279, 12884, 13]]
>>> reward_func(prompts=prompts, completions=completions, completion_ids=completion_ids)
[2.0, 4.0]

Ví dụ 2: Thưởng cho các phần hoàn thành có định dạng cụ thể

Dưới đây là một ví dụ về hàm phần thưởng kiểm tra xem phần hoàn thành có đúng định dạng cụ thể hay không. Ví dụ này được lấy cảm hứng từ hàm phần thưởng định dạng được sử dụng trong bài báo DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. Nó được thiết kế cho định dạng hội thoại, trong đó các prompt và phần hoàn thành bao gồm các tin nhắn có cấu trúc.

import re

def format_reward_func(completions, **kwargs):
    """Reward function that checks if the completion has a specific format."""
    pattern = r"^<think>.*?</think><answer>.*?</answer>$"
    completion_contents = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, content) for content in completion_contents]
    return [1.0 if match else 0.0 for match in matches]

Bạn có thể kiểm tra hàm này như sau:

>>> prompts = [
...     [{"role": "user", "content": "What is the result of (1 + 2) * 4?"}],
...     [{"role": "user", "content": "What is the result of (3 + 1) * 2?"}],
... ]
>>> completions = [
...     [{"role": "assistant", "content": "<think>The sum of 1 and 2 is 3, which we multiply by 4 to get 12.</think><answer>(1 + 2) * 4 = 12</answer>"}],
...     [{"role": "assistant", "content": "The sum of 3 and 1 is 4, which we multiply by 2 to get 8. So (3 + 1) * 2 = 8."}],
... ]
>>> format_reward_func(prompts=prompts, completions=completions)
[1.0, 0.0]

Ví dụ 3: Hoàn thành phần thưởng dựa trên một tham chiếu

Dưới đây là một ví dụ về hàm phần thưởng kiểm tra xem câu trả liệu có chính xác hay không. Ví dụ này được lấy cảm hứng từ hàm phần thưởng độ chính xác được sử dụng trong bài báo DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. Bộ dữ liệu chứa một cột có tên là ground_truth.

import re

def reward_func(completions, ground_truth, **kwargs):
    # Regular expression to capture content inside \boxed{}
    matches = [re.search(r"\\boxed\{(.*?)\}", completion[0]["content"]) for completion in completions]
    contents = [match.group(1) if match else "" for match in matches]
    # Reward 1 if the content is the same as the ground truth, 0 otherwise
    return [1.0 if c == gt else 0.0 for c, gt in zip(contents, ground_truth)]

Bạn có thể kiểm tra hàm này như sau:

>>> prompts = [
...     [{"role": "user", "content": "Solve the equation $2x + 3 = 7$."}],
...     [{"role": "user", "content": "Solve the equation $3x - 5 = 10$."}],
... ]
>>> completions = [
...     [{"role": "assistant", "content": r"The solution is \boxed{2}."}],
...     [{"role": "assistant", "content": r"The solution is \boxed{6}."}],
... ]
>>> ground_truth = ["2", "5"]
>>> reward_func(prompts=prompts, completions=completions, ground_truth=ground_truth)
[1.0, 0.0]

Ví dụ 4: Các hàm phần thưởng đa nhiệm

Dưới đây là một ví dụ về việc sử dụng nhiều hàm phần thưởng trong AsyncGRPOTrainer. Trong ví dụ này, chúng ta định nghĩa hai hàm phần thưởng dành riêng cho từng tác vụ: math_reward_func và coding_reward_func. Hàm math_reward_func cấp phần thưởng cho các bài toán dựa trên tính chính xác của chúng, trong khi hàm coding_reward_func cấp phần thưởng cho các bài toán lập trình dựa trên việc giải pháp có hoạt động được hay không.

from datasets import Dataset
from trl import AsyncGRPOTrainer

# Define a dataset that contains both math and coding problems
dataset = Dataset.from_list(
    [
        {"prompt": [{"role": "user", "content": "What is 2+2?"}], "task": "math"},
        {"prompt": [{"role": "user", "content": "Write a function that returns the sum of two numbers."}], "task": "code"},
        {"prompt": [{"role": "user", "content": "What is 3*4?"}], "task": "math"},
        {"prompt": [{"role": "user", "content": "Write a function that returns the product of two numbers."}], "task": "code"},
    ]
)

# Math-specific reward function
def math_reward_func(prompts, completions, task, **kwargs):
    rewards = []
    for prompt, completion, t in zip(prompts, completions, task):
        if t == "math":
            # Calculate math-specific reward
            correct = check_math_solution(prompt, completion)
            reward = 1.0 if correct else -1.0
            rewards.append(reward)
        else:
            # Return None for non-math tasks
            rewards.append(None)
    return rewards

# Coding-specific reward function
def coding_reward_func(prompts, completions, task, **kwargs):
    rewards = []
    for prompt, completion, t in zip(prompts, completions, task):
        if t == "code":
            # Calculate coding-specific reward
            works = test_code_solution(prompt, completion)
            reward = 1.0 if works else -1.0
            rewards.append(reward)
        else:
            # Return None for non-coding tasks
            rewards.append(None)
    return rewards

# Use both task-specific reward functions
if __name__ == "__main__":
    trainer = AsyncGRPOTrainer(
        model="Qwen/Qwen2.5-0.5B-Instruct",
        reward_funcs=[math_reward_func, coding_reward_func],
        train_dataset=dataset,
    )

    trainer.train()

Trong ví dụ này, math_reward_func và coding_reward_func được thiết kế để hoạt động với một tập dữ liệu hỗn hợp chứa cả các bài toán toán học và lập trình. Cột task trong tập dữ liệu được dùng để xác định hàm thưởng nào sẽ được áp dụng cho từng bài toán. Nếu không có hàm thưởng phù hợp cho một mẫu trong tập dữ liệu, hàm thưởng sẽ trả về None, và AsyncGRPOTrainer sẽ tiếp tục với các hàm và tác vụ hợp lệ. Điều này cho phép AsyncGRPOTrainer xử lý nhiều hàm thưởng với phạm vi áp dụng khác nhau.

Lưu ý rằng AsyncGRPOTrainer sẽ bỏ qua các phần thưởng None được trả về bởi các hàm phần thưởng và chỉ xem xét các phần thưởng được trả về bởi các hàm liên quan. Điều này đảm bảo rằng mô hình được huấn luyện trên các tác vụ liên quan và bỏ qua các tác vụ không có hàm phần thưởng liên quan.

Ví dụ 5: Các hàm phần thưởng bất đồng bộ

Các hàm phần thưởng tùy chỉnh cũng có thể được định nghĩa dưới dạng coroutine async def. Điều này hữu ích nếu phần thưởng của bạn phụ thuộc vào I/O chậm (ví dụ: gọi một dịch vụ từ xa). Khi bạn truyền nhiều hàm phần thưởng async, AsyncGRPOTrainer sẽ thực thi chúng đồng thời để độ trễ của chúng chồng lên nhau.

Dưới đây là một ví dụ tối giản về một hàm thưởng bất đồng bộ mô phỏng một thao tác bị giới hạn bởi I/O:

import asyncio

async def async_reward_func(prompts, completions, **kwargs):
    # Simulate an I/O-bound call (e.g., HTTP request, database lookup)
    await asyncio.sleep(0.01)
    # Simple toy reward: 1.0 if the completion is non-empty, else 0.0
    return [1.0 if completion else 0.0 for completion in completions]

Truyền hàm phần thưởng cho trình huấn luyện

Để sử dụng hàm phần thưởng tùy chỉnh của bạn, hãy truyền nó vào AsyncGRPOTrainer như sau:

from trl import AsyncGRPOTrainer

trainer = AsyncGRPOTrainer(
    reward_funcs=reward_func,
    ...,
)

Bạn có thể truyền nhiều hàm phần thưởng dưới dạng một danh sách; danh sách này có thể bao gồm cả các hàm đồng bộ và bất đồng bộ:

from trl import AsyncGRPOTrainer

trainer = AsyncGRPOTrainer(
    reward_funcs=[reward_func, async_reward_func1, async_reward_func2],
    ...,
)

và phần thưởng sẽ được tính là tổng các phần thưởng từ mỗi hàm.

Tùy chỉnh vòng lặp rollout

Việc tạo dữ liệu, tính phần thưởng, tính lợi thế và lọc được thực hiện trong quy trình rollout, trong AsyncRolloutLoop. Để thay đổi một trong các bước đó, hãy tạo lớp con của vòng lặp, ghi đè phương thức cho bước tương ứng, và đặt lớp con đó làm loop_cls trên một lớp con của trainer. Constructor của trainer không thay đổi.

Ví dụ, một chương trình đào tạo dừng lấy mẫu các prompt mà chính sách đã giải quyết được, với ngưỡng của nó nằm trong một lớp con của cấu hình:

from dataclasses import dataclass

import numpy as np

from trl import AsyncGRPOConfig, AsyncGRPOTrainer
from trl.trainer.async_grpo_trainer import AsyncRolloutLoop


@dataclass
class CurriculumConfig(AsyncGRPOConfig):
    pass_rate_cap: float = 0.9


class CurriculumRolloutLoop(AsyncRolloutLoop):
    def filter_prompt(self, row_index, row, epoch):
        return self.pass_rates.get(row_index, 0.0) < self.config.pass_rate_cap

    async def setup(self):
        self.pass_rates = {}

    def compute_advantages(self, group):
        scored = group.rewards[~np.isnan(group.rewards)]
        if len(scored):
            self.pass_rates[group.row_index] = float(np.mean(scored > 0))
        return super().compute_advantages(group)


class CurriculumTrainer(AsyncGRPOTrainer):
    loop_cls = CurriculumRolloutLoop

Rewards được cộng dồn trên reward_funcs, các môi trường định nghĩa get_reward, và reward do generate_one tùy chỉnh đặt trên Rollout của nó: để chấm điểm bằng judge, hãy truyền một hàm reward bất đồng bộ, hoặc chấm điểm bên trong generate_one. compute_advantages được chạy cho mọi nhóm, bao gồm cả nhóm sau đó bị loại bỏ vì tất cả các reward được chấm điểm của nó đều bằng nhau; việc loại bỏ dựa trên các reward, không phải trên các advantages. Vòng lặp được xây dựng trong một tiến trình được khởi tạo, vì vậy hãy định nghĩa lớp con ở cấp độ module và giữ mã huấn luyện bên dưới if __name__ == "__main__":. Hãy tạo các client trong setup, vốn được chạy trên event loop của vòng lặp, thay vì ghi đè __init__. Một prompt bị filter_prompt từ chối sẽ bị bỏ qua trong epoch đó và được tính vào epoch, giống như một nhóm bị loại bỏ. Nếu nó từ chối mọi prompt của một epoch, vòng lặp sẽ ném ngoại lệ.

Tiếp tục từ một checkpoint

Trainer cơ bản sẽ tiếp tục bằng cách bỏ qua và phát lại các batch đã được huấn luyện, điều này không áp dụng cho hàng đợi rollout trực tiếp, do đó ignore_data_skip luôn là True. Thay vào đó, chỉ số của prompt đầu tiên chưa được huấn luyện sẽ được lưu vào rollout_state.json cùng với mỗi checkpoint và được khôi phục khi tiếp tục, giúp worker tua nhanh đến prompt đó mà không cần phát lại các mẫu. Đây là vị trí đã huấn luyện, không phải vị trí của generator: worker chạy trước tiến độ huấn luyện một khoảng bằng độ sâu của hàng đợi rollout, và các mẫu được đệm đó sẽ bị mất khi phiên chạy kết thúc, do đó việc tiếp tục từ vị trí của generator sẽ bỏ qua các prompt đã được tạo ra nhưng chưa bao giờ được huấn luyện. Datasets dạng streaming (IterableDataset) không thể được định vị lại, do đó worker của chúng sẽ bắt đầu lại từ prompt đầu tiên khi tiếp tục.

Đào tạo Agent

Asynchronous GRPO hỗ trợ đào tạo agent: mô hình gọi công cụ trong quá trình sinh và học từ kết quả.

• Một công cụ là một hàm Python thuần túy (đồng bộ hoặc bất đồng bộ) được cung cấp cho mô hình. Sử dụng tools cho các cuộc gọi không trạng thái (máy tính, tìm kiếm web).
• Môi trường là dạng tổng quát hơn: một đối tượng có trạng thái được tạo mới cho mỗi lần rollout, trong đó các phương thức công khai được phơi bày dưới dạng công cụ, kèm theo hook vòng đời reset và get_reward tùy chọn cho phép nó tự quản lý phần thưởng. Sử dụng environment_factory khi bạn cần trạng thái cho từng lần rollout, hook reset, hoặc phần thưởng do môi trường quản lý.

Chúng có thể kết hợp với nhau — bạn có thể truyền các tools độc lập cùng với một environment_factory.

Công cụ

Tham số tools mong đợi một danh sách các hàm Python (đồng bộ hoặc bất đồng bộ) xác định các công cụ có sẵn cho tác tử:

from trl import AsyncGRPOTrainer

trainer = AsyncGRPOTrainer(
    tools=[tool1, tool2],
    ...,
)

Mỗi công cụ phải là một hàm Python chuẩn với tham số và kiểu trả về có chú thích kiểu (type-hint), cùng với một docstring theo phong cách Google mô tả mục đích, các tham số và giá trị trả về. Để biết thêm chi tiết, hãy xem hướng dẫn Truyền công cụ.

Mô hình gọi công cụ theo tên của nó, vì vậy mỗi công cụ được đăng ký dưới __name__ của nó và phải là một hàm ở cấp module, điều này cũng giúp nó có thể được pickle cho worker rollout. Các công cụ đồng bộ chạy trên một thread pool có kích thước bằng max_concurrent_rollouts, do đó một công cụ chậm sẽ không bao giờ chặn các rollout khác đang chạy. Trong một lượt của trợ lý, các lời gọi công cụ được thực hiện theo thứ tự mà mô hình đã phát ra. Do đó, tối đa max_concurrent_rollouts công cụ đồng bộ có thể chạy cùng lúc, vì vậy một công cụ đọc hoặc ghi trạng thái chia sẻ phải an toàn cho thread.

Vòng lặp gọi công cụ yêu cầu mẫu chat phải giữ nguyên phần đầu (việc thêm một thông điệp công cụ không được làm thay đổi cách các thông điệp trước đó được hiển thị). Đối với các họ model đã biết (ví dụ: Qwen3, DeepSeek-V3), TRL sẽ tự động thay thế bằng mẫu huấn luyện đã vá khi công cụ được bật. Xem Mẫu Chat để biết danh sách đầy đủ.

Sử dụng max_tool_calling_iterations trong AsyncGRPOConfig để giới hạn số lượt gọi công cụ. Mặc định không có giới hạn, và quá trình sinh sẽ dừng lại khi mô hình tạo ra lượt phản hồi không chứa bất kỳ lời gọi công cụ nào.

Ví dụ:

from trl import AsyncGRPOTrainer

def multiply(a: int, b: int) -> int:
    """
    Multiplies two integers.

    Args:
        a: The first integer.
        b: The second integer.

    Returns:
        The product of the two integers.
    """
    return a * b

async def async_add(a: int, b: int) -> int:
    """
    Asynchronously adds two integers.

    Args:
        a: The first integer.
        b: The second integer.

    Returns:
        The sum of the two integers.
    """
    return a + b

trainer = AsyncGRPOTrainer(
    tools=[multiply, async_add],
    ...,
)

Môi trường

Bạn cũng có thể cung cấp các công cụ thông qua environment_factory. Ở chế độ này, AsyncGRPOTrainer tạo một instance môi trường cho mỗi rollout và công khai các phương thức công khai của môi trường dưới dạng các công cụ.

Khi bạn sử dụng environment_factory, môi trường sẽ sở hữu dữ liệu: không có train_dataset bên ngoài. Ở mỗi lượt rollout, reset() tạo ra tác vụ — tự lấy mẫu từ một kho ngữ liệu mà môi trường giữ, hoặc tạo trạng thái theo thủ tục — và trả về prompt. max_steps đặt độ dài huấn luyện. Dưới đây là một ví dụ tối giản về một môi trường tự lấy mẫu một mục tiêu và phơi bày phương thức increment (được phơi bày dưới dạng công cụ).

import random

from trl import AsyncGRPOConfig, AsyncGRPOTrainer

class IncrementEnv:
    # Reserved methods
    def reset(self, **kwargs) -> str | None:  # required; called at the start of each rollout
        self.counter = 0
        self.target = random.randint(1, 6)  # self-sample the task
        return f"Increment the counter by {self.target}."  # returned string becomes the prompt

    def get_reward(self) -> float:  # optional: the environment scores itself from its own state
        return float(self.counter == self.target)

    # Public methods (exposed as tools)
    def increment(self, step: int) -> int:
        """
        Increment the internal counter.

        Args:
            step: Value to add to the counter.

        Returns:
            The updated counter value.
        """
        self.counter += step
        return self.counter

if __name__ == "__main__":
    trainer = AsyncGRPOTrainer(
        model="Qwen/Qwen3-0.6B",
        args=AsyncGRPOConfig(max_steps=1000, chat_template_kwargs={"enable_thinking": False}),
        environment_factory=IncrementEnv,
    )
    trainer.train()

Một lớp môi trường có hai phương thức được dành riêng: reset và get_reward. Khác với các phương thức công khai khác, những phương thức này không được phơi bày cho mô hình dưới dạng công cụ.

• reset (bắt buộc) được gọi khi rollout bắt đầu. Nó có thể trả về None hoặc một chuỗi. Khi trả về một chuỗi, chuỗi đó chính là prompt của người dùng.
• get_reward (tùy chọn, đồng bộ hoặc bất đồng bộ) không nhận đối số và trả về một float: môi trường chấm điểm cho episode vừa chạy dựa trên trạng thái nội bộ của chính nó (trò chơi có kết thúc bằng chiến thắng không? từ đã được đoán đúng chưa?). Hàm này được gọi một lần cho mỗi rollout hoàn thành và đóng vai trò là nguồn phần thưởng.

Cả hai đều chạy trên pool luồng công cụ của worker, do đó một môi trường chậm sẽ chỉ làm chậm rollout của chính nó mà không ảnh hưởng đến các rollout khác đang chạy. Thời lượng của chúng được ghi log dưới dạng env/reset_s và env/reward_s.

Vì môi trường tự lấy mẫu trên mỗi reset(), G thành viên của một nhóm GRPO có thể không chia sẻ cùng một trạng thái ban đầu, khiến đường cơ sở của nhóm hơi nhiễu hơn. Đối với các trạng thái giống nhau trong một nhóm, hãy để reset() suy ra trạng thái một cách tất định từ một khóa chung.

Cung cấp bộ dữ liệu bên ngoài (tùy chọn)

Thay vì để môi trường tự lấy mẫu, bạn có thể cung cấp một train_dataset bên ngoài. Ví dụ, khi các tác vụ của bạn đã nằm trong một dataset, hoặc để định tuyến giữa nhiều môi trường. Với mỗi rollout, trainer sẽ lấy mẫu một hàng, sử dụng cột "prompt" của hàng đó làm prompt, và truyền các cột của hàng vào reset() dưới dạng các tham số keyword (để môi trường đọc tác vụ từ dataset).

from datasets import Dataset
from trl import AsyncGRPOTrainer

# Each row carries the prompt and the task data (here, the counter target to reach).
dataset = Dataset.from_dict(
    {
        "prompt": [[{"role": "user", "content": f"Increment the counter to {i}."}] for i in range(1, 7)],
        "target": list(range(1, 7)),
    }
)

class IncrementEnv:
    def reset(self, target, **kwargs) -> None:  # the row's `target` column arrives as a keyword argument
        self.counter = 0
        self.target = target

    def get_reward(self) -> float:
        return float(self.counter == self.target)

    def increment(self, step: int) -> int:  # exposed as a tool (see the full docstring above)
        self.counter += step
        return self.counter

if __name__ == "__main__":
    trainer = AsyncGRPOTrainer(
        model="Qwen/Qwen3-0.6B",
        train_dataset=dataset,  # the dataset owns the task; the "target" column reaches reset(target=...)
        environment_factory=IncrementEnv,
    )
    trainer.train()

Phần thưởng

Rewards do môi trường và rewards do trainer sở hữu không loại trừ nhau. Truyền reward_funcs cho các rewards tự nhiên để tính từ completion (ví dụ: kiểm tra định dạng trên văn bản) và để môi trường trả về reward phụ thuộc trạng thái của nó thông qua get_reward. Tất cả các nguồn đều được cộng lại với trọng số 1.

def format_reward(completions, **kwargs):  # trainer-owned: scores the completion text
    ...

trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen3-0.6B",
    args=AsyncGRPOConfig(max_steps=1000),
    reward_funcs=format_reward,          # trainer-owned: scores the completion
    environment_factory=IncrementEnv,    # env-owned: scores its internal state via `get_reward`
)

Nhiều môi trường

Để huấn luyện trên một hỗn hợp các tác vụ trong một lần chạy (ví dụ: một tác vụ lập trình và một trò chơi), hãy truyền một từ điển ánh xạ tên môi trường đến các nhà máy thay vì một callable duy nhất. Mỗi rollout sẽ chọn môi trường của nó thông qua trường environment trong bộ dữ liệu, và chỉ các công cụ của môi trường đó được phơi bày. Điều này tránh việc rò rỉ các công cụ không liên quan (công cụ move của trò chơi cho một ví dụ lập trình, và ngược lại). Mỗi môi trường vẫn sở hữu dữ liệu của riêng mình và trả về prompt từ reset(), do đó bộ dữ liệu chỉ cần cột environment để định tuyến giữa chúng.

import random

from datasets import Dataset
from trl import AsyncGRPOConfig, AsyncGRPOTrainer


class CodingEnv:
    def reset(self, **kwargs) -> str:
        n = random.randint(1, 20)
        return f"Compute the {n}th Fibonacci number."

    def run_code(self, code: str) -> str:  # exposed as a tool only for "coding" examples
        """Run the given Python code and return its stdout.

        Args:
            code: The Python source to execute.

        Returns:
            The captured standard output.
        """
        ...

class GameEnv:
    def reset(self, **kwargs) -> str:
        return "Reach the goal tile."

    def move(self, action: str) -> str:  # exposed as a tool only for "game" examples
        """Apply an action to the game and return the new observation.

        Args:
            action: One of "up", "down", "left", "right".

        Returns:
            The observation after the move.
        """
        ...

# The dataset only routes: each row picks the environment for that rollout
if __name__ == "__main__":
    dataset = Dataset.from_dict({"environment": ["coding", "game"] * 500})

    trainer = AsyncGRPOTrainer(
        model="Qwen/Qwen3-0.6B",
        args=AsyncGRPOConfig(chat_template_kwargs={"enable_thinking": False}),
        train_dataset=dataset,
        reward_funcs=reward_func,
        environment_factory={"coding": CodingEnv, "game": GameEnv},
    )
    trainer.train()

Các instance môi trường được tái sử dụng giữa các bước (được reset giữa các episode), do đó một __init__ tốn kém chỉ được thực hiện một lần thay vì ở mỗi bước. Trường environment được trainer sử dụng và không được chuyển tiếp đến reset.

Cuộn đa lượt

Mỗi lượt hội thoại sẽ render toàn bộ cuộc trò chuyện thông qua chat template và sinh ra từ đó. Một chat template không phải lúc nào cũng render một cuộc trò chuyện đang phát triển theo cùng một cách (ví dụ, một số loại bỏ phần suy luận của các lượt trước), do đó các token prompt của một lượt có thể khác với các token đã được giữ lại cho đến thời điểm đó. Vào cuối quá trình rollout, worker sẽ đối chiếu các lượt để tạo thành các mẫu huấn luyện: một lượt có prompt bắt đầu bằng các token đã giữ sẽ được thêm vào cùng một mẫu, còn bất kỳ lượt nào khác sẽ bắt đầu một mẫu mới. Mỗi token được sinh ra đều được huấn luyện trong đúng một mẫu, do đó một nhánh rẽ (fork) chỉ tốn token forward mà không tạo ra tín hiệu huấn luyện.

Mỗi mẫu của một rollout nhận được lợi thế của rollout đó. rollout/fork_frac và rollout/drift_tokens_* cho thấy các lượt đã được điều chỉnh như thế nào.

Tích hợp Môi trường

Để huấn luyện các tác nhân chạy trong môi trường riêng của chúng, chẳng hạn như các tác nhân lập trình trong các sandbox Harbor, hãy xem Training on harnesses trong hướng dẫn OpenEnv.

Huấn luyện Mô hình Ngôn ngữ Thị giác

Các mô hình thị giác-ngôn ngữ (Qwen3.5, Qwen3.6, Qwen3-VL, …) có thể được huấn luyện trên các bộ dữ liệu chỉ chứa văn bản. Truyền id của mô hình như bình thường; không có gì khác thay đổi.

Tháp thị giác bị đóng băng. Bộ dữ liệu chỉ chứa văn bản không bao giờ tạo ra token hình ảnh, do đó tháp thị giác không bao giờ được thực thi trong bước truyền thẳng. Mọi thành phần nằm ngoài mô hình văn bản (tháp thị giác, bộ chiếu đa phương tiện) đều có requires_grad=False: chúng không nhận gradient, không có trạng thái optimizer, và quá trình đồng bộ trọng số sẽ bỏ qua chúng, nên máy chủ giữ nguyên các giá trị đã tải từ checkpoint. Tháp thị giác vẫn được tải, vì vậy nó vẫn chiếm bộ nhớ GPU cho phần trọng số.

Ảnh chưa được hỗ trợ. Các prompt chứa ảnh sẽ không được truyền đến máy chủ vLLM hay bước forward trong quá trình huấn luyện. Huấn luyện đa phương thức cũng cần đường dẫn packing không có padding để xây dựng vị trí M-RoPE 3D từ lưới ảnh, điều này chưa được triển khai ở đây.

AsyncGRPOTrainer

>>> from trl import AsyncGRPOTrainer
>>> from trl.rewards import accuracy_reward
>>> from datasets import load_dataset

>>> dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

>>> trainer = AsyncGRPOTrainer(
...     model="Qwen/Qwen2.5-0.5B-Instruct",
...     reward_funcs=accuracy_reward,
...     train_dataset=dataset,
... )
>>> trainer.train()

>>> class CurriculumTrainer(AsyncGRPOTrainer):
...     loop_cls = CurriculumRolloutLoop

AsyncGRPOConfig

AsyncRolloutLoop

Rollout

TurnRecord

RolloutGroup