cs231n 2025 ver. - Lecture 11

    이번 강의는 2017년 버전의 강의에는 없던 내용이다. GPT, Gemini 등의 LLM이 등장하며 많은 GPU를 사용한 훈련이 필요해지게 되었다. 그 많은 GPU를 훈련에 사용하기 위해 등장한 여러 가지 방법들이 이번 강의에서 다룰 'Large-Scale Distributed Training'이다. 나는 하드웨어 쪽으로는 지식이 부족해서 이해하는 데 조금 어려움이 있었다.

    기존의 LLM들은 자신들의 모델을 어떻게 훈련시켰는 지에 대한 스펙을 공개해 왔으나, GPT 4 이후 그러한 모델 스펙들을 공개하지 않는 트렌드가 시작되었다. 그래서 이번 강의에서는 비교적 최근에 나오고, 스펙이 공개되어 있는 Llama3-405B 모델을 이용해 여러 개념들을 설명한다. (Llama3-405B는 2024년에 나온 Meta의 논문이다.)

1. GPUs

    GPU는 인공지능을 훈련시키는데 필수적인 하드웨어이다. 여기에 최근 들어 LLM이 등장하며 덩치가 큰 모델들을 훈련시키기 위해 엄청난 양의 GPU를 사용하게 되었다. 이번 강의에서는 크게 두 가지 주제에 대해 다루게 될 것이다.

1. GPU 하드웨어
2. 많은 GPU를 가지고 모델을 훈련시키는 알고리즘

먼저 GPU 하드웨어부터 알아보자.

1) GPU Hardware

    GPU = Graphics Processing Unit 이다. 원래는 그래픽을 위해서 디자인 된 하드웨어지만, 현재는 general parallel processor로 사용된다. (데이터를 병렬적으로 처리하기 위한 processor라는 뜻이다.)

NVIDIA의 H100 GPU를 살펴보자. 중간에 보면 Compute core가 있고, 그 위 아래로 HBM Memory가 있는 것을 볼 수 있다. 메모리는 compute core와 분리되어 있는데, 이 둘은 서로 소통하며 데이터를 Memory <-> core로 옮긴다. 옮길 때의 속도는 위의 사진에서도 볼 수 있듯이 3352GB/sec (약 3TB)이다.

Compute Cores

    Compute Core는 아래의 사진과 같이 생겼다.

안에는 50MB의 L2 Cache(나도 잘은 모르지만 메모리의 종류라고 한다.)가 있다. L2 cache는 computing core에 가까워서 더 빠른 access가 가능하다고 한다. 그 옆에는 132개의 Streaming Multiprocessor(SM)들이 있는데, 이들은 독립적인 parallel core이다.

Streaming Multiprocessors (SM)

    SM은 GPU의 memory hierarchy를 따른다. 구체적인 구조는 아래 사진과 같다.

자세히 보면 더 큰 bit의 메모리일 수록 compute core에서 멀리 떨어져 있는 것을 볼 수 있다. bit 수는 크지만 속도는 느리다. 반대로, compute core에 가까워질 수록 메모리의 bit수는 작아지지만, 속도는 빨라진다. H100에는 3개의 hierarchy가 있는데, 차례대로 256KB L1 cache, 50MB L2 cache, 80GB HBM이다.

    여기에 추가로 128개의 FP32 core와 4개의 tensor core들이 있다. FP32 core는 generalized floating point operation을 진행할 수 있는 arithmetic unit이라고 한다. 예를 들면, ax + b와 같은 연산을 FP32 core에서 담당한다는 뜻이다. 덧셈, 나눗셈, 뺄셈, 곱셈 한 번이 1개의 FLoating point Operation, 즉 FLOP이 되는 것이다. ax+b 는 2FLOPs이니 , 하나의 SM에 들어 있는 128개의 FP32 core들은 256 FLOP/cycle per SM을 연산할 수 있다.

    Tensor core는 1개의 SM에 4개씩 들어있는데, 말 그대로 Tensor간의 연산에 특화된 core이다. AX+B 연산을 진행하는데, 위의 사진에서의 예시와 같은 matmul 연산을 진행하게 되는 것이다. 4096 FLOP/cycle per SM의 스펙을 가지고 있고, mixed precision도 사용된다.

