[Inference] Tensor model parallelism

[goliaro]

Description of changes:

Related Issues:

Linked Issues:

• Issue #

Issues closed by this PR:

• Closes #

Before merging:

• Did you update the flexflow-third-party repo, if modifying any of the Cmake files, the build configs, or the submodules?

[jiazhihao]

@gabrieleoliaro The PR looks great to me overall. Do you think we can merge this into the inference branch?

[xinhaoc]

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?

[jiazhihao]

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

[goliaro]

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

[xinhaoc]

thanks!

[goliaro]

@gabrieleoliaro The PR looks great to me overall. Do you think we can merge this into the inference branch?

yes! I'll merge it as soon as CI passes!