Context Parallelism

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

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

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

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

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

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

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

Почему важен контекстный параллелизм#

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

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

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

Контекстный параллелизм в сравнении со смежными методами#

• Tensor parallelism делит операции или матрицы весов внутри отдельных слоев. Контекстный параллелизм вместо этого делит токены вдоль измерения последовательности.
• Pipeline parallelism назначает разные группы слоев модели разным устройствам. Он разделяет глубину модели, а не длину последовательности.
• Data parallelism реплицирует модель и передает каждому репликанту разные обучающие примеры. PyTorch distributed overview рекомендует его, когда полная модель и каждый образец помещаются на одном GPU.
• Sequence parallelism часто разбивает активации на шарды для выбранных операций, связанных с тензорным параллелизмом. Контекстный параллелизм применяет разбиение последовательности более широко для входов сети и активаций.

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

Реальные приложения#

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

• Понимание длинных видео и мультимодальных данных: Видео-трансформеры и large vision models могут представлять кадры, патчи изображений, аудиосегменты и текст как одну длинную последовательность токенов. Распределение этой последовательности помогает моделям анализировать продолжительные записи без агрессивного уменьшения количества кадров или пространственной детализации. AWS Neuron context parallelism overview иллюстрирует, как группы ускорителей могут обмениваться шардами KV для таких рабочих нагрузок с длинным контекстом.

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

Практическое использование и компромиссы#

Следующий минимальный пример использует экспериментальный API контекстного параллелизма PyTorch и scaled dot-product attention controls. Сохрани его как 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 предоставляет более простые облачные и локальные рабочие процессы для разметки датасетов, обучения, развертывания и мониторинга без необходимости настройки контекстного параллелизма вручную.