(
program, rank, nranks, wait_port, current_endpoint, endpoints
)
| 150 | |
| 151 | |
| 152 | def init_communicator( |
| 153 | program, rank, nranks, wait_port, current_endpoint, endpoints |
| 154 | ): |
| 155 | if nranks < 2: |
| 156 | return |
| 157 | endpoints_str = ",".join(endpoints) |
| 158 | other_endpoints = endpoints[:] |
| 159 | other_endpoints.remove(current_endpoint) |
| 160 | block = program.global_block() |
| 161 | if rank == 0 and wait_port: |
| 162 | wait_server_ready(other_endpoints) |
| 163 | if core.is_compiled_with_cuda(): |
| 164 | nccl_id_var = block.create_var( |
| 165 | name=base.unique_name.generate('nccl_id'), |
| 166 | persistable=True, |
| 167 | type=base.core.VarDesc.VarType.RAW, |
| 168 | ) |
| 169 | |
| 170 | block.append_op( |
| 171 | type='c_gen_nccl_id', |
| 172 | inputs={}, |
| 173 | outputs={'Out': nccl_id_var}, |
| 174 | attrs={ |
| 175 | 'rank': rank, |
| 176 | 'endpoint': current_endpoint, |
| 177 | 'other_endpoints': other_endpoints, |
| 178 | }, |
| 179 | ) |
| 180 | |
| 181 | block.append_op( |
| 182 | type='c_comm_init', |
| 183 | inputs={'X': nccl_id_var}, |
| 184 | outputs={}, |
| 185 | attrs={ |
| 186 | 'nranks': nranks, |
| 187 | 'rank': rank, |
| 188 | 'ring_id': 0, |
| 189 | 'endpoints': endpoints_str, |
| 190 | }, |
| 191 | ) |
| 192 | elif core.is_compiled_with_xpu(): |
| 193 | bkcl_id_var = block.create_var( |
| 194 | name=base.unique_name.generate('bkcl_id'), |
| 195 | persistable=True, |
| 196 | type=base.core.VarDesc.VarType.RAW, |
| 197 | ) |
| 198 | |
| 199 | block.append_op( |
| 200 | type='c_gen_bkcl_id', |
| 201 | inputs={}, |
| 202 | outputs={'Out': bkcl_id_var}, |
| 203 | attrs={ |
| 204 | 'rank': rank, |
| 205 | 'endpoint': current_endpoint, |
| 206 | 'other_endpoints': other_endpoints, |
| 207 | }, |
| 208 | ) |
| 209 |
no test coverage detected