suggestions for more examples

[daire-byrne]

I finally made it to the end of the examples (well, the deepspeed one just hangs for me), and now I'm hungry for more!

I think it would be nice to include one on HSDP (FSDP + DDP) to sit between the DDP, FSDP and the 2D (TP + FSDP) sections. And maybe one on context/sequence parallelism too?

I think it is common to have access to some sort of PCIe bases multi GPU server, and being able to use GPUs across sockets is a useful trick for HSDP.

Pipeline parallelism is just too hard to get my head around so you can ignore that one!

Many thanks again for putting this all together.

[daire-byrne]

I had a go at editing 04-fully-sharded-data-parallel/train_llm.py similar to your 2D example:

--- train_llm.py	2025-11-20 15:28:11.843371344 +0000
+++ hsdp_train_llm.py	2025-11-20 16:11:59.667060601 +0000
@@ -59,6 +59,14 @@
     torch.cuda.set_device(device)
     dist.init_process_group(rank=rank, world_size=world_size, device_id=device)
 
+    assert world_size % args.replicas_parallel == 0
+
+    mesh_hsdp = dist.device_mesh.init_device_mesh(
+        "cuda",
+        (args.replicas_parallel, world_size // args.replicas_parallel),
+        mesh_dim_names=("dp_replicate", "dp_shard"),
+    )
+
     logging.basicConfig(
         format=f"[rank={rank}] [%(asctime)s] %(levelname)s:%(message)s",
         level=logging.INFO,
@@ -67,6 +75,7 @@
     LOGGER.debug(os.environ)
     LOGGER.debug(args)
     LOGGER.debug(f"local_rank={local_rank} rank={rank} world_size={world_size}")
+    LOGGER.debug(f"dp_shard={mesh_hsdp['dp_shard'].size()} dp_replicate={mesh_hsdp['dp_replicate'].size()}")
 
     dtype = torch.bfloat16
     torch.manual_seed(args.seed)
@@ -84,6 +93,7 @@
         reshard_after_forward=True,
         offload_policy=CPUOffloadPolicy() if args.cpu_offload else None,
         mp_policy=MixedPrecisionPolicy(param_dtype=dtype, reduce_dtype=torch.float32),
+        mesh=mesh_hsdp,
     )
     for decoder in model.model.layers:
         fully_shard(decoder, **fsdp_config)
@@ -387,6 +397,7 @@
     parser.add_argument("--log-freq", default=10, type=int)
     parser.add_argument("--ckpt-freq", default=500, type=int)
     parser.add_argument("-s", "--seq-length", default=1024, type=int)
+    parser.add_argument("-rp", "--replicas-parallel", default=1, type=int)
     parser.add_argument("--cpu-offload", default=False, action="store_true")
     return parser

Using the arg "-rp 1", it performs just like the previous version (fsdp), but if I increase it, then performance does not look right. I would have though the with local world size = 4 and "-rp 4" it would be more like DDP but it is much slower.

I guess it's a bit more complicated than my quick hack.

It would be nice to add this feature though as HSDP is widely used. You could even add it as an option to the 2d version to make it 3d? :)