비동기 GRPO

이 trainer는 vllm, transformers>=5.2.0 및 kernels가 필요합니다(학습 시 이를 통해 Hub에서 FlashAttention을 로드합니다). 분산 학습의 경우 FSDP2만 지원되며(DeepSpeed ZeRO는 지원되지 않습니다).

현재 vllm과 transformers의 종속성 제약 조건이 충돌합니다. 이를 우회하려면 먼저 vLLM을 설치한 후, transformers를 강제 설치하고, kernels를 TRL의 extra를 통해 추가하세요. 이 extra는 transformers가 허용하는 버전을 고정합니다.

복사됨

pip install vllm
pip install 'transformers>=5.2.0' --no-deps
pip install 'trl[kernels]'

개요

experimental.async_grpo.AsyncGRPOTrainer은 동일한 GRPO 알고리즘을 구현하지만, rollout 생성과 학습을 분리합니다. 백그라운드 워커가 vLLM 서버에서 completions을 지속적으로 스트리밍하는 동안 학습 루프가 이를 소비하므로, 생성과 그래디언트 업데이트가 번갈아 수행되는 대신 겹쳐서 실행됩니다. API는 GRPOTrainer와 동일합니다. GRPO 방법 자체(advantage 계산, KL 추정, 손실 함수 구성, 보상 함수 등)에 대한 자세한 내용은 GRPO Trainer 문서를 참고하세요. GRPOTrainer의 모든 기능이 지원되지는 않으며, 지원되는 파라미터는 experimental.async_grpo.AsyncGRPOConfig를 참조하세요.

이 trainer는 Quentin Gallouédec와 Amine Dirhoussi가 기여했습니다.

GRPOTrainer와의 차이점

표준 GRPOTrainer에서는 생성과 학습이 순차적으로 수행됩니다. 배치 생성, 손실 계산, 가중치 업데이트를 반복하는 방식입니다. 생성이 동일한 GPU에서 실행되는 vLLM colocate mode에서도 한 단계가 완료되어야 다른 단계가 시작됩니다.

experimental.async_grpo.AsyncGRPOTrainer는 이 두 가지 문제를 분리합니다:

• Rollout worker(백그라운드 프로세스) — vLLM 서버에 프롬프트를 전송하고, 보상 함수로 완성 결과를 점수화하며, 이점을 계산한 후 학습 준비가 된 샘플을 큐에 넣습니다.
• 학습 루프 (메인 프로세스) — 큐에서 샘플을 가져오고, 클립된 서로게이트 손실을 계산한 후 모델 가중치를 업데이트합니다.

rollout worker는 trainer에서 생성된 별도 프로세스에서 실행되므로, 보상 계산은 GIL을 두고 학습 루프와 경쟁하지 않습니다. 이는 reward_funcs, tools, environment_factory에 전달할 수 있는 값에 두 가지 영향을 미칩니다(후자의 경우, 계약과 사용 가능한 통합을 다루는 OpenEnv 가이드를 참조하세요):

rollout worker를 별도의 프로세스에서 실행하기 때문에, 해당 프로세스에 전달되는 모든 데이터는 pickled됩니다. 따라서 각 reward function, tool, environment_factory(그리고 이들이 close over하는 모든 것)은 picklable이어야 합니다: 모듈 레벨 함수, functools.partial, 또는 callable class instance를 사용하세요. tool은 예외입니다: 모델은 tool을 이름으로 호출하므로, 각 tool은 __name__ 아래에 등록되어야 하며 모듈 레벨 함수여야 합니다. lambda와 closure는 trainer.train()에서 TypeError를 발생시킵니다. 이는 reward function이 프로세스 내에서 호출되고 closure가 작동하는 GRPOTrainer와 다른 점입니다.

롤아웃 과정도 CUDA_VISIBLE_DEVICES=""로 실행되므로 GPU를 사용할 수 없습니다. GPU 기반 보상 모델(예: AutoModelForSequenceClassification 스코어)은 오류 없이 로드되지만 CPU로 조용히 폴백됩니다(GRPOTrainer에서는 이러한 보상 모델이 트레이너의 GPU를 공유한다는 점에 유의하세요). 보상 함수는 CPU 측에서 가볍게 유지하세요(예: accuracy_reward와 같은 검증기, 형식/길이 검사).

GPU 리워드 모델이 실제로 필요할 경우, 권장되는 접근 방식은 독립적인 추론 엔진(vLLM, TGI, …) 뒤에 배치하여 서빙하는 것입니다. 이 경우 별도의 GPU에서 모델을 서빙하고, 가볍고 picklable한 리워드 함수가 HTTP를 통해 해당 모델을 호출하도록 구성합니다. 이렇게 하면 리워드 모델은 자체 디바이스에 유지되는 반면, rollout 프로세스는 CPU 전용으로 유지되며, 트레이너와 독립적으로 확장됩니다.

