DistributedDataParallel

class torch.nn.parallel.DistributedDataParallel(module, device_ids=None, output_device=None, dim=0, broadcast_buffers=True, process_group=None, bucket_cap_mb=25, find_unused_parameters=False, check_reduction=False, gradient_as_bucket_view=False, static_graph=False) [source]

Реализует распределённую пакетную обработку, основанную на пакете torch.distributed на уровне модуля.

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

Размер пакета должен быть больше, чем количество используемых локально графических процессоров.

См. также: Основы и Использование nn.parallel.DistributedDataParallel вместо multiprocessing или nn.DataParallel. Применяются те же ограничения к входу, что и в torch.nn.DataParallel.

Для создания этого класса необходимо, чтобы torch.distributed был уже инициализирован вызовом torch.distributed.init_process_group().

DistributedDataParallel доказано, что значительно быстрее, чем torch.nn.DataParallel при обучении с использованием нескольких графических процессоров одного узла.

Для использования DistributedDataParallel на узле с N графическими процессорами, необходимо запустить N процессов, гарантируя, что каждый процесс работает исключительно с одним графическим процессором от 0 до N-1. Это можно сделать, либо установив CUDA_VISIBLE_DEVICES для каждого процесса, либо вызвав:

>>> torch.cuda.set_device(i)

где i принимает значения от 0 до N-1. В каждом процессе для построения этого модуля необходимо выполнить следующие действия:

>>> torch.distributed.init_process_group(
>>>     backend='nccl', world_size=N, init_method='...'
>>> )
>>> model = DistributedDataParallel(model, device_ids=[i], output_device=i)

Для запуска нескольких процессов на одном узле можно использовать либо torch.distributed.launch , либо torch.multiprocessing.spawn.

Примечание

Для краткого введения во все функции, связанные с распределенным обучением, обратитесь к Обзору распределенного обучения PyTorch.

Примечание

DistributedDataParallel может использоваться совместно с torch.distributed.optim.ZeroRedundancyOptimizer для уменьшения занимаемой памяти состояний оптимизатора на каждом узле. Подробнее об этом см. рецепт ZeroRedundancyOptimizer.

Примечание

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

Примечание

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

Примечание

Если вы используете torch.save на одном процессе для создания контрольной точки модуля, и torch.load на других процессах для восстановления, убедитесь, что map_location настроено должным образом для каждого процесса. Без map_location, torch.load восстановит модуль на устройствах, с которых модуль был сохранён.

Примечание

Когда модель обучается на M узлах с batch=N, градиент будет M раз меньше по сравнению с той же моделью, обученной на одном узле с batch=M*N, если функция потерь суммируется (а не усредняется, как обычно) по примерам в пакете (потому что градиенты между разными узлами усредняются). Это необходимо учитывать, если требуется получить математически эквивалентный процесс обучения по сравнению с локальным обучением. Однако в большинстве случаев можно считать, что модель, обернутая DistributedDataParallel, DataParallel и обычная модель на одном графическом процессоре, одинаковы (например, использовать одинаковую скорость обучения для эквивалентного размера пакета).

Примечание

Параметры никогда не передаются между процессами. Модуль выполняет операцию all-reduce над градиентами и предполагает, что они будут изменены оптимизатором во всех процессах одинаковым образом. Буферы (например, статистические данные BatchNorm) передаются из модуля в процессе с номером ранга 0 во все другие копии в системе на каждой итерации.

Примечание

Если вы используете DistributedDataParallel совместно с Фреймворком распределённых RPC, вы всегда должны использовать torch.distributed.autograd.backward() для вычисления градиентов и torch.distributed.optim.DistributedOptimizer для оптимизации параметров.

Примечание

DistributedDataParallel в настоящее время имеет ограниченную поддержку отложенного вычисления градиентов с torch.utils.checkpoint(). DDP будет работать как ожидается, если в модели нет неиспользуемых параметров, и каждый слой вычисляется не более одного раза (убедитесь, что вы не передаёте find_unused_parameters=True в DDP). В настоящее время мы не поддерживаем случай, когда слой вычисляется несколько раз, или когда в проверенной модели имеются неиспользуемые параметры.

Пример:

