| 40 | } |
| 41 | |
| 42 | void flushReduceBucket(at::ScalarType scalar_type) override |
| 43 | { |
| 44 | if (!hasKey(reduce_tasks_, scalar_type)) { return; } |
| 45 | |
| 46 | blockCopyEvents(scalar_type); |
| 47 | applyPreDivision(scalar_type); |
| 48 | |
| 49 | // NCCL AllReduce operation |
| 50 | ncclGroupStart(); |
| 51 | for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { |
| 52 | ncclResult_t result = ncclAllReduce(t.getSendBuf().data_ptr(), |
| 53 | t.getSendBuf().data_ptr(), |
| 54 | t.getSendBuf().numel(), |
| 55 | get_nccl_data_type(scalar_type), |
| 56 | getReductionOp(), |
| 57 | nccl_comm_, |
| 58 | rs_stream_); |
| 59 | if (result != ncclSuccess) { throw std::runtime_error("NCCL AllReduce failed"); } |
| 60 | } |
| 61 | ncclGroupEnd(); |
| 62 | |
| 63 | // Copy or accumulate results to gradient buffers |
| 64 | { |
| 65 | at::cuda::CUDAStreamGuard guard(rs_stream_); |
| 66 | for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { |
| 67 | bool acc_grad = has_acc_grad_.at(t.getDSId()); |
| 68 | auto param = param_registry_->getParam(t.getDSId()); |
| 69 | auto grad_buf = param.getGradBuffer().flatten(); |
| 70 | |
| 71 | if (grad_buf.numel() == 0) { continue; } |
| 72 | |
| 73 | int64_t offset = param.getOffset(); |
| 74 | auto recv_buf = t.getSendBuf().flatten().index( |
| 75 | {torch::indexing::Slice(offset, offset + grad_buf.numel())}); |
| 76 | if (acc_grad) { |
| 77 | grad_buf.add_(recv_buf); |
| 78 | } else { |
| 79 | grad_buf.copy_(recv_buf); |
| 80 | } |
| 81 | has_acc_grad_[t.getDSId()] = true; |
| 82 | } |
| 83 | } |
| 84 | |
| 85 | performCleanup(scalar_type); |
| 86 | } |
| 87 | }; |
| 88 | |
| 89 | namespace { |
nothing calls this directly
no test coverage detected