Распределённая RPC-инфраструктура

Основы

Распределённая RPC-инфраструктура упрощает удалённый вызов функций, поддерживает ссылки на удалённые объекты без копирования самих данных и предоставляет API autograd и optimizer для прозрачного выполнения обратного прохода и обновления параметров через границы RPC. Эти функции можно разделить на четыре группы API.

1. Удалённый вызов процедур (RPC) позволяет выполнять функцию на указанном целевом рабочем процессе с заданными аргументами и получать возвращаемое значение либо создавать ссылку на него. Существует три основных API RPC: rpc_sync() (синхронный), rpc_async() (асинхронный) и remote() (асинхронный, возвращает ссылку на возвращаемое значение удалённого вызова). Используйте синхронный API, если выполнение пользовательского кода невозможно продолжить без возвращаемого значения. В противном случае используйте асинхронный API, чтобы получить future, и дождитесь его завершения, когда возвращаемое значение понадобится вызывающей стороне. API remote() полезен, когда требуется создать что-либо удалённо, но получать это на вызывающей стороне не нужно. Представьте, что управляющий процесс настраивает сервер параметров и обучающий процесс. Управляющий процесс может создать таблицу эмбеддингов на сервере параметров, а затем передать ссылку на неё обучающему процессу, но сам никогда не будет использовать эту таблицу локально. В этом случае rpc_sync() и rpc_async() уже не подходят, поскольку они подразумевают, что возвращаемое значение будет передано вызывающей стороне сразу или в будущем.
2. Удалённая ссылка (RRef) служит распределённым указателем на локальный или удалённый объект. Ею можно делиться с другими рабочими процессами, а подсчёт ссылок выполняется прозрачно. У каждой RRef есть только один владелец, и объект существует только у этого владельца. Рабочие процессы, не являющиеся владельцами, но располагающие RRef, могут явно запросить у владельца копию объекта. Это полезно, когда рабочему процессу необходимо получить доступ к некоторому объекту данных, но он сам не является ни создателем (вызывающим remote()), ни владельцем объекта. Распределённый optimizer, о котором мы расскажем ниже, — один из примеров такого сценария.
3. Распределённый Autograd объединяет локальные движки autograd на всех рабочих процессах, участвующих в прямом проходе, и автоматически обращается к ним во время обратного прохода для вычисления градиентов. Это особенно полезно, когда прямой проход должен выполняться на нескольких машинах, например при распределённом параллельном обучении моделей, обучении с сервером параметров и т. д. Благодаря этой функции пользовательскому коду больше не нужно заботиться о том, как передавать градиенты через границы RPC и в каком порядке запускать локальные движки autograd. Это может быть довольно сложно, если в прямом проходе есть вложенные и взаимозависимые вызовы RPC.
4. Распределённый Optimizer принимает в конструкторе Optimizer() (например, SGD(), Adagrad() и т. д.) и список RRef параметров, создаёт экземпляр Optimizer() на каждом уникальном владельце RRef и соответствующим образом обновляет параметры при выполнении step(). Если прямой и обратный проходы распределены, параметры и градиенты находятся на разных рабочих процессах, поэтому на каждом участвующем рабочем процессе требуется optimizer. Распределённый Optimizer объединяет все эти локальные оптимизаторы в один и предоставляет лаконичные конструктор и API step().

RPC

Перед использованием RPC и примитивов распределённого автограда необходимо выполнить инициализацию. Чтобы инициализировать фреймворк RPC, используйте init_rpc(), который инициализирует фреймворк RPC, фреймворк RRef и распределённый автоград.

torch.distributed.rpc.init_rpc(name, backend=None, rank=-1, world_size=None, rpc_backend_options=None) [исходный код]

Инициализирует примитивы RPC, такие как локальный агент RPC и распределённый автоград, после чего текущий процесс сразу становится готовым отправлять и получать RPC-вызовы.

Параметры:

• name (str) – глобально уникальное имя этого узла (например, Trainer3, ParameterServer2, Master, Worker1). Имя может содержать только цифры, латинские буквы, символ подчёркивания, двоеточие и/или дефис и должно быть короче 128 символов.
• backend (BackendType, optional) – тип реализации бэкенда RPC. Поддерживается значение BackendType.TENSORPIPE (по умолчанию). Дополнительные сведения см. в разделе Бэкенды.
• rank (int) – глобально уникальный идентификатор/ранг этого узла.
• world_size (int) – количество рабочих процессов в группе.
• rpc_backend_options (RpcBackendOptions, optional) – параметры, передаваемые конструктору RpcAgent. Это должен быть подкласс RpcBackendOptions, специфичный для агента, содержащий конфигурацию инициализации этого агента. По умолчанию для всех агентов устанавливается тайм-аут 60 секунд и выполняется rendezvous с базовой группой процессов, инициализированной с помощью init_method = "env://"; это означает, что переменные окружения MASTER_ADDR и MASTER_PORT необходимо задать правильно. Дополнительные сведения и список доступных параметров см. в разделе Бэкенды.

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

Возвращаемый тип:

bool

