fix float8 training TP+SP integration tests

[vkuzo]

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:

./test/float8/test_dtensor.sh

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• enable to_mxfp8 cast for DTensor #2420
• make dtensor shared test util more generic #2416
• rename torchao.testing.float8 to torchao.testing.training #2415
• -> fix float8 training TP+SP integration tests #2414

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2414

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (2 Unrelated Failures)

As of commit 5c23c6b with merge base 4e3d019 ():

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests / test (CPU 2.7, linux.4xlarge, torch==2.7.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh) (trunk failure)
  test/quantization/pt2e/test_numeric_debugger.py::TestNumericDebuggerInfra::test_prepare_for_propagation_comparison
• Run Regression Tests / test (CUDA 2.7, linux.g5.12xlarge.nvidia.gpu, torch==2.7.0, cuda, 12.6) / linux-job (gh) (trunk failure)
  test/quantization/pt2e/test_numeric_debugger.py::TestNumericDebuggerInfra::test_prepare_for_propagation_comparison

This comment was automatically generated by Dr. CI and updates every 15 minutes.