반올림 위치 정렬이 BF16 LLM 강화학습의 붕괴를 막는다
Where Rounding Happens Matters: Rounding-Location Alignment Stabilizes LLM Reinforcement Learning
무엇인가
이 논문은 LLM 강화학습에서 추론 시스템과 훈련 시스템이 분리되어 있을 때 생기는 정책 불일치 문제를 다룬다. 가중치와 BF16 텐서가 같아도 구현 차이 때문에 의도한 온폴리시 업데이트가 오프폴리시가 되고 훈련이 불안정해질 수 있다는 점이 출발점이다. 기존에는 추론·훈련 스택 전체를 맞추거나 전역 정밀도를 바꾸거나 학습 알고리즘을 수정했지만, 이 논문은 실행 스택의 일부만 정렬해도 BF16 RL이 안정되는지 묻는다.
어떻게 동작하나
제안 방법은 훈련 쪽의 반올림 위치를 추론 쪽과 맞추는 것이다. 대상은 RoPE, SwiGLU, 잔차 RMSNorm 세 영역이며, 선택된 중간값을 FP32로 유지하다가 해당 영역의 BF16 출력 지점에서 반올림한다. 즉 전체 수치 스택을 통일하는 대신, 추론과 훈련 사이에 어긋나는 반올림 지점을 국소적으로 정렬한다.
무엇과 다른가
중요한 점은 추론 엔진, 행렬곱과 어텐션 커널, 학습 목적함수는 그대로 둔다는 것이다. RoPE는 위치 정보를 넣는 경로, SwiGLU는 피드포워드 계열 활성화, 잔차 RMSNorm은 잔차 스트림의 정규화 지점으로, 각각 BF16 연산 중 어디서 반올림이 일어나느냐가 추론과 훈련에서 달라질 수 있는 지점이다. 이 논문은 이런 위치 정렬만으로 나머지 구현 차이가 남아 있어도 훈련 안정성을 확보할 수 있는지 본다.
어떻게 쓰나
실험은 Qwen 계열 모델 네 개에서 진행됐다. 초록에 따르면 세 영역을 모두 정렬하면 기본 BF16 베이스라인에서 관찰된 훈련 붕괴를 막을 수 있었다. 또한 연구된 실행에서는 RoPE 경로 정렬만으로도 충분했으며, 그 경우에도 수치적 불일치는 남아 있었다. 다만 초록은 데이터셋 이름이나 구체적 지표 수치를 제시하지 않으므로 그 부분은 확인할 수 없다.
전제와 한계
개발자에게 이 결과는 추론과 훈련을 서로 다른 프레임워크로 구성한 RL 파이프라인에서 반올림 위치가 실질적인 디버깅 포인트가 된다는 뜻이다. 가중치와 BF16 설정을 같게 했는데도 온폴리시 데이터가 어긋나거나 학습이 붕괴하면, RoPE·SwiGLU·RMSNorm 주변의 중간 정밀도와 반올림 지점을 추론 엔진과 비교해 볼 필요가 있다. 전체 커널을 교체하거나 전역 FP32로 올리지 않고도 선택적 정렬을 시도할 수 있다는 점이 실무적 시사점이다.
한계는 초록이 밝힌 범위 안에서 해석해야 한다. 실험은 Qwen 계열 네 모델과 연구된 실행에 한정되며, RoPE 경로 정렬만으로 충분했다는 결과도 해당 실행에서의 관찰이다. 세 영역 정렬 후에도 또는 RoPE 정렬 후에도 수치적 불일치가 남을 수 있고, 추론 엔진과 행렬곱·어텐션 커널, 학습 목적함수는 바꾸지 않은 조건에서 얻은 결과다. 따라서 다른 모델·커널·정밀도 조합에서도 같은 효과가 보장된다고 일반화하기는 어렵다.