GPU 후퇴 없는 효율성: TRL 에서 Co-located vLLM 을 통한 효율성 해제

TRL 은 최근 DeepSeekMath 논문에서 소개된 온라인 학습 알고리즘인 GRPO 를 사용하여 LLM 을 훈련할 수 있도록 지원합니다. GRPO 는 모델이 자신의 출력으로부터 학습합니다: 훈련 중에 응답을 생성하고 피드백을 받으며, 이 피드백을 통해 시간이 지남에 따라 자신을 개선합니다.

이는 생성이 훈련 루프에서 중요한 단계 — 그리고 또한 주요 병목 현상임을 의미합니다. 생성 속도를 높이기 위해 TRL 은 vLLM 과 통합됩니다. 이 조합은 GRPO 설정에서 강력한 모델을 더 효율적으로 훈련할 수 있게 합니다. 그러나 한 가지 문제가 있습니다.

TRL v0.18.0 이전에는 vLLM 은 서버 모드에서만 지원되었으며, 다른 GPU 에서 별도의 프로세스로 실행되었습니다. 이는 HTTP 를 통해 훈련 스크립트와 통신하여 모듈러하고 사용하기 쉬웠지만, 또한 GPU 비효율성을 도입했습니다.

다음과 같은 일이 발생합니다:

• 훈련 중에는 모델이 빈번하게 완료를 생성해야 합니다.
• 트레이너는 vLLM 서버에 요청을 보냅니다 (서버는 자체 GPU 에서 실행됨).
• vLLM 이 생성하는 동안 훈련 GPU 는 대기합니다.
• 생성이 끝나면 vLLM GPU 가 대기하고 훈련이 재개됩니다.

이러한 훈련과 생성 간의 "핑퐁"은 다음과 같은 결과를 초래합니다:

• 양쪽 모두의 GPU 시간 낭비
• 추론을 실행하기 위해 추가적인 GPU 를 필요로 하는 수요 증가
• 전체적인 처리량 감소 및 높은 비용

온라인 학습 방법 (예: GRPO) — 생성이 지속적으로 발생하는 경우 — 이 비효율성은 더욱 고통스럽습니다. 하드웨어 비용을 더 지출하지만 예상되는 성능을 얻지 못합니다.

따라서 핵심 질문은 다음과 같습니다: 훈련과 생성을 분리하는 대신 동일한 GPU 를 공유할 수 있을까요?

주된 문제는 훈련과 추론이 별도의 GPU 에서 실행되어 대기 시간과 자원 활용도가 낮아진 것입니다. 자연스러운 해결책? 동일한 GPU 에서 두 가지 모두를 실행합니다. vLLM 을 독립적인 서버로 자체 프로세스와 장치에서 실행하는 대신, vLLM 이 분산 프로세스 그룹 내에서 훈련 코드와 함께 실행할 수 없을까요? 이는 훈련과 추론이 동일한 장치를 공유하고 효율적으로 작업을 전환하며 자원을 낭비하지 않는 단일 분산 작업을 시작할 수 있게 합니다.

이 접근법은 우리가 콜로케이션이라고 부릅니다. 훈련과 추론은 동일한 GPU 에서 콜로케이션되며, 동일한 프로세스 그룹을 통해 조정되어 서로를 기다리지 않고 순차적으로 전환합니다. 추가적인 하드웨어가 필요 없습니다.

기존에는 TRL 이 vLLM 을 외부 HTTP 서버로 의존했으므로 불가능했습니다. 이는 PR #3394 가 vLLM 의 외부 라운처와 훈련 프로세스에 대한 진정한 통합을 추가함으로써 변경되었습니다.

통합 실행: 동일한 프로세스 그룹에 vLLM 을 내장함으로써, 훈련과 추론 작업이 동일한 GPU 를 공유하고 서로를 기다리지 않고 전환할 수 있습니다. 이는 대기 시간을 줄이고 전체적인 효율성을 향상시킵니다.HTTP 통신 생략: REST API 호출 또는 네트워킹 필요 없음 — vLLM 은 훈련 루프와 함께 실행되어 오버헤드 및 지연을 피합니다.Torchrun 호환성: torchrun과 원활하게 작동합니다.

