Хуки коммуникации DDP

Как использовать хук коммуникации?

Чтобы использовать хук коммуникации, пользователю достаточно зарегистрировать хук в модели DDP перед циклом обучения, как показано ниже.

torch.nn.parallel.DistributedDataParallel.register_comm_hook()

С чем работает хук коммуникации?

Хук коммуникации предоставляет гибкий способ выполнять all-reduce градиентов. Поэтому он работает главным образом с градиентами каждой реплики до операции all-reduce, которые распределяются по корзинам, чтобы увеличить перекрытие коммуникации и вычислений. В частности, torch.distributed.GradBucket представляет собой корзину тензоров градиентов для операции all-reduce.

class torch.distributed.GradBucket

Этот класс передаёт в хук коммуникации DDP одномерный тензор градиентов (возвращаемый методом buffer()). Этот тензор можно дополнительно разделить на список тензоров отдельных параметров в этой корзине (возвращаемый методом get_per_parameter_tensors()), чтобы применять поуровневые операции.

torch.distributed.GradBucket.index(self: torch._C._distributed_c10d.GradBucket) → int

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

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

Возвращает:

Индекс корзины, в которой хранятся градиенты нескольких смежных слоёв. Все градиенты распределяются по корзинам.

torch.distributed.GradBucket.buffer(self: torch._C._distributed_c10d.GradBucket) → torch.Tensor

Возвращает:

Одномерный сплющенный буфер torch.Tensor, который можно дополнительно разделить на список тензоров отдельных параметров в этой корзине.

torch.distributed.GradBucket.gradients(self: torch._C._distributed_c10d.GradBucket) → list[torch.Tensor]

Возвращает:

Список torch.Tensor. Каждый тензор в списке соответствует градиенту.

torch.distributed.GradBucket.is_last(self: torch._C._distributed_c10d.GradBucket) → bool

Возвращает:

Указывает, является ли эта корзина последней корзиной, для которой выполняется all-reduce в данной итерации. Это также означает, что эта корзина соответствует первым нескольким слоям прямого прохода.

torch.distributed.GradBucket.set_buffer(self: torch._C._distributed_c10d.GradBucket, buffer: torch.Tensor) → None

Заменяет тензор в корзине входным буфером тензора.

torch.distributed.GradBucket.parameters(self: torch._C._distributed_c10d.GradBucket) → list[torch.Tensor]

Возвращает:

Список torch.Tensor. Каждый тензор в списке соответствует параметру модели.

Регистрация хука коммуникации

torch.distributed.algorithms.ddp_comm_hooks.register_ddp_comm_hook(comm_hook_type, model, state=None) [исходный код]

Регистрирует ddp_comm_hooks в модели DDP.

Регистрирует хуки torch.distributed.algorithms.ddp_comm_hooks в модели DDP. Пользователь может указать тип хука в качестве типа перечисления DDPCommHookType с помощью входного аргумента comm_hook_type. Модели передаётся входное состояние. Используются реализации хуков коммуникации на Python.

Пример::

>>> register_ddp_comm_hook(DDPCommHookType.FP16_COMPRESS, model, state)

torch.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook.post_localSGD_hook(state, bucket) [исходный код]

Запускает алгоритм post-localSGD.

Этот хук коммуникации DDP используется для запуска алгоритма post-localSGD совместно с компонентом усреднения модели (например, PeriodicModelAverager), который выполняется после шага оптимизатора.

Параметры:

• state (PostLocalSGDState) – информация о состоянии для запуска post-localSGD. Пользователям в основном нужно настроить start_localSGD_iter, чтобы определить, когда начинать локальный SGD.
• bucket (dist.GradBucket) – корзина, в которой хранится одномерный сплющенный тензор градиентов, объединяющий несколько тензоров отдельных переменных. Обратите внимание: поскольку хук коммуникации DDP поддерживает только режим «один процесс — одно устройство», в этой корзине хранится ровно один тензор.

