torchfoo.module.make_ddp#

torchfoo.module.make_ddp(m, find_unused_parameters=False)[source]#

Wrap a module with DDP and convert BatchNorm to SyncBatchNorm.

Returns the module unchanged if not in distributed mode.

Warning

SyncBatchNorm only works with GPU modules. On CPU, BatchNorm layers will not be converted and a warning is emitted.

Parameters:

m (Module) – the module to wrap
find_unused_parameters (bool) – forwarded to DistributedDataParallel. When True, DDP traverses the autograd graph to find parameters that did not receive gradients. Defaults to False.

Returns:

The DDP-wrapped module, or the original module if not distributed.