Как работает tensor parallelism для LLM training? Чем отличается от инференса?

Краткий тезис

parallelism|Tensor Parallelism (TP) — это способ распределения параметров одного слоя нейросети между несколькими GPU, чтобы каждая карта хранила только часть тензоров. При обучении после каждого backward необходимо синхронизировать градиенты через AllReduce, что увеличивает объём коммуникации в 2–3 раза по сравнению с инференсом, где выполняется только forward pass. TP критичен для очень больших моделей (от 10+ млрд параметров), когда один слой не помещается в память одной GPU.

1. Что такое Tensor Parallelism (TP)

parallelism|Tensor Parallelism — это стратегия model parallelism (распараллеливания модели), при которой веса одного слоя (например, матрицы weight в линейном слое) разрезаются (шардируются) по одному из измерений. Каждый GPU хранит фрагмент тензора и вычисляет свою часть результата. Затем части объединяются с помощью коммуникационной операции (обычно AllReduce или ReduceScatter + AllGather).

TP отличается от Data Parallel|Data Parallelism (DP), где каждый GPU хранит копию всей модели, и от Pipeline Parallelism (PP), где разные слои распределены по GPU. В TP один слой «размазан» между GPU.

Цель TP — позволить обучать модели, которые не влезают в память одной карты, и уменьшить нагрузку на память за счёт хранения лишь части параметров.

2. Как TP работает на уровне матричных операций

Рассмотрим линейный слой: Y = X @ W + b, где
X — входной тензор формы batch, Вики/Hidden dimension|hidden,
W — вес формы hidden, out_features].

Row‑wise partitioning (разрезание по строкам W)

• W делится на N частей по строкам: W = [W0; W1; ...; W_{N-1}].
• Каждый GPU получает W_i формы hidden, out_features/N].
• Вычисляет Y_i = X @ W_i (локальное умножение).
• Затем AllReduce по оси признаков, чтобы собрать полный Y.

Column‑wise partitioning (разрезание по столбцам W)

• W делится по столбцам: W = [W0 | W1 | ... | W_{N-1}].
• Каждый GPU получает W_i формы [hidden/N, out_features].
• Вход X сначала разрезается по hidden размеру (через AllReduce нет — требуется AllGather).
• После локального умножения результат Y_i остаётся неполным — их конкатенируют вдоль признаков.

На практике для линейных слоёв чаще используют column‑wise, а для внимания — row‑wise, чтобы минимизировать число all‑reduces.

Пример кода (псевдо, PyTorch + NCCL):

