Gather

Gather는 여러 rank가 가진 shard를 하나의 destination rank에 모으는 collective operation이다.

Before Gather

rank 0

12

rank 1

34

rank 2

56

->

Gather

After Gather

rank 0

123456

destination

rank 1

34

rank 2

56

Before와 after

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

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

여기서는 rank 0이 destination이다. 핵심은 전체 결과를 모든 rank가 갖는 것이 아니라, 하나의 rank만 모은 결과를 갖는다는 점이다.

행렬곱에서 왜 필요한가

여러 rank가 C의 row shard를 나눠 계산했다고 하자.

rank 0: C rows 0-1
rank 1: C rows 2-3
rank 2: C rows 4-5

이 결과를 하나의 rank에서 저장하거나 평가해야 한다면 shard들을 한 곳에 모아야 한다.

Gather output shards after local matmul

rank 0

C rows 0-1+C rows 2-3+C rows 4-5

result

full C[6 x 2]

언제 쓰는가

Gather는 분산된 결과 조각을 하나의 rank에 모아야 할 때 쓴다.

different shard per rank -> one rank has full tensor

확인

• Gather 이후 full tensor를 갖는 rank는 하나인가, 모든 rank인가?
• Gather와 scatter는 어떤 관계인가?
• 여러 rank가 계산한 C 조각을 하나의 파일로 저장하려면 왜 gather가 필요할 수 있는가?