How Deep Learning Computes

DDP vs FSDP Lab

마지막 수정:

pytorchdistributedddpfsdpprofiling

DDP와 FSDP는 둘 다 data parallel 축 위에서 작동하지만, 무엇을 복제하고 무엇을 shard하는지가 다르다.

DDP:
  each rank has full params + full grads + full optimizer state
  backward triggers gradient all-reduce

FSDP:
  ranks keep parameter shards
  layer compute triggers parameter all-gather
  backward/update triggers reduce-scatter style communication

실습은 같은 tiny Transformer를 두 wrapper로만 바꿔 실행한다.

cd labs/pytorch-transformer
torchrun --nproc_per_node=2 distributed_train.py --steps 20 --device cuda
torchrun --nproc_per_node=2 fsdp_train.py --steps 20 --device cuda

비교할 값은 세 가지다.

tokens/sec
peak memory per GPU
profiler에서 보이는 communication op

작은 모델에서는 FSDP가 더 느릴 수 있다. 이것은 실패가 아니다. FSDP는 parameter memory를 줄이기 위해 더 많은 all-gather/reduce-scatter를 넣는다. 모델이 작거나 interconnect가 느리면 communication overhead가 memory 이득보다 커질 수 있다.

이 실습의 목표는 “FSDP가 항상 빠르다”를 확인하는 것이 아니라, DDP와 FSDP가 서로 다른 병목을 만든다는 것을 보는 것이다.

확인

  • DDP에서 가장 먼저 profiler로 찾아야 할 collective는 무엇인가?
  • FSDP에서 DDP보다 추가로 생기는 communication은 무엇인가?
  • 작은 모델에서 FSDP가 DDP보다 느릴 수 있는 이유는 무엇인가?
먼저 보면 좋은 카드PyTorch DDP Training Loop단일 GPU training loop를 DistributedDataParallel로 감싸고, sampler, gradient sync, no_sync를 연결한다.먼저 보면 좋은 카드PyTorch FSDP와 ZeRO 기본기DDP가 복제하던 parameter, gradient, optimizer state를 FSDP/ZeRO가 어떻게 shard하는지 PyTorch 관점에서 이해한다.먼저 보면 좋은 카드PyTorch Distributed Profilermulti-GPU 학습에서 step time, GPU memory, all-reduce/FSDP communication, overlap을 측정한다.같은 경로의 다음 카드nanochat Custom Distributed Optimizernanochat이 PyTorch DDP wrapper 대신 custom distributed optimizer로 multi-GPU data parallel training을 구성하는 이유를 이해한다.