TorchTitan - PyTorch 기반 대규모 LLM 사전 학습 치트시트
TorchTitan은 대규모 생성 모델(LLM) 사전 학습을 위한 오픈소스 PyTorch 네이티브 플랫폼입니다. PyTorch의 최신 분산 학습 기술의 깔끔한 참조 구현으로, 4D 병렬화 — 데이터, 텐서, 파이프라인, 컨텍스트 병렬 — 을 모듈식으로 조합합니다. FSDP2, torch.compile, 분산 체크포인팅을 기반으로 탄력적 스케일링을 지원합니다. 타사 스택 대신 네이티브 PyTorch 스케일링을 사용하여 모델을 사전 학습하려는 팀에 가장 좋은 시작점입니다.
분산 사전 학습은 리소스 집약적입니다. 대규모 다중 노드 실행을 수행하기 전에 소규모 모델/스케일에서 구성을 검증하세요.
설치
| 방법 | 명령어 |
|---|
| 소스에서 | git clone https://github.com/pytorch/torchtitan && cd torchtitan && pip install -e . |
| 요구사항 | 최신 PyTorch (nightly 권장), CUDA GPU |
| 데이터셋 | HF 데이터셋 / 토크나이저 사용 |
| 확인 | 1개 GPU에서 디버그 구성 실행 |
병렬화 차원 (4D)
| 차원 | 분할 |
|---|
| 데이터 병렬 (FSDP2) | 배치, 매개변수/최적화기/그래디언트 샤딩 |
| 텐서 병렬 | 개별 레이어 전체 GPU |
| 파이프라인 병렬 | GPU 전체의 모델을 단계로 |
| 컨텍스트 병렬 | GPU 전체의 긴 시퀀스 |
이들은 구성됩니다: 예를 들어, dp × tp × pp × cp는 모델을 대규모 GPU 메시에 매핑합니다.
구성
TorchTitan은 TOML 구성 파일과 명령줄 오버라이드에 의해 구동됩니다.
# 구성으로 학습 시작 (torchrun을 통한 다중 GPU)
CONFIG_FILE=./train_configs/llama3_8b.toml ./run_train.sh
| 구성 섹션 | 제어 |
|---|
[model] | 아키텍처, 크기, 일반 |
[training] | 배치 크기, 단계, 시퀀스 길이, LR |
[parallelism] | data_parallel_shard_degree, tensor_parallel_degree, pipeline_parallel_degree, context_parallel_degree |
[optimizer] | 최적화기 + 하이퍼파라미터 |
[checkpoint] | 분산 체크포인트 빈도/경로 |
[activation_checkpoint] | 재계산 정책 |
주요 기능
| 기능 | 이점 |
|---|
| FSDP2 | 최신 완전 샤드 데이터 병렬 |
torch.compile | 속도를 위한 그래프 컴파일 |
| 비동기 분산 체크포인팅 | 거대한 실행을 효율적으로 저장/복구 |
| Float8 / 낮은 정밀도 | 지원되는 하드웨어에서 더 빠른 학습 |
| 탄력적 스케일링 | GPU 가용성 변경에 적응 |
| 메타 디바이스 초기화 | OOM 없이 거대한 모델 빌드 |
다중 노드 실행
# 노드 전체의 torchrun (개념적)
torchrun --nnodes=4 --nproc_per_node=8 \
--rdzv_backend=c10d --rdzv_endpoint=$MASTER:29500 \
-m torchtitan.train --job.config_file llama3_70b.toml
| 오버라이드 | 예 |
|---|
| 병렬화 차원 | --parallelism.tensor_parallel_degree 8 |
| 단계 | --training.steps 10000 |
| 시퀀스 길이 | --training.seq_len 8192 |
| 컴파일 | --training.compile |
모니터링 및 체크포인트
| 측면 | 방식 |
|---|
| 로깅 | TensorBoard / W&B 통합 |
| 메트릭 | 손실, 토큰/초, MFU (모델 FLOPs 활용도) |
| 체크포인트 | 분산 체크포인트 (DCP) 저장/로드 |
| 복구 | 체크포인트 단계를 가리키도록 구성 |
TorchTitan vs 다른 학습 프레임워크
| 측면 | TorchTitan | Megatron-LM | Nanotron |
|---|
| 출처 | PyTorch 팀 | NVIDIA | Hugging Face |
| 스타일 | PyTorch 네이티브, 조합 가능 | 기능 풍부, 성숙 | 최소한, 해킹 가능한 3D |
| 병렬화 | 4D (dp/tp/pp/cp) | 3D+ (광범위) | 3D |
| 최적 용도 | 네이티브 PyTorch 참조 스케일링 | 최대 스케일 MoE/dense | 작고 해킹 가능한 3D 학습 |
자료