GPU의 성능이 최근 들어 굉장히 빨라졌고, 한 개 이상의 GPU로도 훈련이 가능하다는 점이 굉장히 인상적이다. 10년 간 약 1000배 정도 속도가 빨라졌다. GPU 하드웨어에 대한 부분을 처음 배우다 보니 정리한 내용을 완전히 이해하지는 못하는 것 같다. 추가적으로 이후에 공부를 해봐야 할 부분인 것 같다.
이제 하나의 GPU에서 벗어나 여러 개의 GPU에 대해 알아보자.

2. Zooming out of a Single GPU

    하나의 GPU를 다루는 상황에서는 compute element에서 멀어질수록 memory bandwidth는 느려진다.

• Memory Bandwidth : 시스템의 여러 부분들로 bit를 옮길 수 있는 능력

이러한 경향성은 GPU를 여러 개 사용할 때에도 비슷하게 나타난다. GPU 여러 개를 사용하는 상황을 또 하나의 memory hierarchy로 볼 수도 있다. 하나의 GPU를 넘어선 다음 단위들에 대해 알아보자.

Cluster > Pod > Rack > Server > Single GPU 의 단위로 커져간다. 각 단위가 몇 개의 GPU를 포함하고 있는 지는 위의 사진을 보면 알 수 있다. 참고로 위의 사진은 Llama3에서 사용된 GPU 개수이다. 보통 요새 나오는 모델들은 이러한 세팅으로 몇 달, 만약 더 길어지면 1년까지도 훈련을 시킨다고 한다. 이러한 Cluster에서 GPU를 효율적으로 사용해서 좋은 결과를 내는 것이 중요해지는 것이다.

Cluster

    하나의 cluster는 하나의 큰 컴퓨터로 봐도 된다.

하나의 cluster의 스펙은 위의 사진과 같다. (1PB = $10^15$ byte)

Others

    GPU 말고 다른 processor들로는 Google의 TPU, AMD와 AWS의 칩들이 있지만, TPU를 제외하고 나머지는 잘 사용되지 않는다.

3. How to train on lots of GPUs

    이제 여러 개의 GPU를 어떻게 훈련시켜야 할 지, 그 알고리즘에 대한 부분을 알아보자. 컴퓨터가 하는 일은 크게 두 가지이다.

1. Computation(계산) : input bits -> output bits
2. Communication(소통) : 다른 위치로 bit를 이동시키는 것.

이렇게 두 가지 일을 하는데, 우리가 가지고 있는 자원을 parallel한 방식으로 활용하도록 해야 하고, 동시에 그 자원들이 효율적으로 소통해 하나의 큰 Neural Network를 훈련시킬 수 있도록 해야 하는 것이다.

    여러 GPU를 가지고 훈련을 할 때 가장 중요한 것은 위에서도 언급했다시피 'Parallelization'이다. 그런데 어떤 차원에서 parallelization을 할 것이냐에 따라 DP, CP, PP, TP로 나뉜다.

만약 L개의 layer를 가지고 있는 모델이 있고, 이들이 각각 tensor (batch, sequence, dim) 에서 작동한다고 해보자. 그러면 DP, CP, PP, TP는 아래와 같은 기준으로 나뉘게 된다.

1. Data Parallelism (DP) : batch dimension에서 나누는 것
2. Context Parallelism (CP) : sequence dimension에서 나누는 것.
3. Pipeline Parallelism (PP) : L dimension에서 나누는 것.
4. Tensor Parallelism (TP) : dim dimension에서 나누는 것.

하나씩 살펴보자.

1) Data Parallelism (DP)

    이전에 우리가 loss를 처리하던 방식을 기억해보자. Loss는 주로 minibatch N개의 sample들에 대한 평균을 내서 계산한다. DP의 아이디어는 MN개의 minibatch를 가지고, M개의 GPU가 각각 N개의 sample을 갖게 하자는 것이다. 이렇게 할 수 있는 이유는, gradient가 linear하기 때문이다.

하나의 GPU 안에서는 이전에 하던 것과 동일하게 계산을 진행하고, M개의 GPU에서 나온 gradient를 평균 내는 것이다.

   

1. 먼저, 각 GPU는 optimizer, model의 복사본을 갖고 있어야 한다.
2. 그리고 나서 각 GPU에 N개의 sample이 주어지면, forward pass를 각자 진행한다. 이때, 각 GPU는 각자 다른 데이터를 가지고 있다.

3. 그리고 나서 각 GPU는 backward pass를 통해 gradient를 계산한다. 이때까지 GPU끼리의 소통은 전혀 없고, 모두 각자의 GPU 안에서 계산을 할 뿐이다.
4. 각자 backward pass를 진행하고 나서는 모든 GPU의 gradient 평균을 계산해야 하고, 이때 서로 소통을 하게 된다.