Следующие API позволяют пользователям удалённо выполнять функции, а также создавать ссылки (RRef) на удалённые объекты данных. В этих API при передаче Tensor в качестве аргумента или возвращаемого значения целевой рабочий процесс попытается создать Tensor с такими же метаданными (например, формой, шагом и т. д.). Мы намеренно запрещаем передачу тензоров CUDA, так как это может привести к сбою, если списки устройств в исходном и целевом рабочих процессах не совпадают. В таких случаях приложение может явно переместить входные тензоры на CPU на стороне вызывающего процесса, а при необходимости — на нужные устройства на стороне вызываемого процесса.

torch.distributed.rpc.rpc_sync(to, func, args=None, kwargs=None, timeout=-1.0) [исходный код]

Выполняет блокирующий RPC-вызов функции func на рабочем процессе to. Отправка и получение сообщений RPC выполняются параллельно с исполнением кода Python. Этот метод потокобезопасен.

Параметры:

• to (str or WorkerInfo or int) – имя/ранг/WorkerInfo целевого рабочего процесса.
• func (Callable) – вызываемая функция, например вызываемый объект Python, встроенный оператор (например, add()) или аннотированная функция TorchScript.
• args (tuple) – кортеж аргументов для вызова func.
• kwargs (dict) – словарь именованных аргументов для вызова func.
• timeout (float, optional) – тайм-аут этого RPC в секундах. Если RPC не будет завершён за это время, возникнет исключение, указывающее на превышение времени ожидания. Значение 0 означает бесконечный тайм-аут, то есть ошибка тайм-аута не возникнет. Если значение не указано, используется значение по умолчанию, заданное при инициализации или с помощью _set_rpc_timeout.

Возвращает:

Возвращает результат выполнения func с аргументами args и kwargs.

Пример::

Убедитесь, что на обоих рабочих процессах правильно заданы MASTER_ADDR и MASTER_PORT. Дополнительные сведения см. в документации API init_process_group(). Например:

export MASTER_ADDR=localhost export MASTER_PORT=5678

Затем выполните следующий код в двух разных процессах:

>>> # On worker 0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> ret = rpc.rpc_sync("worker1", torch.add, args=(torch.ones(2), 3))
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

Ниже приведён пример выполнения функции TorchScript с помощью RPC.

>>> # On both workers:
>>> @torch.jit.script
>>> def my_script_add(tensor: torch.Tensor, scalar: int):
>>>    return torch.add(tensor, scalar)

>>> # On worker 0:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> ret = rpc.rpc_sync("worker1", my_script_add, args=(torch.ones(2), 3))
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

torch.distributed.rpc.rpc_async(to, func, args=None, kwargs=None, timeout=-1.0) [исходный код]

Выполняет неблокирующий RPC-вызов функции func на рабочем процессе to. Отправка и получение сообщений RPC выполняются параллельно с исполнением кода Python. Этот метод потокобезопасен. Метод сразу возвращает объект Future, ожидать завершения которого можно позднее.

Параметры:

• to (str or WorkerInfo or int) – имя/ранг/WorkerInfo целевого рабочего процесса.
• func (Callable) – вызываемая функция, например вызываемый объект Python, встроенный оператор (например, add()) или аннотированная функция TorchScript.
• args (tuple) – кортеж аргументов для вызова func.
• kwargs (dict) – словарь именованных аргументов для вызова func.
• timeout (float, optional) – тайм-аут этого RPC в секундах. Если RPC не будет завершён за это время, возникнет исключение, указывающее на превышение времени ожидания. Значение 0 означает бесконечный тайм-аут, то есть ошибка тайм-аута не возникнет. Если значение не указано, используется значение по умолчанию, заданное при инициализации или с помощью _set_rpc_timeout.

Возвращает:

Возвращает объект Future, завершения которого можно дождаться. После завершения из объекта Future можно получить возвращаемое значение func с аргументами args и kwargs.

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

Использование тензоров GPU в качестве аргументов или возвращаемых значений func не поддерживается, поскольку передача тензоров GPU по сети не поддерживается. Перед использованием тензоров GPU в качестве аргументов или возвращаемых значений func необходимо явно копировать их на CPU.

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

API rpc_async копирует хранилища тензоров-аргументов только перед отправкой по сети; в зависимости от типа бэкенда RPC это может выполняться в другом потоке. Вызывающая сторона должна гарантировать, что содержимое этих тензоров останется неизменным до завершения возвращённого объекта Future.

Пример::

Убедитесь, что на обоих рабочих процессах правильно заданы MASTER_ADDR и MASTER_PORT. Дополнительные сведения см. в документации API init_process_group(). Например:

export MASTER_ADDR=localhost export MASTER_PORT=5678

Затем выполните следующий код в двух разных процессах:

>>> # On worker 0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> fut1 = rpc.rpc_async("worker1", torch.add, args=(torch.ones(2), 3))
>>> fut2 = rpc.rpc_async("worker1", min, args=(1, 2))
>>> result = fut1.wait() + fut2.wait()
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

Ниже приведён пример выполнения функции TorchScript с помощью RPC.

