Broadcast

Broadcast는 하나의 source rank가 가진 tensor를 모든 rank에 복사하는 collective operation이다.

Before와 after

예를 들어 rank 0만 값이 있고, rank 1과 rank 2는 빈 tensor를 가지고 있다고 하자.

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

rank 0을 source로 broadcast하면 모든 rank가 같은 값을 갖게 된다.

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

핵심은 source rank의 값이 사라지거나 나뉘는 것이 아니라, 같은 값이 다른 rank들에 복사된다는 점이다.

행렬곱을 여러 GPU에 나눠 계산할 때, 어떤 입력은 여러 rank가 똑같이 필요할 수 있다.

예를 들어:

A[2 x 3] x B[3 x 4] = C[2 x 4]

여기서 B를 column 방향으로 두 조각으로 나눠 rank 0과 rank 1이 각각 계산한다고 하자.

rank 0: B의 왼쪽 columns -> C의 왼쪽 columns 계산
rank 1: B의 오른쪽 columns -> C의 오른쪽 columns 계산

두 rank 모두 자기 B 조각과 곱할 같은 입력 A가 필요하다. 만약 A가 rank 0에만 있다면, rank 1은 계산을 시작할 수 없다.

이때 rank 0의 A를 다른 rank로 broadcast하면 각 rank가 자기 weight shard와 곱할 수 있다.

Broadcast는 한 rank가 가진 값을 다른 모든 rank가 똑같이 필요할 때 쓴다.

one rank has it -> every rank has it

이 카드는 broadcast의 의미만 다룬다. 실제 대규모 학습에서는 데이터가 처음부터 이미 각 rank에 준비되어 있거나, broadcast 대신 다른 collective와 조합되는 경우도 많다.