MIT TLT로 살펴보는 유휴 GPU 활용과 LLM 강화학습 훈련

롤아웃이 멈추면 GPU도 함께 기다린다

LLM 강화학습(RL)에서 계산 시간을 크게 차지하는 구간은 롤아웃이다. 현재 모델이 여러 응답을 생성하고, 그 결과로 보상 신호를 계산하는 단계다. 롤아웃은 전체 RL 훈련 시간의 최대 85%를 소비한다.

분산 환경에서는 응답 길이가 이 병목을 더 키운다. 짧은 응답을 먼저 끝낸 GPU는 긴 응답을 생성하는 다른 GPU가 완료될 때까지 다음 단계로 넘어가지 못한다. 이 대기 구간이 유휴 컴퓨팅 자원으로 남는다.

⤢✕

추측적 디코딩(Speculative Decoding)은 소형 드래프터 모델이 토큰을 먼저 만들고 대형 타겟 모델이 검증하는 방식이다. 다만 RL 훈련에서는 타겟 모델이 수천 번의 업데이트를 거치며 계속 바뀐다. 고정된 드래프터 모델은 몇 차례 업데이트 뒤 타겟 모델과 분리(misalign)되어 더 이상 유용하지 않을 수 있다.

유휴 시간을 드래프터 갱신에 연결한 TLT

MIT 연구팀의 TLT(Taming the Long Tail)는 롤아웃 길이 분포의 긴 꼬리(long-tail), 즉 일부 GPU가 긴 응답을 생성하면서 전체 배치를 지연시키는 상황을 다룬다.

시스템의 첫 축은 적응형 드래프터 훈련기(Adaptive Drafter Trainer)다. 대기 중인 GPU가 드래프터 모델을 자동으로 학습한다. 드래프터는 타겟 모델보다 훨씬 소형이므로 유휴 시간 안에 업데이트할 수 있다. 새 컴퓨팅 자원을 추가하지 않고 기존에 낭비되던 GPU 사이클을 활용하는 구조다. 드래프터가 타겟 모델과 함께 갱신되므로 정렬 문제도 줄일 수 있다.

다른 축인 적응형 롤아웃 엔진(Adaptive Rollout Engine)은 배치의 응답 길이 분포를 보고 디코딩 방식을 고른다. 응답 길이가 균일한 배치에서는 표준 디코딩을 유지하고, 긴 꼬리 분포가 예상되면 드래프터 기반 추측적 디코딩을 활성화한다.

⤢✕

TLT는 테스트한 모든 모델에서 정확도 손실 없이 훈련 속도를 70~210% 향상시켰다. 개선 폭은 모델 크기와 태스크별 롤아웃 길이 분포에 따라 달라지며, 긴 꼬리 분포가 클수록 효과도 커진다. 훈련이 끝난 드래프터 모델은 별도 비용 없이 효율적인 추론에 활용할 수 있다.

데이터·텐서·파이프라인 병렬화가 만드는 대기 구간

대규모 LLM 훈련은 보통 하나의 병렬화 기법만 쓰지 않고 3D 병렬화(3D Parallelism)를 조합한다.

데이터 병렬(Data Parallelism)은 동일한 모델 복사본을 여러 GPU에 배치하고, 각 GPU가 서로 다른 미니배치를 처리하는 방식이다. 그래디언트를 All-Reduce로 동기화해 모델을 갱신한다. 구현이 단순하고 효율적이지만 모델이 단일 GPU 메모리에 들어가야 한다.

모델 병렬(Model Parallelism, Tensor Parallelism)은 레이어 내부를 GPU들에 나눈다. 예를 들어 어텐션 헤드를 여러 GPU가 분담할 수 있다. 단일 GPU 메모리를 넘는 대형 레이어를 다룰 수 있는 대신 GPU 간 통신이 늘어난다.

파이프라인 병렬(Pipeline Parallelism)은 레이어를 여러 스테이지로 나누고 각각을 별도 GPU 그룹에 할당한다. 마이크로배치를 파이프라인으로 흘려 GPU 활용률을 높인다. 그러나 스테이지 간 의존성 때문에 파이프라인 버블(pipeline bubble), 즉 GPU 유휴 시간이 생긴다.

⤢✕

TLT가 특히 효과적인 지점은 파이프라인 병렬에서 남는 버블을 드래프터 학습으로 바꾸는 데 있다. GPipe나 1F1B 같은 스케줄링 알고리즘으로 버블을 줄일 수는 있지만 완전히 없애기는 어렵다. TLT는 그 잔여 시간을 생산적인 연산에 할당한다.

메모리와 인스턴스 비용을 함께 줄이는 방법

대규모 LLM 훈련에서는 유휴 자원 활용 외에도 연산 정밀도와 메모리 배치가 비용을 좌우한다.

혼합 정밀도 훈련(Mixed Precision Training)은 대부분의 연산을 FP32 대신 BF16 또는 FP8로 처리하고, 수치 안정성이 필요한 일부 연산만 FP32로 둔다. 메모리 사용량을 절반 이하로 줄이고 Tensor Core 활용률을 높이는 방식이다.

그래디언트 체크포인팅(Gradient Checkpointing)은 순전파의 중간 활성값을 저장하지 않고, 역전파에서 필요할 때 다시 계산한다. 메모리 사용량을 크게 낮추는 대신 약 30%의 추가 계산이 발생한다. 메모리가 병목일 때 효과적이다.

DeepSpeed의 ZeRO 옵티마이저(Zero Redundancy Optimizer)는 옵티마이저 상태, 그래디언트, 모델 파라미터를 여러 GPU에 분산해 GPU별 메모리 사용량을 GPU 수에 비례해 줄인다. ZeRO-3에서는 파라미터까지 분산해 이론상 무한한 모델 크기를 지원한다.

Flash Attention은 HBM(High Bandwidth Memory)과 SRAM 사이의 데이터 이동을 줄이도록 어텐션 연산을 재구성한다. I/O 효율을 높여 어텐션 계산 속도를 2~4배 향상시키고 메모리 사용량도 낮춘다.

클라우드에서는 스팟 인스턴스(Spot Instances)를 활용할 수 있다. AWS, GCP, Azure의 스팟 GPU 인스턴스는 온디맨드 대비 60~90% 저렴하지만 언제든 회수될 수 있다. 따라서 체크포인트를 빈번히 저장하고, 중단된 지점에서 재개할 수 있는 내결함성(fault-tolerant) 훈련 인프라가 필요하다.

TLT의 유휴 GPU 재활용은 데이터·모델·파이프라인 병렬화, 혼합 정밀도, ZeRO 옵티마이저, Flash Attention 같은 기법과 함께 검토할 수 있다. 분산 훈련에서 없애기 어려운 대기 시간을 무엇에 쓸지까지 스케줄링 대상으로 삼는 접근이다.

Sources

• New method could increase LLM training efficiency | MIT News
• New method could increase LLM training efficiency | MIT EECS
• Taming the Long-Tail: Efficient Reasoning RL Training with Adaptive Drafter | arXiv
• Adaptive drafter model uses downtime to double LLM training speed | TechXplore
• Mist: Efficient Distributed Training of Large Language Models via Memory-Parallelism Co-Optimization | arXiv
• Efficient Training of Large Language Models on Distributed Infrastructures: A Survey | arXiv
• LeMix: Unified Scheduling for LLM Training and Inference on Multi-GPU Systems | arXiv