>>> # On both workers:
>>> @torch.jit.script
>>> def my_script_add(tensor: torch.Tensor, scalar: int):
>>>    return torch.add(tensor, scalar)

>>> # On worker 0:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> fut = rpc.rpc_async("worker1", my_script_add, args=(torch.ones(2), 3))
>>> ret = fut.wait()
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

torch.distributed.rpc.remote(to, func, args=None, kwargs=None, timeout=-1.0) [исходный код]

Выполняет удалённый вызов func на рабочем процессе to и сразу возвращает RRef на результат. Рабочий процесс to становится владельцем возвращённого RRef, а рабочий процесс, вызывающий remote, — его пользователем. Владелец управляет глобальным счётчиком ссылок своего RRef, и объект RRef владельца уничтожается только тогда, когда на него не остаётся ни одной действующей ссылки.

Параметры:

• to (str or WorkerInfo or int) – имя/ранг/WorkerInfo целевого рабочего процесса.
• func (Callable) – вызываемая функция, например вызываемый объект Python, встроенный оператор (например, add()) или аннотированная функция TorchScript.
• args (tuple) – кортеж аргументов для вызова func.
• kwargs (dict) – словарь именованных аргументов для вызова func.
• timeout (float, optional) – тайм-аут этого удалённого вызова в секундах. Если создание этого RRef на рабочем процессе to не будет успешно обработано на этом рабочем процессе за указанное время, при следующей попытке использовать RRef (например, to_here()) возникнет ошибка тайм-аута, указывающая на сбой. Значение 0 означает бесконечный тайм-аут, то есть ошибка тайм-аута не возникнет. Если значение не указано, используется значение по умолчанию, заданное при инициализации или с помощью _set_rpc_timeout.

Возвращает:

Экземпляр пользовательского RRef, ссылающийся на результат. Чтобы получить результат локально, используйте блокирующий API torch.distributed.rpc.RRef.to_here().

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

API remote копирует хранилища тензоров-аргументов только перед отправкой по сети; в зависимости от типа бэкенда RPC это может выполняться в другом потоке. Вызывающая сторона должна гарантировать, что содержимое этих тензоров останется неизменным до подтверждения RRef владельцем; проверить это можно с помощью API torch.distributed.rpc.RRef.confirmed_by_owner().

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

Ошибки, например тайм-ауты API remote, обрабатываются по возможности. Это означает, что при сбое удалённых вызовов, инициированных remote, например из-за ошибки тайм-аута, мы предпринимаем все возможные меры для обработки ошибки. Ошибки обрабатываются и устанавливаются в результирующем RRef асинхронно. Если приложение не использовало RRef до завершения этой обработки (например, с помощью to_here или вызова fork), при последующем использовании RRef ошибки будут корректно возникать. Однако приложение может использовать RRef до обработки ошибок. В этом случае ошибки могут не возникнуть, поскольку обработка ещё не завершена.

Пример:

Make sure that ``MASTER_ADDR`` and ``MASTER_PORT`` are set properly
on both workers. Refer to :meth:`~torch.distributed.init_process_group`
API for more details. For example,

export MASTER_ADDR=localhost
export MASTER_PORT=5678

Then run the following code in two different processes:

>>> # On worker 0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> rref1 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 3))
>>> rref2 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 1))
>>> x = rref1.to_here() + rref2.to_here()
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

Below is an example of running a TorchScript function using RPC.

>>> # On both workers:
>>> @torch.jit.script
>>> def my_script_add(tensor: torch.Tensor, scalar: int):
>>>    return torch.add(tensor, scalar)

>>> # On worker 0:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> rref = rpc.remote("worker1", my_script_add, args=(torch.ones(2), 3))
>>> rref.to_here()
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

torch.distributed.rpc.get_worker_info(worker_name=None) [исходный код]

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

Параметры:

worker_name (str) – строковое имя рабочего процесса. Если None, возвращается идентификатор текущего рабочего процесса (по умолчанию None).

Возвращает:

Экземпляр WorkerInfo для указанного worker_name или WorkerInfo текущего рабочего процесса, если worker_name имеет значение None.

torch.distributed.rpc.shutdown(graceful=True, timeout=0) [исходный код]

Завершает работу агента RPC, а затем уничтожает его. Локальный агент перестаёт принимать ожидающие запросы, а фреймворк RPC завершает работу, останавливая все потоки RPC. Если graceful=True, вызов блокируется, пока все локальные и удалённые процессы RPC не достигнут этого метода; затем он дожидается завершения всех ожидающих задач. В противном случае, если graceful=False, выполняется локальное завершение работы без ожидания, пока другие процессы RPC достигнут этого метода.

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

Для объектов Future, возвращаемых функцией rpc_async(), future.wait() нельзя вызывать после shutdown().

Параметры:

graceful (bool) – следует ли выполнять корректное завершение работы. Если True, метод: 1) ожидает, пока не останется системных сообщений для UserRRefs, и удаляет их; 2) блокируется, пока все локальные и удалённые процессы RPC не достигнут этого метода, а затем ожидает завершения всех ожидающих задач.

Пример::

Убедитесь, что на обоих рабочих процессах правильно заданы MASTER_ADDR и MASTER_PORT. Дополнительные сведения см. в документации API init_process_group(). Например:

export MASTER_ADDR=localhost export MASTER_PORT=5678

Затем выполните следующий код в двух разных процессах:

>>> # On worker 0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> # do some work
>>> result = rpc.rpc_sync("worker1", torch.add, args=(torch.ones(1), 1))
>>> # ready to shutdown
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> # wait for worker 0 to finish work, and then shutdown.
>>> rpc.shutdown()

class torch.distributed.rpc.WorkerInfo

Структура, содержащая сведения о рабочем процессе в системе. Содержит имя и идентификатор рабочего процесса. Этот класс не предназначен для непосредственного создания; экземпляр можно получить с помощью get_worker_info(), а результат передавать в такие функции, как rpc_sync(), rpc_async(), remote(), чтобы не копировать строку при каждом вызове.

property id

Глобально уникальный идентификатор рабочего процесса.

property name

Имя рабочего процесса.

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

torch.distributed.rpc.functions.async_execution(fn) [исходный код]

Декоратор функции, указывающий, что возвращаемое значение функции гарантированно является объектом Future и что эта функция может выполняться асинхронно на стороне получателя RPC. В частности, получатель извлекает Future, возвращённый декорированной функцией, и устанавливает последующие шаги обработки в качестве обратного вызова для этого объекта Future. После завершения установленный обратный вызов считывает значение из объекта Future и отправляет его обратно в качестве ответа RPC. Это также означает, что возвращённый объект Future существует только на стороне получателя и никогда не передаётся через RPC. Этот декоратор полезен, когда выполнение декорированной функции (fn) нужно приостановить и возобновить, например, если она содержит rpc_async() или ожидает других сигналов.

Примечание

Для включения асинхронного выполнения приложения должны передавать в API RPC объект функции, возвращённый этим декоратором. Если RPC обнаруживает атрибуты, установленные этим декоратором, он понимает, что функция возвращает объект Future, и обрабатывает его соответствующим образом. Однако это не означает, что при определении функции данный декоратор обязательно должен быть внешним. Например, при использовании вместе с @staticmethod или @classmethod декоратор @rpc.functions.async_execution должен быть внутренним, чтобы целевая функция распознавалась как статический метод или метод класса. Такая функция по-прежнему может выполняться асинхронно, поскольку статический метод или метод класса сохраняет атрибуты, установленные декоратором @rpc.functions.async_execution.

Пример::

Возвращённый объект Future может быть получен из rpc_async(), then() или конструктора Future. В примере ниже непосредственно используется объект Future, возвращённый функцией then().

>>> from torch.distributed import rpc
>>>
>>> # omitting setup and shutdown RPC
>>>
>>> # On all workers
>>> @rpc.functions.async_execution
>>> def async_add_chained(to, x, y, z):
>>>     # This function runs on "worker1" and returns immediately when
>>>     # the callback is installed through the `then(cb)` API. In the
>>>     # mean time, the `rpc_async` to "worker2" can run concurrently.
>>>     # When the return value of that `rpc_async` arrives at
>>>     # "worker1", "worker1" will run the lambda function accordingly
>>>     # and set the value for the previously returned `Future`, which
>>>     # will then trigger RPC to send the result back to "worker0".
>>>     return rpc.rpc_async(to, torch.add, args=(x, y)).then(
>>>         lambda fut: fut.wait() + z
>>>     )
>>>
>>> # On worker0
>>> ret = rpc.rpc_sync(
>>>     "worker1",
>>>     async_add_chained,
>>>     args=("worker2", torch.ones(2), 1, 1)
>>> )
>>> print(ret)  # prints tensor([3., 3.])

При сочетании с декораторами TorchScript этот декоратор должен быть внешним.

>>> from torch import Tensor
>>> from torch.futures import Future
>>> from torch.distributed import rpc
>>>
>>> # omitting setup and shutdown RPC
>>>
>>> # On all workers
>>> @torch.jit.script
>>> def script_add(x: Tensor, y: Tensor) -> Tensor:
>>>     return x + y
>>>
>>> @rpc.functions.async_execution
>>> @torch.jit.script
>>> def async_add(to: str, x: Tensor, y: Tensor) -> Future[Tensor]:
>>>     return rpc.rpc_async(to, script_add, (x, y))
>>>
>>> # On worker0
>>> ret = rpc.rpc_sync(
>>>     "worker1",
>>>     async_add,
>>>     args=("worker2", torch.ones(2), 1)
>>> )
>>> print(ret)  # prints tensor([2., 2.])

При сочетании со статическим методом или методом класса этот декоратор должен быть внутренним.

