| 24 | namespace tensorflow { |
| 25 | |
| 26 | void NcclBroadcaster::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) |
| 38 | << "NcclBroadcast calling NcclManager::AddBroadcastSend/Recv num_tasks " |
| 39 | << col_params_->group.num_tasks << " current task " |
| 40 | << col_params_->instance.task_names[col_params_->default_rank] |
| 41 | << " num local devices " << num_local_devices << " num global devices " |
| 42 | << num_global_devices << " rank " << col_params_->default_rank |
| 43 | << " device " << col_ctx_->device_name << " instance " |
| 44 | << col_params_->instance.instance_key << " source " |
| 45 | << col_params_->is_source; |
| 46 | if (col_params_->is_source) { |
| 47 | NcclManager::instance()->AddBroadcastSend( |
| 48 | std::move(participant), |
| 49 | {std::move(nccl_collective_key), num_local_devices, num_global_devices, |
| 50 | col_params_->group.runtime_details.communicator_key, |
| 51 | col_params_->source_rank}); |
| 52 | } else { |
| 53 | NcclManager::instance()->AddBroadcastRecv( |
| 54 | std::move(participant), |
| 55 | {std::move(nccl_collective_key), num_local_devices, num_global_devices, |
| 56 | col_params_->group.runtime_details.communicator_key, |
| 57 | col_params_->source_rank}); |
| 58 | } |
| 59 | { |
| 60 | // `WaitForDependencies` may block if the collective instances on which this |
| 61 | // op depends have not yet launched. When this function returns, this op is |
| 62 | // ready to go. |
| 63 | profiler::TraceMe activity("WaitForDependencies", |
| 64 | profiler::TraceMeLevel::kInfo); |
| 65 | col_ctx_->col_exec->WaitForDependencies(*col_params_); |
| 66 | NcclManager::instance()->SignalMultiNodeReady(nccl_collective_key); |
| 67 | } |
| 68 | { |
| 69 | // When all devices at this worker have called `SignalMultiNodeReady`, the |
| 70 | // `NcclManager` will enqueue the NCCL kernel on the NCCL stream. Thus the |
| 71 | // implementation of `Launched` keeps track of the number of devices that |
| 72 | // have launched. |
| 73 | profiler::TraceMe activity("Schedule", profiler::TraceMeLevel::kInfo); |
| 74 | col_ctx_->col_exec->Launched(*col_params_); |
| 75 | } |
| 76 | } |
| 77 | |
| 78 | REGISTER_COLLECTIVE(NcclBroadcast, NcclBroadcaster); |
| 79 |
nothing calls this directly
no test coverage detected