HF Jobs 전반에 걸친 LoRA를 사용한 Async GRPO: 버킷, 프록시, 그리고 NCCL 불필요

TL;DR

AsyncGRPOTrainer는 이제 LoRA 어댑터를 학습시키고 그 어댑터만 vLLM(TRL v1.14)에 동기화할 수 있습니다.

• 랭크-1 어댑터는 몇 메가바이트에 불과하므로, NCCL을 통해 전송하는 대신 모든 Job에 마운트된 스토리지 버킷을 통해 이동할 수 있습니다. 트레이너와 vLLM 복제본은 별도의 머신에서 실행되는 별개의 Hugging Face Job으로 작동합니다.
• 복제본 앞에 배치된 작은 프록시는 인증 헤더를 추가하고, 각 롤아웃을 이미 자신의 KV 프리픽스를 보유한 복제본으로 라우팅하며, 모든 복제본에 어댑터 로드를 브로드캐스트합니다.
• AsyncGRPO 지표는 병목 현상이 어디에 있는지 보여줍니다. 다섯 번의 실행은 동일한 레시피를 사용하여 500 스텝을 처리하는 데 걸리는 시간을 3시간 27분에서 53분으로 단축했습니다.

LoRA 지원이 TRL의 AsyncGRPOTrainer에 최근 추가되었으며, PR #7017과 함께 TRL v1.14를 통해 제공됩니다. 이 비동기 트레이너는 이제 전체 모델 대신 어댑터를 학습시킬 수 있으며, LoRA 어댑터만 vLLM으로 동기화합니다. 본 게시물은 이를 기반으로 구축된 실제 프로젝트에 대해 다루며, 여기서 학습과 추론이 더 이상 같은 머신을 공유하지 않습니다.

LoRA 학습은 Thinking Machines의 블로그 'LoRA Without Regret'에서 보여주었듯이 RL(강화학습)에 특히 적합합니다. 그들은 랭크-1 LoRA가 정책 경사(policy-gradient) RL에 대해 전체 미세 조정과 일치할 수 있음을 보여줍니다. 이는 장점 함수(advantage function)가 에피소드당 ~O(1) 비트의 정보만 제공하기 때문에, 총 정보 비트 관점에서 볼 때 각 단계에서 학습할 것이 많지 않다는 사실에서 비롯됩니다. 랭크-1 어댑터는 이를 흡수할 충분한 용량을 가지고 있습니다.

LoRA 학습에는 시스템적인 결과도 따릅니다. 15억 매개변수(1.5B) 모델의 랭크-1 어댑터는 몇 메가바이트인 반면, 전체 모델은 약 3GB입니다. 업데이트 후마다 전체 정책을 추론 워커로 보내는 대신, 어댑터만 보낼 수 있습니다. vLLM은 또한 여러 개의 어댑터를 한 번에 로드 상태로 유지할 수 있습니다. 이전 롤아웃은 시작한 정책으로 완료되고, 새로운 롤아웃은 최신 것을 사용합니다.

TRL의 AsyncGRPOTrainer

이미 트레이닝과 추론(generation)이 분리되어 있습니다. 트레이너와 vLLM은 서로 다른 머신에서 각자의 속도로 실행될 수 있습니다. 이는 두 프로세스가 파일 시스템을 공유하거나 NCCL 그룹을 형성할 수 있는 단일 노드 또는 클러스터 환경에서는 쉽습니다.

우리가 원하는 것은 Hugging Face Jobs를 사용하여 동일한 구성을 실행하는 것입니다. 본질적으로 HF Job은 하나의 VM에서 실행되는 하나의 컨테이너입니다. 이는 하나의 Job이 트레이너와 vLLM 서버 군(fleet)을 호스팅하기 위해 여러 노드를 스폰할 수 없다는 것을 의미합니다 (적어도 지금은). AsyncGRPOTrainer는 바로 그러한 규모에 맞춰 구축되었기 때문에, 질문은 다음과 같았습니다: 트레이너와 추론 서버가 같은 노드를 공유해야 한다는 요구 사항을 제거한다면 어디까지 할 수 있을까?

음, 전체 가중치(full-weight) 동기화(sync)를 사용한다면 답은 '많지 않다'일 것입니다. 모든 업데이트는 기가바이트 단위로 머신 간에 이동해야 하는데, 이것이 밀집 클러스터에서 NCCL이 하는 역할이지만, Jobs는 노드 간 통신을 할 수 없습니다. 공유된 로컬 디스크도 없고 당연히 공유 localhost도 없습니다.

