Refactoring Distributed test cases to be device agnostic [1/n]

[AnantGulati]

In this series of PR we intend to refactoring distributed test cases to enable to be completely device agnostic.

These changes will include the following approaches to do the same :

• Allowing for multiple device types using instantiate_device_type_test
• Replacing calls to cuda stream with torch.get_device_module(device) wherever it applies
• Skipping set up steps required while using MultiProcessTestCase with DistributedTestBase (Generalization of distributed test cases for non-CUDA devices #138216) wherever applicable
• Replacing explicit calls to distributed backend (NCCL,HCCL,etc) with get_default_backend_for_device (API to retrieve default distributed backend from device #140536).

This should result in significant improvement in usability for all devices

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @mcarilli @ptrblck @leslie-fang-intel @voznesenskym @penguinwu @EikanWang @Guobing-Chen @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov @LucasLLC @MeetVadakkanchery @mhorowitz @pradeepfn @ekr0 @xmfan

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• ROCM Infra failures during checkout of PyTorch

❌ 2 New Failures, 2 Unrelated Failures

As of commit 908427d with merge base 106acf0 ():

NEW FAILURES - The following jobs have failed:

• trunk / linux-focal-rocm6.3-py3.10 / test (default, 1, 2, linux.rocm.gpu.2) (gh)
  Process completed with exit code 1.
• trunk / linux-focal-rocm6.3-py3.10 / test (default, 2, 2, linux.rocm.gpu.2) (gh)
  Process completed with exit code 1.

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• pull / linux-focal-py3.13-clang10 / test (dynamo_wrapped, 2, 3, lf.linux.2xlarge) (gh) (similar failure)
  'test/test_reductions.py::TestReductionsCPU::test_sum_all_cpu_float64'
• trunk / linux-focal-rocm6.3-py3.10 / test (distributed, 1, 1, linux.rocm.gpu.4) (gh) (similar failure)
  Process completed with exit code 1.

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

[ankurneog]

LGTM

[AnantGulati]

@kwen2501 could you please help land this PR
Thanks

[AnantGulati]

The most recent commit changes the way cuda is set in DIstributedTestBase. Here we explicitly pass f"cuda:{self.rank}" to torch.set_device rather than passing the rank directly as a parameter. This was causing rank allocation errors in some test cases

[AnantGulati]

@kwen2501 @fegin Could you please help land this PR

[AnantGulati]

Changing the device setting logic from using instantiate_device_type_test to using get_devtype from common_fsdp to keep the CI green

[AnantGulati]

@kwen2501 Could you please help land this PR

[AnantGulati]

"Try to land this since the failure is unrelated."
@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

Merge failed

Reason: 2 jobs have failed, first few of them are: trunk / linux-focal-rocm6.3-py3.10 / test (default, 1, 2, linux.rocm.gpu.2), trunk / linux-focal-rocm6.3-py3.10 / test (default, 2, 2, linux.rocm.gpu.2)

Details for Dev Infra team

Raised by workflow job

[AnantGulati]

"Try to land this since the failure is unrelated."
@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

Merge failed

Reason: 2 jobs have failed, first few of them are: trunk / linux-focal-rocm6.3-py3.10 / test (default, 1, 2, linux.rocm.gpu.2), trunk / linux-focal-rocm6.3-py3.10 / test (default, 2, 2, linux.rocm.gpu.2)

Details for Dev Infra team

Raised by workflow job

[xmfan]

@pytorchbot merge -i

[pytorchmergebot]

Merge started

Your change will be merged while ignoring the following 4 checks: pull / linux-focal-py3.13-clang10 / test (dynamo_wrapped, 2, 3, lf.linux.2xlarge), trunk / linux-focal-rocm6.3-py3.10 / test (default, 1, 2, linux.rocm.gpu.2), trunk / linux-focal-rocm6.3-py3.10 / test (default, 2, 2, linux.rocm.gpu.2), trunk / linux-focal-rocm6.3-py3.10 / test (distributed, 1, 1, linux.rocm.gpu.4)

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here