Torchtitan — PyTorch native distributed training

2020년대 초까지 분산 학습은 보통 Megatron-LM 이나 DeepSpeed 같은 framework 의 일이었다. PyTorch 위에 layered 인 별도 stack 이 직접 model 을 받아 자기 식으로 분산 — 그러면 사용자가 그 framework 의 model code 를 따라가야 했다. 분산이 어렵다는 인식의 대부분이 거기서 나왔다.

Torchtitan 의 출발점은 단순하다 — “PyTorch 가 직접 분산을 한다.” 사용자는 자기 model 을 PyTorch 로 그대로 짜고, distributed 변환은 한 함수 호출 한 줄로 모델 위에 얹는다. FSDP, TP, SP, PP 가 모두 이 형태로 들어간다.

그래서 강의 끝에 손에 잡혀야 할 자산 — (1) 5 가지 분산 차원 (DP / FSDP / TP / SP / PP) 이 어떻게 한 stack 위에서 결합되는가, (2) DTensor 가 “한 tensor 가 여러 GPU 에 나뉘어 있다” 를 어떻게 표현하는가, (3) 통신 overlap 이 forward / backward 양쪽에서 어떻게 일어나는가, (4) checkpoint 가 분산 환경에서 어떻게 sharded 로 떨어지는가.

강의에서 Tianyu 가 짚은 흐름 — “DDP 는 ‘모델이 한 GPU 에 들어간다’ 의 가정. FSDP 는 ‘들어가지 않는다’ 의 답. 그런데 FSDP 만으로는 큰 모델의 단일 layer 도 한 GPU 에 못 들어가는 경우가 있어 — 그래서 layer 안 weight 도 자르는 TP 가 필요. 이게 다 같은 추상으로 통일된 게 DTensor.”

강의에서 흥미로운 결정 — Torchtitan 은 PP 를 명시적으로 “이 강의에서는 안 다룬다” 고 못 박는다. 이유 — pipeline parallel 은 model 코드에 if/else 분기가 들어가는 경향이 있어 “단순한 model 코드 + 한 줄 분산” 의 철학이 깨진다. 다른 4 개는 model 코드 손 안 대고 wrap 한 줄로 적용 가능.

DTensor 의 핵심 — 하나의 tensor 가 (1) device mesh (이 GPU들에) + (2) placement (이렇게 자르거나 복제) 라는 두 metadata 를 갖는다. operator 가 이 metadata 를 보고 자동으로 통신을 끼워넣는다.

# DTensor — 한 tensor 가 mesh 위에서 sharded
from torch.distributed.tensor import (
    DeviceMesh, Shard, Replicate, distribute_tensor)

# 8 GPU 를 2x4 mesh 로
mesh = DeviceMesh("cuda", [[0,1,2,3],
                            [4,5,6,7]],
                  mesh_dim_names=("dp", "tp"))

# 같은 weight 를 dp 차원에는 복제, tp 차원으로는 분할
W = distribute_tensor(big_weight, mesh,
                      placements=[Replicate(), Shard(0)])

# 일반 tensor 처럼 쓴다 — operator 가 알아서 통신
y = x @ W                        # 자동으로 all-gather + matmul

분산 학습의 잘 안 보이는 어려운 문제 — checkpoint. 1000 GPU 가 sharded 로 weight 를 나눠 가지고 있을 때, 이걸 disk 에 어떻게 적고, 다시 읽을 때 다른 분산 형태(예: 다른 GPU 수, 다른 mesh)로 어떻게 reshape 하는가.

PyTorch 의 답 — DCP (Distributed Checkpoint). 두 개의 design 결정.

• 각 rank 가 자기 자리만 쓴다 — gather 없이. 그래서 checkpoint 가 여러 파일로 떨어진다 — rank 별 한 파일. 디스크 IO 가 분산.
• Logical layout 을 metadata 로 따로 적는다 — 어떤 tensor 가 어떻게 sharded 되어 있는지 metadata 가 별도 파일. 다시 읽을 때 다른 분산 형태로 reshape 가능.

큰 모델의 학습은 메모리 부족이 거의 항상. activation 이 forward 동안 쌓여서 backward 에 쓰이는데, 이걸 다 보존하면 메모리 limit 을 넘긴다. 표준 답 — activation checkpointing (re-compute).

Torchtitan 이 추가한 것 — selective AC. 모든 layer 를 다 re-compute 하면 forward time 의 30% 정도 추가. 하지만 layer 마다 “저장 비용 / re-compute 비용” trade-off 가 다르다 — softmax 처럼 cheap 한 op 는 저장 안 하고, MLP 의 GEMM 결과처럼 expensive 한 op 는 저장.

