NVFP4 강화학습 레시피 정리: 4비트 양자화 안정성 확보와 그래디언트 스파이크 해결법

핵심 메시지

• 최신 GPU에서 4비트 텐서 코어 처리량이 BF16 대비 크게 벌어지면서, 저정밀 연산은 선택이 아니라 하드웨어를 쓰는 기본 조건이 되고 있다.
• 사전학습용 4비트 레시피는 강화학습에 그대로 옮겨지지 않는다. 보상과 그래디언트가 이미 잡음투성이라 양자화 오차가 유용한 갱신을 덮어버린다.
• 텐서 단위 스케일링을 토큰 단위로 바꾸면 배치 구성에 따라 값이 흔들리지 않고, 롤아웃과 학습이 같은 정책을 평가한다는 약속이 지켜진다.
• 블록마다 인코딩을 두 가지로 계산해 오차가 작은 쪽을 고르는 기법은 정확도를 높이지만 연산량이 늘어, 커널을 다시 설계해 느려짐을 8배에서 2.8배로 줄였다.
• 역전파를 어떻게 처리할지는 옵티마이저 선택과 함께 결정해야 한다. 애덤과 짝지었을 때만 그래디언트 스파이크가 사라졌다.

쉽게 이해하기

GPU MODE 강연에 나온 이 발표는 4비트 부동소수점 형식으로 대규모 언어 모델의 강화학습을 돌린 경험을 공개한 것으로, 발표자는 긴 호흡의 에이전트형 강화학습을 연구하는 팀에서 이 레시피를 만들고 오픈소스로 공개했다. 저정밀 연산을 쓰는 이유는 명확하다. 메모리 사용량이 줄어 더 적은 GPU로 모델을 올릴 수 있고, 통신량이 줄며, 무엇보다 최신 GPU의 4비트 텐서 코어 처리량이 BF16과 비교할 수 없이 크기 때문이다. 다만 4비트라고 해서 값을 그냥 4비트로 자르는 것은 아니다. 실제로는 16개 정도의 작은 블록마다 8비트 스케일 값을 하나씩 붙이고, 그 위에 텐서 전체를 아우르는 32비트 스케일을 한 겹 더 얹는다. 문제는 이 두 번째 스케일이다. 추론에서는 보정 과정에서 미리 구한 값을 체크포인트에 저장해 두면 되지만, 강화학습은 정책이 계속 바뀌고 가중치를 동기화해야 한다. 매번 다시 보정하기엔 느리고, 온라인으로 구하면 롤아웃 엔진과 학습 엔진의 배치 구성이 달라 값이 어긋난다.

해결책은 스케일 기준을 텐서에서 토큰 단위로 내리는 것이다. 토큰마다 자기 스케일을 스스로 만들기 때문에 배치가 어떻게 묶이든 값이 변하지 않고, 롤아웃과 학습의 순전파가 같은 규칙을 따른다는 약속이 성립한다. 발표자는 이를 정밀도 문제라기보다 '동작 계약'의 문제로 봐야 한다고 강조했다. 같은 진단은 다른 팀의 기술 보고서에서도 나왔는데, 텐서 단위 스케일이 미래 토큰의 정보를 과거 토큰으로 새게 만들어 그래디언트를 편향시킨다는 지적이었다.

정확도를 더 끌어올린 두 번째 장치는 블록마다 인코딩을 골라 쓰는 방식이다. 4비트로 표현 가능한 값들은 0 근처가 촘촘하고 위로 갈수록 성겨서, 특히 4와 6 사이에 큰 구멍이 있다. 블록의 최댓값을 6에 맞추는 관행적 방식 대신 4에 맞추는 인코딩을 함께 계산해 오차가 작은 쪽을 쓰면 결과가 나아진다. 대신 두 인코딩을 모두 계산하고 역양자화해 오차를 비교해야 해서, 원래 메모리 대역폭에 묶여 있던 커널이 연산에 묶이는 커널로 바뀐다. 발표자는 오차 비교를 16비트 영역에서 정확히 수행하고 스케일 곱을 한 번만 적용하는 식으로 다시 짜서, 기준 구현 대비 8배였던 느려짐을 2.8배까지 줄였다.

역전파에서는 더 미묘한 선택이 남는다. 사전학습 레시피처럼 역전파까지 4비트로 밀면 양자화 오차가 강화학습 신호를 덮어버린다. 그렇다고 고정밀 원본으로 역전파하면 실제로 순전파에서 쓴 양자화된 연산이 아니라 다른 연산의 기울기를 구하는 셈이 된다. 발표팀은 순전파에서 쓴 그대로를 역양자화한 뒤 역전파하는 방식을 택했는데, 이 방법이 항상 이기는 것은 아니었다. 순수 경사하강법에서는 고정밀 역전파가 오히려 자연스러운 평활 효과를 내 더 안정적이었고, 애덤처럼 파라미터마다 모멘텀을 들고 있는 옵티마이저와 짝지었을 때 비로소 그래디언트 스파이크가 사라졌다.

마지막 장치는 지렛대가 큰 레이어만 높은 정밀도로 남기는 것이다. 마지막 몇 개 레이어는 최종 로그 확률에 직접 영향을 주고 역전파의 출발점이기도 하다. 전문가 혼합 모델의 공유 전문가도 라우터 가중치로 희석되지 않고 항상 켜져 있어 영향이 직접적이다. 대형 모델에서는 라우팅되는 전문가가 전체 가중치의 97.5%를 차지해, 나머지를 통째로 BF16으로 남겨도 메모리 부담이 크지 않다. 최종 레시피는 학습 처리량을 일부 내주는 대신 안정성을 사는 쪽으로 정리됐고, 학습 자체는 보통 전체 루프의 병목이 아니라는 판단이 그 근거였다.

주요 인사이트

• 사전학습과 강화학습은 최적화 목표가 다르다. 전자는 대략의 방향만 지키면 되지만, 후자는 드물고 작은 유효 갱신이 양자화 오차에 묻히지 않아야 한다.
• 강화학습에서 저정밀 도입의 핵심 제약은 정확도가 아니라 롤아웃과 학습이 같은 정책을 평가한다는 일관성이다. 배치 구성에 따라 값이 달라지는 스케일링은 이 전제를 깬다.
• 비트 단위로 동일한 결과를 스택 전체에서 보장하는 것은 현실적으로 불가능하지만, 최소한 학습 가중치와 롤아웃 가중치만은 완전히 같아야 한다는 것이 발표팀의 기준선이다.
• 역전파 방식과 옵티마이저는 따로 고를 수 없다. 같은 역전파 설정이 경사하강법에서는 불리하고 애덤에서는 유리하게 뒤집힌 실험 결과가 이를 보여준다.
• 부수 효과로, 이 방식은 보정 과정을 거친 별도 체크포인트를 기다리지 않고 기존 8비트·16비트 체크포인트를 로딩 시점에 4비트로 바꿔 서빙할 수 있게 해 준다.

자주 묻는 질문

관련 AI 소식