각 GPU는 서로 자신들의 gradient 값을 보내고, 다른 GPU들의 gradient 값들을 받는다. 이를 all-reduce operation이라고 하는데, 모든 GPU가 여러 GPU에서 계산된 값을 공유하는 것을 의미한다고 한다.
5. 각 GPU는 자신의 weight들을 업데이트한다. GPU끼리의 통신을 통해 average한 값으로 update하게 되는 것이다.

이때, 3,4번은 parallel하게 진행될 수 있다. 예를 들면, Layer (L+1)의 communication이 진행되는 동시에 Layer L의 backward pass를 진행하는 것이다. 이러한 방식을 사용하면 backward pass가 끝났을 때 all-reduce operation도 끝나 있기 때문에 시간을 절약할 수 있다.

Problem

    DP의 문제는 Model size가 GPU memory에 의해 제약을 받는다는 것이다. 각 GPU가 모두 모델의 복사본을 가져야 하기 때문이다. 뿐만 아니라 각 weight는 weight, grad, Adam optimizer의 beta1, beta2 4개의 숫자가 필요하고, 각 숫자는 2byte를 차지한다. 이런 식이면 1B param만 가지고 있어도 8GB의 용량을 가지게 된다. 10B param이라면 80GB를 차지하게 되는 것이다.

이는 모델의 weight들을 여러 개의 GPU로 분산하는 것으로 해결될 수 있다.

2) DP - Fully Sharded Data Parallelism (FSDP)

    기존 DP의 문제점을 해결하기 위해 등장한 것이 FSDP이다. 모델의 weight들을 여러 개의 GPU에 나눠서 넣어 놓는 방식이다.

각 weight $W_i$는 하나의 GPU가 가지고 있게 되고, 그 GPU가 weight의 grad, optim state도 갖고 있게 된다. 이때의 W는 하나의 layer의 모든 weight matrix를 통틀어서 주르는 표현이다. 만약 위의 예시처럼 4개의 layer가 있고 두 개의 GPU가 있는 상황이라면, $W_1, W_2$는 GPU 1, $W_3, W_4$는 GPU 2가 가지고 있게 된다. 이런 방식이라면 DP가 아닌 것 아니냐고 할 수 있지만, 이렇게 weight를 찢어 놓고 이전에 했던 것과 동일하게 두 GPU가 다른 data를 받게 되니까 여전히 DP라고 할 수 있다. 이제 실제로 어떻게 학습이 진행되는지 알아 보자.

1. Layer i의 forward pass 전에, 그 layer의 weight를 갖고 있는 GPU가 나머지 GPU들에게 weight 값을 broadcasting 해준다. (아래 사진 참조)

2. 모든 GPU에서 layer i에 대한 forward pass를 진행한다. 그리고 나서 broadcast 받은 weight $W_i$를 GPU에서 삭제한다. 이때, layeer i에 대한 forward pass를 진행하는 동안 layer (i+1)의 weight들을 pre-fetch해 놓는다. (아래 사진 참조)

이런 식으로 모든 layer의 forward pass를 진행한다.
3. Layer i의 backward pass 이전에, weight를 갖고 있는 owner GPU가 $W_i$를 broadcasting 해준다. (아래 사진 참조)

4. 모든 GPU는 layer i에 대해 backward pass를 진행하고, local ${dL}\over{dW_i}$를 계산한다. 그리고 나서 $W_i$를 삭제한다.
5. 이후에 모든 GPU는 계산한 local ${dL}\over{dW_i}$ 값을 owner GPU에게 보낸다. 그리고 나서 자신의 $local {dL}\over{dW_i}$는 삭제한다.
6. Owner GPU는 받은 local ${dL}\over{dW_i}$ 값을 이용해 weight들을 업데이트한다. backward 때도 forward pass처럼 동시에 작업을 처리할 수 있다. layer (i+1)의 backward pass를 계산하고 있을 때 $W_i$를 fetch 한다. 또, local gradient를 합치고 $W_i$를 업데이트하면서 Layer (i-1)의 backward pass를 진행한다. (아래 사진 참조)

위의 사진에서 볼 수 있듯이, $W_3$의 통신과 업데이트가 이루어지는 동안 $W_2$의 backward pass와 $W_1$의 fetching이 동시에 일어난다는 것이다.

