Связки коммуникации DDP

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

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

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

torch.nn.parallel.DistributedDataParallel.register_comm_hook()

Что обрабатывает связка коммуникации?

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

class torch.distributed.GradBucket

Этот класс в основном передает сплющенный тензор градиента (возвращаемый buffer()) в связку коммуникации DDP. Этот тензор можно далее разбить на список тензоров параметров в этом пакете (возвращаемый 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

Возвращает

Является ли этот пакет последним пакетом для allreduce в итерации. Это также означает, что этот пакет соответствует первым нескольким слоям в прямом проходе.

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. Каждый тензор в списке соответствует параметру модели.

Связки коммуникации по умолчанию

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

torch.distributed.algorithms.ddp_comm_hooks.default_hooks.allreduce_hook(process_group, bucket) [source]

Эта связка коммуникации DDP просто вызывает allreduce с использованием тензоров GradBucket. После агрегирования тензоров градиента по всем рабочим процессам её обратный вызов then вычисляет среднее значение и возвращает результат. Если пользователь зарегистрирует эту связку, результаты 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) [source]

Эта связка коммуникации DDP реализует простой подход к сжатию градиента, который преобразует тензор GradBucket в формат с плавающей точкой полуточной точности (torch.float16) и затем делит его на размер группы процессов. Он allreduce те float16 тензоры градиента. После того, как сжатые тензоры градиента будут allreduce, цепной обратный вызов 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) [source]

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

Эта связка коммуникации DDP реализует простой подход к сжатию градиента, который преобразует тензор GradBucket в формат с плавающей точкой Brain (torch.bfloat16) и затем делит его на размер группы процессов. Он allreduce те bfloat16 тензоры градиента. После того, как сжатые тензоры градиента будут allreduce, цепной обратный вызов 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) [source]

Эта оболочка преобразует входной тензор градиента заданной связки коммуникации 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) [source]

Предупреждение: Этот API экспериментальный и требует версию NCCL, более позднюю, чем 2.9.6.

Этот обёртка преобразует входной тензор градиента заданного DDP-коммуникационного хука в формат с половинной точностью Brain floating point format <https://en.wikipedia.org/wiki/Bfloat16_floating-point_format> `_ (``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[Тензор]]

Хук связи PowerSGD

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

Состояние PowerSGD

class torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.PowerSGDState(process_group, matrix_approximation_rank=1, start_powerSGD_iter=1000, min_compression_rate=2, use_error_feedback=True, warm_start=True, orthogonalization_epsilon=0, random_seed=0, compression_stats_logging_frequency=10000, batch_tensors_with_same_shape=False) [source]

Хранит как гиперпараметры алгоритма, так и внутреннее состояние для всех градиентов во время обучения. В частности, matrix_approximation_rank и start_powerSGD_iter — основные гиперпараметры, которые должен настраивать пользователь. Для повышения производительности рекомендуется держать двоичные гиперпараметры use_error_feedback и warm_start включёнными.

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

   1.1. Если matrix_approximation_rank слишком низкое, модели потребуется больше шагов обучения, чтобы достичь или никогда не достичь требуемого качества и ухудшится точность.

   1.2. Увеличение matrix_approximation_rank значительно увеличивает вычислительные затраты на сжатие, и точность может не улучшиться сверх определённого порога matrix_approximation_rank.

Для настройки matrix_approximation_rank рекомендуется начинать с 1 и увеличивать в два раза (как экспоненциальный поиск по сетке, 1, 2, 4, …), пока не будет достигнута удовлетворительная точность. Обычно используется небольшое значение 1-4. Для некоторых задач NLP (как показано в приложении D оригинальной статьи) это значение было увеличено до 32.

2. start_powerSGD_iter откладывает сжатие PowerSGD до шага start_powerSGD_iter, а обычный allreduce выполняется до шага start_powerSGD_iter. Эта гибридная схема обычный allreduce + PowerSGD может эффективно повысить точность, даже если используется относительно маленькое значение matrix_approximation_rank. Это связано с тем, что начало этапа обучения обычно очень чувствительно к неточным градиентам, и сжатие градиентов слишком рано может привести к тому, что обучение быстро перейдёт по не оптимальной траектории, что может привести к невосполнимому влиянию на точность.