LoRA를 사용하면 동기화는 단지 몇 메가바이트에 불과합니다. 파일 시스템 부분의 경우, HF Jobs는 Storage Bucket으로 백업되는 볼륨(volumes)을 제공합니다! 이 버킷들은 모든 Job에서 FUSE 파일 시스템으로 마운트될 수 있으며, 노드 간 공유 FS로 작동하기에 충분합니다. 아예 Job 간 네트워크 경로가 필요하지 않습니다.

설정은 상당히 작게 끝났습니다:

• LoRA(그리고 FSDP, 나중에 더 자세히 다룹니다)를 사용하는 AsyncGRPOTrainer 트레이너 Job
• 기본 모델과 트레이너가 마지막으로 게시한 어댑터(adapter)를 서비스하는 두 개의 vLLM Job
• 세 곳 모두에서 동일한 경로로 마운트되는 Storage Bucket (이것이 어댑터를 트레이너에서 서버로 전달하는 방식입니다.)
• 프록시 서버. 왜 필요한지 더 깊이 파고들겠지만, 높은 수준에서는 각 롤아웃(rollout)을 KV 캐시를 보유할 가능성이 가장 높은 복제본(replica)으로 라우팅하고 모든 어댑터 업데이트를 모든 vLLM 복제본에 브로드캐스트하는 프록시가 필요합니다.

AsyncGRPOTrainer의 새로운 어댑터 전용 동기화 경로

작동 방식은 이렇습니다. 트레이너는 vLLM으로 텐서를 전송하지 않습니다. 몇 번의 옵티마이저 스텝마다, 어댑터를 <output_dir>/.vllm_lora/trl-policy-v{N}에 저장하고, 원자적 이름 변경(atomic rename)을 통해 해당 디렉터리를 게시한 다음, 그 경로를 vLLM의 /v1/load_lora_adapter 엔드포인트로 전송합니다. vLLM이 디스크에서 파일을 로드하므로, 롤아웃 워커는 model="trl-policy-v{N}"을 요청할 수 있습니다.

이것이 vLLM에서 런타임 어댑터 로딩이 이미 작동하는 방식입니다. 이 엔드포인트는 텐서가 아닌 경로를 받기 때문에, 트레이너와 서버는 파일 시스템을 공유해야 합니다. Slurm 클러스터의 경우, 이는 네트워크 파일 시스템(network filesystem)입니다. Jobs의 경우, 앞서 언급했듯이 Storage Bucket을 볼륨으로 마운트하여 모든 Job에서 동일한 경로에 접근할 수 있습니다. 내부적으로는 hf-mount를 사용하는데, 이 도구는 해당 버킷을 컨테이너 내부에 POSIX 파일 시스템으로 노출합니다:

# 모든 Job이 같은 절대 경로의 같은 버킷을 받음
hf jobs run ... -v hf://buckets/aminediroHF/asyncgrpo-lora-buckets:/lora ...

TRL이나 vLLM에서 변경해야 할 것은 아무것도 없었습니다. 트레이너는 /lora/<run>/.vllm_lora/에 쓰기하고, 서버는 같은 경로에서 읽습니다. POST 요청에 전송되는 경로는 모든 컨테이너 내부에서 이미 유효합니다.

또한 체크포인트와 최종 어댑터도 버킷에 저장한다는 점에 주목하십시오. HF Jobs는 일시적(ephemeral)이지만, 중단된 트레이너가 훈련을 재개할 수 있습니다. 왜냐하면 최종 어댑터는 항상 버킷에 영구 저장되며 Job이 중지되어도 손실되지 않기 때문입니다.

각 복제본은 하나의 GPU와 기본 vllm/vllm-openai 이미지를 사용합니다. 우리는 런타임 LoRA 로딩을 활성화하고 충분한 어댑터 슬롯만 확보하면 됩니다.

어댑터 슬롯의 개수는 max_staleness에 따라 결정됩니다. AsyncGRPOTrainer에서는 모든 가중치 동기화(weight sync)가 정책 버전을 1 증가시키며, max_staleness는 트레이너가 해당 샘플을 폐기하기 전에 롤아웃 샘플이 현재 정책보다 얼마나 뒤처질 수 있는지를 나타냅니다. max_staleness=4인 경우, 트레이너가 v7에 있을 때도 trl-policy-v3에서 생성된 샘플은 여전히 훈련에 사용됩니다.

. v3에서 시작된 롤아웃은

v3에서도 완료될 수 있어야 합니다.