>>> from torch.distributed import rpc
>>>
>>> # omitting setup and shutdown RPC
>>>
>>> # On all workers
>>> class AsyncExecutionClass:
>>>
>>>     @staticmethod
>>>     @rpc.functions.async_execution
>>>     def static_async_add(to, x, y, z):
>>>         return rpc.rpc_async(to, torch.add, args=(x, y)).then(
>>>             lambda fut: fut.wait() + z
>>>         )
>>>
>>>     @classmethod
>>>     @rpc.functions.async_execution
>>>     def class_async_add(cls, to, x, y, z):
>>>         ret_fut = torch.futures.Future()
>>>         rpc.rpc_async(to, torch.add, args=(x, y)).then(
>>>             lambda fut: ret_fut.set_result(fut.wait() + z)
>>>         )
>>>         return ret_fut
>>>
>>>     @rpc.functions.async_execution
>>>     def bound_async_add(self, to, x, y, z):
>>>         return rpc.rpc_async(to, torch.add, args=(x, y)).then(
>>>             lambda fut: fut.wait() + z
>>>         )
>>>
>>> # On worker0
>>> ret = rpc.rpc_sync(
>>>     "worker1",
>>>     AsyncExecutionClass.static_async_add,
>>>     args=("worker2", torch.ones(2), 1, 2)
>>> )
>>> print(ret)  # prints tensor([4., 4.])
>>>
>>> ret = rpc.rpc_sync(
>>>     "worker1",
>>>     AsyncExecutionClass.class_async_add,
>>>     args=("worker2", torch.ones(2), 1, 2)
>>> )
>>> print(ret)  # prints tensor([4., 4.])

Этот декоратор также работает с вспомогательными функциями RRef, например: . torch.distributed.rpc.RRef.rpc_sync(), torch.distributed.rpc.RRef.rpc_async() и torch.distributed.rpc.RRef.remote().

>>> from torch.distributed import rpc
>>>
>>> # reuse the AsyncExecutionClass class above
>>> rref = rpc.remote("worker1", AsyncExecutionClass)
>>> ret = rref.rpc_sync().static_async_add("worker2", torch.ones(2), 1, 2)
>>> print(ret)  # prints tensor([4., 4.])
>>>
>>> rref = rpc.remote("worker1", AsyncExecutionClass)
>>> ret = rref.rpc_async().static_async_add("worker2", torch.ones(2), 1, 2).wait()
>>> print(ret)  # prints tensor([4., 4.])
>>>
>>> rref = rpc.remote("worker1", AsyncExecutionClass)
>>> ret = rref.remote().static_async_add("worker2", torch.ones(2), 1, 2).to_here()
>>> print(ret)  # prints tensor([4., 4.])

torch.distributed.rpc.backend_registry.backend_registered(backend_name) [исходный код]

Проверяет, зарегистрировано ли имя backend_name в качестве бэкенда RPC.

Параметры:

backend_name (str) – строка для идентификации бэкенда RPC.

Возвращает:

True, если бэкенд зарегистрирован с помощью register_backend, иначе False.

torch.distributed.rpc.backend_registry.register_backend(backend_name, construct_rpc_backend_options_handler, init_backend_handler) [исходный код]

Регистрирует новый бэкенд RPC.

Параметры:

• backend_name (str) – строка бэкенда для идентификации обработчика.
• construct_rpc_backend_options_handler (function) – обработчик, вызываемый при вызове rpc_backend.construct_rpc_backend_options(**dict).
• init_backend_handler (function) – обработчик, вызываемый при вызове функции _init_rpc_backend() с бэкендом. Возвращает агент.

RRef

RRef (Remote REFerence, удалённая ссылка) — это ссылка на значение некоторого типа T (например, Tensor) на удалённом рабочем процессе. Этот дескриптор поддерживает жизнь указанного удалённого значения на узле-владельце, но не означает, что значение в будущем будет передано локальному рабочему процессу. RRef можно использовать при обучении на нескольких машинах, сохраняя ссылки на nn.Modules, размещённые на других рабочих процессах, и вызывая соответствующие функции для получения или изменения их параметров во время обучения. Подробнее см. в разделе Протокол удалённых ссылок.

class torch.distributed.rpc.PyRRef(RRef)

Класс, содержащий ссылку на значение некоторого типа на удалённом рабочем процессе. Этот дескриптор поддерживает жизнь указанного удалённого значения на рабочем процессе. UserRRef будет удалён, когда 1) на него не останется ссылок ни в коде приложения, ни в локальном контексте RRef либо 2) приложение выполнит корректное завершение работы. Вызов методов удалённого RRef приводит к неопределённому поведению. Реализация RRef обеспечивает лишь обнаружение ошибок с наилучшими усилиями, поэтому приложениям не следует использовать UserRRefs после rpc.shutdown().

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

Сериализовать и десериализовать RRef можно только с помощью модуля RPC. Сериализация и десериализация RRef без RPC (например, с помощью Python pickle, torch save() / load(), JIT save() / load() и т. д.) приведёт к ошибкам.

Параметры:

• value (object) – Значение, которое будет обёрнуто этим RRef.
• type_hint (Type, необязательный) – Тип Python, который следует передать компилятору TorchScript в качестве подсказки типа для value.

Пример::

Для простоты в приведённых ниже примерах пропущен код инициализации и завершения RPC. Подробности см. в документации по RPC.

1. Создание RRef с помощью rpc.remote

>>> import torch
>>> import torch.distributed.rpc as rpc
>>> rref = rpc.remote("worker1", torch.add, args=(torch.ones(2), 3))
>>> # get a copy of value from the RRef
>>> x = rref.to_here()

