Megatron-LM
Megatron-LM
Megatron-LM은 NVIDIA가 만든, 대규모 Transformer 모델을 학습하기 위한 GPU 최적화 라이브러리예요. 수십억에서 수조 파라미터까지 모델을 수천 개 GPU에 걸쳐 효율적으로 학습시키는 데 특화돼 있어요. 이 저장소는 두 가지 컴포넌트로 이루어져 있는데, 레퍼런스 예제인 Megatron-LM과 커스터마이즈 가능한 라이브러리인 Megatron Core예요.
Megatron-LM과 Megatron Core
이 저장소는 두 가지 컴포넌트를 담고 있어요.
- Megatron-LM은 Megatron Core에 미리 구성된 학습 스크립트를 더한 레퍼런스 예제예요. 연구팀, 분산 학습 학습자, 빠른 실험에 이상적이에요.
- Megatron Core는 커스텀 학습 프레임워크를 만들기 위한 조합 가능한 라이브러리예요. Transformer 빌딩 블록, 고급 병렬화 전략(TP, PP, DP, EP, CP), 혼합 정밀도 지원(FP16, BF16, FP8, FP4), 모델 아키텍처를 제공해요.
**Megatron Bridge**는 Hugging Face ↔ Megatron 체크포인트 양방향 변환을 프로덕션 레디 레시피와 함께 제공해요.
벤치마크 성능
이 코드베이스는 수천 개 GPU에서 2B에서 462B 파라미터까지의 모델을 효율적으로 학습하며, H100 클러스터에서 최대 47%의 Model FLOP Utilization(MFU)을 달성해요.
- 6,144 H100 GPU: 462B 파라미터 모델 학습 벤치마크 성공.
- 초선형 스케일링: MFU가 41%에서 모델 크기에 따라 47-48%로 상승.
- 엔드투엔드 측정: 처리량에 데이터 로딩, 옵티마이저 스텝, 통신, 로깅이 모두 포함돼요.
약한 스케일링(weak scaling)에서 MFU가 41%(가장 작은 모델)에서 47-48%(가장 큰 모델)로 초선형 상승하는데, 이는 큰 GEMM이 더 높은 산술 강도(arithmetic intensity)를 가져 더 효율적으로 실행되기 때문이에요.
시작하기
PyPI에서 설치하거나 소스에서 빌드할 수 있어요.
uv pip install megatron-core
소스에서 설치하려면 클론하고 editable 모드로 설치해요.
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
uv pip install -e .
소스 빌드는 메모리를 많이 쓸 수 있어요. 메모리가 부족하면
MAX_JOBS로 병렬 컴파일 작업 수를 제한하세요 (예:MAX_JOBS=4 uv pip install -e .).
패키지 구조
Megatron-LM 저장소 구조를 살펴보면 어디에 무엇이 있는지 바로 파악돼요.
Megatron-LM/
├── megatron/
│ ├── core/ # Megatron Core (커널, 병렬화, 빌딩 블록)
│ │ ├── models/ # Transformer 모델
│ │ ├── transformer/ # Transformer 빌딩 블록
│ │ ├── tensor_parallel/ # 텐서 병렬화
│ │ ├── pipeline_parallel/ # 파이프라인 병렬화
│ │ ├── distributed/ # 분산 학습 (FSDP, DDP)
│ │ ├── optimizer/ # 옵티마이저
│ │ ├── datasets/ # 데이터셋 로더
│ │ ├── inference/ # 추론 엔진과 서버
│ │ └── export/ # 모델 내보내기 (예: TensorRT-LLM)
│ ├── training/ # 학습 스크립트
│ ├── post_training/ # 후처리 (양자화, 증류, 프루닝)
│ └── rl/ # 강화학습 (RLHF 포함)
├── examples/ # 바로 쓸 수 있는 학습 예제
└── tools/ # 유틸리티 도구