returnn.torch.distributed

torch.distributed utils

class returnn.torch.distributed.DistributedContext(options: Dict[str, Any])

This class setups some helper functions for torch distributed training

local_rank() → int

local rank

local_size() → int

local size

rank() → int

global rank

size() → int

global size

get_param_sync_step() → int | None

param sync step

maybe_make_distributed_module(module: Module) → DistributedDataParallel | None

Maybe make a wrapped distributed module.

Parameters:

module – original module

Returns:

potentially wrapped module

step_after_param_update(*, module: Module, epoch_step_idx: int)

one train step

returnn.torch.distributed.get_ctx(config=None) → DistributedContext | None

Parameters:

config (Config|None)

Returns:

the global context if Torch distributed is enabled, or None otherwise. If we did not setup the context yet, it will automatically create it.