fix float8 training TP+SP integration tests

Summary:

These tests do not run in CI, and they broke some time ago. The issue
was that each tensor was created on "cuda:0" instead of using the local
rank. For now, fixing by manually specifying the rank. I feel like there
is probably a better way to do this as the rank is supposed to be set
automatically, but leaving that for a future PR.

We should add to CI in the future, saving that for a future PR.

Test Plan:

```bash
./test/float8/test_dtensor.sh
```

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 39bd880
ghstack-comment-id: 2991778315
Pull Request resolved: #2414

Author: vkuzo

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
 
 

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