Симметричная память PyTorch

Зачем нужна симметричная память?

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

Чтобы свести к минимуму влияние на вычисления, также необходимо использовать движки копирования и сетевые интерфейсные карты (NIC) для выполнения обмена данными. Сетевые транспортные протоколы, такие как прямой доступ к памяти по сети (RDMA), повышают производительность, обеспечивая прямой, высокоскоростной обмен данными с низкой задержкой между процессорами и памятью. Такое разнообразие указывает на необходимость более детальных примитивов обмена данными, чем те, что предлагают современные высокоуровневые API коллективных операций. Они позволили бы разработчикам реализовывать алгоритмы, адаптированные под конкретные задачи, например коллективные операции с низкой задержкой, детальное перекрытие вычислений и обмена данными или пользовательские объединённые операции.

Кроме того, современные системы ИИ соединяют GPU высокоскоростными каналами (такими как NVLink, InfiniBand или RoCE), благодаря чему глобальная память GPU становится напрямую доступна другим GPU. Такие соединения дают программистам возможность рассматривать систему как один огромный GPU с большой доступной памятью вместо программирования отдельных «островков GPU».

В этом документе мы покажем, как с помощью PyTorch Symmetric Memory программировать современные системы GPU как «единый GPU» и выполнять детальный удалённый доступ.

Какие возможности открывает PyTorch Symmetric Memory?

PyTorch Symmetric Memory открывает три новые возможности:

• Настраиваемые шаблоны обмена данными: благодаря большей гибкости при написании ядер разработчики могут создавать собственные ядра, реализующие вычисления и обмен данными, адаптированные непосредственно под задачи приложения. Также можно будет легко добавлять поддержку новых типов данных и специальных вычислений, которые могут им требоваться, даже если они пока отсутствуют в стандартных библиотеках.
• Объединение вычислений и обмена данными внутри ядра: возможность инициировать обмен данными на устройстве позволяет разработчикам писать ядра, содержащие инструкции как для вычислений, так и для обмена данными, и объединять вычисления с перемещением данных с минимально возможной детализацией.
• Удалённый доступ с низкой задержкой: сетевые транспортные протоколы, такие как RDMA, повышают производительность симметричной памяти в сетевых средах, обеспечивая прямой, высокоскоростной обмен данными с низкой задержкой между процессорами и памятью. RDMA устраняет издержки, связанные с традиционным сетевым стеком и участием CPU. Кроме того, передача данных переносится с вычислительных ресурсов на NIC, освобождая вычислительные ресурсы для выполнения задач.

Далее мы покажем, как PyTorch Symmetric Memory (SymmMem) позволяет создавать приложения, использующие перечисленные возможности.

Пример «Hello World»

Модель программирования PyTorch SymmMem включает два ключевых элемента:

• создание симметричных тензоров
• создание ядер SymmMem

Для создания симметричных тензоров можно использовать пакет torch.distributed._symmetric_memory:

import torch.distributed._symmetric_memory as symm_mem

t = symm_mem.empty(128, device=torch.device("cuda", rank))
hdl = symm_mem.rendezvous(t, group)

Функция symm_mem.empty создаёт тензор, размещённый в симметричном выделении памяти. Функция rendezvous устанавливает связь с узлами группы и возвращает дескриптор симметричного выделения памяти. Дескриптор предоставляет методы для доступа к сведениям о симметричном выделении памяти, например к указателям на симметричный буфер на узлах-пирах, указателю multicast (если поддерживается) и сигнальным площадкам.

Функции empty и rendezvous необходимо вызывать на всех узлах группы в одном и том же порядке.

Затем для этих тензоров можно вызывать коллективные операции. Например, для выполнения одноэтапной операции all-reduce:

# Most SymmMem ops are under the torch.ops.symm_mem namespace
torch.ops.symm_mem.one_shot_all_reduce(t, "sum", group)

Обратите внимание, что torch.ops.symm_mem — это «пространство имён операций», а не модуль Python. Поэтому его нельзя импортировать с помощью import torch.ops.symm_mem, как нельзя импортировать отдельную операцию с помощью from torch.ops.symm_mem import one_shot_all_reduce. Операцию можно вызвать напрямую, как в примере выше.

Напишите собственное ядро

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

PyTorch Symmetric Memory предоставляет примитивы синхронизации, совместимые с CUDA Graph и работающие с сигнальной площадкой, сопутствующей каждому выделению симметричной памяти. Ядра, использующие симметричную память, можно писать как на CUDA, так и на Triton. Ниже приведён пример выделения симметричного тензора и обмена дескрипторами:

import torch.distributed._symmetric_memory as symm_mem

dist.init_process_group()
rank = dist.get_rank()

# Allocate a tensor
t = symm_mem.empty(4096, device=f"cuda:{rank}")
# Establish symmetric memory and obtain the handle
hdl = symm_mem.rendezvous(t, dist.group.WORLD)

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

hdl.buffer_ptrs
hdl.multicast_ptr
hdl.signal_pad_ptrs

К данным, на которые указывает buffer_ptrs, можно обращаться так же, как к обычным локальным данным; необходимые вычисления также можно выполнять привычными способами. Как и при работе с локальными данными, для повышения эффективности можно и нужно использовать векторизованный доступ.

