fix: keep DD communication inside domain subgroups

[atulcthakur]

ALCHEMI Toolkit Pull Request

Description

In a (pipeline, domain) mesh, each pipeline stage has its own domain
process group. Several DD communication paths either used the global WORLD
group or passed domain-local peer numbers to PyTorch APIs that expect global
ranks.

For a 2×2 mesh with domain groups {0,1} and {2,3}, this caused errors,
incorrect cross-stage communication, or hangs.

Example failure

Consider a four-GPU pipeline with:

• FIRE on global ranks {0,1}
• NVT Nose–Hoover on global ranks {2,3}
• Two domain ranks per pipeline stage

For the NVT stage, domain-local rank 0 is global rank 2. The old controller
broadcast used src=0, which PyTorch interpreted as global rank 0. Because
global rank 0 is not a member of {2,3}, the first NVT step failed with:

ValueError: Global rank 0 is not part of group

The same local-versus-global rank mismatch affected P2P halo exchanges, while
some fixed halo and atom-ownership collectives used all four ranks instead of
the stage’s two-rank domain group.

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Performance improvement
• Documentation update
• Refactoring (no functional changes)
• CI/CD or infrastructure change

Related Issues

Changes Made

• Translate each domain group’s local leader rank to its global rank before
  broadcasting thermostat, barostat, or cell state.
• Gather atom-owner counts only within the current domain group.
• Translate domain-local P2P peer numbers to global ranks before isend and
  irecv.
• Preserve the supplied domain process group in fixed-size halo collectives.
• Publish the active domain group to compiled halo custom operations so their
  forward and backward exchanges do not use WORLD.
• Add four-process Gloo tests using domain groups {0,1} and {2,3}.

Files changed

• nvalchemi/distributed/_dynamics_coordinator.py
• nvalchemi/distributed/sharded_batch.py
• nvalchemi/distributed/_core/gather_primitives.py
• nvalchemi/distributed/_core/particle_halo.py
• nvalchemi/distributed/strategy.py
• test/distributed/test_2d_subgroup_collectives.py
• test/distributed/_core/test_subgroup_communication.py

Testing

• Unit tests pass locally (make pytest)
• Linting passes (make lint)
• New tests added for new functionality meets coverage expectations?

Checklist

• I have read and understand the Contributing Guidelines
• I have updated the CHANGELOG.md
• I have performed a self-review of my code
• I have added docstrings to new functions/classes
• I have updated the documentation (if applicable)

Additional Notes

Tip

This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.

[copy-pr-bot]

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.

[atulcthakur]

/ok to test 1f5e36e

[greptile-apps]

Greptile Summary

This PR confines domain-decomposition communication to each pipeline stage’s domain subgroup.

• Maps subgroup-local broadcast and point-to-point ranks to the global ranks expected by PyTorch.
• Preserves domain process groups in ownership gathers and fixed-size halo collectives.
• Publishes the active domain group for compiled halo forward and backward operations.
• Adds four-process Gloo coverage for two independent domain subgroups.

Important Files Changed

Filename	Overview
nvalchemi/distributed/_core/gather_primitives.py	Preserves process groups in fixed halo collectives and translates subgroup-local P2P peers to global ranks.
nvalchemi/distributed/_core/particle_halo.py	Routes compiled halo custom operations and their gradients through the active domain process group.
nvalchemi/distributed/_dynamics_coordinator.py	Broadcasts replicated controller and cell state from each domain subgroup’s local leader.
nvalchemi/distributed/sharded_batch.py	Restricts atom-owner count gathering to the batch’s mesh process group.
nvalchemi/distributed/strategy.py	Publishes the halo strategy’s domain process group immediately before model execution.
test/distributed/_core/test_subgroup_communication.py	Adds subgroup tests for P2P helpers, fixed halo exchange, and compiled halo autograd.
test/distributed/test_2d_subgroup_collectives.py	Adds two-dimensional mesh tests for controller broadcasts and subgroup-local ownership counts.

_{Reviews (1): Last reviewed commit: "fix: scope DD communication to domain gr..." | Re-trigger Greptile}