분산 LLM 학습 03 - All-Reduce, Ring, 그리고 통신 비용 읽는 법
분산 학습에서 가장 자주 등장하는 collective인 all-reduce를 이해해야 gradient synchronization 비용을 제대로 읽을 수 있다
왜 all-reduce를 먼저 봐야 할까?
distributed training에서 가장 자주 듣는 말 중 하나가 all-reduce다. 이유는 간단하다. data parallel에서는 backward로 계산한 gradient를 모든 rank가 같은 값으로 맞춰야 하기 때문이다.
all-reduce는 두 가지 일을 합친 연산으로 볼 수 있다.
- reduce: 각 rank의 값을 합치거나 평균낸다
- all-gather: 그 결과를 모든 rank에 다시 배포한다
즉, 각 rank가 전체 결과를 갖게 만드는 collective다.
ring all-reduce가 많이 쓰이는 이유
직관적으로는 "중앙 서버 하나에 다 보내고 다시 받으면 되지 않나?"라고 생각할 수 있지만, 그런 구조는 쉽게 병목이 된다. 한 지점에 트래픽이 몰리기 때문이다.
ring all-reduce는 rank들을 고리처럼 연결해서 데이터를 여러 조각으로 나누고, 각 단계마다 이웃 rank와만 통신한다. 이 방식은 모든 링크를 비교적 고르게 사용하므로 대역폭을 잘 활용한다.
핵심 직관은 이렇다.
- 메시지를 chunk로 나눈다
- reduce-scatter 단계에서 부분 합을 만든다
- all-gather 단계에서 결과를 다시 모은다
결국 전체 gradient를 한 번에 움직이지 않고, 조각내서 파이프라인처럼 흘린다.
실제 비용은 무엇으로 결정될까?
통신 시간은 대충 다음 두 항의 합으로 생각하면 된다.
- latency cost: 메시지를 시작하고 단계별 핸드셰이크를 하는 비용
- bandwidth cost: 실제 바이트를 옮기는 비용
작은 tensor가 많으면 latency가 아프고, 큰 tensor 하나가 크면 bandwidth가 아프다. 그래서 gradient bucket을 어떻게 묶느냐가 성능에 큰 영향을 준다.
또한 단순히 "GPU가 8장이다"만 봐서는 부족하다.
- 같은 node 안의 NVLink인가
- PCIe만 쓰는가
- node 사이 Ethernet/InfiniBand인가
이 토폴로지 차이가 통신 비용을 크게 바꾼다.
좋은 관찰 포인트
프로파일링할 때는 다음을 본다.
- backward kernel이 끝난 뒤 NCCL kernel이 길게 이어지는가
- rank마다 step 시간이 고르게 나오는가
- inter-node 구간에서 시간이 갑자기 튀는가
- 작은 all-reduce가 너무 자주 발생하는가
특히 "GPU utilization은 높은데 step time은 잘 안 줄어든다"면 통신이 숨어 있는지 먼저 의심해야 한다.
왜 이 글이 이후 주제의 기반이 되는가
DDP의 bucket, overlap, ZeRO의 sharding, tensor parallel의 collective, pipeline parallel의 stage 간 전달까지 결국은 모두 통신 패턴 문제다. all-reduce를 제대로 이해하면 이후 기법이 왜 그 구조를 택하는지 훨씬 명확해진다.
다음 글에서는 PyTorch DDP가 실제로 gradient synchronization을 어떻게 스케줄링하는지 본다. 이 시점부터 분산 학습은 프레임워크 내부 동작과 성능 튜닝이 직접 연결되기 시작한다.