MoE 언어모델 강화학습, FP4 저정밀도로 5배 빠르게
원제 TRACE: Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models
추천 92댓글 2
Key Point
MoE 모델의 강화학습 훈련에서 롤아웃 단계의 메모리와 속도 병목을 저정밀도 양자화로 5배 이상 해결할 수 있게 되어, 대규모 모델 학습의 실용성이 크게 향상된다.
핵심 요약
- 대규모 언어모델의 강화학습 훈련에서 롤아웃(rollout) 생성 과정이 계산과 메모리 오버헤드를 많이 차지하는 문제가 있어, 저정밀도 롤아웃을 통한 효율화가 필요하다.
- 기존 FP4 강화학습 방법들은 훈련과 롤아웃 경로의 양자화를 독립적으로 최적화하기 때문에 두 경로 간 불일치(discrepancy)가 발생하는 한계가 있다.
- 연구팀은 TRACE(Train-Rollout Quantization Alignment via Compact GuidancE)라는 FP4 양자화 프레임워크를 제안했으며, 이는 롤아웃 측의 양자화 결과를 훈련 측 반올림 결정에 직접 반영하여 불일치를 줄인다.
- TRACE는 심층 레이어의 mantissa와 scale 정보만 선택적으로 유지하는 양자화 정보 캐싱 방식을 도입해 롤아웃 지원으로 인한 저장소와 통신 오버헤드를 감소시킨다.
- 추론(reasoning), 코딩(coding), 장기 강화학습(long-horizon RL) 작업을 포함한 4개의 대규모 MoE 언어모델에서 평가했다.
- TRACE는 FP4 가중치/활성화와 FP4 KV-캐시 롤아웃을 결합하면서도 BF16 롤아웃 수준의 강화학습 성능을 유지하고, 롤아웃 속도를 최대 5.4배 향상시켰다.
- BF16으로 훈련한 모델을 사후 양자화(post-hoc quantization)하는 방식 대비 최종 FP4 성능이 우수하다.