분산 LLM 학습 04 - PyTorch DDP는 내부에서 무엇을 하는가
DDP는 단순 래퍼가 아니라 autograd hook, gradient bucket, process group을 사용해 동기화를 조직하는 런타임이다
DDP를 그냥 "감싸는 것"으로 보면 안 된다
PyTorch DistributedDataParallel은 겉으로 보면 모델을 한 번 감싸는 API처럼 보인다. 하지만 실제로는 backward 흐름 중간에 개입해서 gradient synchronization 시점을 제어하는 런타임에 가깝다.
이 내부를 이해해야 하는 이유는 분명하다.
- 왜 어떤 모델은 overlap이 잘 되고 어떤 모델은 안 되는지
- 왜 unused parameter 설정이 성능과 correctness에 영향을 주는지
- 왜 gradient bucket 크기가 중요해지는지
이 질문들이 전부 DDP 내부 구조와 이어지기 때문이다.
DDP의 핵심 구성
DDP는 크게 다음 요소에 의존한다.
- process group: 어떤 rank들이 collective를 수행할지 정의
- parameter ordering: 파라미터를 bucket으로 묶는 기준
- autograd hook: gradient가 준비되는 시점을 감지
- reducer: 준비된 gradient bucket에 대해 all-reduce를 실행
중요한 점은 backward가 모두 끝난 뒤 한꺼번에 통신하는 것이 아니라, 어떤 bucket이 준비되면 그 bucket부터 all-reduce를 시작할 수 있다는 것이다. 이게 바로 overlap의 출발점이다.
bucket이 왜 필요한가?
gradient를 파라미터 하나씩 동기화하면 collective 호출이 너무 많아진다. 반대로 전체 gradient를 한 번에 모으면 backward가 다 끝날 때까지 아무 통신도 못 한다. bucket은 그 중간 지점이다.
즉 bucket은:
- 호출 수를 줄이고
- 메시지 크기를 어느 정도 확보하면서
- backward와 통신을 겹칠 수 있게 해준다
여기서 bucket 크기가 너무 작으면 latency cost가 늘고, 너무 크면 overlap이 늦어진다.
실무에서 자주 만나는 함정
1. parameter 사용 패턴이 일정하지 않은 모델
조건문이나 branch가 많은 모델은 어떤 step에서 일부 파라미터가 사용되지 않을 수 있다. 이 경우 DDP는 gradient 준비 상태를 추적하는 방식 때문에 추가 비용이나 오류 가능성이 생긴다.
2. no_sync와 gradient accumulation
큰 global batch를 만들기 위해 accumulation을 쓸 때는 매 micro-step마다 all-reduce를 할 필요가 없을 수 있다. 이때 no_sync 같은 패턴을 쓰지만, 잘못 적용하면 메모리 사용량과 step semantics를 동시에 망칠 수 있다.
3. communication overlap이 항상 자동으로 좋아지는 것은 아니다
모델 구조, bucket ordering, kernel 길이, 네트워크 상태에 따라 overlap 효율은 크게 달라진다.
DDP를 볼 때의 관점
DDP는 "gradient synchronization을 언제 시작할지"를 관리하는 시스템으로 보는 편이 좋다. 이렇게 보면 성능 문제를 볼 때도 질문이 더 정확해진다.
- backward 그래프에서 어느 시점에 bucket이 ready 되는가
- all-reduce가 compute와 얼마나 겹치는가
- 특정 bucket이 straggler가 되는가
다음 글에서는 global batch size, gradient accumulation, scaling rule을 함께 묶어서 본다. 여러 GPU를 붙였을 때 왜 학습률과 optimizer 감각까지 다시 잡아야 하는지가 그때 분명해진다.