No GPU left behind: Unlocking Efficiency with Co-located vLLM in TRL

🖼️ 인포그래픽

🖼️ 4컷 인포그래픽

💡 한 줄 요약

TRL의 co-located vLLM은 GRPO의 학습과 생성을 동일한 GPU·분산 프로세스 그룹에서 번갈아 실행해 유휴 시간과 통신 비용을 줄이며, 모델 크기와 텐서 병렬화 설정에 따라 최대 1.73배의 처리량 향상을 보였다.

📌 핵심 요약

• 기존 TRL의 vLLM 서버 모드는 학습 GPU와 생성 GPU를 분리했기 때문에 한쪽이 작업하는 동안 다른 쪽이 쉬는 ‘핑퐁’ 구조와 추가 GPU 비용이 발생했다.
• PR #3394로 도입된 co-locate 모드는 vLLM을 학습 프로세스 안에서 실행하고 동일한 GPU와 torch.distributed 프로세스 그룹을 공유해 별도 HTTP 서버 없이 학습과 생성을 교대로 수행한다.
• 구현은 vLLM의 external_launcher, 텐서 병렬 하위 그룹, 프로세스 간 프롬프트 수집 및 결과 분할을 이용하며 GRPOConfig에서 use_vllm과 vllm_mode="colocate"를 지정해 활성화한다.
• 실험에서는 배치가 커질수록 공동 배치 방식의 이점이 커졌고, 1.5B 모델은 최대 1.43배, 7B 모델은 최대 1.35배의 배치 크기 기반 향상을 기록했다.
• 텐서 병렬화는 작은 1.5B 모델에서는 통신 오버헤드로 성능을 낮췄지만 7B 모델에서는 최대 1.73배 향상을 냈으며, 72B 규모에서는 vLLM Level 2 sleep과 DeepSpeed ZeRO Stage 3로 메모리 경합을 관리했다.

🧩 주요 포인트

1. 기존 TRL의 vLLM 서버 모드는 학습 GPU와 생성 GPU를 분리했기 때문에 한쪽이 작업하는 동안 다른 쪽이 쉬는 ‘핑퐁’ 구조와 추가 GPU 비용이 발생했다.
2. PR #3394로 도입된 co-locate 모드는 vLLM을 학습 프로세스 안에서 실행하고 동일한 GPU와 torch.distributed 프로세스 그룹을 공유해 별도 HTTP 서버 없이 학습과 생성을 교대로 수행한다.
3. 구현은 vLLM의 external_launcher, 텐서 병렬 하위 그룹, 프로세스 간 프롬프트 수집 및 결과 분할을 이용하며 GRPOConfig에서 use_vllm과 vllm_mode="colocate"를 지정해 활성화한다.
4. 실험에서는 배치가 커질수록 공동 배치 방식의 이점이 커졌고, 1.5B 모델은 최대 1.43배, 7B 모델은 최대 1.35배의 배치 크기 기반 향상을 기록했다.
5. 텐서 병렬화는 작은 1.5B 모델에서는 통신 오버헤드로 성능을 낮췄지만 7B 모델에서는 최대 1.73배 향상을 냈으며, 72B 규모에서는 vLLM Level 2 sleep과 DeepSpeed ZeRO Stage 3로 메모리 경합을 관리했다.

🧠 상세 정리

1. GRPO에서 생성이 병목이 되는 이유

TRL은 DeepSeekMath 논문에서 소개된 온라인 학습 알고리즘인 GRPO를 이용한 대규모 언어 모델 학습을 지원한다. GRPO에서는 모델이 학습 도중 직접 응답을 생성하고, 그 응답에 대한 피드백을 받아 점진적으로 개선되므로 생성 단계가 학습 루프의 부가 작업이 아니라 핵심 구성 요소다. 그러나 응답 생성이 반복적으로 일어나는 만큼 생성 속도가 전체 학습 처리량을 제한하는 주요 병목으로 작용한다. TRL은 이를 빠르게 처리하기 위해 vLLM을 통합했지만, 기존 통합 방식은 생성 자체의 속도와 별개로 GPU 배치와 실행 구조에서 새로운 비효율을 만들었다. 글은 이 문제를 해결하기 위해 학습과 추론을 별도 장치에 고정하지 않고 동일한 GPU에서 조정해 실행하는 방식을 제안한다.

2. 기존 서버 모드의 GPU 핑퐁 문제

