| 200 | |
| 201 | template<typename Dtype> |
| 202 | void NCCL<Dtype>::run(int layer) { |
| 203 | CHECK(solver_->param().layer_wise_reduce()); |
| 204 | vector<shared_ptr<Blob<Dtype> > >& blobs = |
| 205 | solver_->net()->layers()[layer]->blobs(); |
| 206 | #ifdef DEBUG |
| 207 | // Assert blobs are contiguous to reduce in one step (e.g. bias often small) |
| 208 | for (int i = 1; i < blobs.size(); ++i) { |
| 209 | CHECK_EQ(blobs[i - 1]->gpu_diff() + blobs[i - 1]->count(), |
| 210 | blobs[i + 0]->gpu_diff()); |
| 211 | } |
| 212 | #endif |
| 213 | if (blobs.size() > 0) { |
| 214 | // Make sure default stream is done computing gradients. Could be |
| 215 | // replaced by cudaEventRecord+cudaStreamWaitEvent to avoid |
| 216 | // blocking the default stream, but it's actually slower. |
| 217 | CUDA_CHECK(cudaStreamSynchronize(cudaStreamDefault)); |
| 218 | |
| 219 | // Reduce asynchronously |
| 220 | int size = 0; |
| 221 | for (int i = 0; i < blobs.size(); ++i) { |
| 222 | size += blobs[i]->count(); |
| 223 | } |
| 224 | if (barrier_) { // NULL in multi process case |
| 225 | barrier_->wait(); |
| 226 | } |
| 227 | NCCL_CHECK(ncclAllReduce(blobs[0]->mutable_gpu_diff(), |
| 228 | blobs[0]->mutable_gpu_diff(), |
| 229 | size, |
| 230 | nccl::dataType<Dtype>::type, |
| 231 | ncclSum, comm_, stream_)); |
| 232 | caffe_gpu_scal(size, (Dtype) 1.0 / Caffe::solver_count(), |
| 233 | blobs[0]->mutable_gpu_diff(), stream_); |
| 234 | } |
| 235 | } |
| 236 | |
| 237 | template<typename Dtype> |
| 238 | void NCCL<Dtype>::on_gradients_ready() { |
no test coverage detected