기술 요약: Jet-RL

문제 정의

강화 학습(RL)은 특히 사고 사슬(Chain-of-Thought, CoT) 생성 과정에서 대규모 언어 모델(LLM)의 복잡한 추론 능력을 향상시키는 데 핵심적인 역할을 하고 있습니다. 그러나 RL 훈련 파이프라인은 계산 효율성이 낮으며, 특히 생성 길이가 8,000 토큰을 초과할 경우 롤아웃(rollout) 단계(응답 생성)가 전체 훈련 시간의 70% 이상을 소비합니다.

이러한 병목 현상을 해결하기 위해, 롤아웃 단계에서는 FP8 양자화를 사용하고 훈련 단계에서는 BF16 정밀도를 유지하는(BF16-train + FP8-rollout) 전략이 널리 채택되고 있습니다. 하지만 본 논문은 이러한 접근 방식에서 두 가지 결정적인 실패 모드를 식별했습니다:

1. 장기 롤아웃 불안정성(Long-Rollout Instability): 생성 길이(예: 8k 토크 이상)가 증가함에 따라, BF16 훈련 분포와 FP8 롤아웃 분포 사이의 수치적 불일치가 누적됩니다. 이는 심각한 훈련 불안정성과 치명적인 정확도 붕괴를 초래합니다.
2. 난이도 높은 작업에서의 실패(Challenging Task Failure): 이 전략은 어려운 추론 작업이나 성능이 낮은 베이스 모델에 적용될 때 실패합니다. 이러한 시나리오에서는 모델의 낮은 신뢰도가 양자화 오차를 증폭시켜, 훈련 곡선이 풀 정밀도(full-precision) 베이스라인에 비해 급격히 발산하게 만듭니다.

근본 원인은 오프-폴리시(off-policy) 특성으로 식별되었습니다. 즉, 액터(actor)는 업데이트 중인 정책(BF16 훈련)과 수치적으로 다른 정책(FP8 롤아웃)에 의해 생성된 데이터로부터 학습하며, 이는 안정적인 RL에 필수적인 온-폴리시(on-policy) 가정을 위반합니다.

방법론: Jet-RL

저자들은 훈련과 롤아웃 사이에 **통합된 정밀도 흐름(unified precision flow)**을 강제함으로써 견고한 온-폴리시 FP8 RL 훈련을 가능하게 하는 프레임워크인 Jet-RL을 제안합니다.

핵심 아키텍처

Jet-RL은 훈련 순전파(forward pass)와 추론 롤아웃 모두가 동일한 FP8 양자화 흐름을 사용하도록 보장함으로써 정밀도 불일치를 제거합니다.

• 통합 정밀도 그래프(Unified Precision Graph): 프레임워크는 모델을 텐서 전파를 나타내는 엣지로 구성된 유향 그래프로 모델링합니다. Jet-RL은 추론 그래프($\mathcal{G}_{infer}$)가 동일한 양자화 입도(granularity)와 정밀도를 공유하는 훈련 순전파 그래프($\mathcal{G}^{fwd}_{train}$)의 부분 그래프가 되도록 보장합니다.
• 마스터 가중치(Master Weights): 순전파와 롤아웃은 FP8로 동작하지만, 표준 혼합 정밀도 훈련과 유사하게 훈련 업데이트를 안정화하기 위해 마스터 가중치는 BF16으로 유지됩니다.

양자화 스킴(Quantization Scheme)

계산 가속화를 유지하면서 안정성을 확보하기 위해, Jet-RL은 선형 레이어(Linear Layers, GEMMs)에 대해 혼합 양자화 전략을 채택합니다.

• 활성화 함수(Activations): 1×128 per-group 양자화를 사용하여 양자화합니다.
• 가중치(Weights): 128×128 per-block 양자화를 사용하여 양자화합니다.
• 역전파(Backward Pass):
  • 역전파를 위해 저장되는 활성화 함수는 FP8로 저장됩니다(이전 연산자와 퓨즈됨).
  • 연산자 간에 전달되는 그래디언트(gradients)는 그래디언트 언더플로우와 수렴 저하를 방지하기 위해 BF16으로 유지됩니다.
  • 역전파 GEMM 연산(DGrad 및 WGrad)은 가속화를 위해 FP8로 수행되며, 하드웨어 요구 사항에 맞추기 위해 특정 커널 구성(예: WGrad를 위한 1×128 및 128×1 양자화)을 활용합니다.