Для настройки start_powerSGD_iter рекомендуется начинать с 10% от общего числа шагов обучения и увеличивать его, пока не будет достигнута удовлетворительная точность. Если в обучении есть этап разогрева, start_powerSGD_iter обычно должно быть не меньше числа шагов разогрева.

3. min_compression_rate — минимальная требуемая скорость сжатия, когда слой сжимается. Из-за вычислительных накладных расходов, связанных со сжатием, тензор стоит сжимать только в том случае, если можно достаточно сэкономить на пропускной способности, где (num_rows + num_cols) * matrix_approximation_rank * min_compression_rate < num_rows * num_cols. Если заданный порог скорости сжатия не может быть удовлетворён, тензор будет непосредственно allreduced без сжатия.

Статистика сжатия регистрируется каждые compression_stats_logging_frequency итераций после начала сжатия PowerSGD.

4. orthogonalization_epsilon может быть очень малым значением (например, 1e-8), добавляемым к каждой нормализованной колонке матрицы на шаге ортогонализации, чтобы предотвратить ошибку деления на ноль, если какая-либо колонка имеет все нули. Если это можно предотвратить (например, с помощью нормализации по батчам), рекомендуется использовать epsilon равное 0 для точности.
5. batch_tensors_with_same_shape управляет тем, следует ли сжимать и распаковывать тензоры с одинаковой формой в пакетной операции для достижения большей параллельности. Обратите внимание, что также необходимо увеличить размер пакета (т.е., аргумент bucket_cap_mb в конструкторе DDP), чтобы в одном пакете появилось больше тензоров с одинаковой формой, однако это может уменьшить перекрытие между вычислениями и обменом, а также увеличить объём памяти из-за объединения тензоров с одинаковой формой. Установите в True в случае, если вычисления сжатия/распаковки являются узким местом.

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

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

Хуки PowerSGD

torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.powerSGD_hook(state, bucket) [source]

Этот DDP-коммуникационный хук реализует алгоритм сжатия градиентов PowerSGD, описанный в статье. После агрегирования тензоров градиента по всем рабочим узлам этот хук применяет сжатие следующим образом:

1. Представляет входной сглаженный тензор градиента 1D как список тензоров на параметр и делит все тензоры на две группы:

   1.1 Тензоры, которые должны быть сжаты перед allreduce, так как сжатие может обеспечить достаточную экономию пропускной способности.

   1.2 Остальные тензоры будут непосредственно allreduced без сжатия, включая все векторные тензоры (для смещений).

2. Обрабатывает несжатые тензоры:

   2.1 Выделяет непрерывную память для этих несжатых тензоров и allreduces все несжатые тензоры как пакет без сжатия;

   2.2 Копирует отдельные несжатые тензоры из непрерывной памяти обратно в входной тензор.

3. Обрабатывает тензоры, которые должны быть сжаты с помощью сжатия PowerSGD:

   3.1 Для каждого тензора M создаёт два тензора низкого ранга P и Q для разложения M, так что M = PQ^T, где Q инициализируется из стандартного нормального распределения и ортогонализируется;

   3.2 Вычисляет каждый P в Ps, который равен MQ;

   3.3 Allreduces Ps как пакет;

   3.4 Ортогонализирует каждый P в Ps;

   3.5 Вычисляет каждый Q в Qs, который приблизительно равен M^TP;

   3.6 Allreduces Qs как пакет;

   3.7 Вычисляет каждый M среди всех сжатых тензоров, который приблизительно равен PQ^T.

Обратите внимание, что этот коммуникационный хук принудительно использует обычный allreduce в течение первых state.start_powerSGD_iter итераций. Это не только даёт пользователю больший контроль над балансом между ускорением и точностью, но и помогает абстрагировать некоторые сложности внутренней оптимизации DDP для будущих разработчиков коммуникационных хуков.

