[QUESTION] Why is CUDA_DEVICE_MAX_CONNECTIONS=1 not required on Blackwell for TP/SP (without FSDP)?
- 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
Assessment
This issue has not been assessed yet.