In practice 상황에서는 backward pass가 끝났을 때 gradient는 모든 GPU에 대해서 소통이 되어 있고, 각 GPU는 weight 업데이트가 끝나 있고, dataloader는 다음 data batch를 CPU를 통해 load하고 있을 것이다. parallelization을 통해 효율적인 자원 사용이 가능해지는 것이다.

2) DP - Hybrid Sharded Data Parallelism (HSDP)

    FSDP도 좋은 방법이지만, 사람들이 사용하는 또 다른 방법으로는 HSDP가 있다.

개념적으로 GPU를 2D grid로 나누는 것이다. parallelism을 진행하는 축이 두 개가 된다고 이해한면 편하다. 2D grid라고 하니 잘 와닿진 않지만, 이러한 방식을 Multidimensional parallelism이라고 한다. 여러 개의 (여기서는 두 가지) parallelism 방식을 사용하는 것을 의미한다. 2D grid이니까 여기서는 두 가지의 parallelism 방식을 사용한 것이다. 위의 사진에서처럼, GPU의 개수는 N=M*K개이다. 이 N개의 GPU를 K개의 group으로 나누고, 각 group은 M개의 GPU를 가진다. 지금부터는 사진의 우측에 있는 예시를 기반으로 설명하겠다. 즉, N=8, M=2, K=4라는 의미이다.
  예시를 보면 각 group은 2개의 GPU를 가지고, 4개의 group이 존재한다. 4개의 group들은 우리가 이전에 배운 DP를 적용한다. 사진을 보면 4개의 group에 weight가 나뉘어 있는 것을 볼 수 있다. 이때, 각 group 내의 GPU 두 개에 모두 같은 weight를 나누어 가진다. 예를 들면 첫째줄의 group은 모두 $W_1$을 가지고, 둘째줄의 group은 모두 $W_2$를 가진다. 그리고 각 group 내에서는 FSDP를 적용한다. group 내에서 forward pass -> backward pass -> local gradient 계산까지 이루어지고, 이들의 전체 gradient 계산을 위해 4개의 group 간의 all-reduce operation이 진행되는 것이다. 그리고 나서 각 group은 계산된 전체 gradient를 가지고 독립적으로 업데이트를 진행한다. 내용이 길지만, 간단히 요약하면 아래와 같다.

1. N개의 GPU를 K개의 group으로 나누고, 각 group은 M개의 GPU를 갖는다. (N=M*K)
2. group 간에는 DP, group 내부에서는 FSDP를 진행한다.
   -> 두 가지 방식을 사용하는 Multidimensional parallelism

    훨씬 복잡해보이는데 이렇게 하는 이유가 무엇일까? 두 가지 방식(FSDP와 DP)이 요구하는 communication 횟수가 다르기 때문이다.

FSDP : forward, backward pass 한 번에 3번 communication이 필요하다.
-> forward pass시의 weight 공유, backward pass 시 weight와 gradient 공유
DP : forward, backward pass 한 번에 1번 communication이 필요하다.
-> 각 GPU가 weight들을 가지고 있기 때문에 all-reduce만 1회 진행하면 된다.

이러한 개념은 우리가 이전에 살펴본 GPU Hierarchy에 사용될 수 있다. 예를 들어, GPU Cluster를 쓴다고 했을 때, GPU server는 FSDP group으로 만들고, GPU server 끼리는 DP group으로 만드는 것이다. GPU server 내의 communication 속도 > GPU server 끼리의 communication 속도일 것이기 때문에, GPU server끼리의 communication은 DP를 사용해 1번만 하도록 설계하고, server 내에서는 비교적 속도가 빠르기 때문에 FSDP를 사용해 3번 communication 하도록 설계한 것이다. 이러한 식으로 GPU Cluster를 어떻게 구성할 지에도 parallelism 개념이 사용될 수 있다.

3) DP - 정리

    지금까지 배운 DP 개념들을 가지고 우리는 하나의 GPU에 들어가지 않는 큰 모델들도 여러 개의 GPU와 parallelism 테크닉들을 통해 학습할 수 있게 되었다. 그러나 DP의 문제점은 Model activation도 용량을 많이 차지할 수 있다는 점이다. 여기서 model activation이란 모델의 중간 출력값을 의미한다.

위의 사진에서도 볼 수 있듯이 Llama3-405B 모델도 activation이 63GB나 차지하는 것을 볼 수 있다. 이에 대한 해결책으로 모든 activation을 메모리에 저장해두지 않고, 필요할 때 다시 계산하는 방식을 생각해 볼 수 있다.