Tools and environment methods can be synchronous or asynchronous. Synchronous tools run on a thread pool sized to max_concurrent_rollouts, so a slow tool never blocks the other in-flight vLLM requests. Within one assistant turn, tool calls run in the order the model emitted them. Up to max_concurrent_rollouts synchronous tools therefore run at once, so a tool that reads or writes shared state must be thread-safe. Tools와 environment methods는 동기(synchronous) 또는 비동기(asynchronous)일 수 있습니다. 동기 tools는 max_concurrent_rollouts 크기의 thread pool에서 실행되므로, 느린 tool이 다른 in-flight vLLM 요청을 차단하지 않습니다. 하나의 assistant turn 내에서 tool calls는 모델이 출력한 순서대로 실행됩니다. 따라서 최대 max_concurrent_rollouts개의 동기 tools가 동시에 실행될 수 있으므로, 공유 상태를 읽거나 쓰는 tool은 thread-safe여야 합니다.

모든 weight_sync_steps 학습 단계마다 업데이트된 가중치가 NCCL을 통해 vLLM 서버로 전송되어, 이후 생성 결과가 최신 정책을 반영하도록 합니다.

생성과 학습이 동시에 실행되므로 학습 샘플은 약간 오래된 버전의 모델에 의해 생성되었을 수 있습니다. max_staleness 매개변수는 샘플이 폐기되기 전에 허용되는 가중치 업데이트 지연 횟수를 제어합니다.

동시성

vLLM 서버로 전송되는 동시 요청의 수는 max_concurrent_rollouts로 제한됩니다. 기본적으로 max(max_staleness, 1) × per_device_train_batch_size × gradient_accumulation_steps × num_processes로 자동 설정되며, 이는 트레이너가 샘플이 스테일(stale) 상태가 되기 전에 소비할 수 있는 최대 샘플 수입니다. 이보다 많은 샘플을 생성하는 것은 낭비입니다. 초과된 샘플은 폐기되므로요.

max_concurrent_rollouts는 상한선이지 실제로 처리 중인 요청의 수가 아닙니다. 워커는 서버의 /metrics 엔드포인트를 몇 초마다 폴링하여, 그 아래에 실시간 상한(rollout/concurrency_cap)을 유지합니다. 특정 동시성 수준을 넘으면 KV 캐시가 차고, 시퀀스가 선점되어 재계산되며, 캐시된 접두사가 제거되고 다시 프리필링되므로 vLLM이 느려집니다. 상한은 서버 시작 시점에 읽어들인 KV cache tokens / max_model_len로 시작합니다. 이는 전체 길이에 맞춰 들어갈 수 있는 요청 수를 의미하며, 첫 단계에서 스래싱이 일어나지 않도록 합니다. 사용량이 0.6 미만이고 대기 중인 작업이 없으면, 상한 분량의 롤아웃이 완료될 때마다 상한이 25%씩 증가하여 상한선까지 올라갑니다. KV 캐시 사용량이 0.8을 넘으면 상한이 30% 줄어듭니다. 실행 중인 요청은 취소되지 않으며, 풀이 새로운 상한 아래로 비워질 때까지 신규 요청 수락만 중단되고, 그 시점까지 추가 감축은 이루어지지 않습니다. 서버가 /metrics를 노출하지 않으면 상한은 상한선 그대로 유지됩니다.

체크포인트 및 재개: ignore_data_skip의 기본값은 True입니다. 기본 Trainer의 skip-and-replay 루프는 라이브 롤아웃 큐에는 적용되지 않습니다. 대신, 아직 학습되지 않은 첫 번째 프롬프트의 인덱스가 각 체크포인트와 함께 rollout_state.json에 저장되고 재개 시 복원되므로, 워커는 샘플을 재생하지 않고 해당 프롬프트까지 빠르게 이동합니다. 이는 생성기의 위치가 아니라 학습된 위치입니다. 워커는 롤아웃 큐 깊이만큼 학습보다 앞서 실행되며, 이러한 버퍼링된 샘플은 실행이 종료되면 손실되므로, 생성기의 위치에서 재개하면 생성되었지만 학습되지 않은 프롬프트가 건너뛰어집니다. 스트리밍 데이터셋(IterableDataset)은 위치를 재설정할 수 없으며, 재개 시 워커는 프롬프트 0부터 다시 시작합니다. bf16=True(기본값) 또는 fp16=True인 FSDP 환경에서는 마스터 가중치가 fp32임에도 불구하고 model.safetensors가 해당 dtype으로 기록되므로, 푸시된 체크포인트의 크기는 서빙되는 모델의 크기와 같고 두 배가 아닙니다. fp32 마스터 가중치와 옵티마이저 상태는 여전히 재개 파일(pytorch_model_fsdp.bin, optimizer.bin)에 저장되며, 재개 시 읽히는 것이 바로 이 파일들이므로 재개에는 영향이 없습니다. FSDP를 사용하지 않으면 체크포인트가 재개 파일이기도 하므로 마스터 dtype을 유지합니다.

빠른 시작

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

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

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

vLLM 서버와 트레이너는 별도의 GPU에서 실행되어야 합니다. CUDA_VISIBLE_DEVICES를 사용하여 GPU를 분할하세요. 예를 들어, GPU 2개를 사용하는 경우 vLLM 서버를 GPU 0에서, 트레이너를 GPU 1에서 다음과 같이 실행할 수 있습니다:

# Terminal 1: vLLM server on GPU 0 (dev mode + NCCL weight transfer are required)
CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3-4B \
    --max-model-len 4096 \
    --logprobs-mode processed_logprobs \
    --api-server-count 8 \
    --weight-transfer-config '{"backend":"nccl"}'

