Skip to content

fix(launcher): honor --nproc-per-node over the local device count - #3775

Open
kabirvashisht4-glitch wants to merge 2 commits into
NVIDIA-NeMo:mainfrom
kabirvashisht4-glitch:kabirvashisht4-glitch/fix/honor-nproc-per-node
Open

kabirvashisht4-glitch wants to merge 2 commits into
NVIDIA-NeMo:mainfrom
kabirvashisht4-glitch:kabirvashisht4-glitch/fix/honor-nproc-per-node

Conversation

@kabirvashisht4-glitch

Copy link
Copy Markdown
Contributor

What does this PR do ?

Makes InteractiveLauncher act on an explicit --nproc-per-node instead of
letting the local GPU probe override it. The probe ran first, so the flag lost
to it in two ways:

num_devices = determine_local_world_size(nproc_per_node="gpu")
assert num_devices > 0, "Expected num-devices to be > 0"

if nproc_per_node == 1 or num_devices == 1:
    return self._run_recipe_in_process(recipe_target, config)
else:
    effective_nproc = nproc_per_node if nproc_per_node is not None else num_devices

1. --nproc-per-node N was silently downgraded on a one-GPU host.
num_devices == 1 short-circuits to the in-process path whatever the user
asked for, so automodel config.yaml --nproc-per-node 4 started one worker
and logged "Launching job locally on a single device" without saying the
requested value was dropped. Running several ranks on one device is a normal way
to exercise FSDP/DDP or PP code paths on a dev box, and torchrun supports it.

2. --nproc-per-node 1 failed with no visible CUDA device.
determine_local_world_size(nproc_per_node="gpu") raises when
torch.cuda.is_available() is false, and it ran before the nproc_per_node == 1
branch — the one path that needs no device count, since it runs the recipe
in-process:

ValueError: Cuda is not available.

Also fires under CUDA_VISIBLE_DEVICES="". The error comes from inside
torch.distributed.run, so it names neither the flag nor an AutoModel
requirement.

Scope, to be explicit: this does not make CPU training work — the recipe
still requires CUDA further down. It makes the documented flag behave as
documented and lets the failure come from the framework rather than from a
device probe the requested single-process path never needed.

Changelog

  • InteractiveLauncher.launch probes determine_local_world_size only when
    nproc_per_node is None; an explicit value is used as-is.
  • The single-process branch keys off the resolved worker count rather than
    nproc_per_node == 1 or num_devices == 1, so an explicit N > 1 reaches
    torchrun even when one device is visible.
  • Three CPU tests in tests/unit_tests/launcher/test_interactive_launcher.py:
    explicit 4 on a one-device host reaches torchrun with 4 workers; an explicit
    count never calls the probe; and --nproc-per-node 1 still runs in-process
    when the probe would raise ValueError("Cuda is not available.").

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you add or update any necessary documentation?

--help already documents the flag as "Number of workers per node for
local/interactive jobs" with no mention of clamping to the device count, so the
code now matches it and no doc change was needed.

Compatibility. Without the flag the behaviour is byte-for-byte the same: the
probe still runs, still asserts > 0, and still picks in-process for one device
or torchrun for more. All 21 pre-existing tests in the file pass unmodified,
including test_interactive_launcher_single_device,
test_interactive_launcher_multi_device, and
test_interactive_launcher_multi_device_explicit_nproc.

Verification (CPU, no GPU needed):

  • pytest tests/unit_tests/launcher/ → 103 passed.
  • Reverting only interactive.py makes all three new tests fail, so they are
    not vacuous.
  • pytest tests/unit_tests/launcher/ tests/unit_tests/_cli/ tests/unit_tests/recipes/ → 1291 passed, 12 skipped.
    (tests/unit_tests/recipes/dllm excluded locally: it fails to import
    transformers.models.diffusion_gemma on main too, unrelated to this change.)
  • ruff format / ruff check clean.

Additional Information

InteractiveLauncher.launch probed the local GPU count before acting on an
explicit --nproc-per-node, so the flag lost to the probe twice over. On a
one-GPU host `num_devices == 1` short-circuited to the in-process path, so
`--nproc-per-node 4` silently started one worker instead of four. And
`determine_local_world_size(nproc_per_node="gpu")` raises when no CUDA
device is visible, so `--nproc-per-node 1` failed inside torch before the
in-process branch that needs no device count could run.

Probe only when no worker count was requested, and pick the single-process
path from the resolved count. Behaviour without the flag is unchanged.

Signed-off-by: kabirvashisht4-glitch <kabirvashisht4@gmail.com>
@kabirvashisht4-glitch
kabirvashisht4-glitch requested a review from a team as a code owner September 1, 2026 03:35
@copy-pr-bot

copy-pr-bot Bot commented Sep 1, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Sep 3, 2026
@akoumpa

akoumpa commented Sep 4, 2026

Copy link
Copy Markdown
Contributor

/ok to test a0989c3

@akoumpa

akoumpa commented Sep 9, 2026

Copy link
Copy Markdown
Contributor

/ok to test e445bb0

@akoumpa

akoumpa commented Sep 9, 2026

Copy link
Copy Markdown
Contributor

/ok to test e445bb0

This branch was successfully deployed

1 active and 2 inactive deployments
public e445bb0f Deployed Sep 9, 2026 by copy-pr-bot[bot] via release / finalize / notify #4402
test e445bb0f Deployed Sep 9, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #10562
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-request waiting-on-maintainers Waiting on maintainers to respond

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants