DP All-Reduce and Overlap

naive DP는 backward가 모두 끝난 뒤 gradient all-reduce를 시작한다.

backward all layers -> all-reduce all gradients -> optimizer update

이 방식은 단순하지만 느리다. all-reduce 동안 GPU가 계산을 하지 못하고 기다리기 때문이다.

Naive DP

backward L4backward L3backward L2backward L1 all-reduce all gradients

Overlap + bucket

compute

backward L4backward L3backward L2backward L1

network

bucket A bucket B bucket C

Backward는 뒤 layer부터 끝난다

모델이 layer 1, 2, 3, 4로 되어 있다고 하자.

forward는 앞에서 뒤로 간다.

layer 1 -> layer 2 -> layer 3 -> layer 4

backward는 반대로 간다.

layer 4 -> layer 3 -> layer 2 -> layer 1

그러면 layer 4의 gradient는 전체 backward가 끝나기 전에 이미 준비된다. 이 gradient를 기다리게 둘 필요가 없다.

layer 4 gradient ready -> layer 4 gradient all-reduce 시작
layer 3 backward 진행 중 -> layer 4 통신도 뒤에서 진행 중

이렇게 계산과 통신을 겹치는 것을 communication-computation overlap이라고 한다.

Bucketing

gradient tensor 하나가 생길 때마다 all-reduce를 호출하면 통신 호출이 너무 많다.

그래서 여러 gradient를 bucket에 모아두고, bucket 단위로 all-reduce한다.

small gradients -> bucket -> one all-reduce

bucket이 너무 작으면 호출 횟수가 많아진다. bucket이 너무 크면 bucket이 찰 때까지 기다려야 하므로 overlap 기회가 줄어든다.

Gradient accumulation과 같이 쓸 때

gradient accumulation 중간 step마다 all-reduce를 하면 낭비다. optimizer update는 마지막에 한 번만 하기 때문이다.

accum step 1: backward, no sync
accum step 2: backward, no sync
accum step 3: backward, all-reduce
optimizer update

PyTorch DDP에서는 보통 no_sync()로 중간 accumulation step의 gradient synchronization을 끈다.

확인

• naive DP는 언제 all-reduce를 시작하는가?
• overlap 방식은 왜 layer 4 gradient부터 먼저 통신할 수 있는가?
• gradient accumulation 중간 step마다 all-reduce하면 왜 낭비인가?