Device Mesh and Tensor Sharding

여러 GPU에 텐서를 나눠 계산하려면 먼저 두 가지를 정해야 한다.

GPU들을 어떤 축으로 볼 것인가?
Tensor의 어떤 차원을 그 축에 나눌 것인가?

첫 번째가 device mesh이고, 두 번째가 tensor sharding이다.

Device mesh

Device mesh는 GPU/TPU들을 이름 붙은 축으로 배열한 것이다.

GPU 4개를 1D mesh로 보면:

mesh = { X: 4 }

X axis: GPU0 - GPU1 - GPU2 - GPU3

같은 GPU 4개를 2D mesh로 보면:

mesh = { X: 2, Y: 2 }

        Y0      Y1
X0      GPU0    GPU1
X1      GPU2    GPU3

여기서 X, Y는 tensor의 차원이 아니라 장치들을 설명하기 위한 축 이름이다.

Tensor sharding

Tensor sharding은 tensor의 차원을 mesh 축에 나눠 놓는 것이다.

예를 들어 전체 행렬이:

A[8 x 8]

이고, row 방향은 mesh X, column 방향은 mesh Y로 나눈다고 하자.

rows -> X
cols -> Y

그러면 GPU 4개는 A의 서로 다른 4개 block을 가진다.

GPU0: rows 0-3, cols 0-3
GPU1: rows 0-3, cols 4-7
GPU2: rows 4-7, cols 0-3
GPU3: rows 4-7, cols 4-7

각 GPU가 들고 있는 local shard의 크기는:

A[4 x 4]

이다.

Global shape과 local shape

분산 텐서를 볼 때는 두 shape을 구분해야 한다.

global shape = 전체 tensor의 논리적 모양
local shape  = 각 GPU가 실제로 들고 있는 조각의 모양

위 예시에서는:

global shape: A[8 x 8]
local shape:  A[4 x 4]

코드를 작성하거나 수식을 볼 때 A는 여전히 8 x 8 행렬처럼 보일 수 있다. 하지만 실제 메모리에서는 GPU마다 4 x 4 조각만 들고 있다.

이 구분이 중요하다. 다음 행렬곱에서 필요한 조각이 내 GPU에 이미 있으면 통신 없이 계산할 수 있지만, 필요한 조각이 다른 GPU에 있으면 why-collective-operations-are-needed에서 본 것처럼 collective operation이 필요해진다.

확인

device mesh는 tensor의 모양인가, GPU들의 논리적 배치인가?
A[8 x 8]을 2x2 mesh에 row/column으로 나누면 각 GPU의 local shape은 무엇인가?
global shape과 local shape을 구분해야 하는 이유는 무엇인가?