MoE 강화학습에서 FP4 롤아웃과 학습 경로의 양자화 불일치를 맞춘다
TRACE: Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models
무엇인가
대규모 언어모델의 사후학습으로 자리 잡은 강화학습은 롤아웃 생성 단계에서 긴 궤적을 반복 생성해야 하기 때문에 계산과 메모리 비용이 크다. 그래서 정책 모델을 저정밀도로 양자화해 생성하는 롤아웃이 유력한 해법인데, FP4(W4A4)는 감속과 메모리 절감 폭이 큰 대신 양자화 공간이 거칠어 학습 경로와 롤아웃 경로 사이에 수치적 불일치가 생기고, 이것이 정책 불일치와 학습 불안정으로 이어진다. Mixture-of-Experts 모델에서는 이 차이가 라우터 점수를 바꿔 다른 전문가를 활성화할 수 있어 문제가 더 커진다. 논문은 기존 FP4 RL 기법들(QUADS, Rollout-ResQ, QaRL, Score Centering 등)이 학습 경로와 롤아웃 경로 각각의 양자화 오차를 독립적으로 줄일 뿐, 두 양자화 실행 경로 사이의 불일치 자체를 직접 줄이지 않는다는 점을 근본적 한계로 지목한다. 각 경로의 오차가 작아도 양자화된 값은 크게 갈라질 수 있다는 것이다.
어떻게 동작하나
제안 방법 TRACE(Train-Rollout Quantization Alignment via Compact GuidancE)의 핵심은 롤아웃 측 양자화 결과로 학습 측 FP4 반올림을 유도하는 것이다. 롤아웃 생성 중 라우팅된 전문가의 활성값과 KV 상태에 대해 FP4 양자화 결과를 기록해 두고, 이후 양자화 인지 학습(QAT) 단계에서 같은 궤적·토큰 위치·양자화 지점에 대응하는 학습 측 활성값의 반올림을 그 기록에 맞춘다. 구체적으로 각 학습 측 활성값에 대해 롤아웃 측 스케일 아래의 인접한 두 FP4 E2M1 코드워드 집합을 잡고, 표준 round-to-nearest 대신 롤아웃 코드워드와의 거리가 더 가까운 쪽을 선택한다. 이 선택은 정의상 RTN보다 롤아웃 코드워드와의 거리를 늘리지 않으므로, 완전 정보를 가정한 국소 수준에서 학습-롤아웃 불일치를 악화시키지 않는다는 보장을 갖는다. 논문은 이 효과를 두 가지 예로 설명한다. BF16에서 60.24와 59.76으로 0.48만 차이 나던 값이 같은 전역·블록 스케일을 적용하면 정규화 값 2.51과 2.49가 반올림 경계의 반대편에 놓여 각각 3과 2로 반올림되고, 불일치가 0.48에서 24로 증폭된다. 반대로 QUADS 예시에서는 2.40과 2.49가 바닐라 FP4에서 둘 다 2로 매핑돼 불일치가 0이었는데, QUADS가 롤아웃 측을 2.50으로 복원하면서 경로별 오차는 줄었지만 학습-롤아웃 불일치가 0.10으로 늘어난다.
무엇과 다른가
여기에 더해 TRACE는 롤아웃 유도 정보를 압축해 통신·저장 비용을 줄인다. 완전한 양자화 정보를 유지하면 데이터량이 감당이 안 되는데, 논문이 든 예에서 Qwen3.5-35B-A3B(히든 2048, KV 헤드 차원 256, KV 헤드 2개, MoE 40층, 풀어텐션 10층, 최대 응답 256K 토큰) 기준으로 양자화 값당 4.5비트를 유지하면 토큰 하나당 활성 유도 정보 약 45KB와 KV 유도 정보 약 5.6KB, 합쳐 약 50KB가 생성된다. RL 스텝당 4,096개 궤적이면 스텝당 최대 약 51TB에 달하고, GPU-CPU 대역폭 300GB/s로도 약 3분, 저장 대역폭 5GB/s에서는 직렬 처리 시 약 3시간이 걸려 RL 스텝 시간을 초과한다. 저자들은 층별로 불일치가 발생한 양자화 값의 99% 이상이 FP4 코드북에서 인접한 한 칸만 차이 난다는 점, 그리고 더 낮은 코드워드로의 반올림 보정이 주로 깊은 층에서 일어난다는 점을 관찰하고, 모델 뒤쪽 절반 층의 가수(mantissa)와 스케일 정보만 통신한다. Qwen3.5-35B-A3B에서는 뒤쪽 20개 층에 대해 1비트 가수와 상각된 스케일 메타데이터를 유지해 토큰당 약 7.5KB(2048×20×1.5=61,440비트)로, 전체 양자화 활성값을 캐싱하는 것보다 캐시 크기를 약 6배 줄인다.
어떻게 쓰나
실험은 네 개의 MoE 모델(Qwen3.5-35B-A3B, Qwen3.5-122B-A10B, Qwen3.8-Flash-Next 125B 총 파라미터/6B 활성, Qwen3.8-2.4T-A95B)에서 추론·코딩·장기 과제 RL로 수행됐다. 추론은 LiveCodeBench v6, AIME 2024/2025, HMMT 2025, 코딩은 DeepSWE v1.1과 Terminal-Bench 2.1, 장기 과제는 GDPval로 평가했다. 비교 대상은 바닐라 NVFP4, NVFP4+QAT, QaRL, QUADS, Score Centering, MXFP4 W4A8/W4A4 변형, 그리고 사후 양자화 기법인 바닐라 NVFP4, 4over6, H-Scale이다. Qwen3.5-35B-A3B 추론 과제에서 TRACE는 QAT·QaRL·QUADS를 네 벤치마크 모두에서 앞서며 평균 점수를 68.8에서 75.3으로 올렸고, HMMT25에서는 59.0에서 70.0으로 향상됐다. 평균 점수는 BF16 롤아웃과 동일한 수준이다. 대형 모델에서도 Qwen3.8-Flash-Next의 Terminal-Bench가 66.4에서 70.6으로 오르는 등 세 모델 모두 QAT를 앞섰고, BF16 롤아웃과도 대등했다(DeepSWE 33.0 대 33.4, Terminal-Bench 70.6 대 68.8, GDPval 90.2 대 90.3).
전제와 한계
학습 동역학에서 QAT·QaRL·QUADS는 학습이 진행될수록 학습-롤아웃 불일치가 커지고 평균 보상이 후반부에 떨어지며 테스트 점수도 급락하는 반면, TRACE는 불일치를 억제해 BF16 롤아웃에 가까운 안정성을 유지한다. 흥미로운 관찰은 적응 현상으로, TRACE의 테스트 점수는 초기에 양자화 열화 때문에 BF16보다 낮지만 약 120 스텝을 지나면 따라잡고 이후에는 대등하거나 일부 구간에서 더 높은 값을 보인다. 즉 불일치가 충분히 통제되면 정책이 저정밀도 실행 환경에 적응한다는 것이다. 효율은 128K 출력 길이에서 BF16 롤아웃 대비 최대 5.4배 디코딩 처리량을 보였고, 엔드투엔드 RL 스텝 시간은 바닐라 결합 FP4 롤아웃 대비 664에서 713으로 7.4%만 증가했다. 롤아웃 시간 분석에서는 모델 순전파가 86%, 가중치 측 유도가 4%, KV 측 유도가 2%, KV 역양자화가 8%를 차지한다.
절제 실험에서는 가중치만 또는 KV만 양자화한 개별 설정보다 결합 FP4가 훨씬 어렵다는 점을 확인했고, MXFP4 형식에서도 W4A8에서 평균 69.0에서 75.1로, 더 공격적인 W4A4에서 67.2에서 73.5로 개선됐다. 유지하는 롤아웃 정보량에 대한 민감도는 40개 층 전체에 1비트 가수를 유지하면 평균 75.8, 뒤쪽 20개 층만 유지하면 소폭 하락, 10개 층 74.1, 5개 층 73.4로 나타났다. Score Centering과의 비교에서는 SC가 학습 붕괴는 막지만 학습-롤아웃 로그확률 차이의 최대값이 1.8~2.1, 최소값이 -4 근처로 BF16 기준에서 멀리 남는 반면, TRACE는 그 차이를 BF16 롤아웃에 가깝게 유지했고 380 스텝에서 평균 75.3 대 74.1로 1.2%포인트 앞섰다. BF16 롤아웃으로 학습한 뒤 사후 양자화한 파이프라인(바닐라 NVFP4 70.4, 4over6 71.0, H-Scale 71.4)과 비교해도 TRACE는 75.3으로 우세하다.
실무 관점에서 이 논문은 FP4 롤아웃을 도입할 때 무엇을 봐야 하는지를 알려준다. 첫째, 롤아웃과 학습의 정책 불일치를 로그확률 차이 같은 지표로 계속 추적해야 하며, 경로별 양자화 오차만 낮추는 것으로는 불충분하다. 둘째, 롤아웃 측 양자화 기록을 전부 저장할 필요가 없다는 실증이 있다. 뒤쪽 절반 층의 가수 1비트와 스케일만으로도 대부분의 성능을 회복하므로, 저장·통신 병목을 6배가량 줄이는 설계가 현실적이다. 셋째, 학습 스택은 VeRL(비동기 RL 조정), Megatron(학습 백엔드), SGLang(롤아웃 생성)의 disagg 구성에 GRPO를 쓰고, R3로 롤아웃 측 전문가 라우팅 결정을 재생하며, 토큰 수준 중요도 비율의 클리핑 상한을 5, max_version_diff를 3으로 두는 조건에서 측정됐다는 점을 감안해야 한다. 롤아웃에서는 MoE 층의 라우팅된 전문가와 소프트맥스 어텐션 층의 KV 캐시만 양자화하고 나머지는 BF16으로 유지한다.
저자들이 명시한 전제와 한계도 분명하다. TRACE는 학습 정책이 롤아웃을 생성한 정책보다 최신일 수 있기 때문에 정책 지연(policy staleness)에서 비롯되는 근본적인 활성값 차이를 없애지는 못하고, 그 차이가 일관되지 않은 FP4 반올림으로 증폭되는 추가 불일치만 줄인다. 또한 국소 지점에서의 불일치 비증가 보장은 전체 네트워크 수준의 불일치가 단조 감소한다거나 하류 비선형 변환을 거친 정책 수준 발산이 줄어든다는 것을 의미하지 않는다. 완전한 롤아웃 유도 정보를 쓰는 변형은 저장·전송 오버헤드 때문에 처리량이 크게 떨어지며, 압축 캐싱이 그 대가를 줄이는 장치라는 점도 함께 밝히고 있다.