[albmalamd]

Linked issue or supporting maintainer

Supported by @jeffdaily

BC-breaking?

No.

Summary (human written only)

On ROCm, the test wrongly treats normal RCCL memory usage as an extra GPU context. This causes failures even when everything works correctly. The change uses AMD SMI to directly check which processes use the GPU. If AMD SMI is unavailable, the unreliable memory check is skipped.

Agent details (if applicable)

The following investigation and implementation details were produced with
Cursor and must be reviewed by the author before submission.

On ROCm, the fallback used when pynvml is unavailable produces both false
positives and false negatives:

• The first RCCL kernel launch allocates approximately 1.1 GiB of queue scratch
  on rank 0's own device. This can exceed the fallback's 70% growth threshold
  even though no other rank created a context there.
• A real extra context created by rank 1 on device 0 can consume far less than
  1.1 GiB, so the same heuristic can miss the condition under test.

The change uses amdsmi_get_gpu_process_list, via
torch.cuda._get_amdsmi_handler, as the ROCm equivalent of the existing NVML
process-count check. If amdsmi is unavailable, the test skips on ROCm instead
of executing the invalid memory heuristic. CUDA behavior is unchanged.

The rejected alternatives were increasing the threshold and warming up RCCL
before taking the baseline. Both would hide the symptom without directly
checking whether another process created a context.

Validation was performed on MI300X GPUs with no foreign VRAM users, using the
ROCm 2.13 nightly wheel:

• upstream stock helper: failed 2/2 because of the memory heuristic;
• patched helper with amdsmi: passed 3/3;
• patched helper with an intentionally injected extra context: failed 2/2 and
  reported two processes on device 0;
• patched helper without amdsmi: skipped as intended.

The equivalent release/2.13 change was also tested directly:

• stock: failed 2/2 with a memory delta byte-identical to the reported CI
  failure;
• patched with amdsmi: passed 5/5;
• intentionally injected extra context: failed 3/3.

Related ROCm CI failure:
https://github.com/ROCm/rockrel/actions/runs/37731484533

Pull Request resolved: pytorch#200937
Approved by: https://github.com/jeffdaily

(cherry picked from commit 7c67769)