Context Parallelism

Параллелизм контекста — это метод распределённых вычислений, при котором длинная входная последовательность делится между несколькими ускорителями. Каждый GPU обрабатывает только часть последовательности и во время вычисления внимания взаимодействует с остальными. Это снижает объём памяти активаций на каждом устройстве и позволяет обучать трансформер на входных данных, которые могут не поместиться в память одного GPU, например на очень длинных документах, продолжительных видео или больших наборах патчей изображения.

В отличие от простого увеличения контекстного окна модели, параллелизм контекста не меняет теоретически допустимый для архитектуры объём информации. Вместо этого он делает обработку такого контекста вычислительно целесообразной, распределяя последовательность по устройствам.

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

Предположим, последовательность содержит 32 000 токенов, а параллелизм контекста использует четыре GPU. Сначала каждое устройство получает примерно по 8 000 токенов и сохраняет соответствующие промежуточные активации.

Большинство операций, например нормализация и полносвязные слои, могут независимо обрабатывать локальные фрагменты последовательности. Сложность связана с механизмом внимания: локальному запросу может понадобиться обратить внимание на ключи и значения, хранящиеся на каждом из остальных устройств.

Поэтому реализации обмениваются блоками ключей и значений, или KV-блоками, между GPU. При использовании кольцевого внимания каждое устройство вычисляет частичное внимание на основе локальных данных, передаёт KV-блок следующему устройству и повторяет процесс, пока не обработает всю последовательность. В учебном руководстве PyTorch по параллелизму контекста этот механизм демонстрируется на примере распределённого внимания с масштабированным скалярным произведением, а в пакете NVIDIA для параллелизма контекста описаны варианты коммуникации на основе all-gather, reduce-scatter и кольцевой схемы.

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

Зачем нужен параллелизм контекста#

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

Параллелизм контекста решает проблему памяти, распределяя активации между устройствами. Он также позволяет распределить вычисление внимания, хотя обмен данными создаёт дополнительные затраты. Эффективные системы совмещают передачу KV с вычислениями с помощью операций, описанных, например, в руководстве по коллективным операциям NCCL.

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

Параллелизм контекста и связанные методы#

• Тензорный параллелизм разделяет операции или матрицы весов внутри отдельных слоёв. Параллелизм контекста вместо этого делит токены по измерению последовательности.
• Конвейерный параллелизм распределяет разные группы слоёв модели по разным устройствам. Он разбивает модель по глубине, а не по длине последовательности.
• Параллелизм данных создаёт копии модели и передаёт каждой копии разные обучающие примеры. В обзоре распределённых вычислений PyTorch его рекомендуют использовать, когда вся модель и каждый пример помещаются на одном GPU.
• Параллелизм последовательности часто распределяет активации для отдельных операций, связанных с тензорным параллелизмом. Параллелизм контекста шире применяет разбиение последовательности к входным данным и активациям сети.

Эти подходы дополняют друг друга. В руководстве NVIDIA по стратегиям параллелизма показано, как параллелизм контекста, тензоров, конвейера и данных формирует многомерную схему распределения устройств.

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

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

• Анализ длинных видео и мультимодальных данных: видеотрансформеры и большие модели компьютерного зрения могут представлять кадры, патчи изображений, аудиофрагменты и текст как одну длинную последовательность токенов. Распределение этой последовательности помогает моделям анализировать длительные записи без значительного сокращения числа кадров или пространственной детализации. В обзоре параллелизма контекста AWS Neuron показано, как группы ускорителей обмениваются фрагментами KV для таких задач с длинным контекстом.

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

Практическое применение и компромиссы#

В следующем минимальном примере используется экспериментальный API PyTorch для параллелизма контекста и настройки внимания с масштабированным скалярным произведением. Сохрани его как cp_example.py и запусти на двух GPU с помощью torchrun --standalone --nproc-per-node=2 cp_example.py.

import os

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.experimental import context_parallel
from torch.nn.attention import SDPBackend, sdpa_kernel

rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])

torch.cuda.set_device(rank)
torch.cuda.manual_seed(0)
dist.init_process_group("nccl")

mesh = init_device_mesh("cuda", (world_size,))
qkv = [torch.randn(1, 4, 4096, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) for _ in range(3)]

with sdpa_kernel(SDPBackend.FLASH_ATTENTION), context_parallel(mesh, buffers=tuple(qkv), buffer_seq_dims=(2, 2, 2)):
    output = F.scaled_dot_product_attention(*qkv, is_causal=True)

output.float().square().mean().backward()
dist.destroy_process_group()

Здесь измерение 2 соответствует измерению последовательности, поэтому каждый процесс получает фрагмент последовательности, а внимание координируется между устройствами в сетке.

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