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

should_sync_now(*, epoch_step_idx: int) → bool

Parameters:

epoch_step_idx – current step index

Returns:

whether to sync the training processes in this step

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

one train step

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

Parameters:

config

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.