NVIDIA / NVIDIA/Megatron-LM

[QUESTION] Why is CUDA_DEVICE_MAX_CONNECTIONS=1 not required on Blackwell for TP/SP (without FSDP)?

Open
#5,105 2 comments 0 reactions 1 assignee Claimed by @youngeunkwon0405 View on GitHub
community-request waiting-on-customer
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

### Background

On pre-Blackwell GPUs (Ampere/Hopper, `get_device_arch_version() < 10`), Megatron-LM asserts that `CUDA_DEVICE_MAX_CONNECTIONS=1` when using tensor parallelism or context parallelism (non-FSDP):

```python
# megatron/training/arguments.py (validate_args)
if (args.tensor_model_parallel_size > 1 or args.context_parallel_size > 1) \
and get_device_arch_version() < 10:
# CUDA_DEVICE_MAX_CONNECTIONS requirement no longer exists since the Blackwell architecture
...
assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') == "1"
```

On Blackwell (`arch >= 10`), this assert is skipped.

The rationale for requiring `=1` on older architectures is documented in `linear_with_grad_accumulation_and_async_allreduce`: async collectives (all-gather / all-reduce / reduce-scatter) should be scheduled before subsequent compute kernels to overlap communication with computation. With the default `CUDA_DEVICE_MAX_CONNECTIONS=8`, the CUDA scheduler may reorder kernels across work queues, breaking this overlap (correctness is unaffected; performance is). Setting `=1` forces FIFO scheduling on a single work queue.

This is related to the earlier discussion in #533.

### Question

**For TP/SP training without FSDP, why is `CUDA_DEVICE_MAX_CONNECTIONS=1` no longer required (or asserted) on Blackwell?**

Specifically:

1. Is there a hardware or driver-level change on Blackwell (SM 10.x) that makes kernel launch ordering / comm-compute overlap work correctly without forcing a single work queue?
2. Should users still **recommend** setting `CUDA_DEVICE_MAX_CONNECTIONS=1` on Blackwell for best TP/SP overlap performance, even though the assert is gone? (`layers.py` still emits a warning when it is not set to `1`, regardless of architecture.)
3. Is the high-priority NCCL stream mechanism (`high_priority_stream_groups` / `ProcessGroupNCCL.Options(is_high_priority_stream=True)`) intended to replace the `=1` workaround for TP groups on Blackwell? Currently it appears to be auto-enabled only for FSDP groups (`dp_cp`, `ep_dp`), not for the `tp` process group.
4. There seems to be an inconsistency: `megatron/training/yaml_arguments.py` still raises `RuntimeError` if `CUDA_DEVICE_MAX_CONNECTIONS != "1"` when sequence parallelism is enabled, with no Blackwell exemption. Is that intentional?

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.