From d0806c8762d2f7cbc880fe9d54b269f1281e59ff Mon Sep 17 00:00:00 2001
From: wz337 <wz337@cornell.edu>
Date: Thu, 27 Jun 2024 15:30:27 -0700
Subject: [PATCH 1/3] update device mesh to include slicing

---
 recipes_source/distributed_device_mesh.rst | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/recipes_source/distributed_device_mesh.rst b/recipes_source/distributed_device_mesh.rst
index dbc4a810434..5e46c68fc99 100644
--- a/recipes_source/distributed_device_mesh.rst
+++ b/recipes_source/distributed_device_mesh.rst
@@ -148,6 +148,26 @@ Then, run the following `torch elastic/torchrun <https://pytorch.org/docs/stable
 
     torchrun --nproc_per_node=8 hsdp.py
 
+How to use DeviceMesh for your custom parallel solutions
+--------------------------------------------------------
+When working with large scale training, one can have a more complicated custom parallel training composition. For example, one may need to slice out submeshes for different parallelism solutions.
+DeviceMesh allows users to slice child mesh from the parent mesh and re-use the NCCL communicators already created when the parent mesh is initialized.
+
+.. code-block:: python
+
+    from torch.distributed.device_mesh import init_device_mesh
+    mesh_3d = init_device_mesh("cuda", (2, 2, 2), mesh_dim_names=("replicate", "shard", "tp"))
+
+    # Users can slice child meshes from the parent mesh.
+    hsdp_mesh = mesh_3d["replicate", "shard"]
+    tp_mesh = mesh_3d["tp"]
+
+    # Users can access the underlying process group thru `get_group` API.
+    replicate_group = hsdp_mesh["Replicate"].get_group()
+    shard_group = hsdp_mesh["Shard"].get_group()
+    tp_group = tp_mesh.get_group()
+
+
 Conclusion
 ----------
 In conclusion, we have learned about :class:`DeviceMesh` and :func:`init_device_mesh`, as well as how

From 7acc2a0f30c731eeb40dfd20aaf55733592ee315 Mon Sep 17 00:00:00 2001
From: Iris Z <31293777+wz337@users.noreply.github.com>
Date: Thu, 27 Jun 2024 17:02:37 -0700
Subject: [PATCH 2/3] Update recipes_source/distributed_device_mesh.rst

Co-authored-by: Svetlana Karslioglu <svekars@meta.com>
---
 recipes_source/distributed_device_mesh.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/recipes_source/distributed_device_mesh.rst b/recipes_source/distributed_device_mesh.rst
index 5e46c68fc99..c2b2c7a99f8 100644
--- a/recipes_source/distributed_device_mesh.rst
+++ b/recipes_source/distributed_device_mesh.rst
@@ -150,7 +150,7 @@ Then, run the following `torch elastic/torchrun <https://pytorch.org/docs/stable
 
 How to use DeviceMesh for your custom parallel solutions
 --------------------------------------------------------
-When working with large scale training, one can have a more complicated custom parallel training composition. For example, one may need to slice out submeshes for different parallelism solutions.
+When working with large scale training, you might have more complex custom parallel training composition. For example, you may need to slice out submeshes for different parallelism solutions.
 DeviceMesh allows users to slice child mesh from the parent mesh and re-use the NCCL communicators already created when the parent mesh is initialized.
 
 .. code-block:: python

From 9a40f62c7c11daa2f9570fe936882765e9d0351b Mon Sep 17 00:00:00 2001
From: Iris Z <31293777+wz337@users.noreply.github.com>
Date: Tue, 9 Jul 2024 14:29:42 -0700
Subject: [PATCH 3/3] Update distributed_device_mesh.rst

---
 recipes_source/distributed_device_mesh.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/recipes_source/distributed_device_mesh.rst b/recipes_source/distributed_device_mesh.rst
index c2b2c7a99f8..08ee247db0a 100644
--- a/recipes_source/distributed_device_mesh.rst
+++ b/recipes_source/distributed_device_mesh.rst
@@ -163,7 +163,7 @@ DeviceMesh allows users to slice child mesh from the parent mesh and re-use the
     tp_mesh = mesh_3d["tp"]
 
     # Users can access the underlying process group thru `get_group` API.
-    replicate_group = hsdp_mesh["Replicate"].get_group()
+    replicate_group = hsdp_mesh["replicate"].get_group()
     shard_group = hsdp_mesh["Shard"].get_group()
     tp_group = tp_mesh.get_group()