使用 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>

namespacec10d{

classBackendDummy:publicBackend{
public:
BackendDummy(intrank,intsize);

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

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

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

classWorkDummy:publicWork{
public:
WorkDummy(
OpTypeopType,
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)){}
boolisCompleted()override;
boolisSuccess()constoverride;
boolwait(std::chrono::millisecondstimeout=kUnsetTimeout)override;
virtualc10::intrusive_ptr<c10::ivalue::Future>getFuture()override;

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

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

namespacec10d{

// 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,
constAllgatherOptions&/* unused */){
for(auto&outputTensorVec:outputTensors){
for(auto&outputTensor:outputTensorVec){
outputTensor.zero_();
}
}

autofuture=c10::make_intrusive<c10::ivalue::Future>(
c10::ListType::create(c10::ListType::create(c10::TensorType::get())));
future->markCompleted(c10::IValue(outputTensors));
returnc10::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,
constAllreduceOptions&opts){
for(auto&tensor:tensors){
tensor.zero_();
}

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

// file name: dummy.hpp
classBackendDummy:publicBackend{
...
<Step1code>
...

staticc10::intrusive_ptr<Backend>createBackendDummy(
constc10::intrusive_ptr<::c10d::Store>&store,
intrank,
intsize,
conststd::chrono::duration<float>&timeout);

staticvoidBackendDummyConstructor()__attribute__((constructor)){
py::objectmodule=py::module::import("torch.distributed");
py::objectregister_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(
constc10::intrusive_ptr<::c10d::Store>&/* unused */,
intrank,
intsize,
conststd::chrono::duration<float>&/* unused */){
returnc10::make_intrusive<BackendDummy>(rank,size);
}

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

# file name: setup.py
importos
importsys
importtorch
fromsetuptoolsimport setup
fromtorch.utilsimport 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}
)

importos

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

importtorch.distributedasdist

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