使用 Cpp 扩展定制进程组后端

Created On: Feb 01, 2022 | Last Updated: Nov 14, 2024 | Last Verified: Nov 05, 2024

作者: Howard Huang、Feng Tian、Shen Li、Min Si

备注

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

前提条件：

• PyTorch 分布式概述

• PyTorch 集体通信包

• PyTorch Cpp 扩展

• 使用 PyTorch 写分布式应用程序

本教程演示如何实现一个自定义 Backend 并将其接入到 PyTorch 分布式包 中，使用 Cpp 扩展。这在您需要针对硬件的专用软件栈或希望尝试新的集体通信算法时非常有用。

基础知识

PyTorch 集体通信支持多个被广泛采用的分布式训练功能，包括 DistributedDataParallel 和 ZeroRedundancyOptimizer。为了使同一集体通信 API 在不同通信后端中工作，分布式包将集体通信操作抽象为 Backend 类。不同的后端可以通过 Backend 的子类来实现，并使用首选的第三方库。PyTorch 分布式包提供了三个默认后端：ProcessGroupNCCL、ProcessGroupGloo 和 ProcessGroupMPI。然而，除这三个后端外，还有其他通信库（如 UCC、OneCCL）、不同类型的硬件（如 TPU、Trainum）以及新兴通信算法（如 Herring、Reduction Server）。因此，分布式包提供扩展 API 以允许定制集体通信后端。

下面的 4 个步骤展示了如何实现一个示例 Backend 后端并在 Python 应用代码中使用它。请注意，本教程旨在展示扩展 API，而不是开发一个可用的通信后端。因此，dummy 后端仅覆盖了部分 API（all_reduce 和 all_gather），并且只是将张量的值设置为 0。

步骤 1：实现 Backend 的子类

第一步是实现一个 Backend 子类，该子类覆盖目标集体通信 API 并运行自定义通信算法。扩展还需要实现 Work 的子类，作为通信结果的未来并允许在应用代码中异步执行。如果扩展使用了第三方库，它可以包含头文件并从 BackendDummy 子类中调用库 API。以下两个代码片段展示了 dummy.h 和 dummy.cpp 的实现。查看 dummy collectives 仓库获取完整实现。

// 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 侧调用的，因此扩展也需要将构造函数 API 暴露给 Python。可以通过添加以下方法实现。在本例中，store 和 timeout 被 BackendDummy 实例化方法忽略，因为这些在此示例实现中未使用。然而，真实扩展应考虑使用 store 来执行 rendezvous 并支持 timeout 参数。

// 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.py 文件，准备路径和命令。然后调用 python setup.py develop 安装扩展。

如果扩展依赖于第三方库，您还可以将 libraries_dirs 和 libraries 指定给 Cpp 扩展 API。查看作为真实示例的 torch ucc 项目。

# 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 后端，就像它是一个内置后端一样。

我们可以通过更改 init_process_group 的 backend 参数来基于后端指定调度。我们可以将 CPU 张量的集体调度指定给 gloo 后端，将 CUDA 张量的集体调度指定给 dummy 后端，只需将 cpu:gloo,cuda:dummy 作为 backend 参数传入即可。

要将所有张量发送到 dummy 后端，我们可以简单地将 dummy 作为 backend 参数指定。

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")