TRL v0.18.0 이전에는 vLLM을 별도 프로세스와 별도 GPU에서 실행하는 서버 모드만 지원했으며, 학습 스크립트는 HTTP를 통해 생성 요청을 전달했다. 이 구조는 구성 요소가 분리돼 사용하기 쉽다는 장점이 있었지만, vLLM이 응답을 생성하는 동안 학습 GPU는 기다리고 학습이 재개되면 vLLM 전용 GPU가 쉬는 순차적 ‘핑퐁’ 실행을 초래했다. 예시 구성에서는 GPU 0~2가 학습에 쓰이고 GPU 3은 vLLM 서버에 전용으로 할당돼, 어느 단계에서든 일부 장치가 유휴 상태가 된다. 그 결과 양쪽 GPU 시간이 낭비되고 추론만을 위한 추가 장치가 필요하며, 전체 처리량은 낮아지고 비용과 운영 복잡성은 커진다. 생성 호출이 끊임없이 반복되는 GRPO 같은 온라인 학습에서는 이러한 비효율이 누적되기 때문에 투입한 하드웨어에 비해 기대한 성능을 얻기 어렵다.

3. 동일한 GPU를 공유하는 공동 배치 구조

공동 배치, 즉 colocation의 핵심은 학습과 추론을 서로 다른 GPU 집합에 격리하지 않고 동일한 장치와 동일한 분산 프로세스 그룹에 배치하는 것이다. 각 GPU는 학습 루프를 실행하는 동시에 같은 프로세스 안에서 vLLM 엔진도 구동하며, 학습과 생성이 자원을 놓고 동시에 경쟁하는 대신 단계에 맞춰 번갈아 GPU를 사용한다. PR #3394는 vLLM의 external launcher를 지원함으로써 외부 HTTP 서버에 의존하던 TRL에 이러한 통합 실행 방식을 추가했다. 이 구조에서는 추론 전용 GPU나 별도 서버 스크립트가 필요하지 않고, REST 요청과 독립 프로세스 사이의 분산 그룹 구성에 따르는 지연과 복잡성도 줄어든다. 글은 이를 임시적인 우회책이 아니라 확장 가능하고 운영에 사용할 수 있는 일급 기능으로 설명한다.

4. 분산 실행과 병렬화 지원

co-located vLLM은 torchrun과 호환되며, 여러 GPU와 여러 노드로 확장할 때 기존 분산 실행 방식에 큰 구성 변경을 요구하지 않는다. 실행 모델은 각 GPU가 동기화된 자체 엔진 인스턴스를 수행하는 SPMD 방식이며, 데이터 병렬화와 텐서 병렬화를 함께 지원한다. 텐서 병렬 크기가 1보다 크면 전체 rank를 지정된 크기의 하위 그룹으로 나눠 각 그룹이 하나의 텐서 병렬 단위로 동작하도록 구성한다. vLLM은 torch.distributed의 프로세스 그룹과 rank 구조를 따르므로 학습 옆에서 초기화해도 분산 실행 구성이 충돌하지 않도록 설계됐다. 독립된 서버 프로세스끼리 별도의 통신 그룹을 맞추는 방식보다 구조가 단순하며, 네이티브 분산 통신과 동일 프로세스 실행을 통해 대규모 학습에 필요한 확장성을 유지한다.

5. 인프로세스 생성의 실제 구현

트레이너는 vLLM의 LLM 객체를 직접 만들고 distributed_executor_backend를 external_launcher로 지정해 엔진을 학습 프로세스 내부에서 시작한다. 초기화에는 텐서 병렬 크기, GPU 메모리 사용률, 최대 동시 시퀀스 수, 프롬프트와 완성 길이를 합산한 최대 모델 길이 등이 전달되며, 같은 텐서 병렬 그룹의 작업자가 동일한 샘플링 결과를 얻도록 그룹 기준 시드도 설정한다. 텐서 병렬 실행에서는 all_gather_object로 그룹 구성원의 프롬프트를 모아 vLLM의 generate를 직접 호출한 뒤, 각 로컬 rank가 담당하는 범위의 completion token ID만 다시 선택한다. 이 과정은 REST API가 아니라 메모리 내 파이썬 호출과 torch.distributed 통신으로 이루어진다. 사용자는 GRPOConfig에서 use_vllm=True와 vllm_mode="colocate"를 지정할 수 있으며, 모델과 학습 작업의 메모리 요구량에 따라 vllm_gpu_memory_utilization을 조절해 과소 활용이나 메모리 부족을 피해야 한다.

6. 1.5B·7B 모델의 처리량 비교

성능 실험은 vLLM이 별도 GPU에서 동작하는 서버 모드와 학습·추론이 동일한 GPU를 공유하는 co-locate 모드를 비교했다. 서버 모드에서는 8개 GPU 중 1개를 추론 서버에 전용으로 할당해 7개만 학습에 사용했지만, 공동 배치 모드에서는 8개 모두가 학습에 참여해 기본 유효 배치 크기부터 달라진다. 연구진은 같은 학습 조건에서 비교하기 위해 서버 모드의 처리량을 8/7로 보정했으며, 두 모드 모두 배치가 커질수록 처리량이 증가했다. 1.5B 모델은 가장 큰 배치에서 공동 배치 방식이 최대 1.43배의 속도 향상을 보였고, 7B 모델도 배치 크기 증가에 더 잘 확장되며 최대 1.35배 향상됐다. 글은 큰 배치가 공동 배치 환경의 공유 GPU 메모리를 더 효과적으로 활용해 처리량 이점을 키운다고 설명한다.

