TRACE, MoE 강화학습의 FP4 학습과 Rollout 경로를 정렬하다

대규모 언어 모델의 포스트트레이닝에서 강화학습은 중요한 단계지만, 비용은 파라미터 업데이트에만 국한되지 않습니다. 모델은 반복적으로 Rollout 샘플을 생성해야 하며, 생성 과정의 연산량과 메모리 사용량, KV 캐시가 주요 병목이 됩니다. 따라서 Rollout을 FP4 같은 저정밀도로 실행하는 방법이 주목받고 있습니다. 그러나 BF16으로 학습한 정책을 학습 후 단순히 FP4로 양자화하는 방식만으로는 강화학습에 필요한 모델 행동을 안정적으로 유지하기 어렵습니다.

핵심 문제: 서로 다른 양자화 실행 경로

기존 FP4 강화학습 방법은 대체로 학습 경로와 Rollout 경로의 양자화 정확도를 각각 최적화합니다. 각 경로의 국소적인 오차가 작더라도 두 경로의 계산 결과가 서로 비슷하다는 보장은 없습니다. 강화학습에서는 현재 정책이 다음 학습 데이터를 직접 생성합니다. 따라서 학습 중 사용한 계산과 실제 샘플 생성에 사용한 계산이 다르면, 정책 업데이트가 실제 Rollout 행동과 어긋날 수 있습니다.

TRACE는 바로 이 차이를 직접 줄이는 데 초점을 둡니다. 이름은 Train-Rollout Quantization Alignment via Compact Guidance의 약자로, 혼합 전문가(MoE) 언어 모델에서 학습과 Rollout의 저정밀 실행 경로를 맞추기 위한 프레임워크입니다.

TRACE의 주요 구성

• Rollout 유도 양자화 인식 학습: Rollout 측에서 얻은 양자화 결과를 학습 과정으로 전달하고, 이를 학습 측 FP4 반올림 결정에 활용합니다. 학습 경로가 독립적으로 양자화 결과를 선택하는 대신 실제 생성에서 사용될 저정밀 동작에 가까워지도록 합니다.
• 생성 파이프라인 전반의 FP4 지원: 가중치와 활성값뿐 아니라 FP4 KV 캐시 Rollout도 지원합니다. 긴 시퀀스 생성에서 큰 메모리를 차지할 수 있는 캐시까지 최적화 범위에 포함합니다.
• 선택적 양자화 정보 캐싱: Rollout 정보를 학습에 전달하면 저장과 통신 비용이 증가할 수 있습니다. TRACE는 모든 정보를 저장하지 않고, 더 깊은 계층의 가수와 스케일 정보를 선택적으로 보존해 추가 비용을 줄입니다.

실험 결과와 의미

논문은 4개의 대규모 MoE 언어 모델을 대상으로 추론, 코딩, 장기 강화학습 과제를 평가했습니다. 가중치, 활성값, KV 캐시를 함께 FP4로 처리하는 Rollout에서도 BF16 Rollout과 비교 가능한 강화학습 성능을 얻었다고 보고합니다. 또한 BF16으로 학습한 정책을 사후 FP4 양자화하는 방식보다 최종 FP4 성능이 우수했으며, 보고된 최대 Rollout 가속은 5.4배입니다.

TRACE의 중요한 관점은 저비트 강화학습의 목표를 바꾼다는 점입니다. 단순히 고정밀 모델의 출력을 최대한 복원하는 것이 아니라, 실제 배포와 샘플 생성에 사용될 저정밀 실행 경로에 맞춰 학습 자체를 조정합니다. MoE 모델은 반복적인 생성과 캐시 이동으로 인해 Rollout 비용이 커질 수 있으므로, 이런 실행 경로 정렬은 실용적인 의미가 있습니다.

다만 제공된 자료에는 모델별, 과제별, 하드웨어별 상세 비교가 포함되어 있지 않습니다. 따라서 5.4배는 논문에서 제시한 최대 가속 수치이며, 모든 환경에서 동일하게 얻어지는 효과로 해석해서는 안 됩니다. 그럼에도 TRACE는 저정밀 강화학습에서 개별 양자화 오차를 줄이는 것뿐 아니라 학습과 Rollout의 실제 계산 경로를 일치시키는 일이 중요하다는 점을 보여줍니다.

출처: Hugging Face Daily Papers