2. Создание RRef из локального объекта

>>> import torch
>>> from torch.distributed.rpc import RRef
>>> x = torch.zeros(2, 2)
>>> rref = RRef(x)

3. Передача RRef другим рабочим процессам

>>> # On both worker0 and worker1:
>>> def f(rref):
>>>   return rref.to_here() + 1

>>> # On worker0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> from torch.distributed.rpc import RRef
>>> rref = RRef(torch.zeros(2, 2))
>>> # the following RPC shares the rref with worker1, reference
>>> # count is automatically updated.
>>> rpc.rpc_sync("worker1", f, args=(rref,))

backward(self: torch._C._distributed_rpc.PyRRef, dist_autograd_ctx_id: SupportsInt | SupportsIndex = -1, retain_graph: bool = False) → None

Выполняет обратное распространение, используя RRef в качестве корня обратного прохода. Если задан dist_autograd_ctx_id, выполняется распределённый обратный проход с использованием указанного ctx_id, начиная с узла-владельца RRef. В этом случае для получения градиентов следует использовать get_gradients(). Если dist_autograd_ctx_id равен None, предполагается, что это локальный граф autograd, и выполняется только локальный обратный проход. В локальном случае узел, вызывающий этот API, должен быть владельцем RRef. Ожидается, что значение RRef — скалярный тензор.

Параметры:

• dist_autograd_ctx_id (int, необязательный) – Идентификатор контекста распределённого autograd, для которого нужно получить градиенты (по умолчанию: -1).
• retain_graph (bool, необязательный) – Если задано False, граф, использованный для вычисления градиента, будет освобождён. Обратите внимание, что почти во всех случаях устанавливать этот параметр в True не требуется; обычно эту задачу можно решить гораздо эффективнее. Как правило, устанавливать значение True нужно для многократного выполнения обратного прохода (по умолчанию: False).

Пример::

>>> import torch.distributed.autograd as dist_autograd
>>> with dist_autograd.context() as context_id:
>>>     rref.backward(context_id)

confirmed_by_owner(self: torch._C._distributed_rpc.PyRRef) → bool

Возвращает, подтверждён ли этот RRef владельцем. OwnerRRef всегда возвращает true, а UserRRef возвращает true только тогда, когда владельцу известен этот UserRRef.

is_owner(self: torch._C._distributed_rpc.PyRRef) → bool

Возвращает, является ли текущий узел владельцем этого RRef.

local_value(self: torch._C._distributed_rpc.PyRRef) → object

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

owner(self: torch._C._distributed_rpc.PyRRef) → torch._C._distributed_rpc.WorkerInfo

Возвращает сведения о рабочем процессе узла, которому принадлежит этот RRef.

owner_name(self: torch._C._distributed_rpc.PyRRef) → str

Возвращает имя рабочего процесса узла, которому принадлежит этот RRef.

remote(self: torch._C._distributed_rpc.PyRRef, timeout: SupportsFloat | SupportsIndex = -1.0) → object

Создаёт вспомогательный прокси для удобного запуска remote, используя владельца RRef в качестве адресата для выполнения функций над объектом, на который ссылается этот RRef. В частности, rref.remote().func_name(*args, **kwargs) эквивалентен следующему:

>>> def run(rref, func_name, args, kwargs):
>>>   return getattr(rref.local_value(), func_name)(*args, **kwargs)
>>>
>>> rpc.remote(rref.owner(), run, args=(rref, func_name, args, kwargs))

Параметры:

timeout (float, необязательный) – Тайм-аут для rref.remote(). Если создание этого RRef не будет завершено успешно в течение заданного времени, при следующей попытке использовать RRef (например, to_here) возникнет ошибка тайм-аута. Если значение не указано, будет использоваться тайм-аут RPC по умолчанию. Описание особенностей тайм-аутов для RRef см. в rpc.remote().

Пример::

>>> from torch.distributed import rpc
>>> rref = rpc.remote("worker1", torch.add, args=(torch.zeros(2, 2), 1))
>>> rref.remote().size().to_here()  # returns torch.Size([2, 2])
>>> rref.remote().view(1, 4).to_here()  # returns tensor([[1., 1., 1., 1.]])

rpc_async(self: torch._C._distributed_rpc.PyRRef, timeout: SupportsFloat | SupportsIndex = -1.0) → object

Создаёт вспомогательный прокси для удобного запуска rpc_async, используя владельца RRef в качестве адресата для выполнения функций над объектом, на который ссылается этот RRef. В частности, rref.rpc_async().func_name(*args, **kwargs) эквивалентен следующему:

>>> def run(rref, func_name, args, kwargs):
>>>   return getattr(rref.local_value(), func_name)(*args, **kwargs)
>>>
>>> rpc.rpc_async(rref.owner(), run, args=(rref, func_name, args, kwargs))

Параметры:

timeout (float, необязательный) – Тайм-аут для rref.rpc_async(). Если вызов не завершится в течение этого времени, будет вызвано соответствующее исключение. Если этот аргумент не задан, будет использоваться тайм-аут RPC по умолчанию.

