NVIDIA/Megatron-LM

해결하는 문제

Megatron-LM은 수천 개의 GPU에 걸쳐 대규모 트랜스포머 모델(20억 ~ 4620억 파라미터 범위)을 학습시키는 과제를 해결합니다. 하드웨어 효율성을 극대화하면서 대규모 모델 학습의 메모리 및 계산 요구 사항을 처리하는 데 필요한 인프라를 제공합니다.

작동 방식

이 프로젝트는 두 가지 주요 부분으로 나뉩니다. GPU에 최적화된 구성 가능한 빌딩 블록(커널, 트랜스포머 컴포넌트 및 옵티마이저)인 Megatron Core와, 이러한 블록을 사전 구성된 학습 스크립트와 결합하여 실험을 용이하게 하는 레퍼런스 구현체인 Megatron-LM입니다.

확장성을 달성하기 위해 다음과 같은 고급 병렬화 전략을 채택합니다:

• Tensor Parallelism (TP), Pipeline Parallelism (PP), Data Parallelism (DP), Expert Parallelism (EP) 및 Context Parallelism (CP).
• 메모리와 속도를 최적화하기 위해 FP16, BF16, FP8, FP4를 포함한 다양한 혼합 정밀도 형식을 지원합니다.
• Mixture-of-Experts (MoE) 아키텍처 및 Falcon-H1 (Transformer-Mamba)와 같은 하이브리드 아키텍처에 대한 특화된 지원을 포함합니다.

대상 사용자

• 맞춤형 학습 파이프라인을 구축하기 위해 고성능의 구성 가능한 라이브러리가 필요한 ML 엔지니어 및 프레임워크 개발자.
• 대규모 언어 모델을 빠르게 실험하고 학습시키기 위한 레퍼런스 구현체를 찾는 연구 팀.

주요 특징

• 극한의 확장성: 6,144개의 H100 GPU를 통해 최대 462B 파라미터 모델 학습 가능.
• 높은 하드웨어 효율성: H100 클러스터에서 최대 47%의 Model FLOP Utilization (MFU) 달성.
• 폭넓은 아키텍처 지원: DeepSeek-V4, MoE 모델 및 하이브리드 Transformer-Mamba 아키텍처 구현 포함.
• 상호 운용성: Megatron Bridge를 통해 Hugging Face와 Megatron 간의 양방향 체크포인트 변환 제공.