<p>Model Parallelism implementation for GPT2 model.
As per my understanding, parallelism implemented as shown in the below picture. Marked blocks are computed parallel.
<a href="https://i.stack.imgur.com/KFhte.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/KFhte.png" alt="enter image description here" /></a></p>
<p>Fig (a)MLP, says</p>
<blockquote>
<p>f and g are conjugate, f is an identity operator in the forward pass
and all-reduce in the backward pass while g is an all-reduce in
forward and identity in backward.</p>
</blockquote>
<p>Similarly self attention block works as shown below picture
<a href="https://i.stack.imgur.com/H00ht.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/H00ht.png" alt="enter image description here" /></a></p>
<p>From the experiment given in the <a href="https://arxiv.org/pdf/1909.08053.pdf" rel="nofollow noreferrer">paper</a>, 1.2 billion parameters fits on single GPU, where as 8 billion parameters requires 8 GPUs (8-ways).</p>
<p>96 is the constant number used as the hidden size per head.
As per Table 2 from <a href="https://arxiv.org/pdf/1909.08053.pdf" rel="nofollow noreferrer">paper</a>, Hidden size might be based on parameter count.</p>
How is model parallelism implemented for GPT2 in MegatronLM?
<p>I am trying to understand the implementation details of <a href="https://github.com/NVIDIA/Megatron-LM#inverse-cloze-task-ict-pretraining" rel="nofollow noreferrer">MegatronLM</a>, which has both model and data parallel. On their <a href="https://nv-adlr.github.io/MegatronLM" rel="nofollow noreferrer">site</a> or in their research <a href="https://arxiv.org/pdf/1909.08053.pdf" rel="nofollow noreferrer">paper</a>, they mentioned how they used intra-layer parallel which is similar to mesh TensorFlow. I am confused with some details.</p>
<p>As shown in the picture below, my understanding is that the computation inside each of the 4 red circles can be parallelized by intra-layer splitting, but MLP must happen after self-attention, so only 2 red circled blocks can be parallelized at the same time. The paper says the model parallel is 8-way. My first question is, <strong>Does this indicate they split each block into 4 intra-layer parts (8/2)?</strong></p>
<p>(8-way divided by 2-blocks)
<a href="https://i.stack.imgur.com/5agbD.jpg" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/5agbD.jpg" alt="enter image description here" /></a></p>
<p>The paper also mentioned</p>
<blockquote>
<p>To have consistent GEMM sizes in the self attention layer, the hidden size per attention head is kept constant at 96 while the number of heads and layers are varied to obtain configurations ranging from 1 billion to 8 billion parameters.</p>
</blockquote>
<p>My second question is <strong>What does the 96 hidden size refer to here?</strong></p>
<p>I am totally new to distributed training, I probably misunderstood something. Any clarification on this topic would be very appreciated! Thanks!</p>