Пример::

>>> from torch.distributed import rpc
>>> rref = rpc.remote("worker1", torch.add, args=(torch.zeros(2, 2), 1))
>>> rref.rpc_async().size().wait()  # returns torch.Size([2, 2])
>>> rref.rpc_async().view(1, 4).wait()  # returns tensor([[1., 1., 1., 1.]])

rpc_sync(self: torch._C._distributed_rpc.PyRRef, timeout: SupportsFloat | SupportsIndex = -1.0) → object

Создаёт вспомогательный прокси для удобного запуска rpc_sync, используя владельца RRef в качестве адресата для выполнения функций над объектом, на который ссылается этот RRef. В частности, rref.rpc_sync().func_name(*args, **kwargs) эквивалентен следующему:

>>> def run(rref, func_name, args, kwargs):
>>>   return getattr(rref.local_value(), func_name)(*args, **kwargs)
>>>
>>> rpc.rpc_sync(rref.owner(), run, args=(rref, func_name, args, kwargs))

Параметры:

timeout (float, необязательный) – Тайм-аут для rref.rpc_sync(). Если вызов не завершится в течение этого времени, будет вызвано соответствующее исключение. Если этот аргумент не задан, будет использоваться тайм-аут RPC по умолчанию.

Пример::

>>> from torch.distributed import rpc
>>> rref = rpc.remote("worker1", torch.add, args=(torch.zeros(2, 2), 1))
>>> rref.rpc_sync().size()  # returns torch.Size([2, 2])
>>> rref.rpc_sync().view(1, 4)  # returns tensor([[1., 1., 1., 1.]])

to_here(self: torch._C._distributed_rpc.PyRRef, timeout: SupportsFloat | SupportsIndex = -1.0) → object

Блокирующий вызов, который копирует значение RRef от владельца на локальный узел и возвращает его. Если текущий узел является владельцем, возвращает ссылку на локальное значение.

Параметры:

timeout (float, необязательный) – Тайм-аут для to_here. Если вызов не завершится в течение этого времени, будет вызвано соответствующее исключение. Если этот аргумент не задан, будет использоваться тайм-аут RPC по умолчанию (60 с).

Дополнительные сведения о RRef

• Протокол удалённых ссылок
  • Общие сведения
  • Предположения
  • Время жизни RRef
    • Обоснование проектных решений
    • Реализация
  • Сценарии протокола
    • Передача RRef от пользователя владельцу в качестве возвращаемого значения
    • Передача RRef от пользователя владельцу в качестве аргумента
    • Передача RRef от владельца пользователю
    • Передача RRef от одного пользователя другому

RemoteModule

RemoteModule — это простой способ удалённо создать nn.Module в другом процессе. Сам модуль находится на удалённом узле, но локальный узел располагает дескриптором этого модуля и может вызывать его так же, как обычный nn.Module. Однако при вызове выполняются RPC-вызовы к удалённой стороне; при необходимости их можно выполнять асинхронно с помощью дополнительных API, поддерживаемых RemoteModule.

class torch.distributed.nn.api.remote_module.RemoteModule(*args, **kwargs) [исходный код]

Экземпляр RemoteModule можно создать только после инициализации RPC.

Он создаёт указанный пользователем модуль на заданном удалённом узле. Он ведёт себя как обычный nn.Module, за исключением того, что метод forward выполняется на удалённом узле. Модуль отвечает за запись операций autograd, чтобы обратный проход передавал градиенты соответствующему удалённому модулю.

На основе сигнатуры метода forward класса module_cls создаются два метода: forward_async и forward. Метод forward_async выполняется асинхронно и возвращает Future. Аргументы методов forward_async и forward совпадают с аргументами метода forward модуля, возвращаемого module_cls.

Например, если module_cls возвращает экземпляр nn.Linear с сигнатурой метода forward: def forward(input: Tensor) -> Tensor:, то созданный RemoteModule будет иметь два метода со следующими сигнатурами:

Параметры:

• remote_device (str) – Устройство на целевом рабочем процессе, на котором нужно разместить этот модуль. Формат: «<workername>/<device>», где поле device должно быть интерпретируемым как тип torch.device. Например: «trainer0/cpu», «trainer0», «ps0/cuda:0». Кроме того, поле device можно не указывать; по умолчанию используется значение «cpu».

• module_cls (nn.Module) –

  Класс модуля, который будет создан удалённо. Например:

  >>> class MyModule(nn.Module):
  >>>     def forward(input):
  >>>         return input + 1
  >>>
  >>> module_cls = MyModule
  

• args (Sequence, необязательный) – аргументы, передаваемые в module_cls.
• kwargs (Dict, необязательный) – именованные аргументы, передаваемые в module_cls.

Возвращает:

Экземпляр удалённого модуля, обёртывающий Module, созданный с помощью предоставленного пользователем module_cls. Он имеет блокирующий метод forward и асинхронный метод forward_async, возвращающий future для вызова forward пользовательского модуля на удалённой стороне.

Пример::

Выполните следующий код в двух разных процессах:

