flagos-ai / flagos-ai/FlagScale
Tensor/context parallel training fails on non-CUDA platforms: CUDA_DEVICE_MAX_CONNECTIONS assertion
- Dominant language
- Python
- Stars
- 537
- Forks
- 174
- Avg merge
- 4d 1h
- Merged PRs (30d)
- 12
Description
### Describe the bug
Training with tensor parallelism or context parallelism fails during argument validation on non-NVIDIA hardware. Reproduced on Moore Threads (MUSA).
```
[default7]:Traceback (most recent call last):
[default7]: File "flagscale/train/megatron/train_gpt.py", line 436, in
[default7]: main()
[default7]: File "flagscale/train/megatron/train_gpt.py", line 416, in main
[default7]: args = parse_and_validate_args(
[default7]: File "flagscale/train/megatron/training/arguments.py", line 124, in parse_and_validate_args
[default7]: validate_args(args, args_defaults)
[default7]: File "flagscale/train/megatron/training/arguments.py", line 1489, in validate_args
[default7]: assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') == "1", \
[default7]:AssertionError: Using tensor model parallelism or context parallelism require setting the environment variable CUDA_DEVICE_MAX_CONNECTIONS to 1
```
### Root cause
Two separate portability problems combine in the block at `arguments.py:1465-1491`:
**1. `CUDA_DEVICE_MAX_CONNECTIONS` is an NVIDIA CUDA driver environment variable.**
Other runtimes (MUSA, CANN, ...) do not read it. Setting it on those platforms has no effect, so the assertion blocks startup without providing any benefit. All three branches of the block — the two `warn_rank_0` paths and the `assert` — discuss only this variable, so the entire block is CUDA-specific, not just the assertion.
**2. `get_device_arch_version()` is not portable.**
The guarding condition is:
```python
if (args.tensor_model_parallel_size > 1 or args.context_parallel_size > 1) \
and get_device_arch_version() < 10:
```
`get_device_arch_version()` (`training/utils/common_utils.py:554`) returns the raw device property `major`:
```python
def get_device_arch_version():
"""Returns GPU arch version (8: Ampere, 9: Hopper, 10: Blackwell, ...)"""
return cur_platform.get_device_properties(cur_platform.device(0)).major
```
That value only carries NVIDIA compute-capability semantics on CUDA. On other platforms it belongs to an unrelated numbering scheme, so comparing it against `10` is not meaningful — it merely happens to be true, which is what lets non-CUDA hardware reach the CUDA-only branch.
### Steps to reproduce
1. Run any FlagScale Megatron training job on non-NVIDIA hardware (e.g. MUSA)
2. Set `tensor_model_parallel_size > 1` or `context_parallel_size > 1`
3. Leave `CUDA_DEVICE_MAX_CONNECTIONS` unset
4. Argument validation aborts with the assertion above
### Expected behaviour
The check should be skipped on platforms where `CUDA_DEVICE_MAX_CONNECTIONS` has no meaning, instead of requiring users to set a variable their runtime ignores.
### Also affects Ascend
`plugin_flagscale/npu_plugin.py:20` overrides `get_device_arch_version()` to return `8`. That avoids reading a device property which may be unavailable, but `8 < 10` still holds, so NPU runs reach the same assertion. This is not MUSA-specific — it applies to every non-NVIDIA backend.
### Workaround
Export the variable anyway, where it is a no-op that exists purely to satisfy the assertion:
```bash
export CUDA_DEVICE_MAX_CONNECTIONS=1
```
For multi-rank jobs this must be set for every rank (e.g. via the `envs` section of the experiment YAML), not just in the launching shell.
### Environment
- FlagScale: commit 73d5f8e
- Megatron-LM-FL: v0.18.2 (commit e15cb69)
- Hardware: Moore Threads (MUSA)
- Python 3.12
### Proposed fix
Gate the whole block on `cur_platform.name() == "cuda"`, matching the existing platform-dispatch idiom in the repo (`train_gr00t_n1_5.py:85`, `train_pi.py:88`). Note `platform.name()` is required rather than `device_name()`, since some backends report `"cuda"` as their device name to reuse PyTorch's CUDA dispatch path (`platform_txda.py` returns `"txda"` from `name()` but `"cuda"` from `device_name()`).
Fix submitted in #1288.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.