--max-model-len을 예상하는 최대 총 시퀀스 길이(프롬프트 + 완성)로 설정하세요. 값이 낮을수록 서버의 GPU 메모리 사용량이 줄어들어 KV 캐시에 더 많은 메모리가 할당되고 처리량이 증가합니다. 좋은 시작점은 설정의 프롬프트 길이에 max_completion_length를 더한 값입니다.

# Terminal 2: training on GPU 1
CUDA_VISIBLE_DEVICES=1 accelerate launch train_async_grpo.py

AMD GPU

Trainer는 추가 설정 없이 AMD GPU(ROCm)에서 실행되며, 위의 명령어는 그대로 작동합니다. 학습은 AMD GPU에서 kernels-community/aiter-flash-attn attention kernel을, NVIDIA Blackwell GPU에서 kernels-community/flash-attn4를, 기타 NVIDIA GPU에서 kernels-community/flash-attn3를 사용하며, 모두 kernels 라이브러리를 통해 Hub에서 다운로드됩니다. attention logits에 상한을 두는 모델(attn_logit_softcapping)은 Gemma 2와 같이 AMD GPU 또는 NVIDIA Blackwell GPU에서 지원되지 않습니다.

비전-언어 모델

비전-언어 모델(Qwen3.5, Qwen3.6, Qwen3-VL, …)은 텍스트 전용 데이터셋으로 학습할 수 있습니다. 모델 ID를 평소와 같이 전달하면 됩니다. 그 외 변경 사항은 없습니다:

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

비전 타워는 동결됩니다. 텍스트 전용 데이터셋은 이미지 토큰을 생성하지 않으므로, 포워드 패스에서 타워가 사용되지 않습니다. 텍스트 타워를 제외한 모든 구성 요소(비전 타워, 멀티모달 프로젝터)는 requires_grad=False 상태입니다. 따라서 그라디언트나 옵티마이저 상태가 할당되지 않으며, 가중치 동기화에서도 완전히 제외됩니다. 서버는 체크포인트에서 로드한 값을 그대로 유지합니다. 타워는 여전히 로드되므로, 가중치를 위해 GPU 메모리를 차지합니다.

이미지는 아직 지원되지 않습니다. 이미지를 포함하는 프롬프트는 vLLM 서버나 학습 forward pass로 전달되지 않습니다. 멀티모달 학습은 이미지 그리드에서 3D M-RoPE 위치를 구성하기 위해 padding-free packing 경로가 필요하지만, 여기서는 구현되어 있지 않습니다.

Top-k 및 top-p 샘플링

Top-k와 top-p는 전체 어휘가 아닌 제한된 후보 집합에 대해 샘플링 분포를 재정규화합니다. 서버의 로그 확률은 이를 반영하며(--logprobs-mode processed_logprobs), 따라서 트레이너도 동일한 집합에 대해 정규화를 수행해야 합니다. 그렇지 않으면 모든 중요도 비율이 편향됩니다. 전체 어휘 기준의 토큰 확률은 실제로 추출된 후보 집합 기준의 확률보다 낮기 때문입니다(DeepSeek-V3.2, 3.1절). top_p < 1 또는 top_k > 0인 경우, 워커는 서버에 각 샘플링된 토큰의 후보 집합을 요청하고 트레이너는 해당 집합에 대해 로그 확률을 재정규화합니다. 이를 위해 서버 플래그가 두 개 더 필요하며, vLLM이 후보 집합의 크기를 top_k로 제한하기 때문에 top_k를 설정해야 합니다:

CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3-4B \
    --max-model-len 4096 \
    --logprobs-mode processed_logprobs \
    --api-server-count 8 \
    --weight-transfer-config '{"backend":"nccl"}' \
    --return-sampling-mask \
    --enable-scale-out

config = AsyncGRPOConfig(top_p=0.95, top_k=512)

--return-sampling-mask은 vLLM이 후보 집합을 기록하도록 합니다(vLLM 0.28 이상, V2 모델 러너 사용 시: 서버가 플래그를 거부하면 VLLM_USE_V2_MODEL_RUNNER=1 설정). --enable-scale-out는 해당 집합을 반환하는 token-in-token-out 엔드포인트를 노출합니다. 워커는 시작 시 서버에서 두 설정을 모두 확인합니다. 구버전 vLLM은 해당 집합을 반환할 수 없으므로, 트레이너는 전체 어휘에 대해 정규화를 유지하고 편향된 비율에 대해 경고를 표시합니다.

Score centering

서버의 log-probability는 트레이너의 것과 정확히 일치하지 않습니다: 커널이 다르고, 정밀도가 다르고, 정책이 몇 단계 뒤처져 있기 때문입니다. 중요도 가중치만으로는 이를 보정할 수 없는데, 샘플링된 토큰의 점수 $\nabla_\theta \log \pi_\theta$ 는 샘플러 하에서 평균이 0이 되지 않으므로, 모든 업데이트가 정책을 샘플러 쪽으로 끌어당깁니다(Score centering). 기본적으로 트레이너는 모든 위치에서 해당 기대값을 뺍니다(score_centering=True), 샘플러의 num_sampler_logprobs 가장 가능성 높은 토큰에 대해 추정하고, 샘플러에 대한 중요도 비율이 [vllm_importance_sampling_clip_min, vllm_importance_sampling_clip_max]. 이를 위해 워커는 서버에서 생성된 모든 토큰의 상위 log-probability를 요청하며, 토큰별로 직렬화합니다. 큰 num_sampler_logprobs 은 서버를 느리게 만들고, 그 --max-logprobs 는 서버가 허용하는 범위를 제한합니다. score_centering=False 는 클립된 목적 함수만으로 학습합니다. sampling/masked_mean 는 경계에 의해 제거된 토큰의 비율이고 sampling/sampler_mass 는 상위 토큰이 커버하는 샘플러의 확률 질량입니다.

