Skip to content

Fix broadcast copying device[0] tensor when not using NCCL #8222

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ssnl merged 4 commits into from
Jun 19, 2018
Merged

Fix broadcast copying device[0] tensor when not using NCCL #8222

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
922402b
Fix broadcast copying device[0] tensor when not using NCCL; Avoids po…
ssnl Jun 6, 2018
e0919aa
use toType
ssnl Jun 6, 2018
115c8d5
revert dense_flat changes
ssnl Jun 7, 2018
3eca907
address comments
ssnl Jun 18, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion aten/src/ATen/DeviceGuard.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ struct DeviceGuard {
}
}

/// Sets the device to the given one if its index is not `nullopt`.
/// Sets the device to the given one.
void set_index(int32_t index) {
if (index == -1) {
return;
Expand Down
11 changes: 8 additions & 3 deletions torch/csrc/cuda/comm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,11 +35,12 @@ std::vector<Tensor> broadcast(const Tensor& tensor, IntList devices) {
"first on devices list");
std::vector<Tensor> tensors;
tensors.reserve(devices.size());
at::DeviceGuard _device_guard;
#ifdef USE_NCCL
if (nccl::is_available({tensor})) {
tensors.push_back(tensor);
for (auto device : devices.slice(1)) {
at::DeviceGuard _device_guard(device);
_device_guard.set_index(device);
tensors.push_back(type.tensor(tensor.sizes()));
}
nccl::broadcast(tensors);
Expand All @@ -48,8 +49,12 @@ std::vector<Tensor> broadcast(const Tensor& tensor, IntList devices) {
{
#endif
auto & gpu_type = type.toBackend(type.is_sparse() ? at::kSparseCUDA : at::kCUDA);
for (auto device : devices) {
at::DeviceGuard _device_guard(device);
if (type.is_cuda()) {
tensors.push_back(tensor);
}
IntList loop_devices = type.is_cuda() ? devices.slice(1) : devices;
for (auto device : loop_devices) {
_device_guard.set_index(device);
tensors.push_back(gpu_type.copy(tensor, true));
}
}
Expand Down