NVIDIA의 Megatron-LM을 462B 파라미터 모델을 학습하는 코드로만 읽으면 중요한 절반을 지나치게 됩니다. 이 저장소는 Megatron-LM과 Megatron Core를 함께 담고 있습니다. Megatron-LM은 분산 학습을 시작하기 위한 레퍼런스 예제와 사전 구성된 스크립트를 제공하고, Megatron Core는 커스텀 훈련 프레임워크에 조립할 수 있는 GPU 최적화 building block을 제공합니다. 연구자가 빠르게 실험하는 입구와, ML 엔지니어가 훈련 시스템을 구성하는 부품이 한 저장소에 공존하는 셈입니다.
구조는 꽤 노골적입니다. megatron/core 아래에 모델과 transformer building block이 있고, tensor_parallel, pipeline_parallel, distributed, optimizer, datasets, inference, export가 각각 분리돼 있습니다. 여기에 post_training과 rl 디렉터리까지 붙습니다. TP·PP·DP·EP·CP 병렬화, FP16·BF16·FP8·FP4 혼합 정밀도, 양자화와 distillation, RLHF 관련 코드가 모델 정의 주변에 흩어진 것이 아니라 훈련 흐름의 일부로 배치돼 있습니다.
벤치마크의 조건도 운영 관점을 뒷받침합니다. 6,144장의 H100에서 462B 모델을 측정했고, 처리량에는 데이터 로딩과 optimizer step, 통신, logging이 모두 포함됐습니다. DP에서는 gradient reduce와 parameter gather를 겹치고, TP 통신도 계산과 겹치도록 구성합니다. 파이프라인 병렬화는 기본적으로 overlap을 사용합니다. 즉, 결과표의 MFU는 커널 하나의 최고 속도가 아니라 여러 단계가 이어지는 훈련 과정의 결과입니다.
그렇다고 47~48%라는 숫자를 모든 환경의 보장처럼 읽을 수는 없습니다. 모델 크기를 키운 weak scaling에서는 큰 GEMM의 높은 arithmetic intensity 덕분에 MFU가 41%에서 47~48%로 상승했습니다. 반대로 GPT-3급 모델의 strong scaling에서는 같은 batch size를 유지한 채 96 H100에서 4,608 H100으로 확장하면서 통신이 더 드러났고, MFU는 47%에서 42%로 떨어졌습니다. GPU 수가 늘어나는 것과 효율이 유지되는 것은 별개의 문제입니다.
실무에서 이 저장소를 검토할 때는 모델 크기보다 반복 훈련의 조건을 먼저 봐야 합니다. 장애 뒤 체크포인트에서 복구할 수 있는지, 병렬화 조합이 클러스터에 맞는지, Hugging Face와 Megatron 사이의 체크포인트 변환이 필요한지 확인해야 합니다. Megatron Bridge가 양방향 변환과 사전 학습·추론 레시피를 별도로 제공하는 이유도 여기에 있습니다. 다만 README의 성능 결과는 수렴까지 학습한 결과가 아니며, 소스 빌드는 메모리를 많이 사용할 수 있습니다. 거대한 모델을 한 번 돌리는 문제와, 그 훈련을 계속 운영하는 문제를 구분해야 이 프로젝트의 설계가 보입니다.