-
Notifications
You must be signed in to change notification settings - Fork 223
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[Inference] Tensor model parallelism #778
Conversation
@gabrieleoliaro The PR looks great to me overall. Do you think we can merge this into the inference branch? |
@@ -127,7 +127,7 @@ Tensor | |||
} | |||
if (bias) { | |||
// q, k, v, o | |||
int dims[1] = {embed_dim * 4}; | |||
int dims[1] = {(qProjSize + kProjSize + vProjSize) * num_heads + oProjSize}; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hi Gabriele, just to make sure, do we have any situation that the embed_dim * 4
is not equal to qProjSize + kProjSize + vProjSize) * num_heads + oProjSize
?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
In Megatron-LM's tensor model parallelism, we partition attention heads across GPUs, so num_heads
of each GPU
can be smaller than the total number of attention heads
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
In the current model we have, the two quantities are indeed always identical. I changed this just to make it easier to remember the data layout
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
thanks!
yes! I'll merge it as soon as CI passes! |
Description of changes:
Related Issues:
Linked Issues:
Issues closed by this PR:
Before merging: