Fix DistributedManager initialization on CPU-only hosts by pjaysaval · Pull Request #1938 · NVIDIA/physicsnemo

pjaysaval · 2026-08-19T15:20:20Z

Description

DistributedManager.setup() failed on hosts without accelerators (e.g. laptop-scale gloo runs, CPU nodes):

it passed the CPU torch.device as the device_id argument of torch.distributed.init_process_group, which recent PyTorch versions reject with ValueError: init_process_group device_id parameter must be an accelerator with an index;
with local_rank=None (the SLURM/OpenMPI initialization paths) it derived the local rank as rank % torch.cuda.device_count(), a division by zero when no accelerator is present.

This PR binds a device to the process group only on CUDA (device_id=None otherwise, the documented default), defaults local_rank to 0 on CPU-only hosts, and additionally guards cleanup() on the process group actually existing so a failed initialization no longer cascades into AssertionError: Process group cannot be None from the atexit handler, obscuring the original error.

Verified with a two-process gloo run launched by torchrun on a CPU-only machine (initialize → broadcast → barrier → cleanup). Adds a regression test that forces the CPU code path via monkeypatched torch.cuda.is_available / device_count, so it exercises the fix even on GPU CI (the existing distributed-manager tests are skipped on CPU by an autouse fixture, which is how this path went unnoticed). Also adds docstrings to two pre-existing undocumented methods in manager.py flagged by the interrogate pre-commit hook.

closes #1936

Checklist

I am familiar with the Contributing Guidelines.
New or existing tests cover these changes.
The documentation is up to date with these changes.
The CHANGELOG.md is up to date with these changes.
An issue is linked to this pull request.
If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

None.

copy-pr-bot · 2026-08-19T15:20:24Z

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.

github-actions · 2026-08-19T15:20:54Z

CODEOWNERS review map

Current for commit f0387e173759. An approval covers every file listed for that owner; one owner is sufficient for shared files.

⏳ @coreyjadams — 2 file(s)

physicsnemo/distributed/manager.py
test/distributed/test_manager.py

⏳ @negin513 — 2 file(s)

physicsnemo/distributed/manager.py
test/distributed/test_manager.py

No CODEOWNER

CHANGELOG.md

Comment /codeowners-info to refresh.

greptile-apps · 2026-08-19T15:23:32Z

Greptile Summary

Fixes direct DistributedManager.setup() use on CPU-only hosts and makes cleanup safe when no process group exists.

Passes a process-group device ID only for CUDA devices.
Avoids modulo zero when setup() itself derives a local rank.
Adds a CPU/gloo regression test, cleanup guard, changelog entry, and missing public-method docstrings.

Important Files Changed

Filename	Overview
physicsnemo/distributed/manager.py	Corrects direct CPU setup and cleanup behavior, but initialize_env() can still divide by zero before reaching the new fallback.
test/distributed/test_manager.py	Exercises setup() under simulated CPU conditions but bypasses the initialize_env() path that retains the failure.
CHANGELOG.md	Documents the intended CPU initialization and cleanup fixes, although the initialization claim is broader than the implementation.

_{Reviews (1): Last reviewed commit: "Fix DistributedManager initialization on..." | Re-trigger Greptile}

greptile-apps · 2026-08-19T15:23:35Z

+                n_devices = torch.cuda.device_count()
+                manager._local_rank = rank % n_devices if n_devices > 0 else 0


CPU fallback misses environment initialization

If a CPU-only process has RANK and WORLD_SIZE set without LOCAL_RANK, initialize_env() evaluates rank % torch.cuda.device_count() before reaching this fallback, causing initialization to fail with ZeroDivisionError before the process group is created.

Knowledge Base Used: Distributed Execution Primitives

Good catch — confirmed: with RANK/WORLD_SIZE set and no LOCAL_RANK on a CPU-only host, initialize_env() divided by the accelerator count before setup() ran. Fixed in c0d5b25 by leaving local_rank as None there and letting setup() derive it (rank modulo the accelerator count, or 0 on CPU-only hosts), with a regression test for the environment-variable initialization path.

DistributedManager.setup() failed on hosts without accelerators: - it passed the CPU torch.device as the device_id argument of torch.distributed.init_process_group, which recent PyTorch versions reject with 'ValueError: init_process_group device_id parameter must be an accelerator with an index'; - with local_rank=None it derived the local rank as rank % torch.cuda.device_count(), a division by zero when no accelerator is present. Bind a device to the process group only on CUDA, and default local_rank to 0 on CPU-only hosts. Also guard cleanup() on the process group actually existing, so a failed initialization no longer cascades into 'AssertionError: Process group cannot be None' at exit. Verified with a two-process gloo run launched by torchrun on a CPU-only machine (initialize, broadcast, barrier, cleanup), and add a regression test that forces the CPU code path via monkeypatched torch.cuda.is_available / device_count. Signed-off-by: Piyoosh Jaysaval <piyoosh.jaysaval@pnnl.gov>

initialize_env() computed rank % torch.cuda.device_count() before calling setup(), so CPU-only hosts with RANK and WORLD_SIZE set but no LOCAL_RANK still failed with ZeroDivisionError. Leave local_rank as None in that case and let setup() derive it (rank modulo the accelerator count, or 0 on CPU-only hosts), and add a regression test for the environment-variable initialization path. Signed-off-by: Piyoosh Jaysaval <piyoosh.jaysaval@pnnl.gov>

peterdsharpe

Looks good to me! Checked this works on my end too. Thanks!

peterdsharpe · 2026-08-20T16:05:25Z

/ok to test c0d5b25

pjaysaval · 2026-09-01T21:28:42Z

Thanks again for the review @peterdsharpe. This has been approved with CI green for a couple of weeks now but appears to still be waiting on a CODEOWNERS review to merge. Is there anything needed from my side? I've held off rebasing onto main to keep the approval and CI state intact, but happy to update the branch if that's preferred.

Resolve the conflict with NVIDIA#1968 (process-group timeout): keep the CPU-safe device_id and pass the new timeout argument. Signed-off-by: Piyoosh Jaysaval <piyoosh.jaysaval@pnnl.gov>

pjaysaval · 2026-09-18T15:40:47Z

@coreyjadams @negin513, this PR needs a review from one of you as CODEOWNERS of physicsnemo/distributed/. It was approved by @peterdsharpe on Aug 19 with CI green and has been waiting since. I have just merged current main into the branch to resolve the conflict with #1968: the CPU-safe device_id is kept and the new timeout argument is passed through unchanged, so the functional diff is the same three hunks as before (CPU device_id, zero-safe local_rank, guarded cleanup()), plus the regression tests. Happy to squash or adjust anything you would like.

negin513 · 2026-10-01T04:53:26Z

/ok to test 93604eb

negin513

LGTM!! :-)

pjaysaval requested review from coreyjadams and negin513 as code owners August 19, 2026 15:20

greptile-apps Bot reviewed Aug 19, 2026

View reviewed changes

pjaysaval added 2 commits August 19, 2026 10:54

pjaysaval force-pushed the fix/distributed-manager-cpu-init branch from f0387e1 to c0d5b25 Compare August 19, 2026 15:54

peterdsharpe approved these changes Aug 20, 2026

View reviewed changes

negin513 mentioned this pull request Sep 9, 2026

Configure the distributed process-group timeout in code or environment #1968

Merged

4 tasks done

Merge branch 'main' into fix/distributed-manager-cpu-init

93604eb

Resolve the conflict with NVIDIA#1968 (process-group timeout): keep the CPU-safe device_id and pass the new timeout argument. Signed-off-by: Piyoosh Jaysaval <piyoosh.jaysaval@pnnl.gov>

negin513 approved these changes Oct 3, 2026

View reviewed changes

Merge branch 'main' into fix/distributed-manager-cpu-init

cd8e046

Navigation Menu

Repository navigation

Fix DistributedManager initialization on CPU-only hosts - #1938

Fix DistributedManager initialization on CPU-only hosts#1938
pjaysaval wants to merge 4 commits into
NVIDIA:mainfrom
pjaysaval:fix/distributed-manager-cpu-init

pjaysaval commented Aug 19, 2026

Uh oh!

copy-pr-bot Bot commented Aug 19, 2026

Uh oh!

github-actions Bot commented Aug 19, 2026

Uh oh!

greptile-apps Bot commented Aug 19, 2026

Uh oh!

greptile-apps Bot Aug 19, 2026

Uh oh!

pjaysaval Aug 19, 2026

Uh oh!

peterdsharpe left a comment

Uh oh!

peterdsharpe commented Aug 20, 2026

Uh oh!

pjaysaval commented Sep 1, 2026

Uh oh!

pjaysaval commented Sep 18, 2026

Uh oh!

negin513 commented Oct 1, 2026

Uh oh!

negin513 left a comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		n_devices = torch.cuda.device_count()
		manager._local_rank = rank % n_devices if n_devices > 0 else 0

Conversation

pjaysaval commented Aug 19, 2026

Description

Checklist

Dependencies

Uh oh!

copy-pr-bot Bot commented Aug 19, 2026

Uh oh!

github-actions Bot commented Aug 19, 2026

CODEOWNERS review map

No CODEOWNER

Uh oh!

greptile-apps Bot commented Aug 19, 2026

Greptile Summary

Important Files Changed

Uh oh!

greptile-apps Bot Aug 19, 2026

Choose a reason for hiding this comment

Uh oh!

pjaysaval Aug 19, 2026

Choose a reason for hiding this comment

Uh oh!

peterdsharpe left a comment

Choose a reason for hiding this comment

Uh oh!

peterdsharpe commented Aug 20, 2026

Uh oh!

pjaysaval commented Sep 1, 2026

Uh oh!

pjaysaval commented Sep 18, 2026

Uh oh!

negin513 commented Oct 1, 2026

Uh oh!

negin513 left a comment

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants