Tensor Parallelism

Définition

Le Tensor Parallelism (parallelisme de tenseurs) est une technique de parallelisme pour les LLMs qui partitionne les matrices de poids individuelles (attention, FFN) sur plusieurs GPUs, chaque GPU ne stockant et ne calculant qu'une partie des operations matricielles. C'est la technique de base pour servir des modeles dont les poids depassent la VRAM d'un seul GPU. Dans le Tensor Parallelism, une multiplication matricielle Y = XW est decomposee sur N GPUs en partitionnant la matrice W par colonne (column-parallel) ou par ligne (row-parallel). Pour le Multi-Head Attention, les h tetes d'attention sont reparties sur les GPUs (h/N tetes par GPU). Pour le FFN, la dimension intermediaire d_ff est divisee. La librairie Megatron-LM (NVIDIA) est la reference pour le Tensor Parallelism dans l'entrainement des LLMs. Pour l'inference, vLLM et TGI implementent le Tensor Parallelism via la biblioteque PyTorch Distributed Communications (NCCL). Le parametre --tensor-parallel-size=N dans ces frameworks active le TP. Le Tensor Parallelism necessite une communication all-reduce entre tous les GPUs apres chaque operation parallelisee, ce qui le rend tres dependant de la bande passante GPU-GPU. NVLink (bande passante 600-900 GB/s entre GPU) est quasi-indispensable pour du TP efficace — sans NVLink, la communication PCIe (64 GB/s max) devient le goulot d'etranglement. En pratique, le Tensor Parallelism est combine avec le Pipeline Parallelism (couches du modele distribuees sur des stages) et le Data Parallelism (repliques du modele pour plusieurs batches) dans les configurations 3D parallelism utilises pour les modeles frontier (GPT-4, LLaMA 405B).