7. 모델 크기에 따라 달라지는 텐서 병렬화 효과

텐서 병렬화의 효과는 모델 크기에 따라 반대 방향으로 나타났기 때문에, 단순히 병렬화 수준을 높이는 것이 항상 유리하지는 않았다. 1.5B 모델에서는 텐서 병렬 크기가 증가할수록 모델 조각 사이의 통신 오버헤드가 커져 공동 배치 모드의 성능이 오히려 낮아졌다. 따라서 비교적 작은 모델에는 지나친 샤딩을 피해야 한다는 것이 해당 실험의 직접적인 결론이다. 반면 7B 모델에서는 텐서 병렬화를 늘릴수록 처리량이 향상됐고, 실험에서 최대 1.73배의 속도 향상을 기록했다. 즉 큰 모델은 샤딩으로 얻는 계산·메모리 분산 효과가 통신 비용을 상쇄할 수 있지만, 작은 모델은 동일한 구성이 불필요한 통신 부담이 될 수 있으므로 모델 규모에 맞춘 설정이 필요하다.

8. 72B 규모의 메모리 관리와 DeepSpeed 구성

Qwen2.5-Math-72B처럼 큰 모델을 여러 GPU와 노드에서 공동 배치하려면 학습과 추론이 같은 장치 메모리를 번갈아 사용할 수 있도록 적극적인 메모리 관리가 필요하다. 이를 위해 GRPO 학습 루프에 vLLM의 sleep API가 추가됐으며, Level 1은 모델 가중치를 CPU 메모리로 옮기고 KV 캐시를 비우는 반면 Level 2는 모델 가중치와 KV 캐시를 모두 GPU에서 내린다. GRPO는 매 단계 뒤 모델이 갱신되므로 이 구성에서는 Level 2 sleep을 사용해 학습에 제공할 GPU 메모리를 최대화하고 학습과 생성 사이의 메모리 경합을 줄였다. 대형 모델 학습에는 기존 TRL과 마찬가지로 DeepSpeed ZeRO Stage 3를 사용해 가중치, 그래디언트, 옵티마이저 상태를 GPU들에 분할하며, 옵티마이저 상태의 CPU 오프로딩도 활성화한다. 또한 통신과 계산의 중첩 및 연속적인 그래디언트 메모리 배치를 사용해 속도와 메모리 접근을 개선하고 단편화를 줄이도록 구성했으며, 제공된 원문은 이 최적화 설명 도중 끝난다.

🧾 핵심 주장 / 시사점

• 공동 배치의 핵심 이득은 생성 엔진 자체의 속도 향상보다 학습 GPU와 추론 GPU를 고정 분리하면서 생기던 상호 유휴 시간을 없애는 데 있다.
• 텐서 병렬화의 최적값은 일률적이지 않으며, 작은 모델에서는 통신 비용이 이득을 압도할 수 있고 큰 모델에서는 샤딩 효과가 처리량 향상으로 이어질 수 있다.
• 대형 모델의 공동 배치는 단순한 프로세스 통합만으로 완성되지 않으며, vLLM 엔진의 완전한 메모리 해제와 ZeRO Stage 3의 상태 분할·CPU 오프로딩을 함께 사용해야 메모리 경합을 제어할 수 있다.

✅ 액션 아이템

• co-locate 모드 적용 대상은 기존 분리형 vLLM 서버 대비 핑퐁 대기시간과 추가 GPU 비용을 정량 비교해 결정한다.
• GRPOConfig에서 use_vllm=True와 vllm_mode="colocate"를 지정하고 external_launcher·텐서 병렬 하위 그룹·프롬프트 수집/결과 분할 연동을 함께 점검한다.
• 1.5B·7B·72B에서 배치 크기별 처리량 곡선을 따로 비교해 텐서 병렬화가 유리한 구간과 성능 저하 구간을 분리 정한다.

❓ 열린 질문

• 배치 크기를 키울수록 공통 이점이 커진다고 했을 때 실제 개선이 확정되는 최소 배치 규모는 어디인가?
• 1.5B에서 텐서 병렬화가 오히려 성능을 낮춘 이유가 통신 오버헤드였는데, 그 임계 조건은 어떤 설정 조합인가?
• 72B에서 vLLM Level 2 sleep과 DeepSpeed ZeRO Stage 3를 썼을 때 메모리 경합이 더는 줄지 않는 한계는 어디인가?