Параметры

• state (PowerSGDState) – Информация о состоянии для настройки скорости сжатия и поддержки обратной связи об ошибках, теплового старта и т. д. Для настройки параметров сжатия необходимо настроить, в основном, matrix_approximation_rank, start_powerSGD_iter и min_compression_rate.
• bucket (dist.GradBucket) — Пакет, хранящий сглаженный тензор градиента 1D, который группирует несколько тензоров на переменную. Обратите внимание, что поскольку DDP comm hook поддерживает только режим одного процесса на одном устройстве, в этом пакете хранится ровно один тензор.

Возвращаемое значение

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

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

Future[Тензор]

Пример::

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

torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.batched_powerSGD_hook(state, bucket) [source]

Этот DDP-коммуникационный хук реализует упрощённый алгоритм сжатия градиентов PowerSGD, описанный в статье. Этот вариант не сжимает градиенты по слоям, а вместо этого сжимает уплощённый входной тензор, объединяющий все градиенты. Поэтому он быстрее, чем powerSGD_hook(), но обычно приводит к намного более низкой точности, если matrix_approximation_rank равно 1.

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

Увеличение matrix_approximation_rank здесь может не обязательно повысить точность, потому что объединение тензоров по параметрам без выравнивания столбцов/строк может разрушить низкоранговую структуру. Поэтому пользователь должен всегда сначала рассмотреть powerSGD_hook(), и рассматривать только этот вариант, когда удовлетворительная точность достигается при matrix_approximation_rank равном 1.

После агрегирования тензоров градиентов по всем рабочим узлам, этот хук применяет сжатие следующим образом:

1. Представляет входной уплощённый одномерный тензор градиента как квадратный тензор M с 0-выми дополнениями;
2. Создаёт два низкоранговых тензора P и Q для разложения M, такие что M = PQ^T, где Q инициализируется из стандартного нормального распределения и ортогонализируется;
3. Вычисляет P, которое равно MQ;
4. Выполняет allreduce для P;
5. Ортогонализирует P;
6. Вычисляет Q, которое приблизительно равно M^TP;
7. Выполняет allreduce для Q;
8. Вычисляет M, которое приблизительно равно PQ^T.
9. Обрезает входной тензор до первоначальной длины.

Обратите внимание, что этот коммуникационный хук применяет обычный allreduce для первых state.start_powerSGD_iter итераций. Это не только даёт пользователю больше контроля над компромиссом между ускорением и точностью, но также помогает абстрагировать некоторые сложности внутренней оптимизации DDP для разработчиков будущих коммуникационных хуков.

Параметры

• state (PowerSGDState) – Информация о состоянии для настройки скорости сжатия и поддержки обратной связи об ошибках, теплого запуска и т.д. Для настройки параметров сжатия, необходимо, в основном, настраивать matrix_approximation_rank и start_powerSGD_iter.
• bucket (dist.GradBucket) – Бакет, который хранит одномерный уплощённый тензор градиента, объединяющий несколько тензоров на переменную. Обратите внимание, что так как DDP comm hook поддерживает только режим одного процесса и одного устройства, в этом бакете хранится ровно один тензор.

Возвращаемое значение

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

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

Future[Tensor]

Пример::

>>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1)
>>> ddp_model.register_comm_hook(state, batched_powerSGD_hook)

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

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

torch.distributed.algorithms.ddp_comm_hooks.debugging_hooks.noop_hook(_, bucket) [source]

Этот DDP-коммуникационный хук возвращает будущее, которое оборачивает вход, поэтому он является «ничего не делающим» и не требует коммуникационных накладных расходов.

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

Пример::

>>> ddp_model.register_comm_hook(None, noop_hook)

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

Future[Tensor]

Сохранение коммуникационных хуков

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

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

classtorch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.PowerSGDState[source]

__getstate__() [source]

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

__setstate__(state) [source]

Принимает предоставленный 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 коммуникационного хука соответствует реализации в оригинальной статье.