따라서 vLLM은 현재 정책과 그 이전 네 개를 서비스해야 합니다. 이것이 트레이너가 max_staleness + 1개의 어댑터 버전을 등록하고 가장 오래된 것을 언로드하는 이유입니다. 각 동기화(sync)는 가장 오래된 것을 언로드하기 전에 새 버전을 로드하므로, 스왑 과정에서 한 슬롯이 더 필요합니다. 이것이 --max-loras 6을 제공합니다.

만약 다섯 개만 사용한다면, vLLM은 매번 동기화 시점에 여전히 진행 중인 롤아웃(rollouts in flight)을 가진 정책을 조용히 제거할 것입니다.

# --expose 8000 reachable at https://<job_id>--8000.hf.jobs
# -v ...:/lora:ro read-only: 서버는 어댑터만 읽습니다
# VLLM_ALLOW_RUNTIME_LORA_UPDATING=1은 /v1/load_lora_adapter를 활성화합니다
...

우리는 vLLM을 v0.27.1로 고정했습니다.

vLLM은 빠르게 변화하며, 위의 플래그들과 런타임 LoRA 엔드포인트가 버전이 노출하는 부분들이므로, 버전을 레시피의 일부로 간주해야 합니다.

트레이너는 최신 어댑터만 유지하고 항상 같은 이름으로 게시하는 또 다른 가능한 설계가 있습니다. 하지만 우리는 그 방법을 사용하지 않았는데, vLLM은 어댑터 이름별로 접두사 캐시(prefix cache)를 키하기 때문입니다. 단일 이름을 사용할 경우, 이전 가중치로 계산된 KV 블록이 스왑 후에도 여전히 일치할 것이므로, 프리필(prefill)을 다시 수행하지 않아 롤아웃이 한 정책 버전에서 접두사를 얻고 다음 버전에서 디코딩을 할 수 있습니다. 트레이너는 이를 알 방법이 없으며, 이는 ratio가 1에서 벗어나는 형태로 나타날 것입니다.

버전별 이름(Versioned names)은 이것을 불가능하게 만듭니다: 이름은 항상 하나의 가중치 세트를 의미하며, 캐시된 접두사는 결코 더 새로운 버전과 일치할 수 없습니다.

우리는 sail/Sanity-Test-R1D-1.5B를 선택했습니다.

이는 FP16을 이용한 훈련-추론 불일치 극복(Defeating the Training-Inference Mismatch via FP16)의 데이터셋입니다 (Qi et al., 2025). 재현 코드는 sail-sg/Precision-RL에 있습니다.

저자들은 DeepSeek-R1-Distill-Qwen-1.5B를 사용하여 각 MATH 문제에 대해 40개의 답변을 생성했습니다. 이들은 성공률이 20%에서 80% 사이인 문제를 유지하여 총 1,460개의 질문을 얻었습니다. 이 데이터셋은 모델이 이미 해결했거나 완전히 희망 없는 문제가 아니기 때문에, 모델이 학습하고 개선할 수 있는 좋은 초기 신호(early signal)를 받을 수 있다는 점에서 RL 검증에 매우 적합합니다.

이는 견고한 end-to-end 테스트로서 훌륭합니다. 만약 vLLM 레플리카 중 하나가 어댑터 이름으로 기본 모델을 조용히 서비스한다면, 우리는 몇십 스텝 내의 곡선에서 그것을 보고 싶습니다. 또한 이 데이터셋은 2시간이 채 걸리지 않아 순환(cycle)할 수 있을 만큼 작습니다.

또한 저희는 하이퍼파라미터로 논문의 LoRA 스크립트인 oat/scripts/lora의 설정을 사용했습니다:

: Qwen/Qwen2.5-Math-1.5B

, LoRA rank 1에 alpha 2, 학습률(learning rate)은 4e-5, 프롬프트당 8개 샘플, 스텝당 128개의 완료(completions), 최대 3,000 생성 토큰, 그리고 4,096 토큰 컨텍스트를 사용했습니다.

트레이너는 TRL이 설치된 동일한 vllm/vllm-openai:v0.27.1 이미지를 사용합니다. 저희가 당시 PR 브랜치를 실행했으며, 이제 같은 코드가 TRL v1.14에 포함되어 있습니다. 학습 스크립트는 일반적인 AsyncGRPOTrainer 스크립트입니다. Job별로 특정한 값은 출력 디렉토리와 서버 URL뿐입니다.