>>> import torch.distributed.autograd as dist_autograd
>>> from torch.nn.parallel import DistributedDataParallel as DDP
>>> import torch
>>> from torch import optim
>>> from torch.distributed.optim import DistributedOptimizer
>>> import torch.distributed.rpc as rpc
>>> from torch.distributed.rpc import RRef
>>>
>>> t1 = torch.rand((3, 3), requires_grad=True)
>>> t2 = torch.rand((3, 3), requires_grad=True)
>>> rref = rpc.remote("worker1", torch.add, args=(t1, t2))
>>> ddp_model = DDP(my_model)
>>>
>>> # Setup optimizer
>>> optimizer_params = [rref]
>>> for param in ddp_model.parameters():
>>>     optimizer_params.append(RRef(param))
>>>
>>> dist_optim = DistributedOptimizer(
>>>     optim.SGD,
>>>     optimizer_params,
>>>     lr=0.05,
>>> )
>>>
>>> with dist_autograd.context() as context_id:
>>>     pred = ddp_model(rref.to_here())
>>>     loss = loss_func(pred, target)
>>>     dist_autograd.backward(context_id, [loss])
>>>     dist_optim.step(context_id)

Примечание

Чтобы загрузить состояние словаря из модели DDP в модель, не являющуюся DDP, необходимо применить consume_prefix_in_state_dict_if_present() для удаления префикса «module.» в словаре состояния DDP перед загрузкой.

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

Конструктор, метод forward и дифференцирование вывода (или функции вывода этого модуля) являются точками синхронизации распределения. Учтите это в случае, если разные процессы могут выполнять разный код.

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

Этот модуль предполагает, что все параметры регистрируются в модели к моменту его создания. Никакие параметры не должны добавляться или удаляться позже. То же самое относится к буферам.

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

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

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

Этот модуль позволяет использовать параметры с несоседними строками. Например, ваша модель может содержать некоторые параметры, у которых torch.memory_format равен torch.contiguous_format , и другие, у которых формат torch.channels_last. Однако соответствующие параметры в разных процессах должны иметь одинаковые шаги.

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

Этот модуль не работает с torch.autograd.grad() (т.е. он будет работать только в том случае, если градиенты будут накапливаться в атрибутах .grad параметров).

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

Если вы планируете использовать этот модуль с бэкендом nccl или бэкендом gloo (который использует Infiniband), вместе с DataLoader, который использует несколько рабочих процессов, измените метод запуска многопроцессорности на forkserver (только Python 3) или spawn. К сожалению, Gloo (который использует Infiniband) и NCCL2 не безопасны для fork, и вы, вероятно, столкнетесь с тупиками, если не измените эту настройку.

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

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

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

Использование DistributedDataParallel совместно с Фреймворком распределённых RPC экспериментально и может измениться.

Параметры:

• module (Модуль) – модуль, который нужно распараллелить

• device_ids (список python:int или torch.device) –

  CUDA устройства. 1) Для модулей с одним устройством, device_ids может содержать ровно один идентификатор устройства, который представляет единственное CUDA устройство, где находится входной модуль, соответствующий этому процессу. В качестве альтернативы, device_ids также может быть None. 2) Для многоустрочных модулей и модулей для CPU, device_ids должно быть None.

  Когда device_ids является None в обоих случаях, как входные данные для прямого прохода, так и сам модуль должны быть размещены на правильном устройстве. (по умолчанию: None)

• output_device (int или torch.device) – Местоположение устройства вывода для модулей CUDA с одним устройством. Для многоустрочных модулей и модулей для CPU, оно должно быть None, а сам модуль определяет местоположение вывода. (по умолчанию: device_ids[0] для модулей с одним устройством)
• broadcast_buffers (bool) – Флаг, который включает синхронизацию (вещание) буферов модуля в начале функции forward. (по умолчанию: True)
• process_group – Группа процессов, которая используется для распределенного всестороннего уменьшения данных. Если None, будет использована группа процессов по умолчанию, которая создана функцией torch.distributed.init_process_group(). (по умолчанию: None)
• bucket_cap_mb – DistributedDataParallel будет группировать параметры в несколько групп, чтобы уменьшение градиента каждой группы могло потенциально перекрываться с обратным вычислением. bucket_cap_mb контролирует размер группы в мегабайтах (МБ). (по умолчанию: 25)
• find_unused_parameters (bool) – Проход по графу autograd от всех тензоров, содержащихся в возвращаемом значении функции обернутого модуля forward. Параметры, которые не получают градиенты в рамках этого графа, предварительно помечаются как готовые к уменьшению. Кроме того, параметры, которые могли быть использованы в функции обернутого модуля forward, но не входили в вычисление потерь и, следовательно, не получали градиенты, предварительно помечаются как готовые к уменьшению. (по умолчанию: False)
• check_reduction – Этот аргумент устарел.
• gradient_as_bucket_view (bool) – Когда установлено True, градиенты будут представлениями, указывающими на разные смещения в allreduce коммуникационных группах. Это может уменьшить максимальное использование памяти, где сохраненный размер памяти будет равен общему размеру градиентов. Кроме того, это избегает накладных расходов на копирование между градиентами и allreduce коммуникационными группами. Когда градиенты являются представлениями, на градиентах нельзя вызывать detach_(). Если возникают такие ошибки, исправьте их, обратившись к функции zero_grad() в torch/optim/optimizer.py в качестве решения. Обратите внимание, что градиенты будут представлениями после первой итерации, поэтому проверку максимальной экономии памяти следует проводить после первой итерации.