vLLM의 API 서버에서 상위 log-probabilities를 직렬화하는 작업은 CPU에서 수행되며, 기본값인 단일 API 서버에서 num_sampler_logprobs=8을 사용하면 생성 처리량이 절반으로 줄어듭니다. vLLM 0.32 이상에서는 서버를 --api-server-count 8 옵션으로 시작하세요. H100 1대에서 1.5B~7B 모델을 측정했을 때, 상위 8개 log-probabilities의 비용이 생성 처리량의 3% 미만으로 줄어듭니다. vLLM 0.31에서는 동일한 플래그를 사용해도 20~40%의 손실이 발생하며, 이는 vllm-project/vllm#58181에서 업스트림으로 수정되었습니다.

LoRA

전체 모델 대신 LoRA 어댑터를 학습하려면 peft_config를 전달하세요:

from peft import LoraConfig

trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen3-4B",
    reward_funcs=reward_len,
    train_dataset=dataset,
    peft_config=LoraConfig(r=32, lora_alpha=64, target_modules="all-linear"),
)

모든 가중치 동기화 시 업데이트된 정책을 서버로 전송해야 합니다. 기본적으로 trainer는 어댑터를 베이스 모델에 병합하여 병합된 전체 가중치를 전송하며, 7B 모델의 경우 동기화마다 약 15 GB가 됩니다. vLLM 서버를 --enable-lora 옵션으로 시작하면 어댑터만 전송하며, r=32일 때 약 160 MB입니다:

CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 vllm serve Qwen/Qwen3-4B \
    --max-model-len 4096 \
    --logprobs-mode processed_logprobs \
    --api-server-count 8 \
    --weight-transfer-config '{"backend":"nccl"}' \
    --enable-lora \
    --max-lora-rank 32 \
    --max-loras 6

VLLM_ALLOW_RUNTIME_LORA_UPDATING=1은 트레이너가 각 새 어댑터를 게시하는 엔드포인트를 노출합니다. --enable-lora를 사용하더라도 --weight-transfer-config를 유지해야 합니다. 트레이너는 시작 시에만 동기화 모드를 선택하는데, 그 시점에는 서버가 이미 실행 중이며 병합 동기화가 폴백으로 사용됩니다.

--max-lora-rank은 1, 8, 16, 32, 64, 128, 256, 320, 512 중 하나여야 합니다. 서버가 제공할 수 있는 최대 랭크를 설정하는 것이므로, 8 환경에서도 r=4 어댑터가 정상적으로 작동합니다. --max-loras는 최소 max_staleness + 2 이상이어야 합니다. 트레이너는 max_staleness + 1개의 어댑터 버전을 등록해 두어, 이전 정책 하에서 시작된 롤아웃이 생성 중간에 정책을 전환하지 않고 해당 정책 하에서 완료될 수 있도록 하며, 각 동기화 시 가장 오래된 버전을 해제하기 전에 다음 버전을 로드합니다. 트레이너는 시작 시 두 값을 확인하고, 서버를 재시작할 때 사용할 값을 알려줍니다.

vLLM은 네트워크가 아닌 디스크에서 어댑터를 로드하므로, 트레이너는 각 버전을 <output_dir>/.vllm_lora/ 아래에 작성하고 서버에 해당 경로를 전달합니다. 서버가 다른 머신에서 실행되는 경우, 해당 디렉토리는 두 머신 모두에서 접근할 수 있는 파일시스템에 위치해야 합니다.

<output_dir>/.vllm_lora/는 서빙 캐시이며, 학습된 어댑터가 저장되는 위치가 아닙니다. 이 디렉터리에는 최신 버전만 보관되며, 각 버전은 스탈니스 윈도우(staleness window)에서 벗어난 직후에 삭제됩니다. 체크포인트가 학습한 어댑터를 보존하므로, 체크포인트 저장을 비활성화하지 마세요. save_strategy="no"를 사용하면 수 시간 동안 실행된 작업이 종료된 후에도 아무것도 남지 않습니다. 또한 output_dir은 서버가 읽을 수 있어야 하고, 작업 종료 후에도 유지되어야 하므로, 느린 공유 파일시스템과의 동기화를 피할 수 있다는 장점에도 불구하고 노드 로컬 스크래치 공간은 사용할 수 없습니다.

Adapter sync가 항상 더 빠르지는 않습니다. vLLM은 LoRA adapter를 서빙할 때 병합된 가중치를 서빙하는 경우보다 토큰당 약 1.33~1.39배 느린 속도로 생성하며, RL에서는 일반적으로 생성 과정이 wall-clock 시간을 지배합니다. 그 대가로 각 sync 시 더 짧은 일시 정지 시간과 약 1%의 sync 대역폭을 얻습니다. 생성 처리량이 병목이라면 병합 sync가 여전히 더 나은 선택인 경우가 많습니다.