# Допустим, N=2, rank 0 и rank 1
W = torch.randn(hidden, out_features)  # полный вес — не хранится
W_local = W[:, rank * out_features//2 : (rank+1)*out_features//2]  # column-wise

# forward
out_local = X @ W_local  # [batch, out_features//2]
# собираем полный выход: AllGather или ReduceScatter? здесь AllGather
out = all_gather(out_local, axis=1)  # [batch, out_features]

# backward (auto diff — градиенты нужно синхронизировать)
# Для backward понадобится AllReduce градиентов W_local

3. Коммуникация при forward pass

В forward‑е TP, как правило, использует AllReduce — усреднение или суммирование частичных результатов, полученных на каждом GPU.

• Если используется row‑wise, результат Y_i — неполный (каждый GPU выдаёт только часть признаков), и для получения полного выхода нужен AllGather (сборка).
• При column‑wise вход разрезается, а выход уже полный — поэтому дополнительная коммуникация для объединения не требуется, но требуется обмен фрагментами входа.

В большинстве реализаций (например, Megatron‑LM) для внимания (Self‑Attention) применяют row‑wise для Q/K/V проекций, а для выходной проекции — column‑wise, чтобы скомпенсировать коммуникацию.

Ключевой момент при инференсе forward выполняется ровно один раз — достаточно одного AllReduce на слой. Коммуникационный объём примерно равен размеру вывода одного слоя.

4. Коммуникация при backward pass (тренировка)

При обучении после каждого backward‑а необходимо вычислить градиенты по весам. Так как веса распределены, градиенты каждого GPU локальны и требуют синхронизации, чтобы получить полные градиенты перед обновлением.

• Градиент локального фрагмента W_local получается через обратное распространение (автоматическое дифференцирование).
• Для каждого слоя, где веса разрезаны, выполняется AllReduce градиентов по всем GPU, работающим в TP‑группе.
• После AllReduce каждый GPU получает одинаковые усреднённые (или суммированные) градиенты для своей части весов.

По сравнению с forward, при обучении возникает дополнительный AllReduce для градиентов. Более того, при backward через сам слой также могут возникать дополнительные all‑reduces (например, для агрегации градиентов по входному тензору при row‑wise).

Итог количество all‑reduce операций на один слой при обучении примерно в 2 раза больше, чем при инференсе (один на forward, один на backward). А с учётом дополнительных all‑reduces внутри внимания — может быть до 3 раз.

5. Сравнение: обучение vs инференс

При инференсе можно использовать FP16/INT8 и была бы возможность Tensor Parallelism без all-reduce (например, просто собрав результат на одной GPU). Но обычно используют тот же механизм для снижения latency, хотя при маленькой загрузке можно обойтись и без TP.

6. Tensor Parallelism и ZeRO-3

ZeRO-3 — это оптимизация памяти, при которой параметры, градиенты и состояния оптимизатора распределены между GPU, а не хранятся полностью. TP разрезает сами вычисления — это другой уровень.

• TP уменьшает память для тензоров весов за счёт хранения только фрагмента.
• ZeRO-3 также шардит веса, но при вычислениях он собирает полный вес перед операцией (через AllGather).
• TP не требует сборки полного веса — каждый GPU работает с фрагментом и собирается только результат.

Сочетание TP + ZeRO-3 возможно, но ведёт к избыточной коммуникации. На практике TP часто используют вместе с Pipeline Parallelism и Data Parallelism (3D Parallelism в Megatron‑LM), а ZeRO-3 применяют в основном для небольших кластеров (2–16 GPU), где TP неэффективен.

7. Другие виды параллелизма (сравнение)

TP даёт минимальный latency (гранулярность — один слой), но максимальный объём коммуникации на один слой. Его стоит применять, когда один слой не помещается в HBM.

8. Практические рекомендации

• Когда использовать TP: модель настолько большая, что даже один слой не влезает в одну GPU (например, GPT‑3 175B с hidden size 12k). Для моделей до 20B обычно достаточно PP + DP.
• Избегать слишком большого числа GPU в TP группе: коммуникация all‑reduce масштабируется как O(N), и при N > 8 пропускная способность становится узким местом. Обычно TP группа — 2–8 GPU на одном узле (NVLink).
• Сочетать TP с Pipeline Parallelism: например, 4 GPU на TP, 8 стадий PP — всего 32 GPU.
• Optimizer state: при обучении нужно либо хранить полные состояния оптимизатора на каждой TP‑группе, либо дополнительно шардить через ZeRO. TP сам по себе не шардит optimizer state.

9. Пет-проект для закрепления

Задача: Написать скрипт, моделирующий TP для матричного умножения на 2 GPU (через PyTorch + torch.distributed) и сравнить время forward+backward при training и только forward при inference.

Инструменты: PyTorch, NCCL, cProfile, nvtop.

Шаги:

1. Инициализировать process group с двумя рангами.
2. Определить функцию tensor_parallel_linear(x, w_fragment, rank, size, mode='train').
   • w_fragment — часть веса (column‑wise).
   • forward: локальное умножение, AllReduce результата.
   • backward: при 'train' — синхронизировать градиенты через AllReduce (можно симулировать вручную, без autograd, для чистоты замера).
3. Запустить torch.manual_seed(0), замерить timeit.
4. Повторить для разных hidden (4096, 8192) и batch size.
5. Построить таблицу: mode, hidden, time, коммуникационный объём.

Ожидаемый результат: Понять, почему training требует в 2–3 раза больше времени на коммуникацию чем inference.

# Эскиз
import torch
import torch.distributed as dist

def tp_linear_forward(x, w_frag):
    out = x @ w_frag
    dist.all_reduce(out)  # sum over ranks
    return out

def tp_linear_backward(grad_out, x, w_frag):
    # симуляция: считаем локальные градиенты, потом all_reduce
    grad_w = x.T @ grad_out
    dist.all_reduce(grad_w)  # синхронить градиенты веса
    grad_x = grad_out @ w_frag.T
    return grad_x, grad_w

# Замер для training и inference

10. Связь с другими вопросами

Навигация

• Предыдущий: 422
• Следующий: 424
• Индекс: 00. Индекс разборов