from peft import LoraConfig
from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer
config = AsyncGRPOConfig(
...

초기화(initialization) 중 TRL은 /server_info를 호출합니다. 만약 lora_config를 발견하면, 어댑터 전용 동기화(adapter-only sync)를 사용합니다. vLLM이 직접 서비스할 수 없는 설정들, 예를 들어 DoRA, modules_to_save, 또는 --max-lora-rank보다 높은 랭크는 경고와 함께 병합 가중치 동기화(merged-weight sync)로 폴백됩니다. 로그에는 Adapter-only vLLM sync enabled가 포함되어야 합니다.

이제 재미있는 부분으로 넘어가겠습니다. 저희는 트레이너와 vLLM Job 사이에 프록시(proxy)가 필요한데, 그 이유는 두 가지입니다:

노출된 Job 포트는 Authorization: Bearer <HF token>을 필요로 합니다.

요청마다 헤더를 추가해야 합니다. 프록시는 이 헤더가 추가되는 곳이므로 TRL은 이를 알 필요가 없습니다. 우리는 여러 GPU에서 생성하는 것보다 더 많은 것을 원합니다. 단일 vLLM 서버의 경우, 일반적으로 그렇게 하는 방법은

--data-parallel-size > 1

이지만, TRL은 좋은 이유로 해당 모드에서 어댑터 전용 동기화(adapter-only sync)를 거부합니다: /v1/load_lora_adapter 호출은 응답하는 DP(Data Parallelism) 랭크에만 도달하기 때문에, 다른 랭크들은 새로운 정책 이름 하에서 계속 기본 모델을 제공하게 됩니다. Jobs의 경우 이 질문 자체가 발생하지 않습니다. 왜냐하면 각 복제본이 자체 머신이기 때문입니다. 따라서 데이터 병렬성은 한 단계 위, 즉 모든 복제본으로 어댑터 로드를 분산시키는 무언가에 존재해야 합니다.

따라서 트레이너 Job에서 127.0.0.1:8000에 작은 프록시를 실행하고 TRL이 마치 단일 vLLM 서버인 것처럼 이 주소를 가리키도록 합니다. 헤더 추가 외에도, 이 프록시는 기능적으로 두 가지 일을 수행합니다:

• 각 완료 요청을 하나의 복제본으로 전송하며, 이 복제본은 해당 프롬프트의 접두사(prefix)가 이미 캐시된 곳으로 선택됩니다 (자세한 내용은 아래 참조).
• 어댑터 로드, 일시 중지 및 재개와 같은 모든 상태 변경 요청을 모든 복제본에 브로드캐스트하여 정책 이름이 어디서든 동일한 의미를 갖도록 합니다.

이것이 왜 중요한지에 대한 간단한 상기입니다. 완료 생성을 하는 것은 매우 다른 워크로드 프로파일을 가진 두 단계로 나뉩니다:

• 프리필(prefill)은 전체 프롬프트를 한 번에 처리하고 모든 프롬프트 토큰의 어텐션 키와 값(attention keys and values)을 계산합니다. - 디코드(decode) 단계는 그 다음 한 번에 하나의 토큰을 생성하며, 각 새로운 토큰은 자신 앞에 있는 모든 토큰의 키와 값에 주의를 기울입니다.

이러한 키와 값들이 바로 KV 캐시입니다. 어텐션이 인과적(causal)이기 때문에, 토큰의 KV는 오직 그 앞의 토큰에만 의존하며 뒤따르는 것에 의존하지 않습니다. 따라서 접두사를 공유하는 두 요청은 해당 접두사의 KV를 공유하므로, 이미 캐시에 가지고 있는 복제본은 프리필의 이 부분을 완전히 건너뛸 수 있습니다. 이제 전체 게임은 그 복제본을 찾는 것이며, 이를 통해 요청이 이미 자신의 접두사를 본 복제본에 착륙하여 혜택을 받을 수 있게 됩니다.

vLLM은 접두사 KV 캐시를 16 토큰 블록 단위로 저장합니다. GRPO 덕분에, 롤아웃 워커는 동일한 프롬프트(G=8인 경우)를 가진 G개의 요청을 전송합니다. 이들이 모두 같은 복제본에 도달하면, 첫 번째 요청이 프리필(prefill)을 계산하고 나머지 일곱 개는 이를 재사용할 수 있습니다. 라운드 로빈 라우팅을 사용하면, 절반은 접두사가 캐시되지 않은 복제본으로 가고, 그 네 개의 요청은 프리필 작업을 다시 수행하여 귀중한 GPU 컴퓨팅 자원을 낭비하게 됩니다.