NVFP4 강화학습 불일치를 업데이트 방향으로 구분해 안정화한다
TRIAGE: Direction-Aware Mismatch Stabilization of Native NVFP4 Reinforcement Learning
무엇인가
대규모 언어모델 강화학습에서 롤아웃 생성은 전체 학습 루프의 병목이다. NVIDIA Blackwell의 NVFP4는 가중치와 활성값을 모두 4비트로 쓰는 W4A4 텐서코어 연산을 네이티브로 지원해 BF16 대비 최대 4배의 행렬 연산 처리량을 준다. 문제는 샘플러와 학습자가 서로 다른 실행 경로로 같은 모델을 돌린다는 점이다. 두 경로의 로그확률 차이 δ = log π_learner − log π_sampler는 동기화된 가중치에서도 존재하고, 저정밀 실행이 이를 더 키운다. 기존 대응은 두 갈래였다. 양자화 인지 학습(QAT)은 학습자를 저정밀 샘플러에 맞추려 하지만 학습자 순전파에서 네이티브 W4A4를 유지하지 못한다. TIS·클리핑·마스킹 같은 중요도 기반 보정은 불일치의 크기만 보고 개입한다. 이 논문은 크기만으로는 그 불일치가 다음 업데이트에서 줄어들지 늘어날지를 알 수 없다고 지적한다.
어떻게 동작하나
저자들은 불일치 에너지 D = ½Σδ²의 1차 전개로 ΔD = η(B_amp − B_con + R_cross)를 유도한다. 여기서 δ·g > 0인 항은 불일치를 증폭시키고, δ·g < 0인 항은 수축시킨다. 바닐라 GRPO에서는 이 기준이 A_i·δ_i,t의 부호로 환원되어, 증폭 영역이 A⁻ = {A_i<0, δ<0}과 A⁺ = {A_i>0, δ>0} 두 개로 갈린다. Qwen3-4B-Base와 Qwen3-30B-A3B-Base의 NVFP4 GRPO 궤적을 25업데이트 윈도우로 나눠 관찰한 결과, 두 증폭 영역의 이론적 대칭이 먼저 깨진다. 이점 가중 점유비 ρ_asym은 4B에서 1.02→2.32, 30B에서 1.34→1.85로 커지는데, 그 시점에 |δ|<0.05인 토큰 비율은 각각 83.0%, 91.9% 이상으로 여전히 대부분이다. 즉 크기 분포가 멀쩡해 보일 때 방향 불균형이 먼저 나타난다. 말기에는 5% 분위 δ가 4B에서 −42.75, 30B에서 −10.94까지 무너진다. 또한 64토큰 세그먼트로 나눠 보면, 꼬리 토큰의 전역 빈도는 4.8~4.9%에서 2.8~3.6%로 오히려 줄지만 상위 5% 세그먼트가 전체 증폭 꼬리 토큰의 약 4분의 1(균등 배치 대비 약 5배)을 담고, 인접 꼬리 과밀 지속성이 우연 대비 2.6배에서 3.8~4.0배로 올라간다. 응답 평균은 이 구조를 가린다. 꼬리 과밀 세그먼트를 가진 응답의 97.2%(4B), 79.9%(30B)가 여전히 |δ̄_i| < 0.05다.
무엇과 다른가
TRIAGE는 이 진단을 업데이트 계수 g_i,t = α_i,t·g^base_i,t + ψ_i,t 두 항으로 옮긴다. 첫째, 세그먼트 단위 진단이다. 세그먼트 S의 가중 평균 갭 δ̄_S와 꼬리 비율 f^tail_S를 계산해 게이트 가중 w_S ∈ [w_min, 1]로 사상한다. 둘째, 방향 선택적 게이팅이다. 위험 세그먼트라도 모든 토큰을 누르지 않는다. A_i<0이면서 δ_i,t<0인, 즉 증폭에 기여하는 토큰에만 w_S를 적용하고 같은 세그먼트의 수축 토큰은 그대로 둔다. 양의 이점 증폭 영역 A⁺는 게이팅하지 않는다. 보상이 양수인 응답의 확률 상승을 누르면 학습 신호 자체를 억제하기 때문이다. 가중 정책 손실을 정규화하면 α_i,t = w_i,t / w̄_i가 되어, 음의 이점 응답 전체를 균일하게 약화하는 대신 증폭 부분집합에서 상대 업데이트 질량을 빼앗아 B_amp와 B_con의 균형을 바꾼다. 셋째, 양의 이점 응답에 대한 유계 보정이다. 게이팅은 곱셈적이라 부호를 뒤집거나 남은 갭을 직접 줄이지 못한다. 그래서 보정 경계 위반량 q_R = [τ_rep − δ̄_R]₊에 슈도-휴버 페널티 φ_β를 걸고, 고정 업데이트 k_rep 이후 λ_k = λ_rep로 활성화한다. 모든 임계값과 스케줄은 실행 전에 고정된다. 핵심은 샘플러와 학습자 양쪽 순전파에서 네이티브 NVFP4 W4A4 실행을 그대로 유지한 채 최적화 목적함수만 수정한다는 점이다.
어떻게 쓰나
실험은 Qwen3-30B-A3B-Base를 주 모델, Qwen3-4B-Base를 재현·절제용으로 쓰고, DAPO 데이터셋에서 GRPO로 학습하며 최대 응답 길이를 20,480으로 둔다. 비교군은 BF16, 네이티브 NVFP4, NVFP4+TIS(절단 임계 C=2.0), NVFP4+TRIAGE 네 가지다. 순수 NVFP4는 4B에서 약 300스텝 후 보상이 급락하고 평균 절대 로그확률 갭이 치솟으며, TIS는 붕괴를 늦출 뿐 막지 못한다. 30B에서도 순수 NVFP4는 약 700스텝 후 불일치가 급증한다. TIS는 1,700스텝 동안 발산하지 않지만 후반 보상 궤적이 BF16과 TRIAGE 아래에 머문다. TRIAGE는 600스텝·1,700스텝 일정 내내 안정적이며, 4B에서 순수 NVFP4 대비 관측 안정 학습 구간을 약 두 배로 늘린다. 다섯 개 수학 추론 벤치마크(GSM8K, MATH-500, AIME 2024/2025, AMC 2023) 평균에서 TRIAGE는 4B 58.49% 대 BF16 58.26%, 30B 70.96% 대 BF16 72.41%로 BF16 수준에 근접한다. 4B의 공통 300스텝 체크포인트에서 TRIAGE는 순수 NVFP4와 NVFP4+TIS를 각각 5.93%p, 2.27%p 앞선다. TIS는 1,700스텝을 완주하고도 붕괴 전 순수 NVFP4의 700스텝 체크포인트(65.87%)에 못 미치는 63.18%를 기록한다.
전제와 한계
시스템 수치는 롤아웃 처리량이 배치 256과 512에서 각각 6,745, 8,717 출력 토큰/초로 BF16의 2.27배, 2.30배다. 롤아웃·학습자 업데이트·가중치 변환과 동기화·오케스트레이션을 포함한 반복 시간은 배치 256에서 208.2초→172.1초, 배치 512에서 254.5초→195.8초로 각각 1.21배, 1.30배 단축된다. TRIAGE 자체의 학습자 추가 비용은 B300 8장에서 동일 배치를 5회 반복 측정해 중앙값 0.84%, 절대 증가 0.225초였다. 전체 실험 규모는 B300 GPU 기준 34,500시간 이상이다. 절제 실험은 4B에서 동일한 TRIAGE 300스텝 체크포인트에서 분기해 수행했는데, 전체 TRIAGE가 가장 높은 보상을 얻는다. 보정(repair)을 제거하면 평균 로그확률 갭은 낮게 유지되지만 이점 기반 업데이트 불균형이 교정되지 않아 보상이 떨어지고, TIS만 단독 적용하면 처음부터 쓸 때와 달리 안정적으로 유지된다. 이는 불일치 역학이 학습 단계에 의존한다는 3장 분석과 일치한다.
실무 관점에서 이 논문이 주는 시사점은 명확하다. 저정밀 롤아웃을 도입할 때 평균 로그확률 갭이나 |δ| 분포만 모니터링하면 위험 구간을 놓친다. 대신 A_i·δ_i,t 부호 조합별 점유비와 응답 내 64토큰 세그먼트 수준의 꼬리 집중도를 함께 봐야 조기 경보를 얻는다. 개입도 응답 전체나 토큰 단위가 아니라 세그먼트 진단과 방향 조건을 결합해 최소한으로 하는 편이 학습 신호 손실이 적다. 다만 TRIAGE는 게이트 임계값·보정 경계 τ_rep·시작 업데이트 k_rep 등 여러 고정 하이퍼파라미터에 의존하므로, 자체 모델과 데이터에서 재조정이 필요하다.
저자들이 밝힌 한계도 분명하다. 첫째, 실험 규모가 제한적이어서 100B 이상 모델이나 훨씬 긴 학습 지평은 평가하지 않았고, 주 실행은 수렴이나 안정적 후반 거동이 관찰되면 종료했다. 둘째, 동기식 RL만 다뤘다. 방향 기준 자체는 비동기 RL에도 적용될 수 있고 정책 노후화가 추가 불일치를 만든다는 점에서 특히 관련이 깊다고 보지만 실험 검증은 없다. 셋째, NVFP4가 롤아웃을 크게 가속해도 전체 반복 속도 향상은 약 1.3배에 그친다. 런타임의 상당 부분이 가속된 GEMM 밖의 연산, 즉 학습자-샘플러 간 파라미터 동기화와 Amax 계산 같은 양자화 관련 리덕션·데이터 변환에 쓰이기 때문이다. 저자들은 더 긴밀하게 통합된 저정밀 커널과 RL 프레임워크 공동 설계, 그리고 증폭 토큰을 일괄 처리하는 대신 크기·국소성·지속성에 적응하는 더 세밀한 보정 규칙을 향후 과제로 제시한다.