data parallel만으로는 해결되지 않는 순간

모델이 너무 커서 한 GPU 메모리에 안정적으로 올라가지 않거나, parameter replication이 너무 비싸면 data parallel만으로는 더 이상 앞으로 가기 어렵다. 이때 등장하는 것이 tensor parallel이다.

핵심 아이디어는 단순하다. 모델 레이어 내부의 큰 연산, 특히 matmul을 여러 GPU에 나눠 계산하는 것이다.

예를 들어 큰 linear layer가 있을 때:

  • output dimension을 나눌 수도 있고
  • input dimension을 나눌 수도 있다

어떤 축을 나누느냐에 따라 필요한 통신 형태가 달라진다.

왜 transformer는 tensor parallel에 잘 맞을까?

transformer는 큰 dense matmul이 반복적으로 등장한다. attention projection과 MLP projection이 대표적이다. 이런 연산은 분할했을 때 계산량이 충분히 크고 패턴이 규칙적이라 tensor parallel의 이점을 얻기 좋다.

반대로 연산이 너무 작거나 irregular하면 통신 비용이 계산 이득을 잡아먹을 수 있다.

column parallel과 row parallel의 감각

큰 linear 연산을 생각해보면 두 가지 전형적인 분할이 나온다.

column parallel

출력 feature 방향으로 weight를 나누면 각 GPU가 출력의 일부를 계산한다. 이후 결과를 모으는 방식이 필요하다.

row parallel

입력 feature 방향으로 weight를 나누면 각 GPU가 부분 합을 계산하고 reduce가 필요해진다.

중요한 것은 어느 쪽이든 "matmul 하나를 쪼개는 대신 통신이 들어온다"는 점이다.

tensor parallel을 볼 때의 판단 기준

  • 연산이 충분히 큰가
  • 통신이 compute에 비해 너무 비싸지 않은가
  • partition한 뒤 activation 흐름이 자연스러운가
  • 다음 레이어와 연결할 때 추가 collective가 얼마나 필요한가

이 판단이 잘 안 되면, tensor parallel은 오히려 구조를 복잡하게만 만들 수 있다.

이 전략의 비용

tensor parallel은 메모리를 줄이는 데 도움을 주지만, 레이어 내부 통신이 늘어난다. data parallel이 step 경계에서 gradient를 동기화한다면, tensor parallel은 forward/backward 안쪽으로 통신이 들어온다.

즉 더 낮은 latency와 좋은 intra-node 연결이 중요해진다.

다음 글에서는 transformer block 안에서 tensor parallel이 구체적으로 어떻게 적용되는지, attention과 MLP를 기준으로 더 세밀하게 본다.