当前位置： 首页 > news >正文

定制PyTorch后端通信（backend）实战

news 来源：原创 2025/6/17 10:29:31

定制PyTorch后端通信（backend）实战

缘起

相关研究中，需要替换PyTorch的通信后端，定制分布式训练中例如all_reduce的算法，包括通信协议的定制。

PyTorch默认的通信后端为‘gloo’和‘nccl’，其他支持的还有mpi（需要基于源码编译），实际上PyTorch 2.7的源代码里还实现了一个‘ucc’的实验性的通信后端，后续可以参考UCC后端的代码进行相关定制。

参见：Distributed communication package - torch.distributed — PyTorch 2.7 documentation

本地环境

$ lsb_release -a   # Ubuntu 23.04

$ python3 --version  # Python 3.11.4

$ python3 -c "import torch; print(torch.__version__)"  # 2.7.0+cu126

实战

对于一个Python和PyTorch的新手，摸索了不知道多久，终于发现解决问题的入口：

1、使用 C++ 扩展定制进程组后端 — PyTorch 教程 2.7.0+cu126 文档 - PyTorch 深度学习库 此链接为中文，标题‘使用c++扩展定制（PyTorch）进程组后端’，正是我要找的东西。

2、https://github.com/H-Huang/torch_collective_extension 一个扩展torch集合通信的例子，该项目给出了两种扩展方式，一种基于custom_backend（推荐的方式），一种基于custom_process_group（旧方式，网上大部分找到的内容都是这种方式），其中基于custom_backend模式的代码与（1）相同，但更全面，默认给出了有关PyTorch所有集合通信的接口示例。

下面主要基于（1）的代码进行验证

代码

在用户目录下创建相关的目录及代码，本例中目录名为‘custom_bankend’，包含四个文件：

  1. dummy.hpp  --  c++头文件
  2. dummy.cpp  --  c++源文件
  3. setup.py  --  Python项目编译和配置脚本
  4. example.py  --  验证测试代码

以下假设用户目录为‘/home/~usrname’

1. dummy.hpp

// file name: dummy.hpp

// 代码源自PyTorch官网实例：https://pytorch.ac.cn/tutorials/intermediate/process_group_cpp_extension_tutorial.html

#pragma once

// 根据实际情况调整路径，参见 setup.py

// #include <torch/python.h>

#include <torch/csrc/api/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.

    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. 注意是在这里指定通信后端的名称（dummy）！

        register_backend("dummy", py::cpp_function(createBackendDummy));

    }

};

class WorkDummy : public Work {

    friend class BackendDummy;

  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

2. dummy.cpp

// file name: dummy.cpp

// 代码源自PyTorch官网实例：https://pytorch.ac.cn/tutorials/intermediate/process_group_cpp_extension_tutorial.html

#include "dummy.hpp"

#include <iostream>

namespace c10d {

bool WorkDummy::isCompleted() {

  return true;

}

bool WorkDummy::isSuccess() const {

  return true;

}

bool WorkDummy::wait(std::chrono::milliseconds /* unused */) {

  return true;

}

c10::intrusive_ptr<c10::ivalue::Future> WorkDummy::getFuture() {

  return future_;

}

// If necessary, pass store/rank/size to the ctor and exchange connection

// information here

BackendDummy::BackendDummy(int rank, int size)

    : Backend(rank, size) {}

// 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));

# 原文此处有误，操作类型应该为：ALLREDUCE

    return c10::make_intrusive<WorkDummy>(OpType::ALLREDUCE, std::move(future));

}

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

}

} // namespace c10d

3. setup.py

# file name: setup.py

# 代码源自PyTorch官网实例：https://pytorch.ac.cn/tutorials/intermediate/process_group_cpp_extension_tutorial.html

# ~/.local/lib/python3.11/site-packages/torch

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/"]

# 在下面的代码行根据本地实际路径替换

sources = ["dummy.cpp"]

include_dirs = ["/home/~~~/.local/lib/python3.11/site-packages/torch/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,

    )

# 不支持 cuda

# 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. example.py

提示：本文件源码改动较大

# 源码初始化有所改动，请密切注意

# 在终端运行命令：

# torchrun --nnodes=1 --nproc-per-node=2 --node_rank=0 --master_addr=127.0.0.1 --master_port=29500 example.py

import os

import torch

import dummy_collectives

import torch.distributed as dist

rank = int(os.getenv('RANK','0'))

world_size = int(os.getenv('WORLD_SIZE','1'))

dist.init_process_group(backend="dummy", rank=rank, world_size=world_size)

# cpu allreduce

x = torch.ones(6)

dist.all_reduce(x)

print(f"cpu allreduce: {x}")

# gpu allrecduce : 实质是计算在 cuda，后端多机之间通信可定制

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