롤아웃 루프 사용자 정의

Generation, rewards, advantages and filtering은 experimental.async_grpo.AsyncRolloutLoop의 rollout 프로세스에서 실행됩니다. 이 중 하나를 변경하려면 루프를 서브클래스로 상속하고 해당 단계의 메서드를 오버라이드한 뒤, trainer 서브클래스의 loop_cls로 해당 서브클래스를 설정합니다. trainer의 생성자는 변경되지 않습니다.

예를 들어, 정책이 이미 해결한 프롬프트의 샘플링을 중단하는 커리큘럼은 설정 서브클래스의 임계값을 통해 정의할 수 있습니다:

from dataclasses import dataclass

import numpy as np

from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer, 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는 reward_funcs, get_reward을 정의하는 환경, 그리고 커스텀 generate_one가 자신의 experimental.async_grpo.Rollout에 설정하는 reward에 대해 합산됩니다. 판정(judge)으로 점수를 매기려면 비동기 보상 함수를 전달하거나, generate_one 내부에서 점수를 매기세요. compute_advantages는 모든 그룹에 대해 실행되며, 점수화된 보상이 모두 동일하여 나중에 제거되는 그룹도 포함됩니다. 제거 여부는 보상을 기준으로 판단하며, 이점(advantages)을 기준으로 하지 않습니다. 루프는 생성된(spawned) 프로세스에서 구성되므로, 서브클래스는 모듈 수준에서 정의하세요. __init__을 오버라이드하는 대신, 루프의 이벤트 루프에서 실행되는 setup에서 클라이언트를 생성하세요. filter_prompt가 거부하는 프롬프트는 해당 에포크에서 건너뛰며, 제거된 그룹과 마찬가지로 해당 에포크에 카운트됩니다. 한 에포크의 모든 프롬프트를 거부하면 루프가 예외를 발생시킵니다.

디자인 철학

이 trainer은 의도적으로 최소한으로 유지되며 범용 솔루션으로 확장될 목적으로 설계되지 않았습니다. 지원하지 않는 기능이 필요하시다면, 저장소를 클론한 후 trainer을 직접 필요에 맞게 수정하는 것을 권장합니다. 커뮤니티의 수요가 상당할 경우에만 새로운 기능이 검토됩니다.

기록된 지표

롤아웃는 그라디언트가 되기까지 여러 단계를 거칩니다. 프롬프트에서 최종 배치 구성에 이르기까지의 각 단계를 먼저 살펴보겠습니다:

Dataset row = PROMPT (message list + reward kwargs)
 └─ GROUP                  (1 prompt, `num_generations` rollouts, 1 advantage baseline)
     └─ ROLLOUT            (1 conversation, keyed by `rollout_id`)
         ├─ TURN           (1 vLLM /v1/completions call + tool messages fed back)
         │   └─ TurnRecord (prompt_ids, output_ids, output_log_probs)
         └─ reconcile      (`_chain_to_sequences` classifies drift per turn: CLEAN/REALIGN/FORK)
             └─ SEQUENCE   (≥1 per rollout; new one per fork; dropped if no trained token)
                 └─ SAMPLE (Sequence + group advantage + reward + metrics)

════════════ Trainer process boundary: `rollout_buffer` (mp.Queue) ════════════

                 └─ SAMPLE (Pulled 1 at a time; dropped if staleness > `max_staleness`)
                     └─ ROW            (Planner inserts it into one of `dp` prefix tries, attention-cost-balanced)
                         └─ MICRO-BATCH (`dp` rows, one per rank)
                             └─ PACKED ROW (Collator lays each trie out: unique tokens, depth as `position_ids`)
                                 └─ FORWARD (`compute_loss`, bs=1, padding stripped, FA3 per trie segment)
                     └─ OPTIMIZER STEP (from `grad_accum` micro-batches)

위에서 아래로 읽으면:

• group은 하나의 프롬프트가 num_generations번 생성된 것을 의미합니다. 이 그룹 내부에서 이점(baseline)이 계산되므로, 그룹은 점수가 매겨지는 단위이기도 합니다.
• rollout는 하나의 대화입니다. 모든 턴에서 전체 메시지 목록을 다시 토큰화하므로, 턴 사이의 드리프트는 마지막에 조정해야 합니다.
• reconciling은 하나의 rollout을 여러 sequences로 fork할 수 있으므로, 하나의 대화에서 여러 개의 학습 샘플이 생성될 수 있습니다.
• sample는 시퀀스와 해당 시퀀스의 advantage, reward, metrics로 구성됩니다. 이것이 프로세스 경계를 넘어 전달되는 데이터입니다.
• 경계를 넘어 trainer는 샘플을 rows(DP rank당 하나)로 계획하고, 각 row를 prefix trie로 접어 공유 토큰이 한 번만 전달되도록 하며, grad_accum개의 micro-batch를 단일 optimizer step으로 누적합니다.

따라서 메트릭은 어떤 단계(또는 entity)에 관한지에 따라 다음과 같이 구분됩니다:

토큰 메트릭의 모호함을 없애기 위해 몇 가지 용어를 정의해 보겠습니다:

