mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 21:49:24 +08:00
1ad0048b64
Headers under torch/csrc/distributed may be referened with relative path, e.g., "<c10d/...>". However, relative path cannot be gracefully handled by Meta internal build when the NCCL PG is hipified to support AMD/RCCL because the "hipified" header files are generated in other directories. Moreover, using absolute path for header inclusion is the state-of-the-art in most components in Pytorch. Thus, this patch refactors all header paths in torch/csrc/distributed to be absolute. See D39835774 for more details about Meta internal complication. **How to test**: commit 9e5d199 removes -I./torch/csrc/distributed in compile options. Thus use it to verify we don't miss any relative path use of torch/csrc/distributed headers. Pull Request resolved: https://github.com/pytorch/pytorch/pull/85780 Approved by: https://github.com/kumpera, https://github.com/huydhn
34 lines
944 B
C++
34 lines
944 B
C++
#include <c10/util/irange.h>
|
|
#include <torch/csrc/distributed/c10d/FileStore.hpp>
|
|
#include <torch/csrc/distributed/c10d/ProcessGroupGloo.hpp>
|
|
|
|
using namespace ::c10d;
|
|
|
|
int main(int argc, char** argv) {
|
|
int rank = atoi(getenv("RANK"));
|
|
int size = atoi(getenv("SIZE"));
|
|
auto store = c10::make_intrusive<FileStore>("/tmp/c10d_example", size);
|
|
ProcessGroupGloo pg(store, rank, size);
|
|
|
|
// Create some tensors
|
|
const auto ntensors = 10;
|
|
std::vector<at::Tensor> tensors;
|
|
for (const auto i : c10::irange(ntensors)) {
|
|
auto x =
|
|
at::ones({1000, 16 * (i + 1)}, at::TensorOptions(at::CPU(at::kFloat)));
|
|
tensors.push_back(x);
|
|
}
|
|
|
|
// Kick off work
|
|
std::vector<c10::intrusive_ptr<Work>> pending;
|
|
for (const auto i : c10::irange(ntensors)) {
|
|
std::vector<at::Tensor> tmp = {tensors[i]};
|
|
pending.push_back(pg.allreduce(tmp));
|
|
}
|
|
|
|
// Wait for work to complete
|
|
for (auto& work : pending) {
|
|
work->wait();
|
|
}
|
|
}
|