torch.distributed.destroy_process_group()

编译&安装配置

注意将下面所有命令中的路径参数替换为自己本地的HOME路径。

$ python3 setup.py build

成功！在当前目录生成‘build’文件夹以及 .so 和 .egg-info文件。

$ python3 setup.py install

报错，如下：

running install

/usr/lib/python3/dist-packages/setuptools/command/install.py:34: SetuptoolsDeprecationWarning: setup.py install is deprecated. Use build and pip and other standards-based tools.

  warnings.warn(

/usr/lib/python3/dist-packages/setuptools/command/easy_install.py:146: EasyInstallDeprecationWarning: easy_install command is deprecated. Use build and pip and other standards-based tools.

  warnings.warn(

error: can't create or remove files in install directory

The following error occurred while trying to add or remove files in the installation directory:

    [Errno 13] Permission denied: '/usr/local/lib/python3.11/dist-packages/test-easy-install-68533.write-test'

The installation directory you specified (via --install-dir, --prefix, or the distutils default setting) was:

    /usr/local/lib/python3.11/dist-packages/

Perhaps your account does not have write access to this directory?  If the installation directory is a system-owned directory, you may need to sign in as the administrator or "root" account.  If you do not have administrative access to this machine, you may wish to choose a different installation directory, preferably one that is listed in your PYTHONPATH environment variable.

For information on other options, you may wish to consult the documentation at:

  https://setuptools.pypa.io/en/latest/deprecated/easy_install.html

Please make the appropriate changes for your system and try again.

目测是权限的问题，目标安装目录‘/usr/local/lib/python3.11/dist-packages/’不允许普通用户写，真正的安装目录应该在用户HOME目录下，即：

‘/home/~usrname/.local/lib/python3.11/site-packages’

通过输出的错误中的提示在命令行中指定安装目录，如下：

$ python3 setup.py install --install-dir=/home/~usrname/.local/lib/python3.11/site-packages

报错，如下：

usage: setup.py [global_opts] cmd1 [cmd1_opts] [cmd2 [cmd2_opts] ...]

   or: setup.py --help [cmd1 cmd2 ...]

   or: setup.py --help-commands

   or: setup.py cmd --help

error: option --install-dir not recognized

似乎是setup.py后跟随的命令不对

$ python3 setup.py --help

输出：

Common commands: (see '--help-commands' for more)

  setup.py build      will build the package underneath 'build/'

  setup.py install    will install the package

Global options: 。。。

百思不得其解，根据命令提示只支持两个指令：build和install

根据参考资料（1）的提示，将命令改为‘develop’，执行：

$ python3 setup.py develop --install-dir=/home/~usrname/.local/lib/python3.11/site-packages

成功，输出：

running develop

/usr/lib/python3/dist-packages/setuptools/command/easy_install.py:146: EasyInstallDeprecationWarning: easy_install command is deprecated. Use build and pip and other standards-based tools.

  warnings.warn(

/usr/lib/python3/dist-packages/setuptools/command/install.py:34: SetuptoolsDeprecationWarning: setup.py install is deprecated. Use build and pip and other standards-based tools.

  warnings.warn(

running egg_info

creating Dummy_Collectives.egg-info

writing Dummy_Collectives.egg-info/PKG-INFO

writing dependency_links to Dummy_Collectives.egg-info/dependency_links.txt

writing top-level names to Dummy_Collectives.egg-info/top_level.txt

writing manifest file 'Dummy_Collectives.egg-info/SOURCES.txt'

/home/~usrname/.local/lib/python3.11/site-packages/torch/utils/cpp_extension.py:576: UserWarning: Attempted to use ninja as the BuildExtension backend but we could not find ninja.. Falling back to using the slow distutils backend.

  warnings.warn(msg.format('we could not find ninja.'))

reading manifest file 'Dummy_Collectives.egg-info/SOURCES.txt'

writing manifest file 'Dummy_Collectives.egg-info/SOURCES.txt'

running build_ext

copying build/lib.linux-x86_64-cpython-311/dummy_collectives.cpython-311-x86_64-linux-gnu.so ->

Creating /home/~usrname/.local/lib/python3.11/site-packages/Dummy-Collectives.egg-link (link to .)

Dummy-Collectives 0.0.1 is already the active version in easy-install.pth

Installed /home/~usrname/src/pytorch/custom_backend

Processing dependencies for Dummy-Collectives==0.0.1

Finished processing dependencies for Dummy-Collectives==0.0.1

验证

$ pip3 show dummy_collectives

Name: Dummy-Collectives

Version: 0.0.1

Summary:

Home-page:

Author:

Author-email:

License:

Location: /home/~usrname/src/pytorch/custom_backend

Editable project location: /home/~usrname/src/pytorch/custom_backend

Requires:

Required-by:

命令执行成功，说明已经完成通信后端 dummy 的注册。

$ python3 -c "import dummy_collectives"

报错，如下：

Traceback (most recent call last):

  File "<string>", line 1, in <module>

ImportError: libc10.so: cannot open shared object file: No such file or directory

通过find命令定位‘libc10.so’

$ find ~ -name libc10.so

根据搜索结果将文件路径添加到环境变量‘LD_LIBRARY_PATH’

本人测试时文件路径为：

/home/~usrname/.local/lib/python3.11/site-packages/torch/lib/libc10.so

$ export LD_LIBRARY_PATH=…

更新环境变量后再次执行：

$ python3 -c "import dummy_collectives"

继续报错，如下：

Traceback (most recent call last):

  File "<string>", line 1, in <module>

ImportError: /home/~usrname/.local/lib/python3.11/site-packages/torch/lib/../../nvidia/cusparse/lib/libcusparse.so.12: undefined symbol: __nvJitLinkGetErrorLog_12_6, version libnvJitLink.so.12

这次是找不到‘libnvJitLink.so.12’

继续通过find命令定位文件，本人测试的文件路径为：

/home/~usrname/.local/lib/python3.11/site-packages/nvidia/nvjitlink/lib/libnvJitLink.so.12

继续将文件路径加入到环境变量‘LD_LIBRARY_PATH’

重要提示：经后续验证，在‘import dummy_collectives’之前执行‘import torch’即可无需配置环境变量‘LD_LIBRARY_PATH’，猜测原因可能是执行‘import torch’时自定加载了相关的.so文件。正确命令如下：

$ python3 -c "import torch; import dummy_collectives"

再次执行测试：

$ python3 -c "import dummy_collectives"

这次没有报错，有一个警告，如下：

/home/~usrname/.local/lib/python3.11/site-packages/torch/distributed/distributed_c10d.py:350: UserWarning: Device capability of dummy unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.

  warnings.warn(

试图消除此警告，通过修改‘.hpp’文件中‘register_backend()’方法，增加形如‘devices="cpu"’的参数，未果，编译不能通过，应该是torch版本的问题，暂时搁置不管，继续验证。

$ torchrun --nnodes=1 --nproc-per-node=1 --node_rank=0 --master_addr=127.0.0.1 --master_port=29500 example.py

成功！输出：

/home/~usrname/.local/lib/python3.11/site-packages/torch/distributed/distributed_c10d.py:350: UserWarning: Device capability of dummy unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.

  warnings.warn(

cpu allreduce: tensor([0., 0., 0., 0., 0., 0.])

cuda allreduce: tensor([0., 0., 0., 0., 0., 0.], device='cuda:0')

got RuntimeError when calling broadcast

提示：可以在上述命令中通过指定‘--nproc-per-node=x’参数启动多个进程同时分布式运行。

重要发现：

如果在import dummy_collectives 之前执行了 import torch，则无需指定之前的环境变量‘LD_LIBRARY_PATH’的操作。

后续工作：基于UDP和RoCE-v2协议定制PyTorch后端的all_reduce操作。

相关文章：

  • MCP终极指南 - 番外篇：抓包分析 Cline 与模型的交互协议
  • ollama常见属性设置
  • AIGC 使用层 | 关于AI prompt
  • Spring常用的注解详细介绍 原理分析，与实际案例
  • 黑客是如何攻击一部手机的？
  • Objective-C与Swift混合编程
  • 使用SVN checkout时报错Unable to connect to a repository at URL
  • LINUX616 问题：vim编辑器；粘滞位（其他需回顾）；配置静态网络报错；more,less,head命令
  • 如何确定驱动480x320分辨率的显示屏所需的MCU主频
  • 微信小程序页面容器弹出层
  • Next.js面试题：API深度解析
  • LangChain 与 Milvus 的碰撞：全文检索技术实践
  • java 设计模式_行为型_23状态模式
  • 论文笔记：CAN LLMS UNDERSTAND TIME SERIES ANOMALIES?
  • Kafka Connect实战：从环境搭建到全流程操作
  • solana 编写智能合约 然后调用它
  • C#/.NET/.NET Core技术前沿周刊 | 第 42 期（2025年6.9-6.15）
  • Debian 编译安装 ruby3.2
  • webpack到vite的改造之路
  • SOME/IP学习随笔
  • 扁平化设计网站 源码/最近疫情最新消息
  • 便宜虚拟主机做网站备份/拼多多seo搜索优化
  • 如何制作网站页面/怎样优化标题关键词
  • 国内外优秀设计网站/深圳网络seo推广
  • 潍坊哪家网站制作公司好/好口碑的关键词优化
  • 佛山定制建站公司推荐/网站收录查询爱站