Возвращает:

Обработчик Future для коммуникации, который обновляет градиенты на месте.

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

Future[Tensor]

Пример::

>>> state = PostLocalSGDState(process_group=process_group, subgroup=subgroup,
                          start_localSGD_iter=10)
>>> ddp_model.register_comm_hook(state, post_localSGD_hook)
>>> # Also need to establish a model averaging module and run model averaging after ``optimizer.step()``.
>>> # Please refer to the examples in ``torch.distributed.algorithms.model_averaging.averagers`` module.

torch.distributed.algorithms.ddp_comm_hooks.ddp_zero_hook.hook_with_zero_step(hook, ddp, zero, shard_buckets=False) [исходный код]

Изменяет hook, чтобы перекрыть шаг оптимизатора ZeroRedundancyOptimizer с обратным проходом DistributedDataParallel.

Этот подход перекрывает вычисления и коммуникацию оптимизатора с коммуникацией обратного прохода. В частности, вычисления обратного прохода выполняются последовательно, а затем начинаются вычисления оптимизатора, перекрываясь с ожидающей коммуникацией обратного прохода (то есть операциями all-reduce) и, возможно, с другой коммуникацией оптимизатора (то есть широковещательными рассылками). Вычисление шага оптимизатора начинается после завершения вычисления последней корзины градиентов.

Этот подход может быть предпочтительнее hook_with_zero_step_interleaved(), если коммуникация относительно медленна по сравнению с вычислениями.

Параметры:

• hook (Callable[[Any, dist.GradBucket], torch.futures.Future]) – хук, который нужно изменить.
• ddp (DistributedDataParallel) – экземпляр DistributedDataParallel для использования.
• zero (ZeroRedundancyOptimizer) – экземпляр ZeroRedundancyOptimizer для использования.
• shard_buckets (bool) – если True, назначение каждой корзины DistributedDataParallel разделяется между несколькими экземплярами ZeroRedundancyOptimizer (то есть несколькими рангами), чтобы добиться приблизительного равномерного распределения; если False, каждая корзина целиком назначается одному экземпляру ZeroRedundancyOptimizer (то есть одному рангу).

Возвращает:

Изменённый хук.

Вызывает исключения:

• ValueError – если zero создан с параметром overlap_with_ddp=False.
• RuntimeError – при использовании любого бэкенда, отличного от NCCL/HCCL, так как в настоящее время Gloo может зависнуть.

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

Callable[[Any, GradBucket], Future[Tensor]]

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

Из-за особенностей текущей реализации перекрытия DistributedDataParallel с ZeroRedundancyOptimizer в первые две или три итерации обучения не выполняются обновления параметров на шаге оптимизатора — в зависимости от того, используется ли static_graph=False или static_graph=True соответственно. Это связано с необходимостью получить информацию о стратегии распределения градиентов по корзинам, используемой DistributedDataParallel. Эта стратегия окончательно определяется только после второго прямого прохода, если используется static_graph=False, или после третьего прямого прохода, если используется static_graph=True.

torch.distributed.algorithms.ddp_comm_hooks.ddp_zero_hook.hook_with_zero_step_interleaved(hook, ddp, zero, shard_buckets=False) [исходный код]

Изменяет hook, чтобы перекрыть шаг оптимизатора ZeroRedundancyOptimizer с обратным проходом DistributedDataParallel

Этот подход перекрывает вычисления и коммуникацию оптимизатора с вычислениями и коммуникацией обратного прохода. В частности, после вычисления градиентов корзины запускаются вычисления оптимизатора с использованием этих градиентов (хотя фактическое вычисление должно дождаться завершения операции all-reduce для этой корзины). В результате операции all-reduce и широковещательные рассылки чередуются в потоке коммуникации.

Этот подход может быть предпочтительнее hook_with_zero_step(), если коммуникация относительно быстра по сравнению с вычислениями.

Параметры:

• hook (Any * dist.GradBucket -> torch.futures.Future) – хук, который нужно изменить.
• ddp (DistributedDataParallel) – экземпляр DistributedDataParallel для использования.
• zero (ZeroRedundancyOptimizer) – экземпляр ZeroRedundancyOptimizer для использования.
• shard_buckets (bool) – если True, назначение каждой корзины DistributedDataParallel разделяется между несколькими экземплярами ZeroRedundancyOptimizer (то есть несколькими рангами), чтобы добиться приблизительного равномерного распределения; если False, каждая корзина целиком назначается одному экземпляру ZeroRedundancyOptimizer (то есть одному рангу).

Возвращает:

Изменённый хук.

Вызывает исключения:

• ValueError – если zero создан с параметром overlap_with_ddp=False.
• RuntimeError – при использовании любого бэкенда, отличного от NCCL, так как в настоящее время Gloo может зависнуть.

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

Callable[[Any, GradBucket], Future[Tensor]]

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

Из-за особенностей текущей реализации перекрытия DistributedDataParallel с ZeroRedundancyOptimizer в первые две или три итерации обучения не выполняются обновления параметров на шаге оптимизатора — в зависимости от того, используется ли static_graph=False или static_graph=True соответственно. Это связано с необходимостью получить информацию о стратегии распределения градиентов по корзинам, используемой DistributedDataParallel. Эта стратегия окончательно определяется только после второго прямого прохода, если используется static_graph=False, или после третьего прямого прохода, если используется static_graph=True.

torch.distributed.algorithms.ddp_comm_hooks.quantization_hooks.quantization_perchannel_hook(process_group, bucket, bucket_size=512) [исходный код]

Применяет логику ``torch.quantize_per_channel`` к DDP с использованием протокола allgather.

По сравнению с поканальным квантованием, основная цель поканального квантования — обработка достаточно больших тензоров. Например, квантование тензора из 6 миллионов элементов по корзинам размером 512 (или 128) элементов может значительно повысить разрешение.

Сначала тензор GradBucket разделяется на несколько частей (каналов) по bucket_size элементов. Затем рабочие процессы выполняют all-gather масштабов и нулевых точек собственных GradBucket перед квантованием. Когда эта информация есть у всех рабочих процессов, первый обратный вызов then, называемый quantize_and_allgather, квантирует собственный тензор градиентов рабочего процесса и использует allgather для передачи данных между всеми рабочими процессами. Последний обратный вызов then, называемый dequantize_and_aggregate, выполняет обратное квантование, сплющивание и локальную агрегацию каждого квантованного тензора градиентов и возвращает среднее значение.

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

Эта функция является экспериментальной и использует протокол allgather, который значительно медленнее протокола allreduce. Работает только со сплющенными градиентами.

Пример::

>>> ddp_model.register_comm_hook(process_group, quantization_perchannel_hook)

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

Future[Tensor]

torch.distributed.algorithms.ddp_comm_hooks.quantization_hooks.quantization_pertensor_hook(process_group, bucket) [исходный код]

Применяет логику torch.quantize_per_tensor к DDP с использованием протокола allgather.

Перед квантованием рабочие процессы выполняют all-gather масштаба и нулевой точки собственных GradBucket. Когда эта информация есть у всех рабочих процессов, первый обратный вызов then, называемый quantize_and_allgather, квантирует собственный тензор градиентов рабочего процесса и использует allgather для передачи данных между всеми рабочими процессами. Последний обратный вызов then, называемый dequantize_and_aggregate, выполняет обратное квантование и локальную агрегацию каждого квантованного тензора градиентов и возвращает среднее значение.

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

Эта функция является экспериментальной и использует протокол allgather, который значительно медленнее протокола allreduce. Работает только со сплющенными градиентами.

Пример::

>>> ddp_model.register_comm_hook(process_group, quantization_pertensor_hook)

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

Future[Tensor]

Хуки коммуникации по умолчанию