4) Activation Checkpointing

    위의 문제를 해결하기 위해 등장한 것이 Activation Checkpointing이다. Neural Network의 layer는 두 가지 function으로 구성된다.

위의 사진에서 볼 수 있듯이, forward pass와 backward pass이다. Layer 4개짜리 Network로 forward pass와 backward pass 시에 사용되는 memory에 대해 알아보자. 이때, 하나의 backward와 forward pass는 compute와 memory에서 모두 O(1)으로 계산한다.

먼저 forward pass를 다 하면, 위와 같이 memory와 compute가 사용된다. backward pass를 진행하면 아래와 같이 진행된다.

따라서 Forward + backward pass를 하면 O(N) Compute, O(N) memory가 되는 것이다. 그러나 모델이 커진다면 메모리가 부족해지는 현상이 나타난다. 따라서 필요할 때 activation을 계산해서 사용하자는 아이디어가 나오게 된 것이다. 먼저 forward pass 시에는 $A_1$을 계산하고, $A_2$를 계산한 후에는 $A_1$을 메모리에서 지우게 된다. 마찬가지로 $A_3$를 계산한 후에는 $A_2$을 지우게 된다. 이런 식으로 메모리에는 계속 한 개의 activation만 남게 되어 forward pass를 다 끝냈을 때에는 아래와 같이 memory와 compute가 나온다.

backward pass시에는 $A_4$가 메모리 상에 있기 때문에 바로 gradient에 대한 계산만 하게 된다. (compute +1)

이제부터가 문제다. 다음 계산을 위해서는 $A_3$가 필요한데, 메모리 상에서 지워버렸기 때문에 다시 계산을 해야 한다. $A_1 \rightarrow A_2 \rightarrow A_3$를 순차적으로 계산해서 3번의 compute를 해야 하고, backward pass 1회 compute하면 4번의 compute가 더해져 총 compute 수는 아래와 같다.

다음으로는 $A_2$가 필요하기 때문에 $A_1 \rightarrow A_2$ 두 번의 compute, 그리고 backward pass 1회를 해서 총 3번의 compute를 더 하게 된다.

같은 방식으로 끝까지 하면 아래와 같은 compute 수가 나오게 된다.

안 좋은 점은 $O(N^2)$ 가 나온다는 점이다. 이를 줄여주기 위해 C layer마다 checkpointing을 해서 거기서부터만 재계산을 하게 만드는 방식을 사용한다.

실제로 현업에서는 $C=\sqrt N$으로 많이 사용한다고 한다.

5) How to train on lots of GPUs

  대부분의 경우, HSDP와 Activation Checkpointing을 사용하면 된다. 구체적인 scaling 방법은 위의 사진과 같다. 그러나 이러한 방법의 문제점은 tuning할 수 있는 것들이 너무 많다는 것이다. 이에 대한 해결책으로 등장한 것이 Model FLOPs Utilization(MFU)이다.

6) Hardware FLOPs Utilization (HFU)

    MFU를 보기 전에 HFU부터 살펴보자. 이전에 H100은 tensor core에서 이론적으로 989.4 TFLOPs/sec의 16bit matrix multiplication이 가능했다. 그러나 이론이 아니라 현실적으로 계산을 한다면 어느 정도 성능을 기대할 수 있을까? HFU는 이론적인 matmul 성능 대비 현실적으로 실현 가능한 성능의 비율이다. H100는 80% 정도의 HFU를 가지는데, 이론적인 성능의 80% 성능을 얻는다는 이야기이다.

그러나 이러한 측정 방식의 문제는 activation checkpointing,data augmentation, optimizer, preprocessing과 같은 계산들을 고려하지 않는다. 따라서 더 이것들을 포함시켜서 사용할 수 있는 측정 단위가 필요하다. 그래서 등장한 것이 MFU이다.

6) Model FLOPs Utilization(MFU)

    MFU는 GPU의 이론적인 최고 FLOPs 중에 얼만큼이 'useful'한 모델 연산에 사용되고 있는지에 대한 비율을 의미한다. 계산하는 방식은 아래의 사진에 나와 있다.

1. 먼저 이론적인 FLOP $\text{FLOP}_{\text{theoretical}}$ 을 계산한다. (= forward, backward pass의 총 matmul FLOPs 수)
2. 사용하는 장비의 이론적인 최고 성능을 찾아본다. ($\text{FLOP/sec}_{\text{theoretical}}$)
3. 이론적으로 forward, backward pass가 걸리는 시간을 계산한다. ($\text{t}_{\text{theoretical}}$)
4. 실제 기기에서 forward, backward pass가 걸리는 시간을 측정한다. ($\text{t}_{\text{actual}}$)
5. MFU = $\text{t}_{\text{theoretical}} / \text{t}_{\text{actual}}$

