fix float8 training TP+SP integration tests (#2414)

vkuzo · web-flow · commit 8b051b49dd29 · 2025-06-24T10:16:36.000-04:00
Update

[ghstack-poisoned]
diff --git a/test/float8/test_dtensor.py b/test/float8/test_dtensor.py
@@ -67,6 +67,8 @@ def setup_distributed():
     device_mesh = init_device_mesh("cuda", (world_size,))
     # seed must be the same in all processes
     torch.manual_seed(1)
+    local_rank = torch.distributed.get_rank()
+    torch.cuda.set_device(local_rank)
     return device_mesh
 
 
diff --git a/test/float8/test_fsdp2_tp.py b/test/float8/test_fsdp2_tp.py
@@ -46,6 +46,8 @@ def setup_distributed():
     )
     # seed must be the same in all processes
     torch.manual_seed(1)
+    local_rank = torch.distributed.get_rank()
+    torch.cuda.set_device(local_rank)
     return device_mesh
 
 

Original file line number	Diff line number	Diff line change
`@@ -46,6 +46,8 @@ def setup_distributed():`
`46`	`46`	`)`
`47`	`47`	`# seed must be the same in all processes`
`48`	`48`	`torch.manual_seed(1)`
	`49`	`+ local_rank = torch.distributed.get_rank()`
	`50`	`+ torch.cuda.set_device(local_rank)`
`49`	`51`	`return device_mesh`
`50`	`52`
`51`	`53`