강의에서 Mark 가 짚은 — “selective AC 의 정책 list 는 hard-coded 가 아니라 offline 분석 결과. 일반적으로 잘 작동.” 즉 어떤 op 을 저장하고 어떤 op 을 re-compute 할지 미리 결정된 표가 있고, 사용자가 평소엔 안 건드린다.

분산 학습의 두 번째 큰 비용 — collective communication (all-reduce, all-gather). NVLink/IB 가 빨라도 모델이 클수록 통신 시간이 꽉 찬다. 답은 계산과 통신의 overlap.

강의에서 Tianyu 가 짚은 사실 — async TP 는 PyTorch 안에 일급 기능으로 들어와 있다. 사용자가 model 코드를 안 바꿔도 parallelize_module(...) 호출에서 async 옵션만 켜면 자동으로 분할 + overlap.

정확한 입장 — Megatron 이 사라지지 않는다. 매우 큰 모델의 깊은 최적화는 Megatron 이 여전히 유리한 영역. Torchtitan 은 “회사가 자기 model 코드를 그대로 쓰면서 분산을 하고 싶을 때”의 선택. 두 stack 이 다른 자리.

강의 시점 기준의 채택 풍경. NDA 로 자세한 회사 이름은 막혀 있지만 패턴은 분명. Llama 류 dense transformer 의 학습이 가장 자연스러운 fit.

• Meta 내부 — Llama 차세대 모델 학습. Mark 가 명시적으로 “Meta 안에서 production train” 이라고 짚음.
• 외부 회사 / 연구실 — startup 들이 Llama-3 변종 fine-tune / pre-train 에 채택. ~50개 회사 정도 채용 중이라는 강의 시점 언급.
• research — popcorn (LLM-for-kernels) 같은 GPU Mode 자체 프로젝트도 Torchtitan 위에서 학습할 가능성. Mark 가 “자체 학습이 필요해진 시점에 우리도 distributed 가 필요하다” 라고 짚음.

손에 새기기 — 실습 시퀀스

1. repo clone + 1 GPU 로 toy run — README 의 train.sh 한 줄. Llama 2 7B 의 초기 학습 step 이 single GPU 에서 도는지.
2. 2 GPU 로 FSDP — 같은 명령에 --training.data_parallel_replicate_degree 1 --training.data_parallel_shard_degree 2. 메모리 사용량이 절반 가까이 줄어드는지.
3. 4 GPU 로 FSDP×TP — TP=2, FSDP=2 의 2D mesh. matmul 안의 weight 가 어떻게 sharded 되는지 DTensor metadata 직접 print.
4. DTensor 직접 만들기 — 작은 tensor 를 distribute_tensor 로 분산하고 redistribute 로 placement 변경. 어떤 collective 가 호출됐는지 NCCL trace 로 확인.
5. checkpoint 떠보기 — 학습 중간에 DCP 로 ckpt. 그 ckpt 를 다른 GPU 수 (예: TP=4, FSDP=2) 로 다시 read. 학습이 이어지는지.
6. activation checkpointing 모드 비교 — full / selective / off 세 모드의 메모리 사용량 + step 시간. selective 가 sweet spot 인지 직접 확인.
7. async TP overlap 측정 — 같은 모델, async TP on/off. Chrome trace (L001) 위에서 collective 와 GEMM 이 실제로 겹치는지.

• MoE 의 본격 지원 — sparse expert 의 분산은 expert parallel 이라는 별도 차원. Torchtitan 이 이걸 어떻게 흡수할지 진행형. expert 별 routing 이 통신 패턴을 비대칭으로 만든다.
• multimodal 학습 — vision encoder + text decoder 의 다른 분산 형태. encoder 는 보통 작아서 DDP, decoder 는 FSDP+TP — 한 모델 안에서 다른 분산 정책.
• pipeline parallel 의 단순화 — 강의 시점에 PP 는 model 코드 분기를 강요하는 stage. PyTorch 가 이걸 단순화할 수 있는가.
• activation checkpointing 의 자동화 — selective AC 의 정책이 hard-coded list. 이 결정을 자동으로 (op cost 분석으로) 만드는 방향이 다음 단계.
• cluster-level resharding — DCP 의 reshape 가 단일 job 안에서는 매끄럽지만, 다른 cluster 의 다른 storage 로 옮길 때의 일관성. 그 운영 시나리오.