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를 결합하는 방법을 보여주는 읽기 쉬운 레시피로 대규모 학습을 배우는 경우.

시작하기(빠른 요약)

  1. 설치 – 리포지토리를 클론하고 PyTorch 나이틀리와 torchtitan 패키지를 설치합니다(또는 사전 빌드된 pip/conda 휠 사용).
  2. 디버그 작업 실행 – NGPU=1 ./run_train.sh로 작은 Llama-3 디버그 모델을 10단계 학습합니다.
  3. 실제 모델 학습 – Hugging Face에서 토크나이저를 다운로드하고 MODULE=torchtitan_recipes.models.llama3 CONFIG=llama3_8b를 설정한 다음 8개 GPU에서 시작합니다.
  4. 사용자 지정 – 병렬화 정도, 데이터 혼합, LoRA 추가 등을 조정하는 자체 레시피 함수를 작성합니다.

문서 및 커뮤니티

  • 문서는 구성, 모델, 분산, 데이터, 체크포인팅, RL, 디버깅, 벤치마크와 같은 주제별 섹션으로 나뉩니다. 각 폴더에는 사용 예제가 포함된 README.md가 있습니다.
  • CI 배지는 CPU, NVIDIA(H100/B200 포함), AMD ROCm 및 RL 워크로드에서 활성 테스트를 보여줍니다.
  • 기여 – 가이드라인은 experiments/ 폴더의 새로운 실험과 풀 리퀘스트를 통한 핵심 개선을 권장합니다.
  • 논문 및 인용 – ICLR-2025 논문에서 설계 선택에 대한 세부 사항을 설명합니다. torchtitan을 기반으로 한 연구를 게시할 때 인용하세요.

위 정보는 모두 리포지토리 README에서 직접 가져온 것입니다. 외부 가정은 추가되지 않았습니다.

관련

  • 프로젝트
  • 프로젝트
  • 프로젝트
  • 프로젝트
  • 프로젝트