How is model parallelism implemented for GPT2 in MegatronLM?

user_11210924 0
Viewed 194

I am trying to understand the implementation details of MegatronLM, which has both model and data parallel. On their site or in their research paper, they mentioned how they used intra-layer parallel which is similar to mesh TensorFlow. I am confused with some details.

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, Does this indicate they split each block into 4 intra-layer parts (8/2)?

(8-way divided by 2-blocks) enter image description here

The paper also mentioned

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.

My second question is What does the 96 hidden size refer to here?

I am totally new to distributed training, I probably misunderstood something. Any clarification on this topic would be very appreciated! Thanks!

transformer-model tensorflow distributed-computing nlp parallel-processing
1 Answers

Model Parallelism implementation for GPT2 model. As per my understanding, parallelism implemented as shown in the below picture. Marked blocks are computed parallel. enter image description here

Fig (a)MLP, says

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.

Similarly self attention block works as shown below picture enter image description here

From the experiment given in the paper, 1.2 billion parameters fits on single GPU, where as 8 billion parameters requires 8 GPUs (8-ways).

96 is the constant number used as the hidden size per head. As per Table 2 from paper, Hidden size might be based on parameter count.

user_142902440
Related