Fix device mesh setup in hybrid Cifar10 example (#192)

haoyuz · facebook-github-bot · commit 35bd72b0e305 · 2025-06-20T13:58:30.000-07:00
Summary: Pull Request resolved: #192 Device mesh shape should be `(replica, WORLD_SIZE // replica)` for consistency across ranks. Reviewed By: tsunghsienlee Differential Revision: D77030402 fbshipit-source-id: bba8555af80864d5333a98c6b06b46bb1ea6c53e
diff --git a/distributed_shampoo/examples/hybrid_shard_cifar10_example.py b/distributed_shampoo/examples/hybrid_shard_cifar10_example.py
@@ -199,7 +199,7 @@ def create_model_and_optimizer_and_loss_fn(
     # initialize device_mesh for hybrid shard data parallel
     device_mesh: DeviceMesh = init_device_mesh(
         "cuda",
-        (args.dp_replicate_degree, WORLD_RANK // args.dp_replicate_degree),
+        (args.dp_replicate_degree, WORLD_SIZE // args.dp_replicate_degree),
         mesh_dim_names=("dp_replicate", "dp_shard"),
     )
 

Original file line number	Diff line number	Diff line change
`@@ -199,7 +199,7 @@ def create_model_and_optimizer_and_loss_fn(`
`199`	`199`	`# initialize device_mesh for hybrid shard data parallel`
`200`	`200`	`device_mesh: DeviceMesh = init_device_mesh(`
`201`	`201`	`"cuda",`
`202`		`- (args.dp_replicate_degree, WORLD_RANK // args.dp_replicate_degree),`
	`202`	`+ (args.dp_replicate_degree, WORLD_SIZE // args.dp_replicate_degree),`
`203`	`203`	`mesh_dim_names=("dp_replicate", "dp_shard"),`
`204`	`204`	`)`
`205`	`205`