Tensor Parallelism

Тензорный параллелизм — это продвинутый метод распределённого обучения в машинном обучении, при котором крупные отдельные математические структуры, или тензоры, распределяются между несколькими аппаратными ускорителями, например GPU или TPU. При обучении огромных моделей глубокого обучения количество параметров может легко превысить объём памяти одного устройства. Вместо того чтобы размещать целый слой нейронной сети на одном GPU, тензорный параллелизм распределяет матрицы весов по нескольким устройствам кластера и делит между ними математические операции, например умножение матриц. Благодаря этому модель использует суммарную память и вычислительную мощность всего оборудования, выполняя параллельные вычисления по схеме одна программа — множество данных (SPMD) и синхронизируя результаты через высокоскоростные соединения, например NVIDIA NVLink.

Как работает тензорный параллелизм#

В основе нейронной сети лежат матричные умножения. Тензорный параллелизм распределяет эти операции между устройствами, разделяя матрицы по строкам или столбцам. Например, в полносвязном слое или механизме внимания трансформера один GPU может вычислять левую половину матрицы, а другой — правую. По завершении параллельных вычислений устройства обмениваются данными — часто с помощью быстрых коллективных операций All-Reduce, — чтобы объединить частичные результаты перед передачей полного тензора следующему слою. В 2025 году новые академические разработки дополнительно оптимизируют этот процесс: частично синхронизируют активации, чтобы сократить накладные расходы на передачу данных, которые обычно становятся узким местом крупных вычислительных кластеров.

Различия между связанными методами параллелизма#

Чтобы понять, как тензорный параллелизм вписывается в общую картину распределённых вычислений, нужно отличать его от других распространённых стратегий:

• Тензорный параллелизм и параллелизм моделей: тензорный параллелизм — это узкая подкатегория параллелизма моделей. Если общий параллелизм моделей подразумевает любое разделение модели между устройствами, то тензорный параллелизм означает именно разделение отдельных тензоров внутри одного слоя.
• Тензорный параллелизм и конвейерный параллелизм: Конвейерный параллелизм — ещё один вид параллелизма моделей, при котором сеть разделяется по глубине: первые слои размещаются на GPU 0, следующие — на GPU 1 и так далее. Из-за этого возникают последовательные зависимости, известные как пузыри конвейера. Тензорный параллелизм разделяет сами слои, позволяя выполнять их одновременно без последовательных задержек, но требует гораздо более высокой пропускной способности сети.
• Тензорный параллелизм и параллелизм данных: при параллелизме данных полная копия модели размещается на каждом GPU, а между устройствами распределяется только обучающий набор данных. Для высокооптимизированных архитектур, например Ultralytics YOLO26, которые легко помещаются в память современных GPU, параллелизм данных через DistributedDataParallel в PyTorch используется по умолчанию. Тензорный параллелизм обычно нужен только тогда, когда параметры одного слоя превышают объём видеопамяти оборудования и возникают ошибки нехватки памяти (OOM).

Примеры применения в реальных условиях#

Тензорный параллелизм незаменим в современной инфраструктуре ИИ, особенно для передовых архитектур, требующих огромных вычислительных ресурсов:

• Обучение больших языковых моделей (LLM): Огромные базовые модели, такие как Llama 3 от Meta, используют тензорный параллелизм, который часто реализуют с помощью платформ вроде NVIDIA Megatron-LM. Примечательно, что DeepSeek V3 обучали без него, используя конвейерный параллелизм и параллелизм экспертов. Поскольку скрытые размерности и головы внимания этих моделей очень велики, для эффективного обучения и низкой задержки при инференсе в реальном времени их необходимо распределять между восемью GPU одного узла.
• Большие модели компьютерного зрения (LVM) и генерация 3D: По мере развития компьютерного зрения в сторону крупных систем мультимодального рассуждения исследователи сочетают тензорный и конвейерный параллелизм в сервисах вроде AWS SageMaker для обучения огромных визуальных трансформеров (ViT). Этот метод позволяет обрабатывать изображения высокого разрешения и генерировать видео, для чего нужны большие непрерывные блоки памяти.

Реализация тензорного параллелизма в PyTorch#

Раньше инженерам приходилось писать сложную пользовательскую логику распределённых вычислений для разделения тензоров. Недавно в PyTorch появился DTensor (Distributed Tensor), который упрощает этот процесс благодаря встроенной поддержке. Ниже показан пример создания тензора, разделённого по строкам, с использованием официального API PyTorch Distributed Tensor:

import torch
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import Shard, distribute_tensor

# Initialize a 1D device mesh across 2 GPUs
mesh = init_device_mesh("cuda", (2,))

# Create a standard PyTorch tensor (e.g., representing a layer's weights)
local_tensor = torch.randn(1024, 1024)

# Distribute the tensor across the GPUs by sharding along the first dimension (row-wise)
# Each GPU now holds a (512, 1024) chunk of the original tensor
distributed_tensor = distribute_tensor(local_tensor, mesh, [Shard(0)])

print(f"Global shape: {distributed_tensor.shape}, Local shape: {distributed_tensor.to_local().shape}")

Для оптимизированных под периферийные устройства задач компьютерного зрения и быстрого развёртывания моделей разработчики обычно используют платформу Ultralytics, которая автоматически обеспечивает оптимальное использование оборудования. Для базовых моделей с миллиардами параметров нужна ручная настройка тензорного параллелизма, но обучать модели, например YOLO26, можно эффективно и из коробки с помощью простых команд CLI. Это обеспечивает максимальную пропускную способность за счёт бесшовного использования встроенных методов параллелизма данных и надёжных рекомендаций по обучению моделей.