Paralelizarea Tensorilor (Tensor Parallelism – TP) în LLM-uri în PyTorch & Deep Learning Practic

Când un model de limbaj (ex: Llama-70B) este prea mare pentru a încăpea în memoria unui singur GPU chiar și în faza de inferență, Tensor Parallelism (TP) împarte fiecare matrice de ponderi pe mai multe acceleratoare interconectate prin NVLink.

În conformitate cu abordarea Megatron-LM:

• Primul strat liniar este partiționat pe coloane (Column Parallel).
• Al doilea strat liniar este partiționat pe rânduri (Row Parallel), urmat de o operație All-Reduce pentru a însuma rezultatele.

from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel, RowwiseParallel

# Exemplu de paralelizare a unui strat Multi-Layer Perceptron (MLP)
parallel_plan = {
    "w1": ColwiseParallel(),  # Proiecție pe coloane
    "w2": RowwiseParallel()   # Proiecție pe rânduri cu All-Reduce
}

parallelize_module(transformer_block.mlp, mesh, parallel_plan)

Tensor Parallelism atinge scalabilitate aproape liniară datorită lățimii uriașe de bandă a magistralelor NVLink (până la 900 GB/s per GPU).