>>> # On worker 0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> from torch import nn, Tensor
>>> from torch.distributed.nn.api.remote_module import RemoteModule
>>>
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> remote_linear_module = RemoteModule(
>>>     "worker1/cpu", nn.Linear, args=(20, 30),
>>> )
>>> input = torch.randn(128, 20)
>>> ret_fut = remote_linear_module.forward_async(input)
>>> ret = ret_fut.wait()
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>>
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

Более практический пример совместного использования с DistributedDataParallel (DDP) приведён в этом руководстве.

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

Возвращает RRef (RRef[nn.Module]), указывающий на удалённый модуль.

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

RRef[Module]

remote_parameters(recurse=True) [исходный код]

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

Как правило, это можно использовать вместе с DistributedOptimizer.

Параметры:

recurse (bool) – если True, возвращает параметры удалённого модуля и всех его подмодулей. В противном случае возвращает только параметры, непосредственно входящие в состав удалённого модуля.

Возвращает:

Список RRef (List[RRef[nn.Parameter]]) параметров удалённого модуля.

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

list[RRef[Parameter]]

Распределённая инфраструктура автоградиента

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

torch.distributed.autograd.backward(context_id: int, roots: List[Tensor], retain_graph=False) → None

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

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

Мы накапливаем градиенты в соответствующем torch.distributed.autograd.context на каждом из узлов. Используемый контекст автоградиента определяется по context_id, переданному при вызове torch.distributed.autograd.backward(). Если для указанного идентификатора не существует допустимого контекста автоградиента, возникает ошибка. Получить накопленные градиенты можно с помощью API get_gradients().

Параметры:

• context_id (int) – Идентификатор контекста автоградиента, для которого нужно получить градиенты.
• roots (list) – Тензоры, представляющие корни вычисления автоградиента. Все тензоры должны быть скалярами.
• retain_graph (bool, необязательно) – Если False, граф, использованный для вычисления градиента, будет освобождён. Обратите внимание, что почти во всех случаях устанавливать этот параметр в True не требуется, и часто этого можно добиться гораздо более эффективным способом. Обычно устанавливать True нужно, если требуется выполнить обратный проход несколько раз.

Пример::

>>> import torch.distributed.autograd as dist_autograd
>>> with dist_autograd.context() as context_id:
>>>     pred = model.forward()
>>>     loss = loss_func(pred, loss)
>>>     dist_autograd.backward(context_id, loss)

class torch.distributed.autograd.context [исходный код]

Объект контекста, охватывающий прямой и обратный проходы при использовании распределённого автоградиента. context_id, созданный в операторе with, необходим для уникальной идентификации распределённого обратного прохода на всех рабочих процессах. Каждый рабочий процесс хранит метаданные, связанные с этим context_id, необходимые для корректного выполнения распределённого прохода автоградиента.

Пример::

>>> import torch.distributed.autograd as dist_autograd
>>> with dist_autograd.context() as context_id:
>>>     t1 = torch.rand((3, 3), requires_grad=True)
>>>     t2 = torch.rand((3, 3), requires_grad=True)
>>>     loss = rpc.rpc_sync("worker1", torch.add, args=(t1, t2)).sum()
>>>     dist_autograd.backward(context_id, [loss])

torch.distributed.autograd.get_gradients(context_id: int) → Dict[Tensor, Tensor]

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

Параметры:

context_id (int) – Идентификатор контекста автоградиента, для которого нужно получить градиенты.

Возвращает:

Отображение, в котором ключом является тензор, а значением — соответствующий градиент этого тензора.

Пример::

>>> import torch.distributed.autograd as dist_autograd
>>> with dist_autograd.context() as context_id:
>>>     t1 = torch.rand((3, 3), requires_grad=True)
>>>     t2 = torch.rand((3, 3), requires_grad=True)
>>>     loss = t1 + t2
>>>     dist_autograd.backward(context_id, [loss.sum()])
>>>     grads = dist_autograd.get_gradients(context_id)
>>>     print(grads[t1])
>>>     print(grads[t2])

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

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

bool

Дополнительная информация об автоградиенте RPC

• Проектирование распределённого автоградиента
  • Общие сведения
  • Запись автоградиента во время прямого прохода
  • Контекст распределённого автоградиента
  • Распределённый обратный проход
    • Вычисление зависимостей
    • Алгоритм режима FAST
    • Алгоритм режима SMART
  • Распределённый оптимизатор
  • Простой сквозной пример

Распределённый оптимизатор

Документацию по распределённым оптимизаторам см. на странице torch.distributed.optim.

Примечания по проектированию

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

• Проектирование распределённого автоградиента

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

• Протокол удалённых ссылок

Учебные материалы

Учебные материалы по RPC знакомят пользователей с инфраструктурой RPC, содержат несколько примеров приложений, использующих API torch.distributed.rpc, а также показывают, как использовать профилировщик для профилирования рабочих нагрузок на основе RPC.

• Начало работы с инфраструктурой распределённых RPC
• Реализация сервера параметров с помощью инфраструктуры распределённых RPC
• Совместное использование Distributed DataParallel и инфраструктуры распределённых RPC (также рассматривается RemoteModule)
• Реализация пакетной обработки RPC