`nvalchemi.distributed`.DistributedManager#

class nvalchemi.distributed.DistributedManager[source]#

Distributed Manager for setting up distributed training environment.

This is a singleton that creates a persistance class instance for storing parallel environment information through out the life time of the program. This should be used to help set up Distributed Data Parallel and parallel datapipes.

Note

One should call DistributedManager.initialize() prior to constructing a manager object

Example

>>> DistributedManager.initialize()
>>> manager = DistributedManager()
>>> manager.rank
0
>>> manager.world_size
1

property broadcast_buffers#: broadcast_buffers in PyTorch DDP

static cleanup(barrier=False)[source]#

Clean up distributed group and singleton

Parameters:: barrier (bool, optional) – Whether to use a global barrier before destroying the process group, by default False

static create_orthogonal_process_group(orthogonal_group_name, group_name, verbose=False)[source]#

Create a process group that is orthogonal to the specified process group.

Parameters:

orthogonal_group_name (str) – Name of the orthogonal process group to be created.
group_name (str) – Name of the existing process group.
verbose (bool) – Print out ranks of each created process group, default False.

static create_process_subgroup(name, size, group_name=None, verbose=False)[source]#

Create a process subgroup of a parent process group. This must be a collective call by all processes participating in this application.

Parameters:

name (str) – Name of the process subgroup to be created.
size (int) – Size of the process subgroup to be created. This must be an integer factor of the parent group’s size.
group_name (Optional[str]) – Name of the parent process group, optional. If None, the default process group will be used. Default None.
verbose (bool) – Print out ranks of each created process group, default False.

property cuda#: If cuda is available

property device#: Process device

property distributed#: Distributed environment

property find_unused_parameters#: find_unused_parameters in PyTorch DDP

static get_available_backend()[source]#: Get communication backend

get_mesh_group(mesh)[source]#

Get the process group for a given mesh.

Creating a group is an expensive operation, so we cache the result manually.

We hash the mesh and use that as the key.

Parameters:: mesh (DeviceMesh)
Return type:: ProcessGroup

property global_mesh#: Returns the global mesh. If it’s not initialized, it will be created when this is called.

group(name=None)[source]#: Returns a process group with the given name If name is None, group is also None indicating the default process group If named group does not exist, PhysicsNeMoUndefinedGroupError exception is raised

group_name(group=None)[source]#: Returns the name of process group

property group_names#: Returns a list of all named process groups created

group_rank(name=None)[source]#: Returns the rank in named process group

group_size(name=None)[source]#: Returns the size of named process group

static initialize()[source]#

Initialize distributed manager

Current supported initialization methods are:

ENV: PyTorch environment variable initialization: https://pytorch.org/docs/stable/distributed.html#environment-variable-initialization
SLURM: Initialization on SLURM systems.: Uses SLURM_PROCID, SLURM_NPROCS, SLURM_LOCALID and SLURM_LAUNCH_NODE_IPADDR environment variables.
OPENMPI: Initialization for OpenMPI launchers.: Uses OMPI_COMM_WORLD_RANK, OMPI_COMM_WORLD_SIZE and OMPI_COMM_WORLD_LOCAL_RANK environment variables.

Initialization by default is done using the first valid method in the order listed above. Initialization method can also be explicitly controlled using the PHYSICSNEMO_DISTRIBUTED_INITIALIZATION_METHOD environment variable and setting it to one of the options above.

static initialize_env()[source]#: Setup method using generic initialization

initialize_mesh(mesh_shape, mesh_dim_names)[source]#

Initialize a global device mesh over the entire distributed job.

Creates a multi-dimensional mesh of processes that can be used for distributed operations. The mesh shape must multiply to equal the total world size, with one dimension optionally being flexible (-1).

Parameters:

mesh_shape (Tuple[int, ...]) – Tuple of ints describing the size of each mesh dimension. Product must equal world_size. One dimension can be -1 to be automatically calculated.
mesh_dim_names (Tuple[str, ...]) – Names for each mesh dimension. Must match length of mesh_shape.

Returns:

The initialized device mesh

Return type:

torch.distributed.DeviceMesh

Raises:

RuntimeError – If mesh dimensions are invalid or don’t match world size
AssertionError – If distributed environment is not available

static initialize_open_mpi(addr, port)[source]#: Setup method using OpenMPI initialization

static initialize_slurm(port)[source]#: Setup method using SLURM initialization

classmethod is_initialized()[source]#

If manager singleton has been initialized

Return type:: bool

property local_rank#: Process rank on local machine

mesh(name=None)[source]#

Return a device_mesh with the given name. Does not initialize. If the mesh is not created already, will raise and error

Parameters:: name (str, optional) – Name of desired mesh, by default None

property mesh_dims#

size)

Type:: Mesh Dimensions as dictionary (axis name

mesh_names()[source]#: Return mesh axis names

mesh_sizes()[source]#: Return mesh axis sizes

property rank#: Process rank

static setup(rank=0, world_size=1, local_rank=None, addr='localhost', port='12355', backend='nccl', method='env')[source]#: Set up PyTorch distributed process group and update manager attributes

property world_size#: Number of processes in distributed environment

nvalchemi.distributed.DistributedManager#

`nvalchemi.distributed`.DistributedManager#