NVIDIA Megatron Core, 비트 단위 결정론적 사전 학습 대규모 구현 및 최적화

비트 단위 결정론 도입

NVIDIA는 Megatron Core에서 대규모 언어 모델의 사전 학습을 위한 비트 단위 결정론(Bitwise-Deterministic)을 구현하고 최적화한다. 이는 독립적인 실행과 체크포인트 재개 실행이 데이터 순서, 아키텍처, 레시피, 병렬 처리, 소프트웨어, 런타임 설정, 하드웨어 등 모든 조건이 고정될 때 동일한 수치적 궤적을 따르도록 보장하는 기술이다. 이로써 대규모 모델 학습 과정에서 발생하는 복잡한 디버깅, 검증, 재현 가능한 학습 재개 문제를 해결한다.

기존에는 수조 개의 매개변수를 가진 모델을 수천 개의 GPU에서 학습할 때, 다차원 병렬 처리, 저정밀도 계산, 분산 체크포인트 등으로 인해 실패를 재현하고 수정 사항을 검증하는 과정이 매우 복잡했다. 비트 단위 결정론은 이러한 문제를 해결하여 신뢰할 수 있는 실패 재현, 손실 급증 디버깅, 시스템 변경 검증, 중단된 학습 실행의 안정적인 복구를 가능하게 한다. 특히 수조 개 매개변수 규모에서는 작은 성능 저하도 수천 GPU-일의 비용으로 이어질 수 있으므로, 결정론적 학습은 효율성을 유지하는 것이 중요하다.

결정론 검증 및 해결

비트 단위 결정론을 검증하기 위해 단순히 출력된 손실 값만으로는 충분하지 않다. 대신 학습 손실, 언어 모델링 손실, 부하 분산 손실, 다중 토큰 예측(MTP) 손실, 그래디언트 노름 등 선택된 수치 메트릭을 각 학습 단계에서 핑거프린트한다. 동일한 구성을 두 번 실행하고 이 핑거프린트를 단계별로 비교하여 첫 불일치 단계를 식별한다. 체크포인트 재개 시에는 매 복원 후 매개변수, 옵티마이저 상태, RNG(Random Number Generator) 상태, 데이터 위치 및 후속 출력이 중단되지 않은 참조 실행과 비트 단위로 동일한지 확인한다.

결정론이 깨졌을 때 이를 해결하기 위한 워크플로우는 Megatron-LM PR #7262에서 제안되었다. 이 워크플로우는 오프라인 비교를 위해 순서가 지정된 랭크별 텐서 핑거프린트를 기록한다. 먼저 동일한 시드, 데이터 순서, 배치 크기, 병렬 처리, 컨테이너, 소프트웨어 스택을 확인한다. 이후 종단 간 비교에서 시작하여 캡처 범위를 학습 반복, 단계, 모듈, 연산, 커널로 점진적으로 좁혀 발산의 원인을 파악한다. `torch.hash_tensor`를 사용하여 GPU 상주 핑거프린트를 생성하며, MXFP8 또는 NVFP4 텐서의 경우 인코딩된 값과 스케일 버퍼를 모두 핑거프린트한다. 오탐(false alarm)을 배제하기 위해 디스패처 사각지대나 초기화되지 않은 텐서와 같은 프로브 아티팩트도 확인한다.

성능 최적화 사례

결정론적 학습은 디버깅에 유용하지만, 상당한 오버헤드가 발생하면 실제 생산 환경에서는 비용이 너무 많이 든다. 따라서 결정론적 및 비결정론적 실행은 처음부터 함께 최적화되어야 한다. NVIDIA는 Mamba 스타일 상태 공간 모델(SSM) 레이어와 트랜스포머 어텐션 레이어를 결합한 수조 개 매개변수 Nemotron 모델을 사용하여 결정론적 사전 학습의 오버헤드를 줄이는 데 성공했다.

최적화 결과, Nemotron 3 Ultra 모델의 단계 시간 결정론 오버헤드는 버퍼 채우기 최적화 후 약 17%에서 2%로, 3,072 GPU에서는 1.5%까지 감소했다. 하이브리드 Triton 프록시의 경우 MoE-MLP 퓨전 복원 후 약 60%에서 37%로, 이후 2%까지 개선되었다. CuTeDSL로 최적화된 가중치-그래디언트 경로는 256 GPU에서 약 3.6%의 오버헤드를 기록했다. 대규모 Nemotron 레시피는 2,432 GPU에서 약 2%의 안정 상태 결정론 오버헤드를 보였으며, 800단계 이상 비트 단위 결정론을 유지했다. 이는 그룹화된 GEMM 에필로그에서 여러 N-타일이 동일한 주소에 누적되어 결과가 도착 순서에 따라 달라지는 문제를 해결하기 위해 각 N-타일에 전용 출력 슬롯을 부여하고 고정된 순서로 결합하는 커널 수준 최적화를 통해 달성되었다.

개발자 영향 및 적용 조건

비트 단위 결정론은 개발자에게 대규모 사전 학습 모델의 디버깅 용이성, 체크포인트 연속성 보존, 예상치 못한 차이 감지, 변경 사항 비교 등 네 가지 주요 이점을 제공한다. 특히, 10,000 GPU에서 100일이 걸리는 비결정론적 실행의 오버헤드를 15%에서 5%로 줄이면 10일, 즉 100,000 GPU-일을 절약할 수 있어 막대한 비용 절감 효과를 가져온다. 또한, 새로운 커널, 퓨전, 정밀도 형식 또는 병렬 처리 구성이 비결정론을 유발할 수 있으므로, Megatron-LM은 `--deterministic-mode`를 통한 레시피 검증, 커널 테스트, 모듈 검증과 같은 다양한 검사 메커니즘을 통해 회귀를 방지한다.

비트 단위 결정론은 동일한 하드웨어 및 소프트웨어 환경 내에서 검증된다는 제약이 있다. GPU 세대, 네트워크 구성 또는 라이브러리 버전 간 비교는 이 검증 범위 밖에 있으며 다른 결과를 생성할 수 있다. 런타임 자동 튜닝, MoE 퓨전, 불완전한 RNG 상태 복원, 저수준 라이브러리 커널 변경, 저정밀도 가중치 재구성, 규모 종속적 통신 경로 등은 일반적인 비결정론의 원인으로 지목된다. 각 원인에 대해 진단 방법과 해결책이 제시되며, Megatron Core 사용자 가이드를 통해 지원되는 레시피를 검토하고 독립 실행 및 체크포인트 재개 실행을 비교하여 재현성을 검증하는 것으로 시작할 수 있다.