싱킹 머신스 랩, 4비트 학습 기법 공개

싱킹 머신스 랩(Thinking Machines Lab)이 11월 13일(미국 시간) 자사 모델 Inkling 2 Small을 4비트 연산으로 학습한 기법을 공개했다. 사전학습 첫 단계부터 마지막 강화학습 롤아웃까지 순전파를 4비트 형식 NVFP4로 돌렸고, 가장 큰 학습 작업에서 전부 8비트 MXFP8로 학습할 때보다 사전학습은 8%, 강화학습은 20% 빨라졌다. 글을 쓴 잭 쿡은 회사 연구 블로그 Connectionism에 "Inkling 2 Small은 우리가 아는 한 네이티브 4비트 순전파로 처음부터 끝까지 학습한 가장 큰 모델"이라고 적었다. 싱킹 머신스 랩은 공식 X 계정에서 "더 적은 비트로 학습하면 빠르지만, 양자화 오차가 학습을 방해할 수 있다"고 문제를 요약했다.

저정밀 학습의 득실은 숫자로 갈린다. 회사가 학습에 쓰는 엔비디아 GB300에서 BF16은 값 하나에 16비트를 쓰고 2.5페타플롭스를 낸다. MXFP8은 8.25비트에 5페타플롭스, NVFP4는 4.5비트에 15페타플롭스다. NVFP4는 BF16보다 3.6배 적은 비트로 저장되면서 최대 6배 빠르게 돈다. 대신 4비트로는 16개 값밖에 나타낼 수 없어 반올림이 거칠어지고, 그 오차가 층과 학습 단계를 거치며 쌓인다. 처음 시도한 완전 NVFP4 학습은 모델의 연산 효율을 회사가 받아들일 수 없을 만큼 떨어뜨렸다.

싱킹 머신스 랩은 학습 단계마다 다른 기준을 적용했다. 수조 개 토큰을 처리하는 사전학습 순전파와 롤아웃마다 수만 개 토큰을 만드는 강화학습 샘플링에서는 속도가 가장 중요하다고 봤다. 가중치를 고치는 역전파는 배치 전체를 평균해 무작위 오차가 희석되지만, 값이 한쪽으로만 반올림되는 편향에는 민감하다. 그래서 4비트는 전문가 혼합(MoE) 층의 순전파에만 쓰고, 역전파는 MXFP8로 계산했다. MoE 층은 연산량이 크면서도 양자화했을 때 정확도 손실이 다른 층보다 적었다는 것이 회사 설명이다.

가장 까다로운 대목은 순전파와 역전파 사이의 어긋남이었다. NVFP4는 연속된 값 16개를 한 묶음으로 같은 배율을 쓰는데, 역전파에서 행렬을 뒤집으면 묶음이 바뀌어 같은 값이 다르게 표현된다. 엔비디아가 제안한 16×16 타일 방식은 묶음을 256개로 키워 이 문제를 없애지만, 큰 값 하나가 나머지를 0으로 만들어 버린다. 회사는 이 방식이 너무 거칠다고 판단하고, 값을 높은 정밀도로 되돌려 뒤집은 다음 확률적 반올림으로 다시 양자화했다. 2.25를 확률적으로 반올림하면 4번 중 3번은 2, 1번은 3이 되어 평균이 원래 값과 같아진다. 블랙웰 GPU는 이 확률적 반올림을 하드웨어로 가속한다.

연산이 빨라지자 이번에는 통신이 막혔다. GB300 NVL72에서 GPU 사이 NVLink 대역폭은 방향마다 초당 900GB인데 NVFP4 연산은 최대 15페타플롭스라, GPU가 1바이트를 받는 동안 최대 1만 6,000번의 부동소수점 연산을 할 수 있다. 회사는 토큰을 전문가에게 보낼 때 NVFP4로 보내고, 결과를 모을 때는 32차원 무작위 아다마르 변환을 거쳐 8비트 형식 S2F6로 줄였다. 모으기 단계의 전송량은 값 하나당 16비트에서 8.25비트로 줄었고, 학습 중 측정 가능한 정확도 저하는 없었다. 이 모든 계산은 회사가 직접 만든 커널 묶음 BBBK(Big Beautiful Blackwell Kernel)가 맡는다.

비용도 공개했다. 마지막 16.6% 층은 MXFP8로 남겨 두었고, 이 방식은 사전학습에서 연산 효율을 약 2% 깎았다. MXFP8로 학습했다면 약 2% 적은 연산으로 같은 손실에 닿았을 것이라는 뜻이다. 강화학습에서는 보상에 눈에 띄는 차이 없이 상당한 속도 향상을 얻었다고 회사는 밝혔다. 처음부터 끝까지 같은 수치 설정을 쓰기 때문에, 학습이 끝난 뒤 양자화 체크포인트를 내보낼 때 따로 보정할 필요도 없다.

보도에 따르면 Inkling 2 Small은 전체 매개변수 5,940억 개 가운데 토큰마다 170억 개를 쓰는 MoE 모델이다. 첫 사전학습에서 토큰 8조 4,000억 개를 처리했고, 강화학습은 GPU 2,048개로 9일 동안 1,700만 회 넘는 롤아웃을 돌렸다. 모델은 회사의 미세조정 서비스 Tinker에서 먼저 쓸 수 있고, 가중치는 외부 레드팀 시험이 끝난 뒤 아파치 2.0 라이선스로 공개할 계획이다.

메탈은 멀티버스 컴퓨팅이 4비트로 압축한 모델이 원본보다 높은 점수를 낸 연구를 보도한 바 있다. 그 연구가 다 학습한 모델을 줄이는 쪽이었다면, 이번 글은 학습 자체를 4비트로 돌리는 쪽이다. 메탈이 확인한 시점까지 싱킹 머신스 랩의 발표 게시물은 조회수 약 40만 8,000회, 좋아요 2,077개를 기록했다. 글은 NVFP4 학습에 관한 엔비디아와 딥시크, 커서의 선행 연구를 출처로 밝히고 그 위에 쌓았다고 적었다.

AI 엔지니어의 눈으로 보면 이 글의 무게는 8%와 20%라는 숫자보다 학습 기간의 변화에 있다. 잭 쿡은 "몇 달 전만 해도 Inkling 2 Small과 비슷한 크기의 모델을 학습하는 데 몇 주가 걸렸다"며 이번에는 GB300 NVL72 랙에서 68시간 만에 사전학습을 마쳤다고 밝혔다. 2020년까지만 해도 기본 형식은 FP32였고, BF16은 엔비디아 암페어와 함께 들어왔다. 회사는 다음 단계로 역전파까지 FP4로 계산하는 방법과 엔비디아 Vera Rubin 시스템에서의 설정을 연구하고 있다. 4비트로 버티는 층이 늘어날수록 같은 랙에서 돌릴 수 있는 실험의 수가 달라진다.