따라서 노드 간 확장 시 최소 설정 변경으로 인해 스케일링이 용이합니다.TP 및 DP 지원: Tensor Parallelism (TP) 과 Data Parallelism (DP) 와 호환되며, 대규모 학습 실행에 적합합니다.SPMD 실행 패턴: 각 GPU 가 엔진의 자체 인스턴스를 동기화하여 실행하는 Single Program, Multiple Data (SPMD) 모델을 사용합니다. 분산 다중 GPU, 다중 노드 환경에 이상적입니다.간소화된 배포: 별도의 서버 스크립트를 유지할 필요가 없으며, vLLM 은 학습 작업 내부에서 직접 시작되고 제어됩니다.향상된 성능: 비활성 GPU 를 피하고 프로세스 간 통신을 제거함으로써 특히 온라인 학습 환경 (예: GRPO) 에서 더 빠른 학습 및 생성 속도를 제공합니다.강력한 프로세스 간 통신: 이는 독립적인 프로세스 간 분산 프로세스 그룹 설정의 복잡성을 피하기 때문에 서버 모드보다 더욱 견고합니다.

이 기능 덕분에 코-로케이션된 학습과 추론은 이제 해킹이 아닌 일류, 확장 가능하고 프로덕션 준비된 상태입니다.

서버 TRL 에서 코-로케이션 TRL 로의 전환은 더 똑똑한 GPU 사용에 관한 것입니다. 아래 다이어그램은 차이를 보여줍니다:

서버 TRL 설정에서 학습과 추론은 별도의 GPU 에서 실행됩니다. 예를 들어:

• GPUs 0~2 는 학습에 사용됩니다.
• GPU 3 은 vLLM 을 별도의 서버로 실행하는 데 전적으로 할당됩니다.

학습 단계 동안 GPU 3 은 비활성 상태입니다.
생성 단계 (추론) 동안, GPUs 0~2 는 비활성 상태이면서 GPU 3 이 출력을 생성합니다.

이는 다음과 같은 결과를 가져옵니다:

• 장치 간 자주 대기하는 비효율적인 GPU 사용
• 추론 전용으로 추가 GPU 를 provision(제공)
• 비용 및 복잡성 증가

반면, 코-로케이션 TRL 설정은 학습과 vLLM 을 동일한 GPU에서 실행합니다. 각 GPU 는:

• 학습 루프를 실행하며,
• 동일한 프로세스 내에서 vLLM 엔진을 시작합니다.

학습과 추론은 GPU 리소스를 번갈아 가며 사용합니다 — 전용 장치 또는 별도의 프로세스가 필요 없습니다.

이 디자인은:

• 비활성 시간을 줄이고,
• 프로세스 간 및 HTTP 통신을 최소화하며,
• 사용 가능한 GPU 메모리와 컴퓨팅 능력을 완전히 활용하며,
• 하드웨어 요구 사항을 증가시키지 않고 더 빠른 성능을 제공합니다.

vLLM 을 서버로 시작하는 대신, 트레이너는 이제 외부 로커를 사용하여 vLLM 을 **프로세스 내 (in-process)**로 시작합니다. 아래에 표시된 바와 같습니다:

