There's no reason to run a second communication for kv layers in MQA. Basically the previous ColumnLinear should already handle that all_gather for us, we just need to run a typical Linear layer. It does mean that we run that kv linear in duplicated fashion, but it's seems to not matter (It's my understanding that sequence parallel wasn't used during the training of starcoder)
The reason why I decided not to build another module, is for backward compatibility. Basically this PR should allow for free throughput improvement, and be able to load previous checkpoints.
I have not tested this PR. cc @loubnabnl
If it does provide the same improvement I've seen on my own codebase, then we need to propagate this PR to our inference models. Typically text-generation-inference and such. (EDIT: this seems already support in text-generation-inference with even an extra merging of qkv)
There's no reason to run a second communication for
kvlayers in MQA. Basically the previous ColumnLinear should already handle that all_gather for us, we just need to run a typical Linear layer. It does mean that we run that kv linear in duplicated fashion, but it's seems to not matter (It's my understanding that sequence parallel wasn't used during the training of starcoder)The reason why I decided not to build another module, is for backward compatibility. Basically this PR should allow for free throughput improvement, and be able to load previous checkpoints.
I have not tested this PR. cc @loubnabnl
If it does provide the same improvement I've seen on my own codebase, then we need to propagate this PR to our inference models. Typically
text-generation-inferenceand such. (EDIT: this seems already support intext-generation-inferencewith even an extra merging of qkv)