pytorch/torchtitan
A PyTorch native platform for training generative AI models
torchtitan – PyTorch 네이티브 대규모 생성형 AI 학습 플랫폼
개요
torchtitan은 PyTorch 팀의 오픈소스 라이브러리로, 대규모 언어 모델(LLM) 및 멀티모달 모델을 완전히 PyTorch에서 학습, 미세 조정, 강화 학습(RL)할 수 있습니다. 코드베이스는 의도적으로 작고 읽기 쉽게 설계되었으며, 전체 학습 파이프라인을 Python '레시피', 즉 구성 객체를 반환하는 함수로 표현합니다. 레시피를 교체하면 모델 크기, 데이터, 병렬화 전략을 변경하거나 모델을 다시 작성하지 않고 사전 학습에서 RL로 전환할 수 있습니다.
핵심 개념 및 중요성
| 개념 | 기능 | 유용성 |
|---|---|---|
| 구성 가능한 병렬화 | 모델은 샤딩 방법을 선언하고, 라이브러리는 DeviceMesh와 DTensor를 통해 Fully-Sharded Data Parallel(FSDP2), Tensor Parallel(TP), Pipeline Parallel(PP), Context Parallel(CP), Expert Parallel(EP)을 스택합니다. |
동일한 코드에서 병렬화의 모든 조합을 실험할 수 있으며, 단일 GPU에서 수백 개 GPU까지 사용자 정의 커널 없이 확장할 수 있습니다. |
| 사전 학습에서 RL까지 원스택 | 동일한 모델 정의와 커널이 고전적인 언어 모델 사전 학습, 지도 미세 조정(SFT), vLLM 생성을 지원하는 TitanRL 강화 학습 루프에도 사용됩니다. | 트레이너와 추론 간에 비트 단위 재현성을 보장하여 RL 기반 정렬 또는 정책 기울기 작업의 디버깅을 단순화합니다. |
| 최첨단 모델 지원 | Llama 3, Qwen 3, DeepSeek V3, FLUX, Kimi, Muse Glimmer 등의 준비된 레시피와 저정밀 형식(MXFP8, NVFP4) 및 MoE(Dist-MoE)를 포함합니다. | 최첨단 모델에서 아키텍처나 양자화 기법을 다시 구현하지 않고 실험을 시작할 수 있습니다. |
| Python 우선 구성 | 레시피는 구성 객체를 반환하는 Python 함수일 뿐입니다. 긴 CLI 플래그 목록을 전달하는 대신 함수를 편집합니다. | 실험을 재현 가능하고 버전 관리하기 쉽게 만들고, 레시피 파일을 동료와 공유하면 정확히 동일한 실행을 얻을 수 있습니다. |
| 광범위한 관찰 가능성 | 내장 메트릭(손실, TFLOPs, MFU), TensorBoard/W&B 로깅, 프로파일링 도구, 결정론 모드, 손실 비교 스크립트, SDC 재생. | 수백 개 GPU 규모에서 중요한 성능 병목 현상이나 수치 버그를 조기에 발견하는 데 도움이 됩니다. |
현재 할 수 있는 작업
- 새 LLM 학습 –
llama3_8b와 같은 레시피를 선택하고 8개 GPU에서./run_train.sh를 실행합니다. 라이브러리가 샤딩, 체크포인팅, 데이터 로딩을 자동 처리합니다. - LoRA 또는 SFT로 미세 조정 – 동일한 레시피 시스템으로 LoRA 어댑터나 채팅 데이터셋 미세 조정 단계를 추가할 수 있습니다.
- RL 기반 정렬 수행 –
TitanRL스택을 사용하여 PPO 스타일 또는 최신 GRPO/DAPO 손실을 실행하고, 학습된 정확한 모델 가중치를 재사용합니다. - 저정밀 학습 실험 – 구성 플래그로 MXFP8 또는 NVFP4를 활성화합니다. 라이브러리는 필요한
torchao커널을 포함합니다. - 수백 개 GPU로 확장 – 문서화된 벤치마크는 H100에서 최대 512개 GPU 실행과 자동 분산 체크포인팅 및 비동기 저장을 보여줍니다.
사용 대상
- 연구 실험실 – 새로운 병렬화나 RL 아이디어를 탐구할 깔끔하고 해킹 가능한 코드베이스가 필요한 경우.
- 엔지니어 – PyTorch 생태계 내에 머물고 싶은(별도 C++/CUDA 프레임워크 불필요) 프로덕션급 LLM 학습 파이프라인을 구축하는 경우.
- 학생 – FSDP, TP, MoE를 결합하는 방법을 보여주는 읽기 쉬운 레시피로 대규모 학습을 배우는 경우.
시작하기(빠른 요약)
- 설치 – 리포지토리를 클론하고 PyTorch 나이틀리와
torchtitan패키지를 설치합니다(또는 사전 빌드된 pip/conda 휠 사용). - 디버그 작업 실행 –
NGPU=1 ./run_train.sh로 작은 Llama-3 디버그 모델을 10단계 학습합니다. - 실제 모델 학습 – Hugging Face에서 토크나이저를 다운로드하고
MODULE=torchtitan_recipes.models.llama3 CONFIG=llama3_8b를 설정한 다음 8개 GPU에서 시작합니다. - 사용자 지정 – 병렬화 정도, 데이터 혼합, LoRA 추가 등을 조정하는 자체 레시피 함수를 작성합니다.
문서 및 커뮤니티
- 문서는 구성, 모델, 분산, 데이터, 체크포인팅, RL, 디버깅, 벤치마크와 같은 주제별 섹션으로 나뉩니다. 각 폴더에는 사용 예제가 포함된
README.md가 있습니다. - CI 배지는 CPU, NVIDIA(H100/B200 포함), AMD ROCm 및 RL 워크로드에서 활성 테스트를 보여줍니다.
- 기여 – 가이드라인은
experiments/폴더의 새로운 실험과 풀 리퀘스트를 통한 핵심 개선을 권장합니다. - 논문 및 인용 – ICLR-2025 논문에서 설계 선택에 대한 세부 사항을 설명합니다. torchtitan을 기반으로 한 연구를 게시할 때 인용하세요.
위 정보는 모두 리포지토리 README에서 직접 가져온 것입니다. 외부 가정은 추가되지 않았습니다.
관련
- 프로젝트
- 프로젝트
- 프로젝트
- 프로젝트
- 프로젝트