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

논문 링크

TRACE: 롤아웃을 따라가는 FP4 양자화 인식 학습 — MoE LLM의 강화학습을 BF16 성능 그대로, 최대 5.4배 빠르게

한 줄 요약 (TL;DR)

MoE LLM의 강화학습(RL)에서 롤아웃 생성을 FP4(W4A4 + FP4 KV) 로 양자화하면 생성 속도가 최대 5.4배 빨라지지만, 학습 경로와 롤아웃 경로의 수치 불일치가 커져 훈련이 불안정해지거나 붕괴한다. 이 논문은 롤아웃 쪽에서 실제로 만들어진 FP4 양자화 결과 를 학습 쪽 라운딩(rounding) 결정을 유도하는 신호로 직접 사용하는 TRACE 를 제안해, BF16 롤아웃과 거의 동일한 성능을 유지하면서 FP4의 속도 이점을 살린다 (근거: §1, §0).

핵심 아이디어

저자들의 진단은 이렇다. 기존 FP4 RL 기법들(QUADS, Rollout-ResQ 등)은 학습·롤아웃 각 경로의 양자화 정확도 를 개별적으로 높이는 데 집중한다. 하지만 두 경로의 양자화 오차를 각각 줄이는 것이 두 경로 사이의 불일치(discrepancy) 를 줄인다는 보장은 없다 (근거: §3.1). TRACE의 두 축은 다음과 같다.

• 롤아웃 유도 QAT(Rollout-Guided QAT) : 롤아웃에서 생성된 FP4 코드워드로 학습 쪽 라운딩 방향을 정해, 경로 간 불일치를 직접 최소화한다 (근거: §3.1).
• 컴팩트 양자화 정보 캐싱 : 불일치의 대부분이 FP4 코드북에서 인접한 한 칸 차이로 나타난다는 관찰에 근거해, 깊은 레이어의 mantissa + scale 정보만 남겨 저장·통신 오버헤드를 약 6배 줄인다 (근거: §3.2).

배경: 그들이 해결한 문제

LLM의 사후학습(post-training)에서 RL은 추론·코딩·장기 과제 성능을 끌어올리는 핵심 수단이 됐다 (근거: §1). 하지만 RL은 긴 트라젝토리를 반복 생성(rollout)해야 하므로 연산·메모리 비용이 크다 . 이 비용을 줄이기 위해 롤아웃을 저정밀도로 양자화하는 접근이 제안됐고, 특히 공격적인 FP4(W4A4) 는 더 큰 가속 여지를 준다 (근거: §1).

문제는 FP4의 거친 양자화 공간이 학습 경로와 롤아웃 경로 사이의 수치 불일치 를 키운다는 점이다. 이 불일치는 정책 불일치(policy mismatch) 로 이어져 RL을 불안정하게 만든다. MoE 모델은 특히 취약 하다. 작은 수치 차이가 라우터 스코어를 바꾸고, 그러면 서로 다른 전문가(expert)가 활성화되며 불일치가 증폭되기 때문이다 (근거: §2). R3(라우팅 재생), PR2(라우터 진화 보정), GSPO(시퀀스 수준 클리핑) 등은 이런 라우팅 불일치를 다루지만, 양자화가 만드는 수치 불일치 자체 는 해결하지 못했다 (근거: §2).

기존 FP4 RL 기법들의 근본 한계를 저자들은 한 문장으로 요약한다. “양자화 정확도 최적화와 학습–롤아웃 불일치 감소 사이의 정렬 실패” — 양자화 오차를 줄여도 두 경로의 양자화된 값은 여전히 크게 벌어질 수 있고, 결국 성능 저하나 훈련 붕괴로 이어진다 (근거: §1).

새로운 접근법: TRACE

TRACE 는 Train-Rollout Quantization Alignment via Compact GuidancE의 약자로, MoE LLM의 FP4 RL 학습 프레임워크다 (근거: §0). 학습 경로와 롤아웃 경로의 양자화를 정렬하는 두 가지 설계 로 구성된다.

