RLHF 비동기 새 APIs
RLHF 비동기 새 APIs (RLHF Async New APIs)
vLLM과 Ray를 사용해 **비동기 강화학습(RL)**을 수행하되, 기본(native) 가중치 동기화 API와 배치 불변(batch-invariant) 생성을 사용하는 예제입니다. 훈련과 추론 워크로드를 서로 다른 GPU로 분리해 Ray가 프로세스 배치와 프로세스 간 통신을 관리하게 합니다.
출처: 문서
원본: https://github.com/vllm-project/vllm/blob/main/examples/rl/rlhf_async_new_apis.py
본문
- Hugging Face Transformer 모델 하나가 훈련용 GPU 1개를 차지하고, vLLM
AsyncLLMEngine이 다른 GPU 1개를 추론용으로 차지합니다. - **배치 불변성(batch invariance)**을 켜면 요청이 몇 개가 배치로 묶이든 생성 출력이 결정적이 됩니다. 검증 단계가 성공하려면 이 설정이 필요합니다. 배치 불변성은 현재 compute capability 9.0 이상의 NVIDIA GPU가 필요합니다(H100, H200 / B100, B200).
- 예제는 단일 노드 2-GPU 클러스터를 가정하지만, Ray는 멀티 노드를 지원합니다. GPU는 vLLM 워크로드 전용으로만 쓰기를 권장합니다. 잔여 GPU 활동이 vLLM 메모리 프로파일링을 방해하고 예상치 못한 동작을 일으킵니다.
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates async reinforcement learning using vLLM and Ray,
with native weight syncing APIs and batch-invariant generation.
The script separates training and inference workloads onto distinct GPUs
so that Ray can manage process placement and inter-process communication.
A Hugging Face Transformer model occupies one GPU for training, and a
vLLM AsyncLLMEngine occupies another GPU for inference.
Batch invariance is enabled so that generation output is deterministic
regardless of how many requests are batched together. This is required
for the validation phase to succeed. Batch invariance currently requires
NVIDIA GPUs with compute capability 9.0 or higher:
- H-series: H100, H200
- B-series: B100, B200
The example performs the following steps:
* Load the training model (Qwen3-1.7B) on one GPU via a Ray actor.
* Initialize the inference engine with a base model (Qwen3-1.7B-Base)
on a separate GPU using vLLM's AsyncLLMEngine with Ray as the
distributed executor backend.
* Set up an NCCL-based weight transfer channel between the trainer
and the inference engine.
* Submit generation requests for a batch of prompts.
* Pause generation once any request reaches a token threshold.
* Broadcast the training model's weights to the inference engine
via the NCCL weight transfer engine, replacing the base weights.
* Resume generation and collect results, noting which tokens were
generated before vs. after the weight swap.
* Validate correctness by launching a fresh vLLM instance loaded
directly with the training model and comparing its output to the
post-swap tokens from the weight-synced engine.
This example assumes a single-node cluster with two GPUs, but Ray
supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
workloads. Residual GPU activity interferes with vLLM memory profiling and
causes unexpected behavior.
"""
import asyncio
import uuid
import ray
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import vllm
from vllm import SamplingParams
from vllm.config import WeightTransferConfig
from vllm.distributed.weight_transfer import (
ModuleSource,
RayVLLMWeightSyncClient,
WeightTransferTrainerFactory,
)
from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo
from vllm.platforms import current_platform
from vllm.utils.network_utils import get_ip, get_open_port
from vllm.v1.executor import Executor
MODEL_NAME_V1 = "Qwen/Qwen3-1.7B-Base"
MODEL_NAME_V2 = "Qwen/Qwen3-1.7B"
PAUSE_TOKEN_THRESHOLD = 10
ATTN_BACKEND = "TRITON_ATTN" if current_platform.is_rocm() else "FLASH_ATTN"
class MyLLM(vllm.AsyncLLMEngine):
"""Configure the vLLM worker for Ray placement group execution."""
def __init__(self, **kwargs):
engine_args = vllm.AsyncEngineArgs(**kwargs)
vllm_config = engine_args.create_engine_config()
executor_class = Executor.get_class(vllm_config)
super().__init__(
vllm_config=vllm_config,
executor_class=executor_class,
log_requests=engine_args.enable_log_requests,
log_stats=not engine_args.disable_log_stats,
)
self._generation_paused = False
self._request_pause_flag = False
async def do_generate(
self, prompt_token_ids: list[int], sampling_params: vllm.SamplingParams
) -> tuple[vllm.RequestOutput, int]:
"""Generate a single request, setting the request pause flag once the
token count reaches the threshold.
Returns (output, pause_token_index). pause_token_index is the number
of tokens generated before the weight change, or -1 if no pause.
"""
pause_token_index = -1
prev_token_count = 0
async for request_output in self.generate(
{"prompt_token_ids": prompt_token_ids},
sampling_params,
request_id=str(uuid.uuid4()),
):
output = request_output
cur_token_count = len(output.outputs[0].token_ids)
if (
cur_token_count >= PAUSE_TOKEN_THRESHOLD
and not self._request_pause_flag
):
self._request_pause_flag = True
if self._generation_paused and pause_token_index == -1:
pause_token_index = prev_token_count
prev_token_count = cur_token_count
return output, pause_token_index
async def pause_after_n_tokens(self):
"""Wait for any request to set the pause flag, then pause."""
while not self._request_pause_flag:
await asyncio.sleep(0)
await super().pause_generation(mode="keep")
await asyncio.sleep(5)
self._generation_paused = True
@ray.remote(num_gpus=1)
class TrainModel:
"""Ray actor that wraps the training model on a dedicated GPU."""
def __init__(self, model_name: str):
from vllm.model_executor.determinism.batch_invariant import (
init_batch_invariance,
)
# need to init all env vars for batch invariance which affect nccl ops
init_batch_invariance()
self.model = AutoModelForCausalLM.from_pretrained(
model_name, dtype=torch.bfloat16
).to("cuda:0")
self.port = get_open_port()
self.master_address = get_ip()
def init_weight_transfer(self, world_size, llm_handle):
"""Build the trainer-side weight-transfer engine and rendezvous."""
self.engine = WeightTransferTrainerFactory.trainer_init(
init_info=NCCLTrainerInitInfo(
master_address=self.master_address,
master_port=self.port,
world_size=world_size,
rank=0, # single-GPU trainer is the sole (sender) rank
packed=True,
),
client=RayVLLMWeightSyncClient(llm_handle),
source=ModuleSource(self.model),
)
def broadcast_weights(self):
"""Push weights to the inference engine (drives start/update/finish)."""
self.engine.send_weights()
@torch.inference_mode()
def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]:
"""Greedy-decode max_new_tokens from the given context."""
input_ids = torch.tensor([token_ids], device="cuda:0")
output = self.model.generate(
input_ids,
max_new_tokens=max_new_tokens,
do_sample=False,
)
new_token_ids = output[0, len(token_ids) :].tolist()
return new_token_ids
# Build platform-specific env vars for Ray
ray_env_vars = {}
if current_platform.is_rocm():
# Workaround for RCCL bug. See https://github.com/ROCm/rocm-systems/issues/5756
ray_env_vars["RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES"] = "1"
else:
# Enable batch invariance for deterministic outputs on NVIDIA
ray_env_vars["VLLM_BATCH_INVARIANT"] = "1"
ray.init(runtime_env={"env_vars": ray_env_vars})
# Launch the training model actor. Ray's resource scheduler will allocate
# 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs.
train_model = TrainModel.remote(MODEL_NAME_V2)
rocm_determinism_kwargs = {}
if current_platform.is_rocm():
# ROCm: To minimize non-determinism, we set fixed seed, no prefix caching, and
# sequential request processing (max_num_seqs=1).
rocm_determinism_kwargs = {
"seed": 0,
"enable_prefix_caching": False,
"max_num_seqs": 1,
}
# Build platform-specific LLM kwargs
llm_kwargs = dict(
model=MODEL_NAME_V1,
enforce_eager=True,
max_model_len=8192,
distributed_executor_backend="ray",
attention_backend=ATTN_BACKEND,
gpu_memory_utilization=0.75,
weight_transfer_config=WeightTransferConfig(backend="nccl"),
)
llm_kwargs.update(rocm_determinism_kwargs)
# Launch the vLLM inference engine.
# With data_parallel_backend="ray", vLLM's CoreEngineActorManager creates
# its own placement groups internally for each DP rank, so we must NOT
# create an outer placement group (it would reserve GPUs and hide them
# from the internal DP resource check).
llm = ray.remote(
num_cpus=0,
num_gpus=0,
)(MyLLM).remote(**llm_kwargs)
PROMPTS = [
"The president of the United States is",
"The capital of France is",
"The largest ocean on Earth is",
"The speed of light in a vacuum is",
"The chemical formula for water is",
"The tallest mountain in the world is",
"The first person to walk on the moon was",
"The Great Wall of China was built to",
"Photosynthesis is the process by which",
"The theory of general relativity was proposed by",
"The boiling point of water at sea level is",
"The largest planet in our solar system is",
"DNA stands for deoxyribonucleic acid and it",
]
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_V1)
batch_prompt_token_ids = [
tokenizer.encode(prompt, add_special_tokens=False) for prompt in PROMPTS
]
# Set up the communication channel between the training process and the
# inference engine.
world_size = 2 # 1 trainer + 1 inference worker
ray.get(train_model.init_weight_transfer.remote(world_size, llm))
N_NEW_TOKENS = 100
# ── Phase 1: concurrent requests with weight sync ───────────────────
print(f"\n{'=' * 50}")
print(f"Prompts ({len(PROMPTS)}):")
for p in PROMPTS:
print(f" - {p!r}")
print(f"{'=' * 50}")
sampling_params = SamplingParams(
temperature=0, max_tokens=PAUSE_TOKEN_THRESHOLD + N_NEW_TOKENS
)
gen_futures = [
llm.do_generate.remote(ptids, sampling_params) for ptids in batch_prompt_token_ids
]
ray.get(llm.pause_after_n_tokens.remote())
ray.get(train_model.broadcast_weights.remote())
ray.get(llm.resume_generation.remote())
results = ray.get(gen_futures)
for i, (output, pause_idx) in enumerate(results):
all_token_ids = list(output.outputs[0].token_ids)
before_text = tokenizer.decode(all_token_ids[:pause_idx])
after_text = tokenizer.decode(all_token_ids[pause_idx:])
print(f"\n Request {i} ({PROMPTS[i]!r}):")
print(f" Old weights ({pause_idx} tokens): {before_text!r}")
n_after = len(all_token_ids) - pause_idx
print(f" New weights ({n_after} tokens): {after_text!r}")
# ── Phase 2: validate with a fresh V2 vLLM instance ────────────────
# This validation relies on batch-invariant (deterministic) generation to
# compare outputs from the weight-synced engine against a fresh V2 instance.
# On NVIDIA, batch invariance is fully supported, so we require 100% exact
# token match. On ROCm, batch invariance is not yet fully implemented
# (see https://github.com/vllm-project/vllm/issues/27433 and
# https://github.com/vllm-project/vllm/issues/33123), so residual
# non-determinism (e.g. GEMM accumulation order, missing kernel overrides)
# can cause single-token divergences that don't indicate a weight-sync
# failure. We relax the pass rate to 90% on ROCm to accommodate this; a
# real regression (broken weight transfer) would cause ~0% pass rate, not 90%+.
MIN_PASS_RATE = 1.0 if not current_platform.is_rocm() else 0.9
print(f"\n{'=' * 50}")
print("VALIDATION: comparing weight-synced vLLM with fresh V2 instance")
if current_platform.is_rocm():
print(f" (ROCm mode: requiring >= {MIN_PASS_RATE:.0%} exact match rate)")
print(f"{'=' * 50}")
ray.get(llm.shutdown.remote())
ray.kill(llm)
ray.kill(train_model)
llm_v2_kwargs = dict(
model=MODEL_NAME_V2,
enforce_eager=True,
max_model_len=8192,
gpu_memory_utilization=0.75,
distributed_executor_backend="ray",
attention_backend=ATTN_BACKEND,
)
llm_v2_kwargs.update(rocm_determinism_kwargs)
llm_v2 = ray.remote(
num_cpus=0,
num_gpus=0,
)(MyLLM).remote(**llm_v2_kwargs)
val_futures = [
llm_v2.do_generate.remote(
list(output.prompt_token_ids) + list(output.outputs[0].token_ids)[:pause_idx],
SamplingParams(
temperature=0, max_tokens=len(output.outputs[0].token_ids) - pause_idx
),
)
for output, pause_idx in results
]
val_results = ray.get(val_futures)
num_pass = 0
num_total = len(results)
for i, ((output, pause_idx), (val_output, _)) in enumerate(zip(results, val_results)):
expected = list(output.outputs[0].token_ids)[pause_idx:]
actual = list(val_output.outputs[0].token_ids)
match = actual == expected
if match:
num_pass += 1
print(f" [PASS] {PROMPTS[i]!r}")
else:
print(f" [FAIL] {PROMPTS[i]!r}")
print(f" weight-synced vLLM: {tokenizer.decode(expected)!r}")
print(f" V2 vLLM: {tokenizer.decode(actual)!r}")
for j, (e, a) in enumerate(zip(expected, actual)):
if e != a:
print(
f" first divergence at output token {j}: "
f"expected {e} ({tokenizer.decode([e])!r}) vs "
f"actual {a} ({tokenizer.decode([a])!r})"
)
break
ray.get(llm_v2.shutdown.remote())
ray.kill(llm_v2)
pass_rate = num_pass / num_total
print(f"\n Result: {num_pass}/{num_total} prompts passed ({pass_rate:.0%})")
print(f" Required: >= {MIN_PASS_RATE:.0%}")
assert pass_rate >= MIN_PASS_RATE, (
f"Validation pass rate {pass_rate:.0%} ({num_pass}/{num_total}) "
f"is below the required {MIN_PASS_RATE:.0%} threshold. "
f"See failures above for details."
)
print("=" * 50)
예제가 수행하는 단계: Ray 액터로 훈련 모델(Qwen3-1.7B)을 한 GPU에 로드하고, 별도 GPU에 vLLM AsyncLLMEngine(Ray executor, base 모델)으로 추론 엔진을 초기화합니다. NCCL 기반 가중치 전송 채널을 구성해 프롬프트 배치의 생성 요청을 보내고, 어느 요청이든 토큰 임계값에 도달하면 생성을 멈춘 뒤 NCCL 가중치 전송 엔진으로 훈련 모델 가중치를 브로드캐스트(base 가중치 대체)하고, 생성을 재개해 결과를 수집합니다. 마지막으로 훈련 모델을 직접 로드한 새 vLLM 인스턴스를 띄워 그 출력과 가중치 동기화 엔진의 스왑 후 토큰을 비교해 정확성을 검증합니다. NVIDIA에서는 100% 정확 토큰 일치를 요구하고, ROCm에서는 미구현 배치 불변성으로 인한 잔여 비결정성을 감안해 통과율 90%로 완화합니다.
더 알아보기 (Learn more)
- RLHF HTTP NCCL — HTTP 컨트롤 플레인 + NCCL 데이터 플레인 예제
- Routed Experts E2E — RDT 라우팅 전문가 예제
- 원본 예제 파일