torchtnt.utils.distributed.PGWrapper

class torchtnt.utils.distributed.PGWrapper(pg: Optional[ProcessGroup])

A wrapper around ProcessGroup that allows collectives to be issued in a consistent fashion regardless of the following scenarios:

pg is None, distributed is initialized: use WORLD as pg pg is None, distributed is not initialized: single process app pg is not None: use pg

__init__(pg: Optional[ProcessGroup]) → None

Methods

__init__(pg)
all_gather_object(obj_list, obj)
barrier()
broadcast_object_list(obj_list[, src])
get_rank()
get_world_size()
scatter_object_list(output_list, input_list)