• generated 토큰은 모델이 생성한 결과입니다(completions/*).
• forwarded 토큰은 forward 패스가 처리하는 모든 토큰을 의미하며, 프롬프트, 도구 결과, 생성된 토큰을 포함합니다.
• trained 토큰은 completion_mask == 1인 부분집합에 대해 손실이 계산되는 토큰을 의미합니다. 예를 들어, 포크(fork) 또는 리얼라인(realign)은 이미 생성된 토큰을 컨텍스트로 강등시킬 수 있으므로 trained ≠ generated입니다.

배치 번호의 구성 방식

row는 하나의 DP rank가 하나의 micro-batch에서 forward하는 단위로, 여러 샘플이 하나의 prefix forest에 패킹되어 공유하는 토큰은 한 번만 forward됩니다. 자세한 내용은 Tree packing을 참고하세요. 플래너는 어떤 샘플이 어떤 row에 배치될지 결정하여(어떤 rank도 straggler가 되지 않도록 attention 비용을 균형 있게 조절) 데이터 콜레이터가 패킹을 수행합니다.

모든 지표에서 “step”은 항상 하나의 FULL optimizer step을 의미하며, micro-batch를 의미하지 않습니다. gradient_accumulation_steps개의 micro-batch가 하나의 step을 구성하며, _per_step로 이름이 붙은 모든 지표는 optimizer step당 값을 나타내며, 이는 global_step, logging_steps 및 Trainer 어휘의 나머지 부분과 일치합니다. 지표가 micro-batch당 값을 나타내는 경우 명시되어 있으며(batch/microbatches_per_step), 그렇지 않으면 행(row)당 양(batch/row_*, batch/samples_per_row)입니다.

따라서 한 단계에는 고정된 수의 row-slots가 포함됩니다:

row-slots per step = gradient_accumulation_steps x world_size

batch 메트릭이 서로 비교할 수 있는 계단 형태로 구성되도록 합니다:

sample                                     sample/forwarded_tokens_mean
  └─ packed into a ROW                     batch/samples_per_row, batch/row_tokens_mean
      └─ one row per DP rank               = one micro-batch
          └─ gradient_accumulation_steps micro-batches = one STEP
                                           batch/samples_per_step, batch/forwarded_tokens_per_step

즉, batch/samples_per_step ≈ row-slots x batch/samples_per_row이며, 토큰도 마찬가지입니다. 두 값은 스텝 안의 마이크로 배치 간 변동성까지만 일치합니다. 스텝별 지표는 합계이고, 행별 지표는 평균이기 때문에 정확한 일치보다는 몇 퍼센트 이내의 오차를 예상해야 합니다.

샘플이 긴 경우 batch/row_fill_frac을 주목해야 합니다. 10k 토큰 샘플은 32k 예산을 효율적으로 채우지 못합니다(3개는 들어가고 4개는 절대 들어가지 않으므로, 패커는 보통 2개만 채우고 행이 약 77%만 차 있는 상태가 됩니다). 반면 1k 토큰 샘플은 예산을 거의 완벽하게 채웁니다. 이는 버그가 아니라 양자화(quantization)의 특성이며, token_budget이 이를 조절하는 레버입니다. 다만 각 시퀀스당 어텐션이 O(L²)이므로, 긴 시퀀스로 구성된 행이 더 차 있다고 해서 메모리 비용이 선형으로 증가하지는 않습니다.

트리 패킹

행(row)은 prefix trie이며 단순한 연결(concatenation)이 아닙니다: 여러 샘플이 공유하는 토큰은 한 번만 전달(forward)되고, 어텐션은 각 샘플이 정확히 자신의 히스토리만 보이도록 합니다. 하나의 프롬프트에 대한 모든 롤아웃은 해당 프롬프트를 반복하며, 멀티턴 롤아웃의 모든 턴은 대화를 다시 전송하므로, 디코더의 작업량은 packing ratio(원시 토큰 수 ÷ 전달된 토큰 수)만큼 감소합니다. 이는 레이아웃 변경일 뿐입니다: 학습된 모든 토큰은 각자의 타깃, 어드밴티지, 이전 로그-확률(old log-probability)을 유지하며, 해당 샘플을 단독으로 전달했을 때의 로그-확률과 일치합니다.

발생 위치:

planner    TreeRow.add            the trie grows sample by sample; a sample is priced by its novel tokens
   │                              (TreeRow.fits), so `token_budget` bounds unique tokens, not raw ones
collator   TreeRow.pack           DFS layout: unique tokens, depth as position_ids, one loss term per
   │                              original trained token (token_index), and per segment its queries and
   │                              the K/V indices of its root-to-segment path, all built on CPU
attention  tree_attention         gather K/V along the paths, one FA3 varlen call, causal

하나의 프롬프트에 대한 두 번의 롤아웃에서 작동 방식:

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

2. Trie (one per row)   S0 [p1 p2] ─┬─ S1 [p3 a1 a2]          nodes are keyed on (token, next token),
                                   └─ S2 [p3 b1 b2]          so the fork happens one token early and
                                                              every node predicts exactly one token;
                                                              chains without a fork become segments

3. Packed row           idx           0  1 | 2  3  4 | 5  6  7
   (DFS order)          input_ids    p1 p2 | p3 a1 a2 | p3 b1 b2  8 tokens: MLPs, norms and projections
                        position_ids  0  1 |  2  3  4 |  2  3  4  run once per unique token
                        segment      S0    | S1       | S2

4. Attention            seg   queries      K/V gathered            kv_index
   (one FA3 call)       S0    [p1 p2]      [p1 p2]                 0 1
                        S1    [p3 a1 a2]   [p1 p2 | p3 a1 a2]      0 1 2 3 4
                        S2    [p3 b1 b2]   [p1 p2 | p3 b1 b2]      0 1 5 6 7

                        FA3 aligns the causal mask bottom-right      K:  p1 p2 p3 a1 a2
                        when queries are fewer than keys, so a       Q: p3   ■  ■  ■  ·  ·
                        segment's queries see their ancestors and       a1   ■  ■  ■  ■  ·
                        their own causal past, and nothing else.        a2   ■  ■  ■  ■  ■

5. Loss                 p2 (idx 1) predicts p3 for both rollouts; token_index points every trained
                        token of every rollout at the position that predicts it

몇 가지 결과:

• 공통 접두사는 공유됩니다: 프롬프트의 롤아웃, 그리고 프롬프트 간 공통 시스템 프롬프트 또는 도구 스키마입니다. 플래너는 샘플을 해당 group_id를 이미 보유하고 있는 행, 즉 해당 프롬프트가 이미 있는 행으로 전송합니다. 해당 행이 가득 차면, 랭크가 아직 비어 있거나 가장 부하가 높은 행의 부하를 높이지 않는 한 샘플은 다른 행으로 이동합니다. 가장 부하가 높은 행은 마이크로 배치의 시간을 결정합니다. 그렇지 않으면 마이크로 배치가 닫힙니다.
• 공유된 것이 없는 경우는 단순한 패킹입니다. 각 샘플은 자기 자신의 경로가 되는 하나의 세그먼트가 되며, 이는 일반적인 패딩 없는 FA3 varlen 패킹입니다.
• K/V만 중복됩니다: 공유 접두사는 하위 세그먼트당 한 번만 수집되며, 이는 언패킹된 학습이 보유할 양과 비슷하고 GQA를 사용하면 작습니다.
• Gated DeltaNet 하이브리드가 지원됩니다(Qwen3.5, Qwen3.6, Qwen3-Next, …): 이 모델의 컨볼루션은 각 토큰의 이전 토큰들을 읽으며, 청크형 델타 규칙은 각 세그먼트를 부모의 최종 상태에서 시작합니다(trl/experimental/async_grpo/tree_gated_delta_rule.py). 이를 위해 transformers>=5.15.0 이상이 필요하며, 해당 버전에서 이러한 레이어는 트리 패킹이 패치하는 모듈 레벨 함수를 호출합니다. 다른 비-어텐션 레이어를 포함하는 모델은 초기화 시 거부됩니다.

구매하는 것: 아무것도 공유하지 않으면 패딩 없는 패킹의 일반 속도와 거의 동일한 속도로 실행되며, 속도 향상은 패킹 비율에 따라 대략적으로 증가합니다(Qwen3-0.6B/1.7B, H100에서 비율 3.4일 때 약 2.8~3배). 비율은 공유 접두사 × 행당 샘플 수이므로 batch/samples_per_row가 주요 지표입니다: 1에 가까우면 아무것도 공유할 수 없으며, 레버는 token_budget입니다.

처리량으로 판단하세요. batch/trained_tokens_per_step / perf/fwd_bwd_s입니다. 패킹된 마이크로 배치에는 더 많은 샘플이 포함되므로 perf/fwd_bwd_s는 증가하고, 포워딩된 토큰으로 나누는 perf/mfu_*는 감소합니다. 둘 다 회귀(regression)가 아닙니다. 행당 샘플 수가 증가하면 스텝당 추출되는 롤아웃도 증가하므로, 생성에 병목이 있는 실행은 그 영향이 더 커집니다.

행을 더 길게 만들거나 줄이는 모든 작업은 공유를 제거합니다. 먼저 fork_threshold_tokens를 확인하세요. 한 턴이 재토큰화 시 해당 토큰 수만큼 드리프트가 발생해야만 독립된 행이 되며, 에이전틱 드리프트는 일반적으로 수십 토큰 수준으로 기본값인 1024에 비해 매우 작아 포크가 발생하지 않고 대화는 랭크당 한 행으로 수렴합니다. rollout/fork_frac는 현재 어떤 상태에 있는지 나타냅니다.

롤아웃 큐

워커는 점수가 매겨진 샘플을 큐에 넣고, 트레이너는 해당 큐에서 데이터를 가져옵니다. 이 하나의 큐를 설명하는 네 가지 지표는 각각 다른 질문에 답합니다:

마지막 두 지표(perf/rollout_wait_s와 rollout/backpressure_s)는 서로 거울 이미지이며, 동시에 큰 값을 가질 수 없습니다. 큐 크기와 함께 읽으면 어느 쪽이 병목인지 알 수 있습니다:

• 큐가 거의 비어 있고 perf/rollout_wait_s가 높으면 생성 병목 상태입니다. 트레이너가 기아 상태에 빠져 있으므로 rollout/generated_tok_s, rollout/concurrency, rollout/score_queue_size를 확인하세요.
• 큐가 거의 가득 차고 rollout/backpressure_s가 높으면 → trainer-bound 상태입니다. 생성이 스로틀링되고 그 출력이 큐에서 오래 머물고 있으므로 sample/staleness_mean가 상승하는지 확인하세요.
• 둘 다 0에 가까움 → 균형 잡힘!

Completions

completion은 vLLM이 하나의 프롬프트에 대해 생성한 결과로, 롤아웃의 모든 턴에 걸쳐 집계됩니다.

롤아웃

rollout는 완전한 대화 하나를 의미합니다. 프롬프트가 완료될 때까지 생성되는 과정이며, 그 과정에서 발생한 모든 도구 왕복을 포함합니다. num_generations개의 rollout은 동일한 프롬프트를 공유하며 하나의 그룹을 형성합니다.

모든 rollout이 동일한 보상을 받았거나(또는 점수화 가능한 rollout이 두 개 미만인 경우) 해당 그룹은 큐에 추가되지 않고 제외됩니다. 기본 advantage를 사용할 경우 모든 위치에서 advantage가 0이 되어 forward와 backward를 수행해도 아무것도 이동하지 않으므로, 이 경우 그룹을 제외하는 것이 효율적입니다. 배치에는 영향이 없으며, 유지된 그룹에서 여전히 samples_per_step만큼 채워지므로, 많은 그룹이 제외되면 한 step에 더 많은 생성이 필요합니다. 제외된 그룹도 epoch에 포함되며, epoch은 프롬프트를 한 번만 순회하는 것을 의미합니다.

도구

모델이 도구를 실행할 때만 기록됩니다. tools/<name>/*는 도구 이름별로 반복되며, 이는 rewards/<func>/*가 보상 함수별로 반복되는 방식과 동일합니다. 하나의 실패한 도구가 전체 평균에서는 보이지 않기 때문입니다. 내장된 도구 루프와 OpenEnv harness trainer(HarnessAsyncGRPOTrainer, harness에서 학습 참조)은 동일한 이름을 기록하므로 하나의 대시보드에서 둘 다를 읽을 수 있으며, path 컬럼은 각 메트릭을 생성하는 주체를 나타냅니다.

환경

environment_factory가 있는 경우에만 기록됩니다. 두 호출 모두 워커의 도구 스레드 풀에서 실행되므로, 느린 환경은 자신의 롤아웃만 지연시키고 진행 중인 다른 롤아웃에는 영향을 주지 않습니다.

큐에서 도착하는 샘플

각 샘플은 RolloutSample 하나입니다. 그룹 내 모든 generation이 완료되고 점수가 매겨진 후, 워커가 _score_group에서 조율된 시퀀스와 그룹 advantage, 보상 및 샘플별 메트릭을 조립하여 생성합니다. 워커는 이들을 rollout_buffer에 푸시하며, 아래에 설명된 모든 항목은 트레이너 측에서 학습을 위해 이를 가져오는 시점에 측정됩니다.

보상 및 오프폴리니스

배치

마이크로 배치 1개는 world_size개의 행으로 구성됩니다(패킹이 1개의 시퀀스로 평탄화한다는 점을 기억하세요). 따라서 DP 랭크당 1개씩 할당되어, 랭크 i는 행 i를 포워드로 전달합니다. 이러한 마이크로 배치 gradient_accumulation_steps개가 하나의 옵티마이저 스텝을 구성합니다. 따라서 스텝 1개는 총 gradient_accumulation_steps × world_size개의 행을 처리합니다. 아래에 나열된 모든 _per_step 메트릭은 전체 스텝과 모든 랭크에 대한 합계이며, batch/row_* 메트릭은 행에 대한 평균입니다.

성능

Throughput은 트레이너가 초당 처리하는 토큰 수를 측정합니다. Model FLOPs Utilization (MFU)는 모델의 forward pass와 backward pass에 사용된 학습 장치의 이론적 연산 용량의 백분율을 추정합니다. 이 추정에서 사용되는 용량은 하드웨어와 학습 정밀도에 따라 달라집니다.

Throughput과 MFU는 두 가지 시간 기준에 따라 보고됩니다:

• _fwd_bwd는 순전파와 역전파 과정에서의 성능을 측정합니다.
• _wall_clock는 옵티마이저 업데이트, 가중치 동기화, 롤아웃 대기 시간을 포함한 전체 옵티마이저 스텝에 걸친 성능을 측정합니다.

이 차이는 계산 성능과 대기 및 조정 비용을 구별하는 데 도움이 됩니다. 두 MFU 지표 모두 학습 장치만 설명합니다. 외부 롤아웃 또는 티처 서버의 계산 또는 하드웨어 용량은 포함되지 않습니다. Wall-clock MFU는 해당 서비스를 위한 트레이너의 대기 시간을 반영하며, 서비스 자체의 사용률을 반영하지 않습니다.

MFU는 선택된 정밀도에서 로컬 학습 장치의 피크 컴퓨트 용량이 알려진 경우에만 보고됩니다. 그렇지 않으면 MFU는 생략되며, 학습은 계속되고 처리량 및 타이밍 지표는 계속 사용할 수 있습니다.

AsyncGRPOConfig

AsyncGRPOTrainer

>>> from trl.experimental.async_grpo 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

AsyncRolloutLoop

롤아웃

TurnRecord

RolloutGroup