fix: skip CUDA_DEVICE_MAX_CONNECTIONS check on non-CUDA platforms

[kolinweiwei]

PR Category

Hardware

PR Types

Bug Fixes

PR Description

Training with tensor parallelism or context parallelism fails during argument validation on non-NVIDIA hardware (reproduced on Moore Threads / MUSA):

File "flagscale/train/megatron/training/arguments.py", line 1489, in validate_args
    assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') == "1", \
AssertionError: Using tensor model parallelism or context parallelism require setting the environment variable CUDA_DEVICE_MAX_CONNECTIONS to 1

Root cause

Two issues combine in the block at arguments.py:1465-1491:

1. CUDA_DEVICE_MAX_CONNECTIONS is NVIDIA-specific. It is a CUDA driver environment variable. Other runtimes (MUSA, CANN, ...) do not read it, so setting it on those platforms has no effect — the assertion blocks startup without buying anything. All three branches of the block (including the two warn_rank_0 paths) talk exclusively about this variable, so the entire block is CUDA-only, not just the assertion.

2. get_device_arch_version() is not portable. It returns the raw device property major, which only carries NVIDIA compute-capability semantics (8: Ampere, 9: Hopper, 10: Blackwell) on CUDA. On other platforms that value belongs to an unrelated numbering scheme, so get_device_arch_version() < 10 is not a meaningful test there — it merely happens to be true, which is what lets non-CUDA hardware fall into the CUDA-only branch.

Fix

Gate the whole block on cur_platform.name() == "cuda". This matches the existing platform-dispatch idiom in the repo (train_gr00t_n1_5.py:85, train_pi.py:88, ...), and cur_platform is already imported in arguments.py:60.

platform.name() is used rather than device_name() on purpose: some backends deliberately report "cuda" as their device name to reuse PyTorch's CUDA dispatch path (e.g. platform_txda.py returns "txda" from name() but "cuda" from device_name()), so device_name() would reintroduce the same misclassification.

Behaviour on CUDA is unchanged — the condition only gains a platform check.

Notes

The workaround until this lands is to export CUDA_DEVICE_MAX_CONNECTIONS=1 on non-CUDA platforms as well, where it is a no-op that exists purely to satisfy the assertion.

This also affects Ascend: plugin_flagscale/npu_plugin.py overrides get_device_arch_version() to return 8, which avoids reading a device property that may be unavailable, but 8 < 10 still holds — so NPU runs reach the same assertion.

Related: #1287, #1286

Fixes #1289

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

BAAI3315 seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}