| 253 | |
| 254 | |
| 255 | def prepare_distributed_context(place=None): |
| 256 | if place is None: |
| 257 | place = ( |
| 258 | base.CUDAPlace(paddle.distributed.ParallelEnv().dev_id) |
| 259 | if paddle.distributed.ParallelEnv().nranks > 1 |
| 260 | else base.CUDAPlace(0) |
| 261 | ) |
| 262 | |
| 263 | place = _get_paddle_place(place) |
| 264 | strategy = paddle.distributed.parallel.ParallelStrategy() |
| 265 | strategy.nranks = paddle.distributed.ParallelEnv().nranks |
| 266 | strategy.local_rank = paddle.distributed.ParallelEnv().local_rank |
| 267 | strategy.trainer_endpoints = ( |
| 268 | paddle.distributed.ParallelEnv().trainer_endpoints |
| 269 | ) |
| 270 | strategy.current_endpoint = ( |
| 271 | paddle.distributed.ParallelEnv().current_endpoint |
| 272 | ) |
| 273 | |
| 274 | if strategy.nranks < 2: |
| 275 | return |
| 276 | |
| 277 | global _parallel_context_initialized |
| 278 | |
| 279 | if not _parallel_context_initialized and isinstance(place, base.CUDAPlace): |
| 280 | |
| 281 | def _init_context(): |
| 282 | communicator_prog = base.Program() |
| 283 | init_communicator( |
| 284 | communicator_prog, |
| 285 | strategy.local_rank, |
| 286 | strategy.nranks, |
| 287 | True, |
| 288 | strategy.current_endpoint, |
| 289 | strategy.trainer_endpoints, |
| 290 | ) |
| 291 | exe = base.Executor(place) |
| 292 | exe.run(communicator_prog) |
| 293 | |
| 294 | if in_dynamic_mode(): |
| 295 | base.disable_dygraph() |
| 296 | _init_context() |
| 297 | base.enable_dygraph(place) |
| 298 | |
| 299 | else: |
| 300 | assert "Only support CUDAPlace for now." |
| 301 | |
| 302 | _parallel_context_initialized = True |
| 303 | return strategy |
| 304 | |
| 305 | |
| 306 | def _update_input_info(inputs): |