Хуки коммуникации по умолчанию — это простые хуки без состояния, поэтому входное состояние в register_comm_hook — это либо группа процессов, либо None. Входной аргумент bucket представляет собой объект torch.distributed.GradBucket.

torch.distributed.algorithms.ddp_comm_hooks.default_hooks.allreduce_hook(process_group, bucket) [исходный код]

Выполняет allreduce с использованием тензоров GradBucket.

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

Если пользователь регистрирует этот хук коммуникации DDP, ожидается, что результаты DDP будут такими же, как и без зарегистрированного хука. Поэтому поведение DDP не изменится, и пользователь может использовать этот хук в качестве примера или изменить его, чтобы регистрировать полезную информацию либо выполнять другие задачи, не влияя на поведение DDP.

Пример::

>>> ddp_model.register_comm_hook(process_group, allreduce_hook)

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

Future[Tensor]

torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_hook(process_group, bucket) [исходный код]

Сжимает тензор, приводя GradBucket к типу torch.float16 и деля на размер группы процессов.

Этот хук коммуникации DDP реализует простой способ сжатия градиентов: тензор GradBucket приводится к формату чисел с плавающей запятой половинной точности (torch.float16), а затем делится на размер группы процессов. Выполняется all-reduce этих тензоров градиентов float16. После выполнения all-reduce сжатых тензоров градиентов связанный обратный вызов decompress приводит их обратно к входному типу данных (например, float32).

Пример::

>>> ddp_model.register_comm_hook(process_group, fp16_compress_hook)

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

Future[Tensor]

torch.distributed.algorithms.ddp_comm_hooks.default_hooks.bf16_compress_hook(process_group, bucket) [исходный код]

Предупреждение: этот API является экспериментальным и требует версии NCCL новее 2.9.6.

Этот хук коммуникации DDP реализует простой способ сжатия градиентов: тензор GradBucket приводится к формату чисел с плавающей запятой половинной точности Brain floating point (torch.bfloat16), а затем делится на размер группы процессов. Выполняется all-reduce этих тензоров градиентов bfloat16. После выполнения all-reduce сжатых тензоров градиентов связанный обратный вызов decompress приводит их обратно к входному типу данных (например, float32).

Пример::

>>> ddp_model.register_comm_hook(process_group, bf16_compress_hook)

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

Future[Tensor]

Кроме того, предоставлена оболочка хука коммуникации для поддержки fp16_compress_hook() или bf16_compress_hook() в качестве оболочки, которую можно комбинировать с другими хуками коммуникации.

torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_wrapper(hook) [исходный код]

Преобразует входной тензор в torch.float16, а результат хука — обратно к входному типу данных.

Эта оболочка преобразует входной тензор градиентов заданного хука коммуникации DDP в формат чисел с плавающей запятой половинной точности (torch.float16), а результирующий тензор заданного хука преобразует обратно к входному типу данных, например float32. Таким образом, fp16_compress_hook эквивалентен fp16_compress_wrapper(allreduce_hook).

Пример::

>>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1, start_powerSGD_iter=10)
>>> ddp_model.register_comm_hook(state, fp16_compress_wrapper(powerSGD_hook))

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

Callable[[Any, GradBucket], Future[Tensor]]

torch.distributed.algorithms.ddp_comm_hooks.default_hooks.bf16_compress_wrapper(hook) [исходный код]

Предупреждение: этот API является экспериментальным и требует версии NCCL новее 2.9.6.

Эта оболочка преобразует входной тензор градиентов заданного хука коммуникации DDP в формат чисел с плавающей запятой половинной точности Brain floating point (torch.bfloat16), а результирующий тензор заданного хука преобразует обратно к входному типу данных, например float32.

Таким образом, bf16_compress_hook эквивалентен bf16_compress_wrapper(allreduce_hook).

Пример::

>>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1, start_powerSGD_iter=10)
>>> ddp_model.register_comm_hook(state, bf16_compress_wrapper(powerSGD_hook))

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

