[bilalsar]

Motivation

In short: test_all_gather_parity and test_reduce_scatter_parity put their tensors on the wrong GPU, and the torch 2.15 default nccl backend now rejects that.

Both test_layers.NcclTensorParallelLayerTest tests fail on ROCm 10.3.0-20261027 (torch 2.15.0+rocm10.3.0a20261027) on every host where Apex_UT reaches them. They passed on 10.0.0-20260729. All runs use apex 11867be, the release/1.15.0 pin in related_commits.

Host	GPUs	test_all_gather_parity	test_reduce_scatter_parity	Log
ctr-navi4x-aj63-ws13	2x RX 9060 XT (gfx1200)	ERROR	ok	job
ctr-navi4x-at48-ws008	4x Radeon (gfx1201)	ERROR	ERROR	log
quanta-ccs-aus-d10-19	8x MI325X	ERROR	ERROR	log
smci350-rck-g03-d13-03	8x MI350X	ERROR	ERROR	log

• Prefilled ticket: Apex_UT test_layers all_gather/reduce_scatter parity FAIL (from this triage report)

  File "/workspace/apex/tests/L0/run_transformer/test_layers.py", line 71, in test_all_gather_parity
    all_gather(chunks, tensor, group=parallel_state.get_tensor_model_parallel_group())
RuntimeError: Expected tensor on cuda:1 but found tensor on cuda:0

  File "/workspace/apex/tests/L0/run_transformer/test_layers.py", line 116, in test_reduce_scatter_parity
RuntimeError: Expected tensor on cuda:2 but found tensor on cuda:0

Technical Details

In short: allocate on the rank's current device instead of cuda:{tensor_model_parallel_rank}.

tests/L0/run_transformer/test_layers.py: both tests allocated on cuda:{tensor_model_parallel_rank} (lines 55 and 103). Whenever the tensor-parallel size is smaller than the world size, a rank whose tensor-parallel rank differs from its global rank passes a tensor on the wrong GPU. test_all_gather_parity hits this at tensor-parallel size 1 (global rank 1 passes a cuda:0 tensor to a group bound to cuda:1). test_reduce_scatter_parity starts at tensor-parallel size 2, so it hits it only with 4 or more GPUs (global rank 2 passes a cuda:0 tensor). The torch 2.15 default nccl backend rejects tensors that are not on the group's bound device; the legacy ProcessGroupNCCL accepted them.

The tensors now go on torch.cuda.current_device(), which DistributedTestBase sets to the rank's device. The values still depend on the tensor-parallel rank, so what the parity checks compare is unchanged.

Effect on other suites: only these two tests change. The UCC variants of both tests are skipped without torch_ucc. No library code changes.

Related: the test_pipeline_parallel_fwd_bwd out-of-memory errors in the same Apex_UT runs are a separate root cause: #345.

Test Result

Full Apex_UT rerun through the frameworks-qa-ci Functional Test Runner on ctr-navi4x-aj63-ws13 with the same image (rocm/frameworks-internal:rocm10.3.0a20261027_ubuntu26.04_py3.14_pytorch_2.15.0_31b8278, wrapper fqa-ci-wrapper:9f4114214363), with apex checked out at this PR's commit d6c652f.

	Failing: 11867be (pin)	Passing: this PR d6c652f
Job	run 40688983363	run 40766959214
test_all_gather_parity	ERROR	ok
test_reduce_scatter_parity	ok	ok
Apex_UT totals (303)	218 passed, 7 errors, 78 skipped	219 passed, 6 errors, 78 skipped

The 6 errors left in the passing run are the test_pipeline_parallel_fwd_bwd out-of-memory errors, fixed by #345. No new errors.

apex-test-layers-fail-expected-tensor-on-cuda1.log (run 40688983363)

ERROR: test_all_gather_parity (test_layers.NcclTensorParallelLayerTest.test_all_gather_parity)
RuntimeError: Process 1 exited with error code 10 and exception:
  File "/workspace/apex/tests/L0/run_transformer/test_layers.py", line 71, in test_all_gather_parity
    all_gather(chunks, tensor, group=parallel_state.get_tensor_model_parallel_group())
RuntimeError: Expected tensor on cuda:1 but found tensor on cuda:0
...
Ran 102 tests in 714.453s
FAILED (errors=7, skipped=56)
|   303    |   218    |    0     |    7     |    78    |   225    |     FAIL     |

apex-test-layers-pass-current-device.log (run 40766959214)

Checking out Apex commit: d6c652f3011b58a59af48192840924414f62da96
HEAD is now at d6c652f Use the rank's own device in test_layers all_gather/reduce_scatter parity
test_all_gather_parity (test_layers.NcclTensorParallelLayerTest.test_all_gather_parity) ... ok
test_reduce_scatter_parity (test_layers.NcclTensorParallelLayerTest.test_reduce_scatter_parity) ... ok
...
Ran 102 tests in 725.479s
FAILED (errors=6, skipped=56)
|   303    |   219    |    0     |    6     |    78    |   225    |     FAIL     |

How the PR commit was run: the job was dispatched from the frameworks-qa-ci branch test/apex-ut-commit-override, which only lets apex_ut take an APEX_COMMIT_OVERRIDE env var instead of the related_commits SHA. Everything else is the same as the nightly job.

Not run: a run with this fix on a host with 4 or more GPUs, which is where test_reduce_scatter_parity fails.

Not covered: CI picks the fix up once related_commits points at an apex commit that contains it (it pins release/1.15.0, so this needs a backport there too).

Submission Checklist

• Change is limited to the failing tests
• Validated on a GPU CI job (failing and passing job links above)

[bilalsar]

Validation, commit d6c652f: fixed.

test_layers.NcclTensorParallelLayerTest.test_all_gather_parity: ERROR on 11867be → ok on d6c652f. Apex_UT 218 passed / 7 errors → 219 passed / 6 errors (the 6 left are the separate test_pipeline_parallel_fwd_bwd OOM). Host ctr-navi4x-aj63-ws13 (2x gfx1200), image rocm10.3.0a20261027_ubuntu26.04_py3.14_pytorch_2.15.0_31b8278.

• Failing: https://github.com/AMD-ROCm-Internal/frameworks-qa-ci/actions/runs/40688983363/job/122818736888
• Passing: https://github.com/AMD-ROCm-Internal/frameworks-qa-ci/actions/runs/40766959214/job/123093727527