| 24 | namespace tensorflow { |
| 25 | |
| 26 | void NcclGatherer::Run(StatusCallback done) { |
| 27 | auto* compute_stream = col_ctx_->op_ctx->op_device_context()->stream(); |
| 28 | auto* gpu_info = col_ctx_->op_ctx->device()->tensorflow_gpu_device_info(); |
| 29 | const int num_global_devices = col_params_->group.group_size; |
| 30 | const int num_local_devices = col_params_->instance.num_devices_per_task.at( |
| 31 | col_params_->instance.task_names[col_params_->default_rank]); |
| 32 | string nccl_collective_key = |
| 33 | NcclCollectiveKey(col_ctx_->exec_key, col_ctx_->step_id); |
| 34 | auto participant = absl::make_unique<NcclManager::Participant>( |
| 35 | compute_stream->parent(), compute_stream, gpu_info, col_ctx_->input, |
| 36 | col_ctx_->output, col_params_->default_rank, std::move(done)); |
| 37 | VLOG(1) << "NcclGatherer calling NcclManager::AddToAllGather num_tasks " |
| 38 | << col_params_->group.num_tasks << " current task " |
| 39 | << col_params_->instance.task_names[col_params_->default_rank] |
| 40 | << " num local devices " << num_local_devices |
| 41 | << " num global devices " << num_global_devices << " rank " |
| 42 | << col_params_->default_rank << " device " << col_ctx_->device_name |
| 43 | << " instance " << col_params_->instance.instance_key; |
| 44 | NcclManager::instance()->AddToAllGather( |
| 45 | std::move(participant), |
| 46 | {std::move(nccl_collective_key), num_local_devices, num_global_devices, |
| 47 | col_params_->group.runtime_details.communicator_key, |
| 48 | /*source_rank=*/-1}); |
| 49 | { |
| 50 | // `WaitForDependencies` may block if the collective instances on which this |
| 51 | // op depends have not yet launched. When this function returns, this op is |
| 52 | // ready to go. |
| 53 | profiler::TraceMe activity("WaitForDependencies", |
| 54 | profiler::TraceMeLevel::kInfo); |
| 55 | col_ctx_->col_exec->WaitForDependencies(*col_params_); |
| 56 | NcclManager::instance()->SignalMultiNodeReady(nccl_collective_key); |
| 57 | } |
| 58 | { |
| 59 | // When all devices at this worker have called `SignalMultiNodeReady`, the |
| 60 | // `NcclManager` will enqueue the NCCL kernel on the NCCL stream. Thus the |
| 61 | // implementation of `Launched` keeps track of the number of devices that |
| 62 | // have launched. |
| 63 | profiler::TraceMe activity("Schedule", profiler::TraceMeLevel::kInfo); |
| 64 | col_ctx_->col_exec->Launched(*col_params_); |
| 65 | } |
| 66 | } |
| 67 | |
| 68 | REGISTER_COLLECTIVE(NcclGather, NcclGatherer); |
| 69 |
nothing calls this directly
no test coverage detected