DeviceMesh

As explained in Deep Learning with PyTorch, Second Edition

A PyTorch abstraction for arranging devices into a communication mesh.

DeviceMesh is a PyTorch abstraction for arranging devices into a communication mesh. It defines named dimensions and creates process groups for coordinating different kinds of parallel work, such as replicated models and partitioned layers. section 16.6 section 16.12

A two-dimensional mesh

A 2 × 4 mesh can represent two model replicas, each split across four devices: from torch.distributed.device_mesh import init_device_mesh mesh_2d = init_device_mesh( "cpu", (2, 4), mesh_dim_names=("replicate", "model_parallel") ) The named dimensions expose separate communication groups. The replicate dimension can synchronize model parameters, while the model_parallel dimension can coordinate a partitioned forward pass. section 16.6

A DeviceMesh helps organize:

• Higher-dimensional parallelism that combines data, model, pipeline, tensor, or context parallelism. section 16.12 section 16.6

Figure 16.9 Visualization of a 2 × 4 device mesh, where the rows represent model replicas and the columns represent model partitions. The dashed arrows represent how each device is communicating.

Visualize how rows represent model replicas, columns represent model partitions, and dashed links show communication between devices. section 16.6

How Deep Learning with PyTorch, Second Edition covers it

How deep it goes

Explains named mesh dimensions, generated process groups, parameter synchronization, and model-parallel communication.

Examples

Uses a 2 × 4 mesh with replicate and model_parallel dimensions to synchronize replicas and coordinate model partitions.

What you can do after

Configure distributed communication with torch.distributed by selecting mesh dimensions and their process groups.

How it connects

Connects data parallelism, model parallelism, pipeline parallelism, tensor parallelism, and fully sharded data parallelism.

liveBookmanning.com

Read Deep Learning with PyTorch, Second Edition on liveBook Get the book at manning.com

Related concepts

• torch.distributedused by · torch.distributed is PyTorch’s package for initializing processes and coordinating communication in distributed computation across devices or machines.
• process groupimplements · A process group is a defined collection of processes that share a communication or coordination context for distributed computation.
• distributed trainingused by · Distributed training is the coordinated training of a machine-learning model across multiple computing devices, processes, or machines by distributing computation or data among them.

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with PyTorch, Second Edition. The section links point to the parts of the book it draws on.