위와 같은 방식으로 MFU를 계산하고, 보통 MFU > 30% 는 괜찮은 성능, 그리고 MFU > 40% 는 굉장히 좋은 성능을 의미한다고 한다.

7) 다양한 Parallelism (CP, PP, TP)

  &nsbp 지금까지 우리는 DP를 집중적으로 살펴봤다. 이제 CP, PP, TP에 대해서도 알아보자.

Context Parallelism (CP)

    CP는 sequence dimension에서 쪼개는 방법이다. 주로 transformer에서 많이 사용한다고 한다. Transformer는 L-length sequence를 단위로 작동하는데, 이때 여러 개의 GPU를 사용해서 하나의 긴 sequence를 처리한다.

위의 사진과 같이 여러 GPU가 sequence의 서로 다른 부분을 나눠 갖는 것이다. 이때, normalization과 residual connection은 weigh가 없기 때문에 넘어간다. MLP는 weight가 있기 때문에 각 GPU가 weight들을 복사해서 갖고 있고, DP와 같이 서로 gradient를 주고 받는다. 중요한 부분은 Attention이다. Attention은 조금 더 깊게 파고 들어야 한다.

  먼저 Attention의 QKV projection 부분이다. MLP layer이기 때문에 DP에서와 같이 gradient를 주고 받는다. 그러나 Attention operator는 Parallelize하기가 어렵다. 그래서 사용하는 두 가지 방법이 있는데, Ring Attention과 Ulysses라는 방법이다. Ring Attention은 Attention matrix를 여러 블럭으로 나눈 다음 각 GPU에게 줘서 parallel하게 작업을 처리하게 하는 방식이다. Ulysses는 조금 더 쉬운 방법인데, head에 대한 parallelism을 하는 것이다. 나머지 부분들은 모두 sequence dimension에서 parallelize 한다.

CP는 sequence length가 길 때 중요해진다. Llama3-405B pretraining을 예시로 보자.

Llama3는 two-stage training을 하는데, stage 1에서는 CP를 사용하지 않고, stage 2에서는 sequence length가 길어지며 16-way CP를 사용한다. 하나의 sequence를 16개의 GPU가 나눠서 parallel하게 처리한다는 뜻이다. 따지고 보면 각 GPU의 batch size는 1/16이라고 할 수 있다.

Pipeline Parallelism (PP)

    PP는 모델의 layer를 잘라서 parallelization 하는 방법이다. parallelization을 생각했을 때 직관적으로 떠오르는 방법이기도 하다.

그러나 이 방법의 문제점은 GPU가 노는 시간이 많다는 것이다. 이전 layer에서 계산이 완료될 때까지는 노는 GPU들이 생길 수 밖에 없다. 이러한 문제는 batch를 microbatch 여러 개로 쪼개서 GPU가 노는 시간을 최대한 줄이는 것으로 해결할 수 있다.

Tensor Parallelism (TP)

    TP는 각 Layer의 weight을 여러 GPU에 나누어서 처리하는 방법이다. 아래 그림을 보면 더 직관적인 이해가 가능하다.

그러나 이러한 방법의 문제점은 forward pass 이후 여러 GPU에 퍼져 있는 Y의 부분들을 모아야 한다는 것이 있다. 이를 해결하기 위한 일종의 트릭이 있다.

2개의 연속적인 TP layer가 있을 때, 첫 번째 layer의 weight matrix(W)는 column 방향으로 나누고, 두 번째 layer의 weight matrix(U)는 row 방향으로 나누는 것이다. 이렇게 하면 첫 번째 layer를 거친 후 output을 합치기 위해 GPU간의 communication을 할 필요가 없고, 두 layer 모두 계산한 후 마지막에만 합치면 된다는 장점이 있다.

어떤 Parallelism이 가장 좋은 선택일까?

    이 질문에 대한 답은 "모두"이다. TP, CP, PP, DP를 모두 한 번에 사용하고, GPU를 4D grid로 배열한다. 실제로 내가 이렇게까지 큰 모델을 다룰 일은 없을 수도 있겠지만, Llama3는 아래와 같은 방식으로 여러 개의 테크닉들을 한 번에 사용했다고 한다.