논문 리뷰 · 2026년 10월 8일

TRACE: MoE 언어 모델의 FP4 강화학습을 위한 롤아웃 유도 양자화 인식 학습

TRACE: Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models

롤아웃 쪽 FP4 반올림 결과로 학습 쪽 반올림을 유도해 학습-롤아웃 불일치를 줄였고, MoE 모델 RL에서 BF16 수준 성능으로 최대 5.4배 롤아웃 가속을 얻었다.

문제

LLM 강화학습에서는 롤아웃 생성 비용이 크다. FP4 저정밀 롤아웃은 매력적이지만, 학습 경로와 롤아웃 경로 사이에 수치 불일치가 커져 정책 불일치와 학습 붕괴를 일으킨다. MoE에서는 이 차이가 전문가 라우팅까지 바꿀 수 있다. 기존 방법은 경로별 양자화 오차를 각각 줄일 뿐, 두 양자화 경로 사이의 차이를 직접 줄이지는 않는다.

방법

(1) 롤아웃 유도 QAT: 롤아웃 중 라우팅 전문가의 활성값과 KV 캐시의 FP4 양자화 결과를 기록한다. 학습 단계에서는 인접한 두 FP4 코드워드 가운데 롤아웃 코드워드에 더 가까운 쪽으로 반올림한다. 일반 RTN 반올림 대비, 각 지점에서 국소 불일치가 늘지 않음이 보장된다. (2) 효율적 캐싱: 불일치는 대부분 인접 코드워드 한 칸 차이이고 보정은 깊은 층에 집중된다. 그래서 후반부 층의 가수(mantissa)와 스케일 정보만 저장·전송한다. 실험은 VeRL·Megatron·SGLang 환경에서 GRPO와 R3 라우팅 재생을 써서 했다.

결과

저자들은 FP4 W/A와 FP4 KV 캐시를 함께 쓰는 롤아웃에서 BF16 롤아웃과 비슷한 RL 성능과 학습 양상을 유지했다고 보고했다. 평가는 Qwen 계열 MoE 4종으로 했고, 과제는 추론, 코딩(DeepSWE, Terminal-Bench), 장기 과제(GDPval)다.
BF16 롤아웃 대비 롤아웃 속도가 최대 5.4배 빨랐다.
Qwen3.8-Flash-Next의 Terminal-Bench 2.1 학습에서 일반 FP4 QAT는 학습-롤아웃 log-prob 차이가 커지며 점수가 급락했다. TRACE는 불일치를 억제했고, BF16과의 격차가 학습 중 점차 줄거나 역전됐다.
BF16 롤아웃으로 학습한 뒤 FP4로 사후 양자화하는 파이프라인(vanilla NVFP4, 4over6, H-Scale)보다 최종 FP4 성능이 높았다.
불일치가 난 양자화 값의 99% 이상이 FP4 코드북에서 인접 한 칸 차이였다. 이것이 가수 정보만 전송하는 설계의 근거다. 전체 정보를 보관하면 RL 한 스텝당 최대 약 51TB가 필요하다는 추정도 함께 제시됐다.

한계

보장은 지점별 국소 불일치에 한정되며, 네트워크 전체나 정책 수준의 차이가 단조롭게 준다는 뜻은 아니라고 저자들이 밝혔다. 비동기 학습의 정책 staleness로 생기는 활성값 차이 자체는 없애지 못한다. 양자화 대상은 라우팅 전문가와 KV 캐시뿐이고 나머지 모듈은 BF16으로 유지했다. 평가 모델이 모두 Qwen 계열 MoE이고, 대형 두 모델은 50스텝만 학습했다.

의미

대규모 MoE의 RL 후학습에서 가장 비싼 롤아웃 단계를 FP4로 돌리면서도 학습 안정성을 지키는 방법을 제시한다. 핵심은 '양자화 오차 최소화'가 아니라 '두 경로의 일치'를 목표로 삼아야 한다는 관점이다.

핵심 용어

학습-롤아웃 불일치(Train-Rollout Mismatch): RL에서 응답을 생성하는 추론 엔진과 그래디언트를 계산하는 학습 엔진이 같은 입력에 서로 다른 확률을 내는 현상이다. 정책 불일치와 학습 불안정을 일으킨다.
양자화 인식 학습(QAT): 순전파에서 저정밀 양자화를 흉내 내며 학습해, 모델이 양자화 오차에 적응하도록 하는 기법이다.
NVFP4 / FP4 E2M1: 지수 2비트, 가수 1비트로 이루어진 4비트 부동소수점 형식이다. 블록 스케일과 함께 쓰며 표현 가능한 값이 매우 적다.
Round-to-Nearest(RTN): 값을 가장 가까운 양자화 코드워드로 반올림하는 표준 방식이다. 경계 근처의 작은 차이가 큰 차이로 증폭될 수 있다.

원문

저자: Xin Wang, Hao Yu, Zhengyang Zhuge, Bochao Mao, Zheng Li, Junda Feng, Yuyan Luo, Yi Zhang, 외 · 게시 2026-10-06 · 라이선스 arXiv 비독점 배포 라이선스 · 본문 기준 분석

이 리뷰는 AI가 논문을 읽고 작성한 요약이며, 정확한 내용은 원문을 기준으로 확인해 주세요.