Callable[[Any, GradBucket], Future[Tensor]]

Коммуникационный хук PowerSGD

PowerSGD (Vogels et al., NeurIPS 2019) — это алгоритм сжатия градиентов, который обеспечивает очень высокую степень сжатия и ускоряет распределённое обучение, ограниченное пропускной способностью сети. Для работы этому алгоритму необходимо хранить гиперпараметры и внутреннее состояние. Поэтому коммуникационный хук PowerSGD — это хук с сохранением состояния, и пользователь должен предоставить объект состояния, определённый ниже.

Отладка коммуникационных хуков

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

torch.distributed.algorithms.ddp_comm_hooks.debugging_hooks.noop_hook(_, bucket) [исходный код]

Возвращает Future, оборачивающий входное значение, и ничего не делает, не создавая накладных расходов на коммуникацию.

Этот хук следует использовать только для анализа потенциального ускорения при оптимизации allreduce, а не для обычной синхронизации градиентов. Например, если после регистрации этого хука удаётся добиться ускорения обучения менее чем на 10%, это обычно означает, что allreduce не является узким местом производительности в данном случае. Такая инструментация может быть особенно полезна, если трассировки GPU трудно получить или их анализ усложнён такими факторами, как перекрытие allreduce с вычислениями или рассинхронизация между рангами.

Пример::

>>> ddp_model.register_comm_hook(None, noop_hook)

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

Future[Tensor]

Создание контрольных точек коммуникационных хуков

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

В PowerSGDState реализованы __setstate__ и __getstate__; его можно использовать в качестве примера.

классtorch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.PowerSGDState[исходный код]

__getstate__() [исходный код]

Возвращает Dict[str, Any], который будет сериализован с помощью pickle и сохранён.

process_group не сериализуется и исключается из возвращаемого состояния.

__setstate__(state) [исходный код]

Принимает предоставленный state и устанавливает его для этого экземпляра PowerSGDState.

Для process_group устанавливается значение по умолчанию.

Ниже приведён простой сквозной пример сохранения и повторной загрузки состояния и хука PowerSGD.

import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

from torch.nn.parallel import DistributedDataParallel
from torch.distributed.algorithms.ddp_comm_hooks import powerSGD_hook as powerSGD

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(24,24)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(24,12)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def run_demo(demo_fn, world_size):
    mp.spawn(
        demo_fn,
        args=(world_size,),
        nprocs=world_size,
        join=True)

def demo_serialization(rank, world_size):
    setup(rank, world_size)

    CHECKPOINT = tempfile.gettempdir() + "/checkpoint.pt"

    model = SimpleModel().to(rank)
    ddp_model = DistributedDataParallel(model, device_ids=[rank])

    powersgd_hook = powerSGD.powerSGD_hook
    powersgd_state = powerSGD.PowerSGDState(process_group=None)

    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    ddp_model.register_comm_hook(powersgd_state, powersgd_hook)

    state = {
        'state_dict': ddp_model.state_dict(),
        'comm_hook': powersgd_hook,
        'comm_hook_state': powersgd_state}

    if rank == 0:
        torch.save(state, CHECKPOINT)

    dist.barrier()
    map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
    checkpoint = torch.load(CHECKPOINT, map_location=map_location)

    new_ddp_model = DistributedDataParallel(SimpleModel().to(rank), device_ids=[rank])
    new_ddp_model.load_state_dict(checkpoint['state_dict'])
    powersgd_hook = checkpoint['comm_hook']
    powersgd_state = checkpoint['comm_hook_state']

    new_ddp_model.register_comm_hook(powersgd_state, powersgd_hook)

    if rank == 0:
        os.remove(CHECKPOINT)

    cleanup()

if __name__ == "__main__":
    n_gpus = torch.cuda.device_count()
    assert n_gpus >= 2, f"Requires at least 2 GPUs to run, but got {n_gpus}"
    world_size = n_gpus
    run_demo(demo_serialization, world_size)

Благодарности

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