self.llm = LLM(
model=model.name_or_path,
tensor_parallel_size=args.vllm_tensor_parallel_size,
...

코-로케이션 vLLM 은 torch.distributed 프로세스 그룹과 랭크 구조를 존중합니다. 이는 vLLM 을 학습과 함께 충돌 없이 초기화할 수 있게 하며, TP/DP 설정을 원활하게 작동시킵니다:

if self.vllm_tensor_parallel_size > 1:
# TP 를 위한 랭크의 서브그룹 생성, 각 그룹은 `vllm_tensor_parallel_size` 개의 랭크를 포함합니다.
self.tp_group, _ = torch.distributed.new_subgroups_by_enumeration(
...

코-로케이션 vLLM 은 이제 REST API 에 의존하지 않으며, 메모리 내에서 직접 실행되고 네이티브 Python 호출을 통해 통신합니다:

if self.vllm_tensor_parallel_size > 1:
orig_size = len(prompts_text)
gathered_prompts = [None for _ in range(self.vllm_tensor_parallel_size)]
...

이 설정을 사용하려면 GRPO 구성에서 vllm_mode="colocate" 를 설정하면 됩니다:

training_args = GRPOConfig(
...,
use_vllm=True,
...

참고: 모델 크기와 학습에 대한 전체 GPU 메모리 요구 사항에 따라, 과소 활용 또는 메모리 부족 오류를 피하기 위해 GRPOConfig 의 vllm_gpu_memory_utilization 파라미터를 조정해야 할 수 있습니다.

colocation 의 영향을 측정하기 위해 우리는 전통적인 server mode(vLLM 이 별도의 GPU 에서 독립 서버로 실행되는 경우) 과 새로운 co-locate mode(학습과 추론이 동일한 GPU 를 공유하는 경우) 를 비교하는 일련의 실험을 수행했습니다.

server mode에서는 vLLM 추론 서버가 1 개 GPU 를 완전히 전용으로 사용하므로, 학습에는 7 개 GPU 만 사용됩니다.

co-locate mode에서는 모든 8 개 GPU 가 학습에 사용되며, 이는 기본값으로 유효 배치 크기를 증가시킵니다.

공정한 비교를 보장하기 위해 우리는 server mode 의 Throughput 을 8/7 배로 정규화했습니다. 이 조정은 co-locate mode 의 더 큰 학습 능력을 반영하며, 두 설정을 동일한 학습 조건 하에서 비교할 수 있도록 합니다.

• 배치 크기가 증가함에 따라 두 설정 모두 성능이 향상됩니다.
  Co-located setup 은 최대 배치 크기에서 1.43× speedup까지 도달합니다.- co-locate mode 에서 더 큰 배치는 공유 GPU 메모리를 더 잘 활용합니다.

• co-located setup 에서 TP 를 증가시키면 성능이 감소합니다. - 더 많은 sharding 은 더 많은 통신 오버헤드를 유발하며, 이는
  작은 모델에는 불리하지 않습니다. Takeaway: 작은 모델의 경우 co-locate mode 에서 과도한 sharding 을 피하세요.

• 다시 말하지만, co-located mode 는 배치 크기와 함께 더 잘 확장됩니다. - 테스트된 최대 배치 크기에서 gains 은 1.35× speedup까지 도달합니다.

• 1.5B 모델과 반대되는 경향입니다.

• 7B 모델의 경우,
  TP 를 증가시키면 성능이 향상되어 최대 1.73× speedup까지 도달합니다. 큰 모델은 co-locate setup 에서 sharding 을 이득으로 합니다.

Qwen2.5-Math-72B와 같은 큰 모델을 학습할 때는 많은 GPU 와 노드에서 효율적이며 확장 가능하고 안정적인 학습을 위해 올바른 전략을 사용해야 합니다. 우리의 설정에서는 co-located vLLM과 몇 가지 주요 최적화를 결합하여 이를 효율적으로 실행했습니다.

co-located training 을 사용할 때, 학습과 추론이 동일한 장치에서 원활하게 실행되도록 GPU 메모리 관리는 매우 중요합니다. 이를 지원하기 위해 우리는 GRPO 학습 루프에 vLLM 의 sleep() API 를 추가했습니다.

sleep() 함수는 vLLM 엔진을 일시적으로 중지하고 GPU 메모리를 방출합니다. 이는 두 가지 수준을 지원합니다:

Level 1: GPU 에서 모델 가중치를 언로드 (CPU 메모리에 유지) 하고 KV cache 를 정리합니다. 같은 모델이 곧 재사용될 때 유용합니다.Level 2: 모델 가중치와 KV cache 를 모두 완전히 언로드합니다. 모델이 변경되거나 곧 재사용되지 않을 경우 가장 좋습니다.

GRPO 에서 모델은 각 단계마다 업데이트되므로, 우리는 Level 2 sleep을 사용합니다.

Level 2 sleep 의 이점:

학습용 최대 GPU 메모리 확보

• 학습과 생성 사이의 메모리 경쟁 방지
• 큰 모델 (예: Qwen2.5-72B) 일컬어도 colocation 을 효율적으로 유지합니다

이 작은 추가는 원활하고 확장 가능한 co-located training 을 가능하게 하는 데 큰 차이를 만듭니다.

Qwen2.5-72B 와 같은 큰 모델을 학습하기 위해 우리는 DeepSpeed ZeRO Stage 3에 의존하며, 이는 plain TRL 에서 사용된 동일한 설정입니다.

ZeRO 는 GPU 간 메모리 분배를 통해 큰 모델을 확장합니다. Stage 3 은 더 나아가 다음을 분할합니다:

• 모델 가중치
• 그래디언트
• 옵티마이저 상태

이는 단일 GPU 에 맞지 않는 모델에 필수적입니다. ZeRO Stage 3 에서 각 GPU 는 모델의 일부만 처리합니다.

추가 옵션은 다음과 같습니다:

"offload_optimizer": {"device": "cpu"}

옵티마이저 상태를 CPU 로 이동하여 GPU 메모리를 방출합니다 - co-locate setup 에서 필수적입니다."overlap_comm": true

계산과 통신 오버레이를 활성화하여 학습 속도를 높입니다."contiguous_gradients": true

기억 블록 단위로 그래디언트를 할당하여 메모리 접근성을 개선하고 분할을 줄입니다.

이러한 최적화들은 72B 모델을 효율적으로 훈련하는 데 도움이 되며, 제한된 메모리 제약 하에서 코로케이션 (colocation) 이 안정적으로 유지되도록 보장합니다.

TRL 에서 권장한 대로 우리는 Accelerate라는 경량 라이브러리를 사용합니다. 이는 분산 훈련을 단순화하며 다음 기능을 처리합니다:

• 다중 GPU 및 다중 노드 작업 시작
• 데이터 병렬성 (Data parallelism)
• 그래디언트 누적 (Gradient accumulation)
• 분산 데이터 로딩

이로 인해 설정은 깔끔하고 확장 가능하며 유지 관리가 쉽습니다.

4 개 GPU 가 더 적게 사용되더라도, 코로케이션 설정은 일반 TRL 보다 약 1.26 배 빠릅니다.
이는 sleep() 를 사용한 스마트한 GPU 공유 및 메모리 정리의 효과성을 강조합니다.

코로케이션 및 일반 설정의 훈련 보상 (reward) 플롯은 거의 동일하며, 이는 다음과 같습니다:

우리는 Math500 벤치마크에서 세 가지 모델을 평가했습니다: Base 모델, Co-locate 훈련된 모델, Plain 훈련된 모델. 두 훈련된 모델 모두 Base 를 상회하며, co-locate 모델은 일반 훈련된 모델과 동등한 성능을 보이며 이는 코로케이션이 다운스트림 성능에 영향을 미치지 않는 것을 확인합니다.

우리는 co-located vLLM 으로 GRPO 훈련을 확장하는 작업을 통해 몇 가지 중요한 과제를 직면했으며, 대규모 모델을 훈련할 때 효율성, 유연성 및 시스템 설계에 대해 중요한 교훈을 얻었습니다.

vLLM ≥ 0.8.0 의 Tensor Parallelism 버그.Tensor Parallelism (TP) 는 vLLM 버전 0.8.0 이상에서 external_launcher 와 함께 작동하지 않습니다. 이는 Issue #15895 에서 추적되었습니다. 중단점을 식별하기 위해 우리는 vLLM 개발자 블로그 게시글에 설명된 방법을 따랐으며, 이 게시물은 각 커밋에 대한 wheels 를 제공합니다. 바이섹팅 (bisect) 후 중단 커밋을 cc10281 로 식별했습니다. 근본 원인은 결정성 (determinism) 이었습니다. 새 버전은 랜덤 시드를 명시적으로 설정해야 하므로, 시드를 설정하면 문제가 사라졌습니다.

Level 2 Sleep Buffer 버그.초기에는 load_weights 를 사용하여 가중치를 다시 로드할 때 level 2 sleep 이 제대로 작동하지 않았습니다. 이는 Issue #16564 에서 추적되었습니다. 문제는 모델 버퍼 (예: BatchNorm 의 running mean/var) 가 sleep 에서 깨어날 후 복원되지 않았다는 것입니다. PR #16889 를 통해 문제가 해결되었으며, 이는 level 2 sleep 에서 깨어날 때 버퍼를 명시적으로 복원하는 로직을 추가했습니다. 우리는 이제 원래 버퍼의 복사본을 유지하며 새로운 가중치를 로드할 후 수동으로 재적용합니다.

종료 시 Segmentation Fault.vLLM sleep 은 훈련이 끝날 때 프로세스를 닫을 때 segmentation fault 를 일으키는 문제가 여전히 열려 있습니다. 이는 Issue #16993 에서 보고되었습니다. 이 크래시 (crash) 는 종료 중 발생하며 훈련 자체를 파괴하지 않으므로 우리는 블로그에서 공유된 모든 데모 및 실험을 완료할 수 있었습니다. 그러나 우리는 sleep() 를 TRL upstream 으로 완전히 통합하기 전에 공식적인 수정 사항을 기다리고 있습니다.

이러한 과제는 블록터가 아니었으며, vLLM 이 메모리와 병렬성을 어떻게 관리하는지에 대한 깊은 이해를 필요로 하는 신중한 디버깅 및 버전 관리를 요구했습니다.

Co-located inference 는 GPU 활용도를 크게 개선합니다. 훈련 및 생성이 동일한 GPU 를 공유하도록 허용함으로써 대기 시간을 제거하고 하드웨어 요구 사항을 줄이며, 더 적은 GPU 와도 높은 트루풋을 달성할 수 있습니다.

vLLM 의 sleep() 기능은 대규모 코로케이션에 필수적입니다. 이는 메모리 사용을 세밀하게 제어할 수 있게 하며, 훈련이 생성 단계 사이에서 GPU 메모리를 완전히 반환할 수 있게 합니다. 이는 Qwen2.5-72B 와 같은 모델의 핵심 활성화 요소입니다.