はじめに

LLMはモデルがサイズがデカすぎるため、複数のGPUにモデルを分割配置して学習を行います。MegatronLMはモデルを分割する手法の一つであり、Transoformerを分割することに特化しています。今回はそんなMegatronLMを説明したいと思います。

モデルの分割

皆さんもご存知の通りLLMのモデルは非常にサイズがでかいです。例えばllama2-70bのモデルは140GBです。GPUのメモリサイズは多くても100GBほどですので、1つのGPUには乗り切らないです。学習際にはモデル以外にactivationなどをメモリに保持しないといけないので、1つのGPUで学習することは当然できないです。それを可能にしてくれるのがモデル分割手法たちです。

モデル分割手法はレイヤー間でモデルを分割する手法とレイヤー内でモデルを分割する手法があります。レイヤー間の分割は非常に単純でTransformer blockの切れ目でモデルを分割するといった感じです。レイヤー内のモデルの分割はTransoformer block内のAttentionやFFNを分割する手法です。MegatronLMはそんなレイヤー内のモデル分割手法であり、AttentionのFFNを分割します。

MegatronLM

MegatronLMではTransformer blockのAttentionとFFNを分割しますが、分割方法はそれぞれ少し異なっています。

Attentionの分割

Attentionの分割は正確にとMulti-Head Attentionを分割します。簡単にいうとhead数が8でGPU数が２だとするとhead1~4をGPU1が実行し、head5~8をGPU2が実行して、最後にお互いに結果を送りあうという感じです。これを順に説明してきます。

まず、Muti-Head Attentionは次の式でかけます。$n$はhead数で、$I$は入力、$W_i^Q,W_i^K,V_i^V$はそれぞれ、$QKV$を計算する重みを列方向に$i$個に分割した時の$i$番目です。

\begin{align}
MutiHeadAttention(Q,K,V) & = Concat(h_1,h_2,...h_n)W^O\\
h_i & = Attention(Q_i,K_i,V_i)\\
Attention(Q,K,V) & = Softmax(QK^T)V\\
Q_i & = IW_{i}^Q\\
K_i & = IW_{i}^K\\
V_i & = IW_{i}^V\\
\end{align}

この計算を複数のGPUで実行します。でもそんなには難しくないです。$n$が2でGPU数も2の時を考えると下のように書けるはずです。$W_i^O$は$W^O$を行方向に分割した$i$番目のことで、よりわかりやすくするためにConcatの部分を式変形しています。

\begin{align}
MutiHeadAttention(Q,K,V) & = Concat(h_1,h_2)W^O\\
& = 
\begin{bmatrix}
h_1 & h_2
\end{bmatrix}
\begin{bmatrix}
W_1^O \\
W_2^O
\end{bmatrix}
\\
& = h1W_1^O + h2W_2^O
\end{align}

これの$h1W_1^O$をGPU1で行って、$h2W_2^O$をGPU２で行います。その後、お互いに自分が計算した結果を相手に送ってあげて、相手から受けとた結果を自分が計算した結果に足してあげれば、無事に計算ができることがわかると思います。

###　FFNの分割
FFNの分割はAttentionの分割を理解していれば非常に簡単です。まず、FFNは次のように書けると思います。

\begin{align}
FFN(I) = GeLU(IW)
\end{align}

これのWを分割してあげます。分割する方法は列方向で2分割した場合は下記のように$W_1,W_2$と分割できます。

\begin{align}
FFN(I) = GeLU(I
\begin{bmatrix}
W_1\\
W_2
\end{bmatrix}
)
\end{align}

これをさらに変形してあげると次のようになるはずです。

\begin{align}
FFN(I) &= GeLU(I
\begin{bmatrix}
W_1\\
W_2
\end{bmatrix}
)
\\
& = GeLU(Concat(IW_1, IW_2)) \\
& = Concat(GeLU(IW_1), GeLU(IW_2))
\end{align}

Concatは行列を列方向にくっつけるだけですので、GeLUの中にあった、Concatを外に出してあげて大丈夫です。これの$GeLU(IW_1)$をGPU1で計算し、$GeLU(IW_2)$をGPU2で計算し、お互いに送り合ってConcatしてあげれば、無事に計算できるはずです。

これで、無事にAttentionもFFNも分割して行えることがわかりました。最後まで読んでいただきありがとうございます。