mapping_networks.distributed.setup_ddp

mapping_networks.distributed.setup_ddp(rank: int, world_size: int, backend: str = 'gloo', init_method: str = 'env://', **kwargs: Any) → None

Initialize the default distributed process group.

Call this once at the start of each worker process before constructing the model. Use cleanup_ddp() at the end of training to destroy the group.

Parameters:

  • rank – Global rank of this process (0-based).

  • world_size – Total number of processes.

  • backend – Distributed backend — "gloo" (CPU/cross-platform), "nccl" (CUDA multi-GPU, recommended for GPU training), or "mpi". Defaults to "gloo".

  • init_method – URL for rendezvous. Defaults to "env://" which reads MASTER_ADDR / MASTER_PORT environment variables.

  • **kwargs – Additional keyword arguments forwarded to torch.distributed.init_process_group().

Raises:

RuntimeError – If torch.distributed is not available.