Симметричная память особенно удобна для написания ядер на Triton. Если раньше Triton устранял препятствия для написания эффективного кода CUDA, то теперь в ядра Triton легко добавлять обмен данными. Приведённое ниже ядро демонстрирует написанную на Triton операцию all-reduce с низкой задержкой.

@triton.jit
def one_shot_all_reduce_kernel(
    buf_tuple,
    signal_pad_ptrs,
    output_ptr,
    numel: tl.constexpr,
    rank: tl.constexpr,
    world_size: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    ptx_utils.symm_mem_sync(
        signal_pad_ptrs, None, rank, world_size, hasSubsequenceMemAccess=True
    )

    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE

    while block_start < numel:
        offsets = block_start + tl.arange(0, BLOCK_SIZE)
        mask = offsets < numel
        acc = tl.zeros((BLOCK_SIZE,), dtype=tl.bfloat16)

        for i in tl.static_range(world_size):
            buffer_rank = buf_tuple[i]
            x = tl.load(buffer_rank + offsets, mask=mask)
            acc += x

        tl.store(output_ptr + offsets, acc, mask=mask)
        block_start += tl.num_programs(axis=0) * BLOCK_SIZE

    ptx_utils.symm_mem_sync(
        signal_pad_ptrs, None, rank, world_size, hasPreviousMemAccess=True
    )

Синхронизация в начале и конце приведённого выше ядра гарантирует, что все процессы видят согласованные данные. Основная часть ядра представляет собой обычный код Triton, который Triton оптимизирует в фоновом режиме, обеспечивая эффективный доступ к памяти за счёт векторизации и развёртывания циклов. Как и любые ядра Triton, его легко изменить, добавив вычисления или изменив алгоритм обмена данными. Дополнительные утилиты и примеры использования симметричной памяти для реализации распространённых шаблонов в Triton см. в https://github.com/meta-pytorch/kraken/blob/main/kraken.

Одностороннее чтение

Симметричная память также предоставляет небольшой API для одностороннего get, позволяющий копировать данные из симметричного выделения узла-пира в локальный тензор:

src = symm_mem.empty(1024, device=device)
hdl = symm_mem.rendezvous(src, group)

if dist.get_rank(group) == 0:
    dst = torch.empty((512,), device=device)
    # Copy the last 512 elements of the peer's allocation into dst.
    symm_mem.get(dst, hdl, peer=1, offset=512)

hdl — это дескриптор симметричной памяти, возвращённый функцией rendezvous; удалённый источник — это выделение узла-пира, соответствующее этому дескриптору. Число копируемых элементов определяется по dst, поэтому передайте представление (например, dst[:n]), чтобы заполнить только часть тензора; значение offset задаётся в элементах типа данных dst и по умолчанию равно 0. dst может быть обычным тензором CUDA или другим симметричным тензором; он должен находиться на том же устройстве, что и hdl, и занимать непрерывную область памяти. Копирование выполняется в текущем потоке CUDA.

Масштабирование

В больших языковых моделях эксперты распределяются более чем по 8 GPU, поэтому требуется возможность доступа между узлами. В этом помогают NIC с поддержкой RDMA. Кроме того, программные библиотеки, такие как NVSHMEM или rocSHMEM, скрывают различия между доступом внутри узла и между узлами, предоставляя примитивы, немного более высокого уровня, чем доступ по указателю, например put и get.

PyTorch предоставляет плагины NVSHMEM, расширяющие возможности ядер Triton для обмена данными между узлами. Как показано в приведённом ниже фрагменте кода, команду put для обмена данными между узлами можно инициировать внутри ядра.

import torch.distributed._symmetric_memory._nvshmem_triton as nvshmem
from torch.distributed._symmetric_memory._nvshmem_triton import requires_nvshmem

@requires_nvshmem
@triton.jit
def my_put_kernel(
    dest,
    src,
    nelems,
    pe,
):
    nvshmem.put(dest, src, nelems, pe)

Декоратор requires_nvshmem используется для указания, что ядру требуется библиотека устройств NVSHMEM в качестве внешней зависимости. При компиляции ядра Triton будет искать библиотеку устройств NVSHMEM в системных путях. Если она доступна, Triton включит необходимые инструкции устройства для использования функций NVSHMEM.

Использование пула памяти

Пул памяти позволяет PyTorch SymmMem кэшировать выделения памяти, для которых уже была установлена связь, экономя время при создании новых тензоров. Для удобства в PyTorch SymmMem добавлен API get_mem_pool, возвращающий пул симметричной памяти. Возвращённый MemPool можно использовать с менеджером контекста torch.cuda.use_mem_pool. В приведённом ниже примере тензор x будет создан в симметричной памяти:

    import torch.distributed._symmetric_memory as symm_mem

    mempool = symm_mem.get_mem_pool(device)

    with torch.cuda.use_mem_pool(mempool):
        x = torch.arange(128, device=device, dtype=torch.float32)

    torch.ops.symm_mem.one_shot_all_reduce(x, "sum", group_name)

Аналогичным образом можно поместить вычислительную операцию в контекст MemPool — результирующий тензор также будет создан в симметричной памяти.

    dim = 1024
    w = torch.ones(dim, dim, device=device)
    x = torch.ones(1, dim, device=device)

    mempool = symm_mem.get_mem_pool(device)
    with torch.cuda.use_mem_pool(mempool):
        # y will be in symmetric memory
        y = torch.mm(x, w)

Начиная с torch 2.11, бэкенды CUDA и NVSHMEM поддерживают MemPool. Поддержка MemPool в бэкенде NCCL находится в разработке.

Симметричные ядра NCCL

В NCCL 2.27+ добавлено семейство ядер устройств, внутреннее название которых — «SymK». Они специально разработаны для симметричных буферов, зарегистрированных как окна. Поскольку каждому рангу заранее известны адреса буферов всех узлов-пиров, эти ядра обходятся без универсального механизма proxy/ring и вместо этого используют варианты LL (низкая задержка), multimem/NVLS и TMA. NCCL выбирает вариант для каждого вызова в зависимости от размера сообщения, поэтому одна и та же операция dist.all_reduce использует ядро, оптимизированное для задержки, для небольших сообщений и ядро, оптимизированное для пропускной способности, для больших.

Симметричные ядра запускаются через стандартный API коллективных операций — dist.all_reduce, dist.all_gather_into_tensor, dist.reduce_scatter_tensor — без изменения места вызова. Важно, чтобы буферы были зарегистрированы в NCCL как симметричные окна. Это можно сделать двумя способами.

import torch
import torch.distributed as dist

device = torch.device("cuda", rank)

# `device_id` eagerly initializes the NCCL communicator. `register_mem_pool`
# requires a communicator that already exists, and raises otherwise.
dist.init_process_group(backend="nccl", device_id=device)
pg = dist.group.WORLD

backend = dist.get_backend_impl(pg, device)

# A MemPool backed by `ncclMemAlloc` / `ncclMemFree`.
pool = torch.cuda.MemPool(backend.mem_allocator)

# `symm=True` registers each segment with `ncclCommWindowRegister` using
# `NCCL_WIN_COLL_SYMMETRIC`, which is what makes the symmetric kernels
# eligible. The default `symm=False` performs ordinary user-buffer
# registration, which does not.
backend.register_mem_pool(pool, symm=True)

with torch.cuda.use_mem_pool(pool):
    x = torch.ones(1024 * 1024, dtype=torch.bfloat16, device=device)

# Dispatches to a NCCL symmetric kernel.
dist.all_reduce(x, op=dist.ReduceOp.SUM)

# De-register before the pool is torn down.
backend.deregister_mem_pool(pool)

import torch.distributed as dist
import torch.distributed._symmetric_memory as symm_mem

symm_mem.set_backend("NCCL")

x = symm_mem.empty(1024 * 1024, dtype=torch.bfloat16, device=device)
symm_mem.rendezvous(x, group=dist.group.WORLD.group_name)

dist.all_reduce(x, op=dist.ReduceOp.SUM)

NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=TUNING python train.py

AllReduce [Symmetric]: 2097152 Bytes -> Kernel AllReduce_RSxLDMC_AGxSTMC nchannels 16 nthreads 512 nWorks 1

Коллективные операции с использованием движка копирования

Коллективные операции с использованием движка копирования (CE) — это оптимизация коллективных операций NCCL, при которой перемещение данных выполняют движки копирования GPU (движки DMA), а не потоковые мультипроцессоры CUDA (SM). Это освобождает SM для вычислений и позволяет эффективнее перекрывать обмен данными и вычисления при распределённом обучении.

Чтобы использовать коллективные операции CE, необходимо:

1. Настроить группу процессов NCCL с политикой zero-CTA
2. Настроить симметричную память с бэкендом NCCL
3. Выделять тензоры с помощью симметричной памяти
4. Зарегистрировать тензоры в симметричной памяти посредством установления связи

После настройки стандартные функции коллективных операций, такие как all_gather_single() и all_to_all_single(), будут автоматически использовать движки копирования при работе с тензорами в симметричной памяти.

Пример

import torch
import torch.distributed as dist
import torch.distributed._symmetric_memory as symm_mem

# Initialize process group with zero-CTA policy for CE collectives
opts = dist.ProcessGroupNCCL.Options()
opts.config.cta_policy = dist.ProcessGroupNCCL.NCCL_CTA_POLICY_ZERO
device = torch.device("cuda", rank)
dist.init_process_group(backend="nccl", pg_options=opts, device_id=device)

# Set up symmetric memory with NCCL backend
symm_mem.set_backend("NCCL")
group_name = dist.group.WORLD.group_name

# Allocate tensors using symmetric memory
numel = 1024 * 1024
inp = symm_mem.empty(numel, device=device)
out = symm_mem.empty(numel * world_size, device=device)

# Register tensors for symmetric memory operations
symm_mem.rendezvous(inp, group=group_name)
symm_mem.rendezvous(out, group=group_name)

# Perform collective operation using copy engines
# This now runs on DMA engines instead of SMs
work = dist.all_gather_single(out, inp, async_op=True)
work.wait()

Преимущества

• Перенос нагрузки с SM: обмен данными выполняется на движках копирования, освобождая SM для вычислений
• Более эффективное перекрытие: позволяет эффективнее перекрывать вычисления и обмен данными
• Прозрачный API: используется тот же API коллективных операций, но с тензорами в симметричной памяти

Требования и ограничения

• Требуется NCCL версии 2.28 или новее
• На GPU должен быть включён доступ peer-to-peer (P2P)
• Тензоры должны быть выделены с помощью torch.distributed._symmetric_memory() и участвовать в установлении связи
• Группа процессов NCCL должна быть настроена с параметром NCCL_CTA_POLICY_ZERO или для переменной среды NCCL_CTA_POLICY должно быть задано значение 2
• В NCCL 2.28 коллективные операции CE не могут выполняться в потоке по умолчанию, поэтому для активации внутреннего потока ProcessGroupNCCL необходимо использовать флаг async_op=True либо создать отдельный поток самостоятельно

Редукция с повышенной точностью

Если тензоры выделены в симметричной памяти, реализация симметричных ядер NCCL включает внутреннюю редукцию с повышенной точностью. Например, при входных данных BF16 NCCL автоматически накапливает значения внутри в FP32, а затем преобразует результат в BF16 (BF16 на входе → накопление в FP32 → BF16 на выходе). Это повышает точность редукции без изменения вызова коллективной операции.

Область применения

• Поддерживаемые операции: только reduce_scatter и all_reduce
• Область действия: внутри домена NVLink начиная с torch 2.9 (NCCL 2.27); NVLink + сеть для reduce_scatter начиная с torch 2.11 (NCCL 2.29)
• Точность: BF16/FP16 на входе → внутреннее накопление в FP32 → BF16/FP16 на выходе

Пример

import torch.distributed as dist
import torch.distributed._symmetric_memory as symm_mem

# Allocate tensors using NCCL symmetric memory
symm_mem.set_backend("NCCL")
inp = symm_mem.empty(1024, 1024, device=device, dtype=torch.bfloat16)
symm_mem.rendezvous(inp, group_name)

# reduce_scatter and all_reduce on symmetric memory tensors
# automatically benefit from FP32 internal accumulation
dist.all_reduce(inp)

Установление связи при масштабировании

По умолчанию rendezvous обменивается метаданными через TCPStore. Каждый ранг группы симметричной памяти выполняет одну операцию записи в хранилище и N-1 операций чтения (где N — размер группы, обычно 8–72 для доменов NVLink). При большом размере мира TCPStore (пропускная способность около 200 тыс. запросов в секунду) становится узким местом: например, для групп NVLink из 72 рангов и общего числа рангов 10 тыс. однократное установление связи через TCPStore занимает около 3,6 с; при 100 тыс. рангов время увеличивается примерно до 36 с.

Чтобы вместо этого использовать allgather NCCL группы процессов, задайте use_pg_for_symm_mem_rendezvous в параметрах группы процессов:

opts = dist.ProcessGroupNCCL.Options()
opts.use_pg_for_symm_mem_rendezvous = True
pg = dist.new_group(ranks, pg_options=opts)

t = symm_mem.empty(size, device=device)
hdl = symm_mem.rendezvous(t, group=pg)

Если группа процессов используется только для симметричной памяти и впоследствии не будет использоваться для обычных коллективных операций (например, группа параллелизма экспертов), после установления связи можно освободить коммуникатор NCCL с помощью abort(). Дескриптор симметричной памяти останется пригодным к использованию, поскольку он зависит только от отображённой памяти, а не от коммуникатора:

opts = dist.ProcessGroupNCCL.Options()
opts.use_pg_for_symm_mem_rendezvous = True
ep_pg = dist.new_group(ep_ranks, pg_options=opts)

t = symm_mem.empty(size, device=device)
hdl = symm_mem.rendezvous(t, group=ep_pg)

# Release the NCCL communicator since ep_pg won't be used for collectives.
# The symm_mem handle is still usable — it only needs the mapped memory.
ep_pg.abort()

Справочник API

torch.distributed._symmetric_memory.empty(*size: _int, dtype: _dtype | None = None, device: _device | None = None) → Tensor [исходный код]

torch.distributed._symmetric_memory.empty(size:Sequence[_int], *, dtype:_dtype|None=None, device:_device|None=None) → Tensor

Аналог torch.empty(). Возвращённый тензор можно использовать с помощью torch._distributed._symmetric_memory.rendezvous() для создания тензора с симметричной памятью между участвующими процессами.

Примечание

Это синхронное выделение памяти на хосте. Вместе с rendezvous() эта операция предназначена для выполнения при инициализации: выделите тензор с симметричной памятью один раз и используйте его повторно, вместо выделения памяти на критических путях выполнения.

Параметры:

size (int...) – последовательность целых чисел, задающая форму выходного тензора. Можно передать переменное число аргументов или коллекцию, например список или кортеж.

Именованные аргументы:

• dtype (torch.dtype, необязательно) – требуемый тип данных возвращаемого тензора. По умолчанию: если None, используется глобальное значение по умолчанию (см. torch.set_default_dtype()).
• device (torch.device, необязательно) – требуемое устройство возвращаемого тензора. По умолчанию: если None, используется текущее устройство для типа тензора по умолчанию (см. torch.set_default_device()). device будет CPU для типов тензоров CPU, текущим устройством CUDA для типов тензоров CUDA и текущим устройством XPU для типов тензоров XPU.

torch.distributed._symmetric_memory.rendezvous(tensor, group) → _SymmetricMemory [исходный код]

Создаёт тензор с симметричной памятью между участвующими процессами. Это коллективная операция.

Примечание

Это операция инициализации, блокирующая хост: при первой встрече с тензором выполняется обмен дескрипторами и сопоставление между процессами, а также синхронизация хоста с устройством. Её нельзя поставить в очередь на поток CUDA или захватить в граф CUDA. Выполните rendezvous для буфера один раз и повторно используйте возвращённый дескриптор, не вызывая эту операцию на критических путях выполнения; последующие вызовы для того же тензора возвращают кэшированный дескриптор.

Параметры:

• tensor (torch.Tensor) – локальный тензор, используемый для создания тензора с симметричной памятью. Он должен быть выделен с помощью torch._distributed._symmetric_memory.empty(). Форма, тип данных и тип устройства должны быть одинаковыми во всех участвующих процессах.
• group (Union[str, torch.distributed.ProcessGroup]) – группа, определяющая участвующие процессы. Это может быть имя группы или объект группы процессов.

Тип возвращаемого значения:

_SymmetricMemory

torch.distributed._symmetric_memory.get(dst, hdl, peer, offset=0) [исходный код]

Копирует dst.numel() элементов, начиная с offset, из симметричного выделения памяти процесса peer в локальный dst, используя односторонний доступ к симметричной памяти.

hdl — это дескриптор симметричной памяти, возвращённый функцией torch.distributed._symmetric_memory.rendezvous(); удалённым источником служит выделенная память процесса peer, связанная с этим дескриптором. Число копируемых элементов определяется по dst; чтобы заполнить только часть тензора, передайте представление (например, dst[:n]). offset задаётся в элементах типа данных dst. dst может быть обычным тензором CUDA или тензором с симметричной памятью; он должен находиться на том же устройстве, что и hdl, и использовать непрерывную область памяти. Копирование выполняется в текущем потоке CUDA.

Параметры:

• dst (Tensor) – локальный тензор назначения.
• hdl (SymmetricMemory) – дескриптор, удалённое выделение памяти узла-соседа которого является источником.
• peer (int) – ранг, откуда выполняется копирование.
• offset (int, optional) – смещение в элементах в выделенной памяти узла-соседа, с которого начинается чтение. По умолчанию — 0.

torch.distributed._symmetric_memory.is_nvshmem_available() → bool [исходный код]

Проверяет, доступны ли NVSHMEM (CUDA) или rocSHMEM (ROCm) в текущей сборке и могут ли они использоваться во время выполнения. В ROCm версия rocSHMEM VERSION должна быть не ниже 3.3.0 (см. rocshmem/rocshmem.hpp).

Тип возвращаемого значения:

bool

torch.distributed._symmetric_memory.set_backend(name) [исходный код]

Задаёт бэкенд для выделения симметричной памяти. Это глобальная настройка, влияющая на все последующие вызовы torch._distributed._symmetric_memory.empty(). Обратите внимание: бэкенд нельзя изменить после выделения тензора с симметричной памятью.

Параметры:

backend (str) – бэкенд для выделения симметричной памяти. На данный момент поддерживаются только “NVSHMEM”, “CUDA”, “NCCL”.

torch.distributed._symmetric_memory.get_backend(device) [исходный код]

Возвращает бэкенд для выделения симметричной памяти для указанного устройства. Если бэкенд не найден, возвращает None.

Параметры:

device (torch.device или str) – устройство, для которого требуется получить бэкенд.

Тип возвращаемого значения:

str | None

torch.distributed._symmetric_memory.get_mem_pool(device) [исходный код]

Возвращает пул симметричной памяти для указанного устройства. Если пул не найден, создаёт новый.

Выделение тензоров из этого пула должно быть симметричным для всех рангов. Выделенные тензоры можно использовать в симметричных операциях, например в операциях, определённых в torch.ops.symm_mem.

Параметры:

device (torch.device или str) – устройство, для которого требуется получить пул симметричной памяти.

Возвращает:

пул симметричной памяти для указанного устройства.

Тип возвращаемого значения:

torch.cuda.MemPool

Пример:

>>> pool = torch.distributed._symmetric_memory.get_mem_pool("cuda:0")
>>> with torch.cuda.use_mem_pool(pool):
>>>     tensor = torch.randn(1000, device="cuda:0")
>>> tensor = torch.ops.symm_mem.one_shot_all_reduce(tensor, "sum", group_name)

torch.distributed._symmetric_memory.is_symm_mem_tensor(tensor) → bool [исходный код]

Возвращает True, если tensor был выделен в симметричной памяти (то есть с помощью torch.distributed._symmetric_memory.empty() или _SymmetricMemory.empty_strided_p2p()).

Это неколлективная проверка со сложностью O(1).

Параметры:

tensor (torch.Tensor) – тензор для проверки.

Тип возвращаемого значения:

bool

torch.distributed._symmetric_memory.set_signal_pad_size(size) [исходный код]

Задаёт размер сигнальной области для будущих выделений симметричной памяти.

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

Предупреждение

Эту функцию необходимо вызвать до любых выделений симметричной памяти. После выполнения выделений изменить размер нельзя.

Параметры:

size (int) – размер сигнальной области в байтах. Размер должен быть пропорционален числу запущенных блоков и размеру мира.

Пример:

>>> # Set a larger signal pad size before any allocations
>>> torch.distributed._symmetric_memory.set_signal_pad_size(1024 * 1024)  # 1MB

torch.distributed._symmetric_memory.get_signal_pad_size() [исходный код]

Возвращает текущий размер сигнальной области для выделений симметричной памяти.

Возвращает заданный пользователем размер, если он был установлен с помощью set_signal_pad_size(); в противном случае возвращает размер по умолчанию.

Возвращает:

размер сигнальной области в байтах.

Тип возвращаемого значения:

int

Пример:

>>> size = torch.distributed._symmetric_memory.get_signal_pad_size()
>>> print(f"Signal pad size: {size} bytes")

Справочник по операциям

torch.distributed._symmetric_memory.reduce_scatter_offset(input, out, group, *, dim, offsets, dst_ranks, red_op='sum') → None [исходный код]

Одновременно выполняет редукцию N блоков двумерного тензора input из буфера симметричной памяти, направляя каждый блок определённому рангу-получателю. Только dst_ranks[i] записывает результат редукции для блока i; результат записывается в непрерывный выходной тензор той же формы, что и блок i.

Аргумент dim задаёт измерение, по которому выполняется шардинг:

• dim=0 (шардинг по строкам): блок i охватывает input[offsets[i-1] : offsets[i], :]. Каждый out[j] имеет форму (size_j, input.size(1)).
• dim=1 (шардинг по столбцам): блок i охватывает input[:, offsets[i-1] : offsets[i]]. Каждый out[j] имеет форму (input.size(0), size_j).

Блоки задаются с помощью offsets — включительной префиксной суммы размеров блоков вдоль dim (по соглашению первый блок начинается с индекса 0). Смещения блоков могут быть равномерными или неравномерными; в случае неравномерных смещений должно выполняться следующее условие: для каждого j принадлежащий ему блок с номером j должен иметь одинаковый размер на всех рангах (чтобы out[j] имел одинаковую форму); значения j могут различаться.

Параметры:

• input (Tensor) – двумерный тензор, выделенный в симметричной памяти (внутреннее измерение должно быть непрерывным).
• out (list[Tensor]) – выходные тензоры для блоков, принадлежащих этому рангу. Длина должна равняться числу блоков, принадлежащих этому рангу (то есть количеству i, для которых dst_ranks[i] == my_rank). Каждый out[j] должен быть непрерывным и иметь тот же тип данных, что и input.
• group (str) – имя ProcessGroup, в котором выполняется операция.
• dim (int) – измерение, вдоль которого определяются блоки (0 или 1).
• offsets (list[int] | None) – включительная префиксная сумма размеров блоков вдоль dim, длина N. Если не задано, input.size(dim) делится на блоки одинакового размера в соответствии с размером group.
• dst_ranks (list[int] | None) – ранг-получатель для каждого блока. Если не задано, блоки распределяются по рангам по кругу.
• red_op (str) – операция редукции; в настоящее время поддерживается только 'sum'.

Пример:

>>> # Each rank holds a Grouped GEMM gradient buffer in symmetric memory.
>>> # The buffer has W experts laid out as equal column blocks; each expert
>>> # is reduced to a specific rank (dst_ranks[i] == i % world_size).
>>> buf = symm_mem.empty(H, W * C, dtype=torch.bfloat16, device="cuda")
>>> symm_mem.rendezvous(buf, group=group_name)
>>> offsets = [i * C for i in range(1, W + 1)]  # inclusive prefix-sum
>>> dst_ranks = [i % world_size for i in range(W)]
>>> n_owned = sum(r == rank for r in dst_ranks)
>>> out = [torch.empty(H, C, dtype=torch.bfloat16, device="cuda") for _ in range(n_owned)]
>>> symm_mem.reduce_scatter_offset(buf, out, group_name, dim=1, offsets=offsets, dst_ranks=dst_ranks)

torch.ops.symm_mem.multimem_all_reduce_(input: Tensor, reduce_op: str, group_name: str) → Tensor

Выполняет операцию мультимемного all-reduce над входным тензором. Для этой операции требуется аппаратная поддержка мультимемных операций. На графических процессорах NVIDIA требуется NVLink SHARP.

Предупреждение

Все коллективные операции symm_mem для заданной группы необходимо запускать в одном потоке CUDA. Ядра синхронизируют ранги с помощью общей сигнальной области, индексируемой по идентификатору блока, без изоляции по потокам; одновременный запуск из разных потоков для одной группы приведёт к взаимоблокировке. Чтобы использовать коллективные операции symm_mem из нескольких потоков, сериализуйте их в одном выделенном потоке с помощью stream.wait_stream() / current_stream.wait_stream().

Параметры:

• input (Tensor) – входной тензор для выполнения all-reduce. Должен быть симметричным.
• reduce_op (str) – выполняемая операция редукции. В настоящее время поддерживается только “sum”.
• group_name (str) – имя группы, в которой выполняется all-reduce.

torch.ops.symm_mem.multimem_all_gather_out(input: Tensor, group_name: str, out: Tensor) → Tensor

Выполняет операцию мультимемного all-gather над входным тензором. Для этой операции требуется аппаратная поддержка мультимемных операций. На графических процессорах NVIDIA требуется NVLink SHARP.

Предупреждение

Все коллективные операции symm_mem для заданной группы необходимо запускать в одном потоке CUDA. Подробности см. в описании multimem_all_reduce_().

Параметры:

• input (Tensor) – входной тензор для выполнения all-gather.
• group_name (str) – имя группы, в которой выполняется all-gather.
• out (Tensor) – выходной тензор для хранения результата операции all-gather. Должен быть симметричным.

torch.ops.symm_mem.one_shot_all_reduce(input: Tensor, reduce_op: str, group_name: str) → Tensor

Выполняет одноэтапную операцию all-reduce над входным тензором.

Предупреждение

Все коллективные операции symm_mem для заданной группы необходимо запускать в одном потоке CUDA. Подробности см. в описании multimem_all_reduce_().

Параметры:

• input (Tensor) – входной тензор для выполнения all-reduce. Должен быть симметричным.
• reduce_op (str) – выполняемая операция редукции. В настоящее время поддерживается только “sum”.
• group_name (str) – имя группы, в которой выполняется all-reduce.

torch.ops.symm_mem.one_shot_all_reduce_out(input: Tensor, reduce_op: str, group_name: str, out: Tensor) → Tensor

Выполняет одноэтапную операцию all-reduce на основе входного тензора и записывает результат в выходной тензор.

Предупреждение

Все коллективные операции symm_mem для заданной группы необходимо запускать в одном потоке CUDA. Подробности см. в описании multimem_all_reduce_().

Параметры:

• input (Tensor) – входной тензор для выполнения all-reduce. Должен быть симметричным.
• reduce_op (str) – выполняемая операция редукции. В настоящее время поддерживается только “sum”.
• group_name (str) – имя группы, в которой выполняется all-reduce.
• out (Tensor) – выходной тензор для хранения результата операции all-reduce. Может быть обычным тензором.

torch.ops.symm_mem.two_shot_all_reduce_(input: Tensor, reduce_op: str, group_name: str) → Tensor

Выполняет двухэтапную операцию all-reduce над входным тензором.

Предупреждение

Все коллективные операции symm_mem для заданной группы необходимо запускать в одном потоке CUDA. Подробности см. в описании multimem_all_reduce_().

Параметры:

• input (Tensor) – входной тензор для выполнения all-reduce. Должен быть симметричным.
• reduce_op (str) – выполняемая операция редукции. В настоящее время поддерживается только “sum”.
• group_name (str) – имя группы, в которой выполняется all-reduce.

torch.ops.symm_mem.all_to_all_vdev(input: Tensor, out: Tensor, in_splits: Tensor, out_splits_offsets: Tensor, group_name: str) → None

Выполняет операцию all-to-all-v с помощью NVSHMEM; информация о разбиении задаётся на устройстве.

Параметры:

• input (Tensor) – входной тензор для выполнения all-to-all. Должен быть симметричным.
• out (Tensor) – выходной тензор для хранения результата операции all-to-all. Должен быть симметричным.
• in_splits (Tensor) – тензор, содержащий разбиение данных для отправки каждому узлу. Должен быть симметричным и иметь размер (group_size,). Разбиение задаётся в единицах элементов первого измерения.
• out_splits_offsets (Tensor) – тензор, содержащий разбиения и смещения полученных от каждого узла данных. Должен быть симметричным и иметь размер (2, group_size). Строки содержат (в указанном порядке): выходные разбиения и выходные смещения.
• group_name (str) – имя группы, в которой выполняется all-to-all.

torch.ops.symm_mem.all_to_all_vdev_2d(input: Tensor, out: Tensor, in_splits: Tensor, out_splits_offsets: Tensor, group_name: str[, major_align: int = None]) → None

Выполняет двумерную операцию all-to-all-v с помощью NVSHMEM; информация о разбиении задаётся на устройстве. В моделях смеси экспертов эту операцию можно использовать для диспетчеризации токенов.

Параметры:

• input (Tensor) – входной тензор для выполнения all-to-all. Должен быть симметричным.
• out (Tensor) – выходной тензор для хранения результата операции all-to-all. Должен быть симметричным.
• in_splits (Tensor) – тензор, содержащий разбиения данных для отправки каждому эксперту. Должен быть симметричным и иметь размер (group_size * ne), где ne — число экспертов на ранг. Разбиение задаётся в единицах элементов первого измерения.
• out_splits_offsets (Tensor) – тензор, содержащий разбиения и смещения полученных от каждого узла данных. Должен быть симметричным и иметь размер (2, group_size * ne). Строки содержат (в указанном порядке): выходные разбиения и выходные смещения.
• group_name (str) – имя группы, в которой выполняется all-to-all.
• major_align (int) – необязательное выравнивание главного измерения выходного фрагмента для каждого эксперта. Если не задано, предполагается, что выравнивание равно 1. Все корректировки выравнивания отражаются в выходных смещениях.

Ниже показан пример перемешивания 2D AllToAllv (world_size = 2, ne = 2, общее число экспертов = 4):

Source: |       Rank 0      |       Rank 1      |
        | c0 | c1 | c2 | c3 | d0 | d1 | d2 | d3 |

Dest  : |       Rank 0      |       Rank 1      |
        | c0 | d0 | c1 | d1 | c2 | d2 | c3 | d3 |

где каждый c_i / d_i — это срез тензора input, предназначенный для эксперта i; его длина задаётся входным разбиением. Иными словами, перемешивание 2D AllToAllv выполняет транспонирование: на входе данные упорядочены по рангам, а на выходе — по экспертам.

Если major_align не равно 1, выходные смещения c1, c2, c3 будут выровнены вверх до этого значения. Например, если длина c0 равна 5, а длина d0 — 7 (всего 12), и major_align задано равным 16, выходное смещение c1 будет равно 16. То же относится к c2 и c3. Это значение не влияет на смещение младшего измерения, то есть d0, d1, d2 и d3. Примечание: поскольку cutlass не поддерживает пустые сегменты, если выровненная длина равна 0, мы задаём её равной major_align. См. pytorch/pytorch#152668.

torch.ops.symm_mem.all_to_all_vdev_2d_offset(Tensor input, Tensor out, Tensor in_splits_offsets, Tensor out_splits_offsets, str group_name) → None

Выполняет перемешивание 2D AllToAllv; информация о входных разбиениях и смещениях задаётся на устройстве. Входные смещения не обязательно должны быть точными префиксными суммами входных разбиений, то есть между фрагментами допускаются отступы. Однако эти отступы не будут передаваться узлам-получателям.

В моделях смеси экспертов эту операцию можно использовать для объединения токенов, обработанных экспертами на параллельных рангах. Эту операцию можно рассматривать как обратную операцию all_to_all_vdev_2d (которая перемещает токены к экспертам).

Параметры:

• input (Tensor) – входной тензор для выполнения all-to-all. Должен быть симметричным.
• out (Tensor) – выходной тензор для хранения результата операции all-to-all. Должен быть симметричным.
• in_splits_offsets (Tensor) – тензор, содержащий разбиения и смещения данных для отправки каждому эксперту. Должен быть симметричным и иметь размер (2, group_size * ne), где ne — число экспертов. Строки содержат (в указанном порядке): входные разбиения и входные смещения. Разбиение задаётся в единицах элементов первого измерения.
• out_splits_offsets (Tensor) – тензор, содержащий разбиения и смещения полученных от каждого узла данных. Должен быть симметричным и иметь размер (2, group_size * ne). Строки содержат (в указанном порядке): выходные разбиения и выходные смещения.
• group_name (str) – имя группы, в которой выполняется all-to-all.

torch.ops.symm_mem.tile_reduce(in_tile: Tensor, out_tile: Tensor, root: int, group_name: str[, reduce_op: str = 'sum']) → None

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

Параметры:

• in_tile (Tensor) – входной двумерный тензор для редукции. Должен быть выделен в симметричной памяти.
• out_tile (Tensor) – выходной двумерный тензор для хранения результата редукции. Должен быть симметричным и иметь ту же форму, тип данных и устройство, что и in_tile.
• root (int) – ранг процесса в заданной группе, который получит результат редукции.
• group_name (str) – имя группы процессов симметричной памяти, в которой выполняется редукция.
• reduce_op (str) – выполняемая операция редукции. В настоящее время поддерживается только "sum". По умолчанию используется "sum".

Эта функция выполняет редукцию тензоров in_tile всех участников группы и записывает результат в out_tile на корневом ранге. В операции должны участвовать все ранги; они должны передавать одинаковые значения group_name и тензоры одной формы.

Пример:

>>> 
>>> # Reduce the bottom-right quadrant of a tensor
>>> tile_size = full_size // 2
>>> full_inp = symm_mem.empty(full_size, full_size)
>>> full_out = symm_mem.empty(full_size, full_size)
>>> s = slice(tile_size, 2 * tile_size)
>>> in_tile = full_inp[s, s]
>>> out_tile = full_out[s, s]
>>> torch.ops.symm_mem.tile_reduce(in_tile, out_tile, root=0, group_name)

torch.ops.symm_mem.multi_root_tile_reduce(in_tiles: list[Tensor], out_tile: Tensor, roots: list[int], group_name: str, [reduce_op: str = 'sum']) → None

Одновременно выполняет редукцию нескольких фрагментов, причём каждый фрагмент сводится на отдельном корневом ранге.

: param list[Tensor] in_tiles: Список входных тензоров. : param Tensor out_tile: Выходной тензор для хранения результата редукции фрагмента. : param list[int] roots: Список корневых рангов, каждый из которых соответствует входному фрагменту из in_tiles в том же порядке. Один ранг не может быть корневым более одного раза. : param str group_name: Имя группы для коллективной операции. : param str reduce_op: Выполняемая операция редукции. В настоящее время поддерживается только “sum”.

Пример:

>>> 
>>> # Reduce four quadrants of a tensor, each to a different root
>>> tile_size = full_size // 2
>>> full_inp = symm_mem.empty(full_size, full_size)
>>> s0 = slice(0, tile_size)
>>> s1 = slice(tile_size, 2 * tile_size)
>>> in_tiles = [ full_inp[s0, s0], full_inp[s0, s1], full_inp[s1, s0], full_inp[s1, s1] ]
>>> out_tile = symm_mem.empty(tile_size, tile_size)
>>> roots = [0, 1, 2, 3]
>>> torch.ops.symm_mem.multi_root_tile_reduce(in_tiles, out_tile, roots, group_name)