tensorplay.distributed.device_mesh

Creating a mesh

init_device_mesh() creates a DeviceMesh from a device type and a tuple of mesh dimensions:

from tensorplay.distributed.device_mesh import init_device_mesh

# requires an initialized process group, e.g. via dist.init_process_group()
mesh = init_device_mesh("cuda", mesh_shape=(2, 4),
                        mesh_dim_names=("dp", "tp"))

Semantics follow the SPMD model: mesh_shape (and, when given, mesh_dim_names) must be identical across all ranks, and the function blocks until every rank has joined. Each mesh dimension can be given a name via mesh_dim_names; the names are reflected in mesh_dim_names and can be used to address that dimension when sharding a distributed tensor.

Querying the mesh

A DeviceMesh answers the usual questions about rank placement:

• get_rank() — the global rank at a given coordinate (or the local rank of this process).

• get_local_rank() — the ordinal of this process within a mesh dimension.

• get_coordinate() — the coordinates of this process in the mesh.

• get_group() / get_all_groups() — the underlying process groups, per mesh dimension.

Where to go next

• distributed tensors — DTensor values that live on a mesh.

• tensor parallelism — sharding module weights across the mesh.

• the distributed package — the process-group model and collectives the mesh is built on.