2x slowdown using TP

[jph00]

Thanks for this very interesting lib @BlackSamorez! I just tried running your kaggle nb locally, using 2 x A6000s. I used the 'meta-llama/Llama-2-7b-hf' model. The call to model.generate with max_length=200 takes 12 seconds when using tensor_parallel with the 2 GPUs.

However, if I remove tensor_parallel and instead just use a single GPU, generation is over twice as fast, taking 5 seconds.

Is this slowdown expected, or am I doing something wrong?