1. 롤아웃 유도 QAT : 롤아웃 생성 중 FP4 routed-expert 활성화와 FP4 KV 상태의 양자화 결과를 기록하고, 이어지는 QAT 단계에서 이 정보로 학습 쪽 라운딩을 유도한다 (근거: §3.1).
2. mantissa 전용 통신 : 전체 양자화 값 대신, 라운딩 방향을 결정하는 데 필요한 최소 정보(깊은 레이어의 mantissa·scale)만 캐싱해 전송한다 (근거: §3.2).

수식으로 보면 더 명확하다. 동일한 트라젝토리·토큰 위치·양자화 지점에서 짝지어진 학습·롤아웃 활성화 간 양자화 불일치는 다음과 같이 정의된다 (근거: §3.1).

핵심 통찰은, FP4 라운딩에서 작은 BF16 차이가 반올림 경계를 넘으면 크게 증폭 된다는 점이다. 예를 들어 학습·롤아웃 활성화가 각각 $60.24$, $59.76$으로 BF16에서 불과 $0.48$ 차이인데, 같은 글로벌·블록 스케일로 정규화하면 $2.51$과 $2.49$가 되어 반올림 경계 양쪽에 걸리며 각각 $3$과 $2$로 라운딩된다. 불일치가 $0.48$에서 $24$로 약 50배 증폭되는 것이다 (근거: §3.1).

작동 원리: 구체적인 예시로 살펴보기

왜 “정확도 개선"만으로는 부족한가

기존 기법(예: QUADS)이 오히려 불일치를 키우는 예가 Fig. 3에 나온다 (근거: §3.1). QUADS는 BF16 학습 활성화를 유지하면서 롤아웃 쪽 FP4 오차를 잔차 $R(X)=X-Q(X)$로 보정한다. 학습·롤아웃 활성화가 각각 $2.40$, $2.49$인 경우를 보자.

• Vanilla FP4 : 두 값을 모두 $2$로 매핑 → 양자화 불일치 $0$.
• QUADS : 학습 쪽은 $2.40$ 유지, 롤아웃 쪽은 $2.50$으로 재구성 → 경로별 오차는 줄지만 불일치는 $0.10$으로 증가 .

즉, 각 경로의 오차를 줄이는 것과 경로 간 불일치를 줄이는 것은 별개의 목표 다.

TRACE의 라운딩 선택

TRACE는 각 학습 활성화 $X_{\mathrm{train}}$에 대해, 롤아웃 스케일 하에서 정규화된 값의 이웃한 두 FP4 E2M1 코드워드 $\{q_{-}, q_{+}\}$를 후보로 두고, 표준 round-to-nearest(RTN) 대신 롤아웃 쪽 실제 코드워드 $q_{\mathrm{rollout}}$에 더 가까운 쪽을 고른다 (근거: §3.1).

이 선택은 같은 후보 집합에서 다음 보장을 준다 (근거: §3.1).

즉 TRACE는 RTN 대비 각 지점의 국소적 불일치를 증가시키지 않는다 . 이 국소 보장이 전체 네트워크 불일치의 단조 감소를 의미하지는 않지만, “불일치한 반올림이 만드는 불필요한 증폭"을 제거한다는 것이 핵심이다 (근거: §3.1).

왜 mantissa 하나만으로 충분한가

위 가이던스에 필요한 롤아웃 정보를 전부 전송하려 하면 비용이 감당하기 어렵다. Qwen3.5-35B-A3B(은닉 차원 $2048$, KV 헤드 차원 $256$, MoE 레이어 $40$개, 최대 응답 $256$K 토큰) 기준으로, 토큰당 약 $45$ KB의 활성화 가이던스와 $5.6$ KB의 KV 가이던스가 발생한다. RL 한 스텝에 $4096$개 트라젝토리를 쓰면 스텝당 약 51 TB 의 양자화 정보가 된다. GPU→CPU $300$ GB/s로는 약 3분이지만, 저장 매체 $5$ GB/s로 쓰고 읽으면 약 3시간 이 걸려 RL 스텝 시간을 훌쩍 넘는다 (근거: §3.2).

저자들의 관찰은 두 가지다 (근거: §3.2).

