Tensor Parallelism
Tensor Parallelism (TP)
Tensor Parallelism — метод распределения модели, при котором отдельные тензоры (матрицы весов) разрезаются по горизонтали между GPU. Каждая GPU хранит и вычисляет свою часть каждого слоя.
Что такое Tensor Parallelism
Tensor Parallelism (TP) — способ разделения модели, при котором каждый тензор (матрица весов) горизонтально разрезается между GPU. Все GPU работают над одним и тем же слоем, но каждая обрабатывает свою часть тензора.
Принцип работы
Для линейного слоя Y = XA:
Разрезаем A по столбцам на 2 GPU:
A = [A1 | A2]
GPU 0: Y1 = X × A1
GPU 1: Y2 = X × A2
Результат: Y = [Y1 | Y2] (конкатенация)
Для Attention: Q, K, V матрицы разрезаются по головам. Каждая GPU вычисляет attention для своих голов.
TP vs другие виды параллелизма
| Вид | Что разрезается | Коммуникация | Задержка |
|---|---|---|---|
| Data Parallelism | Батч данных | All-Reduce (конец шага) | Низкая |
| Tensor Parallelism | Тензоры слоя | All-Reduce (каждый слой!) | Высокая |
| Pipeline Parallelism | Слои модели | Point-to-point (между стадиями) | Средняя |
Когда использовать
TP требует очень быстрого интерконнекта (NVLink), т.к. коммуникация происходит на каждом слое. Поэтому:
- Внутри ноды (8 GPU + NVLink): TP отлично работает
- Между нодами (InfiniBand): TP слишком медленный, используйте PP или ZeRO
Типичная конфигурация для 8 нод × 8 GPU:
TP = 8 (внутри ноды, через NVLink)
PP = 8 (между нодами, через InfiniBand)
DP = 1 (или больше, если позволяет batch size)
Влияние на batch size
TP не увеличивает эффективный batch size (в отличие от DP). Это важно: если вам нужен большой batch — используйте DP или Gradient Accumulation.
Реализации
- Megatron-LM (NVIDIA) — оригинальная реализация TP для Transformer
- DeepSpeed — через Megatron-DeepSpeed
- PyTorch —
torch.distributed.tensor.parallel - vLLM — TP для инференса (
--tensor-parallel-size N)
# vLLM: TP для инференса
from vllm import LLM
llm = LLM("meta-llama/Llama-3-70B-Instruct", tensor_parallel_size=4)