• static_graph (bool) –

  Когда установлено True, DDP знает, что обученный граф статичен. Статический граф означает 1) Множество используемых и неиспользуемых параметров не будет меняться в течение всего цикла обучения; в этом случае неважно, установили ли пользователи find_unused_parameters = True или нет. 2) Способ обучения графа не будет меняться в течение всего цикла обучения (то есть нет потока управления, зависящего от итераций). Когда static_graph установлено в True, DDP будет поддерживать случаи, которые не поддерживались в прошлом: 1) Многократные обратные вычисления. 2) Многократное кэширование активаций. 3) Кэширование активаций, когда у модели есть неиспользуемые параметры. 4) Есть параметры модели, которые находятся за пределами функции forward. 5) Потенциально улучшить производительность, когда есть неиспользуемые параметры, так как DDP не будет искать граф в каждой итерации для определения неиспользуемых параметров, когда static_graph установлено в True. Чтобы проверить, можно ли установить static_graph в True, один из способов — проверить данные протоколирования DDP в конце предыдущего обучения модели, если ddp_logging_data.get("can_set_static_graph") == True, то чаще всего можно установить и static_graph = True.

  Пример::

  >>> model_DDP = torch.nn.parallel.DistributedDataParallel(model)
  >>> # Training loop
  >>> ...
  >>> ddp_logging_data = model_DDP._get_ddp_logging_data()
  >>> static_graph = ddp_logging_data.get("can_set_static_graph")
  

Переменные:

module (Модуль) – модуль, который нужно распараллелить.

Пример:

>>> torch.distributed.init_process_group(backend='nccl', world_size=4, init_method='...')
>>> net = torch.nn.parallel.DistributedDataParallel(model)

join(divide_by_initial_world_size=True, enable=True, throw_on_early_termination=False) [источник]

Менеджер контекста, который используется совместно с экземпляром torch.nn.parallel.DistributedDataParallel, чтобы иметь возможность тренироваться с неравномерными входами на участвующих процессах.

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

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

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

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

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

Параметры:

• divide_by_initial_world_size (bool) – Если True, разделит градиенты на начальный world_size DDP-обучения был запущен. Если False, вычислит эффективный размер мира (количество рангов, которые еще не исчерпали свои входы) и разделит градиенты на него во время allreduce. Установите divide_by_initial_world_size=True для обеспечения того, чтобы каждый образец входных данных, включая неравномерные входные данные, имели равный вес с точки зрения того, насколько они влияют на глобальный градиент. Это достигается путем постоянного деления градиента на начальный world_size даже при столкновении с неравномерными входами. Если вы установите это в False, мы разделим градиент на оставшееся количество узлов. Это гарантирует соответствие обучению на меньшем world_size числе узлов, хотя это также означает, что неравномерные входные данные будут в большей степени способствовать глобальному градиенту. Обычно вы захотите установить это в True для случаев, когда последние несколько входных данных вашей задачи обучения неравномерны. В крайних случаях, когда есть большая разница в количестве входов, установка этого значения в False может обеспечить лучшие результаты.
• enable (bool) – Включает или выключает обнаружение неравномерных входов. Передайте enable=False для отключения в случаях, когда вам известно, что входы равномерны на участвующих процессах. По умолчанию True.
• throw_on_early_termination (bool) – Выводить ошибку или продолжать обучение, когда по крайней мере один ранг исчерпал входные данные. Если True, то выведет ошибку при первом ранге, достигшем конца данных. Если False, то продолжит обучение с меньшим эффективным размером мира, пока все ранги не присоединятся. Обратите внимание, что если этот флаг указан, то флаг divide_by_initial_world_size будет проигнорирован. По умолчанию False.

Пример:

>>> import torch
>>> import torch.distributed as dist
>>> import os
>>> import torch.multiprocessing as mp
>>> import torch.nn as nn
>>> # On each spawned worker
>>> def worker(rank):
>>>     dist.init_process_group("nccl", rank=rank, world_size=2)
>>>     torch.cuda.set_device(rank)
>>>     model = nn.Linear(1, 1, bias=False).to(rank)
>>>     model = torch.nn.parallel.DistributedDataParallel(
>>>         model, device_ids=[rank], output_device=rank
>>>     )
>>>     # Rank 1 gets one more input than rank 0.
>>>     inputs = [torch.tensor([1]).float() for _ in range(10 + rank)]
>>>     with model.join():
>>>         for _ in range(5):
>>>             for inp in inputs:
>>>                 loss = model(inp).sum()
>>>                 loss.backward()
>>>     # Without the join() API, the below synchronization will hang
>>>     # blocking for rank 1's allreduce to complete.
>>>     torch.cuda.synchronize(device=rank)

join_hook(**kwargs) [source]

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

Параметры:

kwargs (dict) – dict, содержащий любые ключевые аргументы для изменения поведения хука присоединения во время выполнения; все Joinable экземпляры, использующие один и тот же менеджер контекста присоединения, получают одинаковые значения для kwargs.

Хук поддерживает следующие ключевые аргументы:

divide_by_initial_world_size (bool, optional):

Если True, то градиенты делятся на начальный размер мира, с которым был запущен DDP. Если False, то градиенты делятся на эффективный размер мира (т.е. количество не присоединённых процессов), что означает, что неравномерные входы вносят больший вклад в глобальный градиент. Обычно это следует устанавливать в True если степень неравномерности невелика, но в крайних случаях можно установить в False для потенциально лучших результатов. Значение по умолчанию True.

no_sync() [source]

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

Пример:

>>> ddp = torch.nn.parallel.DistributedDataParallel(model, pg)
>>> with ddp.no_sync():
>>>   for input in inputs:
>>>     ddp(input).backward()  # no synchronization, accumulate grads
>>> ddp(another_input).backward()  # synchronize grads

register_comm_hook(state, hook) [source]

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

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

Параметры:

• state (object) –

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

  Хранится локально на каждом рабочем узле и общий для всех тензоров градиента на рабочем узле.

• hook (Callable) –

  Вызываемая функция со следующим сигнатурами: hook(state: object, bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:

  Эта функция вызывается, когда ведро готово. Хук может выполнить любые необходимые операции и вернуть Future, указывающее на завершение любых асинхронных задач (например, allreduce). Если хук не выполняет никакой коммуникации, он всё равно должен вернуть завершённое Future. Future должно содержать новое значение тензоров ведра grad. После того как ведро готово, редуктор c10d вызовет этот хук и использует тензоры, возвращённые Future, и скопирует градиенты в отдельные параметры. Обратите внимание, что возвращаемый тип future должен быть единственным тензором.

  Мы также предоставляем API get_future для получения Future, связанного с завершением c10d.ProcessGroup.Work. get_future в настоящее время поддерживается для NCCL и также поддерживается для большинства операций GLOO и MPI, за исключением операций точка-точка (send/recv).

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

Тензоры ведра grad не будут предварительно разделены на world_size. Пользователь отвечает за деление на world_size в случае операций, таких как allreduce.

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

Хук коммуникации DDP может быть зарегистрирован только один раз и должен быть зарегистрирован перед вызовом backward.

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

Объект Future, возвращаемый хуком, должен содержать единственный тензор, имеющий ту же форму, что и тензоры внутри ведра grad.

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

get_future API поддерживает бэкэнды NCCL и частично GLOO и MPI (нет поддержки операций точка-точка, таких как send/recv) и вернёт torch.futures.Future.

Пример::

Ниже приведен пример хука noop, возвращающего тот же тензор.

>>> def noop(state: object, bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:
>>>     fut = torch.futures.Future()
>>>     fut.set_result(bucket.buffer())
>>>     return fut

>>> ddp.register_comm_hook(state=None, hook=noop)

Пример::

Ниже приведен пример алгоритма Parallel SGD, где градиенты кодируются перед allreduce, а затем декодируются после allreduce.

>>> def encode_and_decode(state: object, bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:
>>>     encoded_tensor = encode(bucket.buffer()) # encode gradients
>>>     fut = torch.distributed.all_reduce(encoded_tensor).get_future()
>>>     # Define the then callback to decode.
>>>     def decode(fut):
>>>         decoded_tensor = decode(fut.value()[0]) # decode gradients
>>>         return decoded_tensor
>>>     return fut.then(decode)

>>> ddp.register_comm_hook(state=None, hook=encode_and_decode)