• 불일치가 생긴 양자화 값의 99% 이상 이 FP4 코드북에서 인접한 한 칸 만 차이난다 (Fig. 6).
• 라운딩 보정은 주로 깊은 레이어 에서 발생한다.

따라서 학습 쪽 활성화와 롤아웃 스케일만 있으면 후보 코드워드가 거의 확정되므로, 전체 값이 아니라 mantissa 코드 정보만 보내도 원하는 이웃 코드워드를 식별할 수 있다. TRACE는 후반부 레이어 절반의 mantissa + scale 만 전송한다. Qwen3.5-35B-A3B에서 1비트 mantissa + amortized scale을 후반 20개 레이어에 적용하면 토큰당 $2048 \times 20 \times 1.5 = 61{,}440$ 비트 ≈ $7.5$ KB 로, 전체 활성화 캐싱 대비 약 6배 줄어든다 (근거: §4.2).

데이터 흐름

  flowchart LR
    A["롤아웃 생성<br/>(SGLang, FP4)"] --> B["FP4 활성화·KV<br/>양자화 결과 기록"]
    B --> C["후반부 레이어만<br/>mantissa + scale 캐싱"]
    C --> D["CPU/저장소로<br/>비동기 전송"]
    D --> E["QAT 시 재로드<br/>→ 라운딩 유도"]
    E --> F["Megatron<br/>정책 업데이트"]

성능 검증: 주요 결과

실험은 4개의 MoE 모델(추론·코딩·장기 과제 RL)에 걸쳐, 공통 인프라(VeRL + Megatron + SGLang, GRPO, R3)에서 진행됐다 (근거: §4).

추론 RL (Qwen3.5-35B-A3B, NVFP4 W/A + KV 공동 양자화) 에서 TRACE는 모든 벤치마크에서 기존 기법을 앞서며 BF16 수준을 회복한다 (근거: §4.1, Tab. 1).

평균 점수는 QUADS 대비 68.8 → 75.3(+6.5) , HMMT25는 59.0 → 70.0(+11.0) 로 뛰며, BF16(74.9)와 사실상 동률이다 (근거: §4.1).

대형 모델 확장성 은 더 인상적이다 (근거: §4.1, Tab. 2).

Flash-Next에서는 BF16(68.8)을 오히려 1.8 포인트 상회 하기도 한다. 흥미로운 점은 훈련 동역학인데, TRACE는 초기에는 양자화로 BF16보다 낮다가 약 120스텝 이후 BF16을 따라잡고 이후 동등하거나 더 높은 수치를 보인다 — 즉 정책이 FP4 환경에 적응(adaptation) 한다 (근거: §4.1, Fig. 7). 반면 QAT·QaRL·QUADS는 훈련이 진행될수록 불일치가 커지며 후반부에 보상과 테스트 점수가 급락한다.

효율성 측면에서, TRACE는 vanilla FP4와 거의 같은 생성 처리량을 유지하면서 BF16 대비 128K 출력 길이에서 최대 5.4배 디코딩 처리량을 낸다. 엔드투엔드 RL 스텝 시간은 vanilla FP4 대비 664 → 713(단위: 상대 시간), 약 7.4% 오버헤드에 그친다. 롤아웃 시간 분해를 보면 모델 forward가 86% , weight/KV 가이던스가 각각 4%/2% , KV 역양자화가 8% 다 (근거: §4.2).

절제 실험 도 일관된다 (근거: §4.3). MXFP4 포맷에서 W4A8은 69.0 → 75.1 , 더 공격적인 W4A4는 67.2 → 73.5 로 개선된다. 민감도 분석에서는 기본 설정(후반 20개 레이어, 1비트 mantissa)이 평균 75.3 , 커버리지를 10개/5개 레이어로 줄이면 각각 74.1/73.4 로 완만하게 떨어진다. Score Centering과의 비교(74.1 vs 75.3), BF16 학습 후 사후 PTQ와의 비교(vanilla 70.4 / 4over6 71.0 / H-Scale 71.4 vs TRACE 75.3 )에서도 우위를 보인다 (근거: §4.4, §4.5).

우리의 관점: 강점, 한계, 그리고 이 연구가 중요한 이유

강점

