Skip to content

Commit b06ba6d

Browse files
committed
test_mla_helix: Add reduce_output=False to disable NCCL all-reduce
The MLA's o_proj layer performs an all-reduce by default, which fails in the MPI worker context because NCCL groups aren't initialized. Setting reduce_output=False disables this all-reduce operation.
1 parent 57fd9a3 commit b06ba6d

File tree

1 file changed

+1
-0
lines changed

1 file changed

+1
-0
lines changed

tests/unittest/_torch/modules/test_mla_helix.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -632,6 +632,7 @@ def _run_mla_distributed(
632632
dtype=scenario.dtype,
633633
config=config,
634634
enable_unit_test=True,
635+
reduce_output=False, # Disable all-reduce to avoid NCCL in MPI workers
635636
).cuda()
636637

637638
# ========================================================================

0 commit comments

Comments
 (0)