pytorch/torchtitan

torchtitan – 대규모 생성형 AI 모델 학습을 위한 PyTorch 네이티브 플랫폼

무엇인가요 – torchtitan은 PyTorch 팀이 개발한 오픈소스 라이브러리로, Llama 3.1과 같은 대규모 언어 모델(LLM)을 오직 네이티브 PyTorch 프리미티브만을 사용하여 스케일링 가능하게 학습할 수 있게 해줍니다. FSDP2, 텐서 병렬, 파이프라인 병렬, 컨텍스트 병렬 등 분산 학습 기술의 컬렉션을 체크포인트, 양자화, 로깅 유틸리티와 통합하여, 깔끔하고 확장 가능한 Python API를 통해 제공합니다.

왜 중요한가요 – 현대의 LLM(8 B에서 405 B 파라미터)을 학습하려면 복잡한 커스텀 CUDA 커널, 제3자 라이브러리, 무거운 오케스트레이션 도구의 스택이 필요합니다. torchtitan은 최신 PyTorch nightly 릴리스와 직접 작동하는 최소 코드베이스를 제공함으로써 이 복잡성을 줄이려 합니다. 연구자는 모델 코드를 다시 작성하지 않고도 새로운 병렬 전략을 실험할 수 있으며, 프로덕션 팀은 동일한 라이브러리를 사용해 프로토타입에서 멀티노드 학습으로 쉽게 전환할 수 있습니다.

핵심 기능 (README에 기재된 내용)

카테고리	기능
병렬 처리	• 다차원 조합 가능한 병렬 처리: FSDP2(파라미터별 샤딩), 텐서 병렬(비동기 TP 포함), 파이프라인 병렬(제로 버블), 초장문 시퀀스용 컨텍스트 병렬.
메모리 및 속도 최적화	• 메타 장치 모델 초기화, 선택적/전체 활성화 체크포인트, 분산 비동기 체크포인트, BF16 옵티마이저 상태, Float8 및 MXFP8 양자화 (Blackwell GPU 지원).
학습 워크플로우	• `torch.compile` 지원, 토큰 수 플래그를 통한 그래디언트 누적, 유연한 워밍업-안정-감소형 LR 스케줄러, C4 사전 구성된 내장 데이터로더, 사용자 정의 데이터셋 플러그인.
관측성	• TensorBoard / Weights & Biases 메트릭 (손실, 메모리, 처리량, MFU), 랭크별 구조화된 로깅, 프로파일링 도구 (CPU/GPU, 플라이트 레코더).
상호운용성	• 분산 체크포인트 형식(DCP)은 `torchtune`에서 직접 로드 가능하여 파인튜닝에 사용 가능. TorchFT 통합, 토크나이저 다운로드 및 체크포인트 변환용 헬퍼 스크립트.
확장성	• 잘 문서화된 확장 포인트, 빠른 프로토타이핑용 `experiments` 폴더 (예: TitanRL RL 스택), 새로운 모델 추가를 위한 명확한 가이드라인.
지원되는 하드웨어	• NVIDIA GPU (H100까지 테스트 완료, 512GPU 실행 가능), AMD ROCm 포크, 모든 CUDA/ROCm 버전에 대응하는 나이트리 빌드.

일반적인 워크플로우

설치 – pip install torchtitan (또는 나이트리 PyTorch + 나이트리 torchtitan 웨일 사용). 단일 노드 또는 Slurm/ParallelCluster 클러스터에서 즉시 작동.
데이터 준비 – 내장 C4 로더 사용 또는 설정 파일로 사용자 정의 데이터셋 지정.
모델 선택 – Llama 3.1 (8 B, 70 B, 405 B)이 포함됨. 새로운 모델 추가는 torchtitan/models/README.md의 짧은 가이드를 따름.
설정 – Python 설정 파일과 CLI 플래그(--module, --config)를 통해 모든 옵션 노출. 병렬 처리, 체크포인트, 양자화 옵션을 자유롭게 활성화/비활성화 가능.
실행 – 1~N 노드에서 torchrun (또는 제공된 run_train.sh) 실행. torchtitan은 요청된 병렬 처리를 자동으로 조합하고 학습을 시작.
모니터링 – 메트릭은 TensorBoard/W&B에 표시. 로그에는 디버깅용 랭크별 추적 정보 포함.
체크포인트 및 파인튜닝 – 저장된 DCP 체크포인트는 torchtune으로 직접 로드하여 하류 SFT 또는 RLHF에 사용 가능.

시작하기 (빠른 예제)

# 1️⃣ 나이트리 PyTorch + torchtitan 웨일 설치
pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu130
pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu130

# 2️⃣ Llama‑3.1 토크나이저 다운로드
python scripts/download_hf_assets.py \
    --repo_id meta-llama/Llama-3.1-8B \
    --assets tokenizer \
    --hf_token <your‑HF‑token>

# 3️⃣ 8GPU 실행 (단일 노드)
MODULE=llama3 CONFIG=llama3_8b ./run_train.sh

CONFIG=llama3_8b 파일은 torchtitan/config/에 있으며, 이미 FSDP2 + 텐서 병렬 + 활성화 체크포인트 등을 활성화하고 있습니다.

누구에게 적합한가요?

연구자 – 새로운 스케일링 법칙, 병렬 조합, 양자화 기술을 탐색하는 사람. 클린룸 구현으로 내부 코드를 쉽게 읽고 수정할 수 있습니다.
엔지니어 – 프로덕션 수준의 LLM 사전 학습 파이프라인을 구축하고 싶은 사람. 외부 프로퍼리터리 스택을 피하고 PyTorch 생태계 내에서 머무르고 싶은 사람.
학생 – 분산 학습 개념을 배우고 싶은 사람. experiments 폴더에는 실행 가능한 데모(예: TitanRL)가 포함되어 있어, 동일한 모델 코드가 RL 학습과 추론에 재사용되는 방식을 보여줍니다.

커뮤니티 및 지원

포럼 – 전용 PyTorch 포럼 카테고리 (pytorch‑forum → distributed → torchtitan).
논문 – ICLR 2025 논문 "TorchTitan: 프로덕션 준비 LLM 사전 학습을 위한 원스톱 PyTorch 네이티브 솔루션" (arXiv 2410.06511).
기여 – 핵심 변경 및 실험적 아이디어 추가에 대한 가이드라인은 CONTRIBUTING.md 및 torchtitan/experiments/README.md에 기재.
제3자 포크 – AMD 최적화 포크 존재 (AMD‑AGI/torchtitan‑amd).

라이선스

torchtitan은 BSD‑3‑Clause 라이선스 하에 배포됩니다. 코드 자체는 허용적인 라이선스이지만, 다운로드하는 외부 데이터나 모델 가중치의 라이선스를 존중해야 합니다 (예: Meta‑Llama 토크나이저 및 체크포인트).

결론 – torchtitan은 생산성에 맞는, 오직 PyTorch만을 사용하는 대규모 LLM 학습 스택을 제공하며, 가독성, 확장성, 최신 분산 학습 연구에 중점을 둡니다. 무거운 외부 생태계를 도입하지 않고 대규모 생성형 모델의 학습이나 실험을 원한다면, 이 라이브러리에서 시작하는 것이 가장 좋습니다.