目录
目录

使用 cpp 扩展自定义进程组后端

创建时间： Feb 01， 2022 |上次更新时间： 2024-11-14 |上次验证： Nov 05， 2024

作者： Howard Huang， Feng Tian， Shen Li， Min Si

注意

编辑在 github 中查看和编辑本教程。

先决条件：

本教程演示了如何使用 cpp 扩展实现自定义并将其插入 PyTorch 分布式包中。当您需要专门的软件时，这很有帮助 堆栈，或者当您想尝试新的 集体通信算法。Backend

基本

PyTorch 集合通信为几个广泛采用的分布式 训练功能，包括 DistributedDataParallel 和 ZeroRedundancyOptimizer。 为了使相同的集合通信 API 与 不同的通信后端，分布式包 abstracts collective communication 操作复制到 Backend 类中。不同的后端可以 然后作为 using preferred 的子类实现 第三方库。PyTorch 分布式带有三个默认后端，即 、 和 。然而 除了这三个后端之外，还有其他通信库 （例如，UCC、OneCCL）、不同类型的硬件 （例如 TPU、Trainum）和新兴 通信算法（例如，Herring、Reduction Server）。 因此，分布式包公开了扩展 API 以允许自定义 集体通信后端。BackendProcessGroupNCCLProcessGroupGlooProcessGroupMPI

以下 4 个步骤显示了如何实现虚拟后端 并在 Python 应用程序代码中使用它。请注意，本教程的重点是 演示扩展 API，而不是开发一个功能齐全的 通信后端。因此，后端只覆盖了 API（ 和 ），并简单地设置 tensor 的值 设置为 0。Backenddummyall_reduceall_gather

第 1 步：实现 的子类Backend

第一步是实现一个子类，该子类将 以集合通信 API 为目标并运行自定义通信算法。 该扩展还需要实现一个子类，该子类 作为通信结果的 future 并允许在 应用程序代码。如果扩展使用第三方库，它可以 包括 headers 并从子类调用库 API。下面的两个代码片段显示了 和 的实现。请参阅 dummy collectives repository 了解完整实现。BackendWorkBackendDummydummy.hdummy.cpp

// file name: dummy.hpp
#include <torch/python.h>

#include <torch/csrc/distributed/c10d/Backend.hpp>
#include <torch/csrc/distributed/c10d/Work.hpp>
#include <torch/csrc/distributed/c10d/Store.hpp>
#include <torch/csrc/distributed/c10d/Types.hpp>
#include <torch/csrc/distributed/c10d/Utils.hpp>

#include <pybind11/chrono.h>

namespace c10d {

class BackendDummy : public Backend {
  public:
    BackendDummy(int rank, int size);

    c10::intrusive_ptr<Work> allgather(
        std::vector<std::vector<at::Tensor>>& outputTensors,
        std::vector<at::Tensor>& inputTensors,
        const AllgatherOptions& opts = AllgatherOptions()) override;

    c10::intrusive_ptr<Work> allreduce(
        std::vector<at::Tensor>& tensors,
        const AllreduceOptions& opts = AllreduceOptions()) override;

    // The collective communication APIs without a custom implementation
    // will error out if invoked by application code.
};

class WorkDummy : public Work {
  public:
    WorkDummy(
      OpType opType,
      c10::intrusive_ptr<c10::ivalue::Future> future) // future of the output
      : Work(
          -1, // rank, only used by recvAnySource, irrelevant in this demo
          opType),
      future_(std::move(future)) {}
    bool isCompleted() override;
    bool isSuccess() const override;
    bool wait(std::chrono::milliseconds timeout = kUnsetTimeout) override;
    virtual c10::intrusive_ptr<c10::ivalue::Future> getFuture() override;

  private:
    c10::intrusive_ptr<c10::ivalue::Future> future_;
};
} // namespace c10d

// file name: dummy.cpp
#include "dummy.hpp"

namespace c10d {

// This is a dummy allgather that sets all output tensors to zero
// Modify the implementation to conduct real communication asynchronously
c10::intrusive_ptr<Work> BackendDummy::allgather(
        std::vector<std::vector<at::Tensor>>& outputTensors,
        std::vector<at::Tensor>& inputTensors,
        const AllgatherOptions& /* unused */) {
    for (auto& outputTensorVec : outputTensors) {
        for (auto& outputTensor : outputTensorVec) {
            outputTensor.zero_();
        }
    }

    auto future = c10::make_intrusive<c10::ivalue::Future>(
        c10::ListType::create(c10::ListType::create(c10::TensorType::get())));
    future->markCompleted(c10::IValue(outputTensors));
    return c10::make_intrusive<WorkDummy>(OpType::ALLGATHER, std::move(future));
}

// This is a dummy allreduce that sets all output tensors to zero
// Modify the implementation to conduct real communication asynchronously
c10::intrusive_ptr<Work> BackendDummy::allreduce(
        std::vector<at::Tensor>& tensors,
        const AllreduceOptions& opts) {
    for (auto& tensor : tensors) {
        tensor.zero_();
    }

    auto future = c10::make_intrusive<c10::ivalue::Future>(
        c10::ListType::create(c10::TensorType::get()));
    future->markCompleted(c10::IValue(tensors));
    return c10::make_intrusive<WorkDummy>(OpType::ALLGATHER, std::move(future));
}
} // namespace c10d

