How Deep Learning Computes

PyTorch Distributed Training

이 경로는 PyTorch 단일 GPU 학습 루프를 multi-GPU 학습으로 확장하는 실습 경로입니다.

핵심 질문은 다음입니다.

같은 training step을 여러 GPU로 나눌 때
무엇을 복제하고, 무엇을 나누고, 무엇을 통신하는가?

먼저 process group과 rank 개념을 잡고, torchrun으로 여러 process를 띄웁니다. 그다음 DDP에서 model replica들이 서로 다른 micro-batch를 처리하고 gradient all-reduce로 동기화하는 방식을 구현합니다.

이후 모델/optimizer/gradient/activation memory가 커질 때 FSDP와 ZeRO 관점으로 넘어가고, DDP/FSDP를 같은 tiny Transformer에서 비교합니다. nanochat의 custom distributed optimizer는 “DDP wrapper만이 multi-GPU 학습의 전부는 아니다”라는 관점에서 읽습니다. 마지막에는 통신과 계산이 실제로 겹치는지 profiler로 확인합니다.

이 path의 산출물은 단순히 “여러 GPU에서 돌아간다”가 아니라, single-GPU baseline 대비 throughput, memory, communication overhead를 설명하는 것입니다.

  1. PyTorch Single-GPU Training Loop — forward, loss, backward, optimizer update, mixed precision을 포함한 단일 GPU 학습 루프를 만든다.
  2. Rank와 Process Group — collective operation을 부르기 전에 rank, world size, process group, root rank가 무엇인지 이해한다.
  3. torchrun과 Process Group 시작하기 — PyTorch distributed 학습에서 여러 process를 띄우고 rank, world size, device를 초기화하는 흐름을 익힌다.
  4. Global Batch, Micro-Batch, Gradient Accumulation — global batch size를 mbs, grad_acc, dp의 곱으로 분해해 본다.
  5. Data Parallelism — 같은 모델 복사본 여러 개가 서로 다른 micro-batch를 처리하고 gradient를 맞추는 방식을 이해한다.
  6. AllReduce — 여러 rank의 값을 합산한 뒤 모든 rank가 같은 합산 결과를 갖게 하는 collective operation.
  7. PyTorch DDP Training Loop — 단일 GPU training loop를 DistributedDataParallel로 감싸고, sampler, gradient sync, no_sync를 연결한다.
  8. DP All-Reduce and Overlap — DP에서 gradient all-reduce를 backward 계산과 겹쳐 통신 대기 시간을 줄이는 방법을 이해한다.
  9. Parallelism vs Sharding — 병렬화는 일을 나누는 전략이고, 샤딩은 큰 텐서를 조각내 저장하는 방법임을 구분한다.
  10. ZeRO — DP에서 중복되는 optimizer state, gradient, parameter를 단계적으로 shard해 메모리를 줄이는 방법을 이해한다.
  11. PyTorch FSDP와 ZeRO 기본기 — DDP가 복제하던 parameter, gradient, optimizer state를 FSDP/ZeRO가 어떻게 shard하는지 PyTorch 관점에서 이해한다.
  12. DDP vs FSDP Lab — 같은 Transformer 학습 루프를 DDP와 FSDP로 실행해 throughput, memory, communication pattern을 비교한다.
  13. nanochat Custom Distributed Optimizer — nanochat이 PyTorch DDP wrapper 대신 custom distributed optimizer로 multi-GPU data parallel training을 구성하는 이유를 이해한다.
  14. Activation Checkpointing — 일부 activation만 저장하고 나머지는 backward 때 재계산해서 메모리를 줄인다.
  15. GPU Interconnect Topology — 여러 GPU가 어떤 링크로 연결되어 있는지에 따라 같은 collective도 비용이 달라진다는 점을 이해한다.
  16. Scaling Efficiency and Topology — GPU 수를 늘릴 때 compute, communication, topology가 scaling efficiency를 어떻게 제한하는지 이해한다.
  17. PyTorch Distributed Profiler — multi-GPU 학습에서 step time, GPU memory, all-reduce/FSDP communication, overlap을 측정한다.