fix: set CUDA device before init_process_group; use_orig_params for FSDP v1#65
Draft
sparticlesteve wants to merge 2 commits into
Draft
Conversation
Without torch.cuda.set_device(local_rank) before init_process_group, all ranks default to CUDA device 0 at communicator creation time. NCCL reads the current device when initialising comm '0'; if every rank reports device 0 the intra-node P2P device-key mapping is wrong, which causes a deadlock when FSDP v2 later creates a second communicator that expects device keys consistent with comm '0'. Also add use_orig_params=True to the FSDP v1 wrapper so that optimisers and named_parameters() see the original parameter names rather than flat parameter views. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Allows the dataset directory to be specified at the command line rather than being hardcoded relative to the script location. Falls back to the original ./dataset/ default when --datadir is not provided. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
torch.cuda.set_device(local_rank)beforeinit_process_groupinsetup_ddp(). We observed consistent hangs with FSDP v2 +SHARD_GRAD_OPon 4× A100 nodes at NERSC Perlmutter that were resolved by adding this call. Our best guess is that without it, ranks may share CUDA device 0 at NCCL communicator init time, causing a mismatch when FSDP v2 creates per-rank device-specific communicators — but we haven't fully traced the root cause. Addingset_devicebeforeinit_process_groupis good practice regardless.use_orig_params=True: Needed fortorch.autograd.grad()to work through FSDP-sharded parameters (e.g. force prediction via-∇E). Note this does not resolve all FSDP v1 issues we observed with energy-force models on Perlmutter — FSDP v2 is the recommended path for force prediction. Still worth including as a correctness improvement for workflows that useautograd.grad()under FSDP v1.--datadirargument: Add optional--datadirCLI argument togfm_mlip_all_mpnn.py, defaulting to the original./dataset/path.Testing
We have an FSDP correctness test suite covering loss and gradient correctness for both FSDP v1 and v2, with and without force prediction. On 4× A100 (Perlmutter, SLURM job 51850250): 13 passed, 1 skipped. The skipped test is gradient correctness for FSDP v1 + force prediction: the two-phase backward (
autograd.gradfor forces followed byloss.backward) conflicts with FSDP v1's post-backward hooks, leaving gradients unavailable — a fundamental incompatibility rather than a gap we expect to close. Loss correctness for FSDP v1 + forces is still verified. FSDP v2 gradient correctness with force prediction is verified and passes.We are working on adapting this test suite to HydraGNN's testing conventions before upstreaming it in a follow-up PR.
🤖 Generated with Claude Code