第 2 步：公开扩展 Python API

后端构造函数是从 Python 端调用的， 因此，扩展还需要向 Python 公开构造函数 API。这可以 通过添加以下方法完成。在此示例中，被 instantiation 方法忽略，因为 这些未在此虚拟实现中使用。但是，实际扩展 应考虑使用 to perform rendezvous 并支持该参数。storetimeoutBackendDummystoretimeout

// file name: dummy.hpp
class BackendDummy : public Backend {
    ...
    <Step 1 code>
    ...

    static c10::intrusive_ptr<Backend> createBackendDummy(
        const c10::intrusive_ptr<::c10d::Store>& store,
        int rank,
        int size,
        const std::chrono::duration<float>& timeout);

    static void BackendDummyConstructor() __attribute__((constructor)) {
        py::object module = py::module::import("torch.distributed");
        py::object register_backend =
            module.attr("Backend").attr("register_backend");
        // torch.distributed.Backend.register_backend will add `dummy` as a
        // new valid backend.
        register_backend("dummy", py::cpp_function(createBackendDummy));
    }
}

// file name: dummy.cpp
c10::intrusive_ptr<Backend> BackendDummy::createBackendDummy(
        const c10::intrusive_ptr<::c10d::Store>& /* unused */,
        int rank,
        int size,
        const std::chrono::duration<float>& /* unused */) {
    return c10::make_intrusive<BackendDummy>(rank, size);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("createBackendDummy", &BackendDummy::createBackendDummy);
}

第 3 步：构建自定义扩展

现在，扩展源代码文件已准备就绪。然后，我们可以使用 cpp 扩展来构建它。为此，请创建一个文件来准备路径和 命令。然后致电安装扩展。setup.pypython setup.py develop

如果扩展依赖于第三方库，您还可以指定 and 到 cpp 扩展 API。请参阅 torch ucc 项目作为真实示例。libraries_dirslibraries

# file name: setup.py
import os
import sys
import torch
from setuptools import setup
from torch.utils import cpp_extension

sources = ["src/dummy.cpp"]
include_dirs = [f"{os.path.dirname(os.path.abspath(__file__))}/include/"]

if torch.cuda.is_available():
    module = cpp_extension.CUDAExtension(
        name = "dummy_collectives",
        sources = sources,
        include_dirs = include_dirs,
    )
else:
    module = cpp_extension.CppExtension(
        name = "dummy_collectives",
        sources = sources,
        include_dirs = include_dirs,
    )

setup(
    name = "Dummy-Collectives",
    version = "0.0.1",
    ext_modules = [module],
    cmdclass={'build_ext': cpp_extension.BuildExtension}
)

第 4 步：在应用程序中使用扩展

安装后，你可以方便地在调用 init_process_group 时使用后端，就像它是一个内置的后端一样。dummy

我们可以通过更改 的参数来指定基于 backend 的调度。我们 可以将带有 CPU 张量的集合调度到后端，并将带有 CUDA 张量的集合调度到后端 指定为 backend 参数。backendinit_process_groupgloodummycpu:gloo,cuda:dummy

要将所有张量发送到 backend，我们只需指定 backend 参数即可。dummydummy

import os

import torch
# importing dummy_collectives makes torch.distributed recognize `dummy`
# as a valid backend.
import dummy_collectives

import torch.distributed as dist

os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '29500'

# Alternatively:
# dist.init_process_group("dummy", rank=0, world_size=1)
dist.init_process_group("cpu:gloo,cuda:dummy", rank=0, world_size=1)

# this goes through gloo
x = torch.ones(6)
dist.all_reduce(x)
print(f"cpu allreduce: {x}")

# this goes through dummy
if torch.cuda.is_available():
    y = x.cuda()
    dist.all_reduce(y)
    print(f"cuda allreduce: {y}")

    try:
        dist.broadcast(y, 0)
    except RuntimeError:
        print("got RuntimeError when calling broadcast")

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源