NVIDIA / NVIDIA/Megatron-LM

[Bug] validate_args() eagerly probes the CUDA device and crashes GPU-less launcher/driver processes (Found no NVIDIA driver)

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

Description

**Describe the bug**

`megatron/training/arguments.py::validate_args()` — pure argument validation — eagerly queries the CUDA device in three places that only MoE + tensor/context parallelism reach. When `validate_args` runs in a process that does not (yet) own a GPU, these probes raise `RuntimeError: Found no NVIDIA driver on your system` during argument parsing, before any training starts.

Under `torchrun` this is masked: `validate_args` runs inside `initialize_megatron` in a process that is about to become a GPU rank (`torch.cuda.set_device(args.local_rank)` follows), so it happens to have a device. It breaks when arguments are validated in a **GPU-less** process — e.g. a launcher/driver that parses and validates args before any device is assigned. This is common in Ray-based orchestration, where a CPU-only driver validates args once and the GPU worker actors consume the already-validated args; there the job dies at argument validation with a confusing "no driver" error.

The three probes (current `main`):

1. `--moe-grouped-gemm` compute-capability check — [arguments.py L1296-1298](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/arguments.py#L1296-L1298): `dc = torch.cuda.get_device_capability(); assert dc[0] >= 8`.
2. TP/CP>1 `CUDA_DEVICE_MAX_CONNECTIONS` note — [arguments.py L1371-1372](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/arguments.py#L1371-L1372): `if (TP>1 or CP>1) and get_device_arch_version() < 10:`.
3. FSDP arch check — [arguments.py L1402](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/arguments.py#L1402): `if args.use_torch_fsdp2 or args.use_megatron_fsdp and get_device_arch_version() >= 10:`.

`get_device_arch_version()` ([common_utils.py L478-480](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/utils/common_utils.py#L478-L480)) does `torch.cuda.get_device_properties(torch.device("cuda:0")).major`, so it hard-requires a device.

**Steps/Code to reproduce bug**

On a GPU-less process (`torch.cuda.is_available() == False`), with Megatron importable — these are the exact calls `validate_args` makes for a MoE + TP/CP>1 config:

```python
import torch
print(torch.cuda.is_available()) # False
torch.cuda.get_device_capability() # RuntimeError: Found no NVIDIA driver ...
from megatron.training.utils import get_device_arch_version
get_device_arch_version() # RuntimeError: Found no NVIDIA driver ...
```

I ran this on a GPU-less driver process and both calls raise `RuntimeError: Found no NVIDIA driver`, so `validate_args` itself raises there for such a config.

**Expected behavior**

`validate_args` — pure argument validation — should complete without a live GPU rather than crashing with `Found no NVIDIA driver`. A device may not be assigned yet at parse time. The device-dependent requirements these probes gate belong on a device-owning process, not in argument parsing that may run device-less.

**Additional context**

Fix: guard each probe with `torch.cuda.is_available()`, checking the feature flag first so `is_available()` is only called for configs that actually probe the device. When a GPU is present the checks are byte-for-byte unchanged; when it is not, argument validation completes (the device-dependent checks are then only enforced if a device-owning process also runs `validate_args`, with a warning so the operator sets the env either way). A PR implementing this is linked. I searched open/closed issues and PRs (`Found no NVIDIA driver` + `validate_args`, `get_device_capability`/`get_device_arch_version` GPU-less) and found no existing report of this specific behavior.

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.