Hugging Face TRL: 효율적인 GRPO 학습을 위한 공동 배치 vLLM

Hugging Face는 TRL 라이브러리 내에서 vLLM을 공동 배치(co‑locate)하도록 지원을 추가했습니다. 이를 통해 학습과 추론이 동일한 GPU를 공유할 수 있어, 생성 중에 학습 GPU가 유휴 상태가 되고 추론 GPU가 학습 중에 유휴 상태가 되는 "핑‑퐁" 비효율성을 없애고 처리량을 높이며 하드웨어 요구량을 감소시킵니다.

서버 모드 추론의 비효율성

TRL v0.18.0 이전에는 vLLM이 서버 모드에서만 지원되었습니다. 이 구성에서는 vLLM이 전용 GPU에서 별도 프로세스로 실행되며, 학습 스크립트와는 HTTP를 통해 통신합니다. 이러한 아키텍처는 생성이 지속적으로 일어나는 Group Relative Policy Optimization (GRPO) 같은 온라인 학습 알고리즘에서 큰 병목을 초래합니다:

  • GPU 활용도 저하: vLLM 서버가 완성을 생성하는 동안 학습 GPU가 유휴 상태가 됩니다.
  • 자원 낭비: 추론 전용 GPU를 별도로 프로비저닝해야 하므로 비용이 증가합니다.
  • 통신 오버헤드: REST API 호출에 의존함으로써 네트워크 지연이 발생합니다.

공동 배치 설계 및 구현

공동 배치는 vLLM이 동일한 분산 프로세스 그룹 내에서 학습 코드와 함께 실행되도록 합니다. 외부 서버 대신, 트레이너는 외부 런처를 사용해 프로세스 내부에서 vLLM을 시작합니다.

주요 기술 기능

  • 통합 실행: 학습과 추론 작업이 동일한 GPU 자원을 번갈아 사용해 유휴 시간을 줄입니다.
  • 직접 메모리 통신: HTTP 호출을 제거하고, vLLM이 네이티브 Python 호출을 통해 학습 루프와 통신합니다.
  • 분산 호환성: 시스템은 torchrun, Tensor Parallelism (TP), Data Parallelism (DP)과 호환되며, Single Program, Multiple Data (SPMD) 실행 패턴을 활용합니다.
  • 배포 간소화: vLLM이 학습 작업 내부에서 직접 제어되므로 별도 서버 스크립트가 필요 없습니다.

설정

사용자는 GRPOConfig에서 vllm_mode="colocate"를 지정하여 이 기능을 활성화할 수 있습니다:

training_args = GRPOConfig(
    ...,
    use_vllm=True,
    vllm_mode="colocate",
)

성능 벤치마크

서버 모드(8 GPU 중 1 GPU를 vLLM 전용으로 사용)와 공동 배치 모드(전체 8 GPU를 학습에 사용) 간의 실험 결과는 일관된 처리량 향상을 보여줍니다. 공정성을 위해 서버 모드 처리량은 8/7 배율로 정규화되었습니다.

모델별 결과

  • 1.5B 모델: 가장 큰 배치 크기에서 최대 1.43배 속도 향상을 달성했습니다. 다만 Tensor Parallelism (TP)을 늘리면 통신 오버헤드로 성능이 감소했습니다.
  • 7B 모델: 다양한 배치 크기에서 최대 1.35배, TP를 늘렸을 때는 최대 1.73배 속도 향상을 기록했으며, 큰 모델일수록 샤딩 효과가 더 크게 나타났습니다.
  • 72B 모델 (Qwen2.5‑Math‑72B): 공동 배치 설정이 일반 TRL 대비 약 1.26배 빠르게 동작했으며, GPU 4대를 적게 사용하면서도 성능이 유지되었습니다.

대형 모델(72B+) 확장

Qwen2.5‑Math‑72B와 같은 대형 모델을 학습하려면 학습과 생성 간 메모리 충돌을 방지하기 위한 고급 메모리 관리가 필요합니다.

vLLM 슬립 모드

Hugging Face는 GRPO 루프에 vLLM의 sleep() API를 통합했습니다. Level 2 슬립을 사용하면 모델 가중치와 KV 캐시를 GPU에서 완전히 언로드하여 학습 단계에 사용할 자유 메모리를 최대화하고 메모리 충돌을 방지합니다.

DeepSpeed 및 Accelerate 최적화

안정성과 효율성을 유지하기 위해 다음 최적화가 적용됩니다:

  • DeepSpeed ZeRO Stage 3: 모델 가중치, 그래디언트, 옵티마이저 상태를 GPU에 걸쳐 파티셔닝합니다.
  • CPU 오프로드: 옵티마이저 상태를 CPU("offload_optimizer": {"device": "cpu"})로 이동시켜 GPU 메모리를 확보합니다.
  • 메모리 관리: contiguous_gradients를 활성화해 메모리 단편화를 줄이고, overlap_comm을 사용해 학습 속도를 높입니다.
  • Accelerate 통합: 다중 GPU/다중 노드 실행 및 데이터 병렬을 위해 사용됩니다.

도전 과제 및 교훈

공동 배치 vLLM 구현 과정에서 여러 기술적 난관이 발견되었습니다:

  • vLLM ≥ 0.8.0에서의 결정성: 외부 런처와 함께 사용되는 Tensor Parallelism에서 무작위 시드를 명시적으로 설정해야 하는 버그가 있었습니다.
  • 버퍼 복원: Level 2 슬립 버그(Issue #16564)로 인해 슬립 후 모델 버퍼(예: BatchNorm의 running mean/var)를 명시적으로 복원해야 했습니다.
  • 종료 안정성: vLLM 슬립이 학습 프로세스 종료 시 세그멘테이션 오류를 일으키는 알려진 이슈가 존재합니다(Issue #16993). 하지만 이는 학습 결과에 영향을 주지 않습니다.

모델 품질에 대한 결론

Math500 및 AIME24 벤치마크 결과, 공동 배치가 모델 성능을 저해하지 않음을 확인했습니다. 공동 배치와 일반 설정의 보상 곡선은 거의 동일하며, 공동 배치로 학습된 모델은 일반 방식으로 학습된 모델과 동등한 성능을 보입니다. 이는 효율성 향상이 정확도 손실을 동반하지 않음을 입증합니다.

Sources