[PyTorch] Distributed tests set device in NCCL init and manually shut down NCCL

[timmoon10]

Description

This attempts to debug NCCL errors in some distributed tests.

#3654 also manually tears down NCCL in the TE ops distributed test.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Set device in NCCL init in distributed tests.
• Manually tear down NCCL in TE ops distributed test.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[timmoon10]

/te-ci pytorch L1

[greptile-apps]

Confidence Score: 5/5

[Low impact] The PR appears safe to merge; the previous cleanup issue is fixed and no new blocking issue was found.

Summary

This PR gives NCCL an explicit CUDA device in two distributed tests and adds explicit cleanup to the fusible-ops test.

• The latest change moves the final barrier and CUDA synchronization onto the success path.
• Group destruction remains in finally, including when a test raises.
• The previous cleanup finding is fixed. No new actionable issues were found.

_{Reviews (2) · Last reviewed commit: "Apply suggestion from @timmoon10" · Reviewed by Greptile}

[timmoon10]

/te-ci pytorch L1

[vthumbe1503]

Would be good to understand why device_id needs to be set manually.

@timmoon10 and I had offline discussion regarding this. It might be because this particular test doesnt use torchrun like other tests do, and so the device is not auto-set to "cuda:rank" like other tests

[ptrendx]

I do actually have a diff already that changes this test to torchrun to avoid this. Can make PR today. @timmoon10

[vthumbe1503]

LGTM