AllGather

distributedcollectiveall-gather

AllGather는 여러 rank가 가진 shard를 모은 뒤, 모든 rank가 full tensor를 갖게 하는 collective operation이다.

Before와 after

Before all-gather
rank 0: [1, 2]
rank 1: [3, 4]
rank 2: [5, 6]

After all-gather
rank 0: [1, 2, 3, 4, 5, 6]
rank 1: [1, 2, 3, 4, 5, 6]
rank 2: [1, 2, 3, 4, 5, 6]

Gather는 destination rank 하나만 full tensor를 갖는다. AllGather는 모든 rank가 full tensor를 갖는다.

행렬곱에서 왜 필요한가

각 rank가 C의 column shard를 계산했다고 하자.

rank 0: C left
rank 1: C middle
rank 2: C right

다음 계산에서 모든 rank가 full C를 입력으로 필요로 한다면, 각자 가진 C shard만으로는 부족하다. 이때 AllGather로 모든 rank가 full C를 갖게 만들 수 있다.

언제 쓰는가

AllGather는 shard된 tensor를 모든 rank에서 full tensor로 복원해야 할 때 쓴다.

different shard per rank -> every rank has full tensor

확인

Gather와 AllGather의 차이는 무엇인가?
AllGather 이후 rank 1은 rank 0의 shard를 갖는가?
다음 layer가 full C를 필요로 한다면 왜 AllGather가 필요할 수 있는가?