• 문제 정의가 날카롭다. “양자화 정확도"와 “학습–롤아웃 불일치"를 분리해, 후자를 직접 목적 으로 삼은 첫 시도라는 점이 이 논문의 가장 큰 기여다 (근거: §1, §3.1).
• 이론적 보장과 실용적 절충의 균형. 국소 불일치 비증가 보장(§3.1)과, 관찰 기반의 mantissa 전용 통신(§3.2)으로 정확도-비용 트레이드오프를 설득력 있게 다룬다.
• 스케일과 과제의 폭. 35B부터 2.4T-A95B까지, 추론·코딩·장기 과제를 아우르며, 단일 설정에 과적합되지 않음을 보여준다 (근거: §4).
• “적응"이라는 발견. FP4 환경에 정책이 적응해 BF16을 따라잡는 현상은 단순 보존을 넘는, 저정밀 RL의 가능성을 시사한다 (근거: §4.1).

한계와 비판적 질문

• 명시적 한계 섹션이 비어 있다. 부록의 “Limitations” 서브섹션이 주석 처리되어 있어, 저자가 인정한 한계가 거의 없다 (근거: Appx.). 이는 공정성 평가를 어렵게 만든다.
• 정책 스테일니스(staleness)의 근본 불일치는 해결하지 못한다. TRACE는 “라운딩이 만드는 추가 불일치"만 줄이며, 학습 정책과 롤아웃 정책이 달라 생기는 활성화 차이 자체는 남는다 (근거: §4). R3, PR2 등 라우팅 정렬 기법과의 상호작용은 여전히 열려 있다.
• 정확한 코드워드가 아닌 재구성에 의존. mantissa 전용 캐싱은 “완전 복원"이 아니라 근사적 참조를 사용하므로(§3.2, Alg. 1), 극단적 분포에서는 가이던스 품질이 떨어질 수 있다.
• FP4 하드웨어 의존성. NVFP4·MXFP4 커널과 분리(disagg) 배포·비동기 통신 파이프라인에 크게 의존해, 다른 정밀도·인프라로의 일반화에는 추가 검증이 필요하다 (근거: §4, Appx. D).
• 비용-편익의 경계가 불명확. 7.4% 스텝 오버헤드와 51 TB 수준의 데이터 이동을 피하기 위한 복잡한 엔지니어링이, BF16 대비 절약된 시간을 실제로 얼마나 넘어서는지에 대한 총소유비용(TCO) 관점의 분석은 부족하다.

왜 중요한가

RL 사후학습은 이제 최신 LLM 개발의 표준이 됐고, 그 비용의 대부분은 롤아웃 생성 에 있다. TRACE는 이 비용 병목을 FP4로 돌파하면서도 품질을 지키는 방법을 제시해, 실무적으로 즉시 재현 가능한 레시피를 제공한다. 동시에 “양자화 오차 vs 경로 간 불일치"라는 개념적 구분은 저정밀 학습 전반에 재사용될 수 있는 프레임이다.

다음 단계는?: 앞으로의 길

저자가 명시한 향후 계획은 없지만, 위 한계를 감안하면 다음 단계가 자연스럽게 도출된다.

• 스테일니스 정렬 통합 : R3·PR2의 라우팅 정렬과 TRACE의 양자화 정렬을 결합해, 정책 불일치의 두 원천(라우팅 + 양자화)을 함께 다루는 통합 프레임.
• 적응형 캐싱 : 레이어·스텝별로 가이던스 비트 수를 동적으로 조절해, 통신 오버헤드를 성능 손실 없이 더 줄이는 방향.
• 다른 정밀도·아키텍처로의 확장 : FP8, INT4, dense 모델, 그리고 더 긴 컨텍스트(1M 토큰급)에서의 검증.
• TCO 및 에너지 분석 : “속도 향상"을 넘어 총 연산·전력 비용 관점에서 저정밀 RL의 실질 이득을 정량화.

논문 원문의 표

arXiv e-print 의 LaTeX 원본에서 기계적으로 옮긴 표입니다. 숫자는 논문의 값이며 모델을 거치지 않았습니다.

표 1. 표 1

표 2. 표 2

표 3. 표 3

표 4. 표 4