구현

Jet-RL은 추론 엔진으로 vLLM을, RL 훈련 프레임워크로 VeRL을 사용하여 구현되었습니다. 또한 최첨단 FP8 GEMM 커널인 DeepGEMM과 양자화, 전치(transposition), 퓨즈된 활성화/RMSNorm 연산을 위한 커스텀 Triton 커널을 활용합니다.

주요 기여

1. 문제 식별: 본 논문은 BF16-train + FP8-rollout 전략이 장기 롤아웃 및 까다로운 작업에서 오프-폴리시 수치적 불일치로 인해 근본적으로 결함이 있으며, 훈련 붕괴를 초래한다는 것을 입증했습니다.
2. Jet-RL 프레임워크: 훈련과 롤아웃 간의 온-폴리시 일관성을 강제하는 통합 FP8 정밀도 흐름을 제안하여, 정밀도 불일치 문제를 해결하고 정밀도 교정(calibration)의 필요성을 제거했습니다.
3. 성능 및 안정성: Jet-RL이 다양한 모델과 작업에 걸쳐 BF16 베이스라인에 근접한 수렴 및 정확도를 유지하면서 상당한 속도 향상을 달라는 것을 입증했습니다.

실험 결과

저자들은 GSM8K, MATH, DeepMATH 데이터셋에 대해 Llama3.1-8B, Qwen2.5-7B, Qwen3-8B-Base 모델을 대상으로 Jet-RL을 평가했습니다.

정확도 및 수렴

• 안정성: Qwen2.5-7B에서 수렴에 실패하고 16k 롤아웃에서 상당한 성능 저하를 보인 BF16-train + FP8-rollout 베이스라인과 달리, Jet-RL은 모든 설정에서 성공적으로 수렴했습니다.
• 정확도 격차: Jet-RL은 BF16 훈련 대비 성능 격차를 약 1% 수준으로 줄였습니다(예: Qwen3-8B-Base에서 2.9% 저하 대비 1.1% 저하). 일부 사례(Llama3.1-8B)에서는 Jet-RL이 BF16 베이스라인보다 2.0% 더 높은 성능을 보이기도 했습니다.
• 강건성: Jet-RL은 나이브한 FP8 롤아웃 전략이 붕괴되었던 어려운 작업(DeepMATH)과 긴 컨텍스트(16k)에서도 성능을 유지했습니다.

효율성

• 롤아웃 속도 향상: FP8 양자화는 롤아웃 단계(tokens/s)에서 **1.07× ~ 1.33×**의 속도 향상을 제공했으며, 더 큰 모델(32B)에서 가장 높은 이득을 보였습니다.
• 훈련 속도 향상: 통합 FP8 흐름은 액터 업데이트 단계를 1.54×, 레퍼런스 모델 추론을 1.80× 가속화하여, 훈련 단계에서 **1.41×**의 속도 향상을 가져왔습니다.
• 엔드 투 엔드(End-to-End): 프레임워크는 8B 모델 훈련에서 **1.16×**의 엔드 투 엔드 속도 향상을 달성했으며, 더 큰 모델에서는 더 높은 이득이 기대됩니다.

의의 및 주장

본 논문은 Jet-RL이 대규모 RL 훈련을 가속화하기 위해 FP8 연산을 적용할 수 있는 신뢰할 수 있고 효율적인 경로를 제공한다고 주장합니다. 정밀도 불일치로 인한 근본적인 불안정성을 해결함으로써, Jet-RL은 FP8을 단순한 추론뿐만 아니라 전체 RL 훈련 루프에 사용할 수 있게 합니다. 이를 통해 복잡한 추론 모델 훈련에 필요한 수렴 안정성이나 정확도를 희생하지 않으면서도 상당한 가속화(롤아웃 최대 33% 향상, 훈련 최대 41% 향상)를 실현할 수 있습니다. 이 연구는 저정밀도 하드웨어로 RL 훈련을 확장하기 위해서는 통합된 정밀도 흐름이 필수적임을 확립합니다.