Интеграция распределённого обучения

Общие сведения

Распределённое обучение позволяет ускорителям масштабировать вычислительные нагрузки на несколько устройств и узлов, координируя коллективный обмен данными (например, allreduce, broadcast, allgather) через бэкенд ProcessGroup. PyTorch поставляется со встроенными бэкендами, такими как NCCL для CUDA и Gloo для CPU, однако фреймворк предоставляет механизм регистрации, позволяющий сторонним поставщикам ускорителей подключать собственную библиотеку коллективного обмена данными без изменения основного кода.

Интеграцию можно разделить на три уровня:

1. Реализация бэкенда на C++ – подкласс c10d::Backend, реализующий операции коллективного обмена данными, обмена точка-точка и синхронизации.
2. Привязки Python – предоставляют класс бэкенда C++ в Python с помощью pybind11.
3. Регистрация бэкенда – регистрирует бэкенд в torch.distributed.Backend.register_backend(), чтобы init_process_group мог обнаружить и создать его экземпляр.

Перед началом работы

Это руководство посвящено только интеграции бэкенда ProcessGroup — регистрации пользовательского бэкенда коллективного обмена данными в torch.distributed. В нём не рассматривается интеграция полного стека с высокоуровневыми API, такими как DDP, FSDP или другие стратегии распределённого обучения.

Перед тем как следовать этому руководству, убедитесь, что у вас есть:

• Импортируемый пакет расширения torch_xxx, регистрирующий ваше устройство с помощью PrivateUse1. Сведения о регистрации устройств, операторах и хуках среды выполнения приведены в предыдущих главах этого руководства.
• Библиотека коллективного обмена данными (CCL) с реализациями базовых коллективных операций, таких как allreduce и broadcast, для вашего устройства. Библиотека CCL может предоставляться поставщиком (например, NCCL для NVIDIA, HCCL для Huawei) или быть пользовательской реализацией.

Проектирование

В этом разделе описаны интерфейсы и понятия, связанные с регистрацией бэкенда.

Реализация

В этом разделе описаны практические шаги по реализации и регистрации бэкенда на примере эталонной реализации OCCL. Реализация состоит из трёх шагов:

1. Реализовать бэкенд на C++
2. Создать привязки Python
3. Зарегистрировать бэкенд в Python

 1  class DummyWork : public Work {
 2   public:
 3    DummyWork();
 4
 5    virtual ~DummyWork();
 6    bool isCompleted() override;
 7    bool isSuccess() const override;
 8    bool wait(std::chrono::milliseconds timeout) override;
 9    void synchronize() override;
10    void abort() override;
11    c10::intrusive_ptr<c10::ivalue::Future> getFuture() override;
12
13   protected:
14    friend class ProcessGroupOCCL;
15
16   private:
17    c10::intrusive_ptr<c10::ivalue::Future> future_;
18  };

 1void initProcessGroupBindings(py::module& m) {
 2  py::class_<c10d::ProcessGroupOCCL, c10d::Backend, c10::intrusive_ptr<c10d::ProcessGroupOCCL>>(m, "ProcessGroupOCCL")
 3      .def(
 4          py::init([](const c10::intrusive_ptr<::c10d::Store>& /*store*/,
 5                      int rank,
 6                      int size,
 7                      std::chrono::milliseconds /*timeout*/) {
 8            return c10::make_intrusive<::c10d::ProcessGroupOCCL>(rank, size);
 9          }),
10          py::arg("store"),
11          py::arg("rank"),
12          py::arg("size"),
13          py::arg("timeout") = std::chrono::milliseconds(30 * 60 * 1000));
14}

 1if torch.distributed.is_available():
 2    try:
 3        from torch_openreg._C import ProcessGroupOCCL
 4
 5        def _create_occl_backend(store, rank, size, timeout):
 6            return ProcessGroupOCCL(store, rank, size, timeout)
 7
 8        torch.distributed.Backend.register_backend(
 9            "occl", _create_occl_backend, devices=["openreg"]
10        )
11    except Exception as e:
12        raise RuntimeError("Failed to register 'occl' process group backend.") from e

Использование

После регистрации бэкенд без проблем интегрируется с torch.distributed:

import torch
import torch.distributed as dist

# Import triggers autoload, which registers the "occl" backend
import torch_openreg

# Initialize process group – OCCL is auto-selected for openreg devices
dist.init_process_group(
    backend="occl",
    init_method="env://",
    world_size=2,
    rank=0,
)

# Use standard distributed APIs
tensor = torch.randn(4, device="openreg")
dist.all_reduce(tensor)

dist.destroy_process_group()

Кроме того, имя бэкенда можно не указывать, если задан device_id: PyTorch выберет бэкенд по сопоставлению устройств и бэкендов:

dist.init_process_group(
    device_id=torch.device("openreg:0"),
    init_method="env://",
    world_size=2,
    rank=0,
)

dist.init_process_group(
    backend="cpu:gloo,openreg:occl",
    init_method="env://",
    world_size=2,
    rank=0,
)

Тестирование

Основные моменты тестирования:

• Убедитесь, что после импорта бэкенд появляется в dist.Backend.backend_list.
• Проверьте, что init_process_group / destroy_process_group выполняется успешно.
• Убедитесь, что коллективные операции принимают тензоры на зарегистрированном устройстве и возвращают завершённые объекты Work.
• Для выполнения тестов в нескольких процессах используйте MultiProcessTestCase из torch.testing._internal.common_distributed.

Пример эталонной реализации см. в наборе тестов OCCL.