fix(distributed): synchronize convergence across ranks

[atulcthakur]

ALCHEMI Toolkit Pull Request

Description

DomainParallel treated convergence indices as Boolean values. Because tensor([0]) means “system 0 converged” but converts to False, convergence could be missed.
Ranks could also disagree locally and enter different collectives, causing a hang.
The DD run and pipeline paths also failed to stop or retire the converged system.

Fix

• Preserve convergence as system-index tensors.
• Reduce a per-system convergence mask across every domain rank.
• Stop completed DD runs and retire converged pipeline systems.
• Avoid convergence collectives when no convergence hook is configured.

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

• nvalchemi/distributed/domain_parallel.py

  • Fix distributed convergence reduction in step().
  • Add convergence-based early exit to run().
  • Fix pipeline completion and completion-reason checks.

• test/distributed/test_domain_parallel_convergence.py

  • Test preservation of tensor([0]).
  • Test rank disagreement without hanging.
  • Test that fixed-step dynamics adds no convergence collective.
  • Test run() early exit.
  • Test pipeline-stage retirement.

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.

[greptile-apps]

Greptile Summary

This PR fixes three inter-related correctness bugs in DomainParallel's convergence handling and adds a dedicated regression-test suite to cover them.

• Falsy-zero bug: system index tensor([0]) evaluated to False in the old boolean-reduction path; the new code keeps convergence as a per-system index tensor and reduces a (num_graphs,) integer mask with ReduceOp.MIN.
• Rank-desync / hang: convergence was previously only checked when the local rank already had a non-None result, so ranks could skip the collective and deadlock; convergence_due is now evaluated uniformly so every rank always enters the same all_reduce.
• Missing early-exit paths: run() never broke on convergence, and pipeline stages never retired a converged system; both are now handled with the correct numel() == num_graphs predicate.

Important Files Changed

Filename	Overview
nvalchemi/distributed/domain_parallel.py	Three convergence bugs fixed: tensor([0]) falsy-boolean, rank-disagree hang, and pipeline retirement; logic is correct throughout.
test/distributed/test_domain_parallel_convergence.py	New regression-test suite covering all fixed paths; one grouped-pipeline test is appropriately skipped with a Gloo-reliability rationale.

_{Reviews (4): Last reviewed commit: "test: skip unreliable grouped Gloo pipel..." | Re-trigger Greptile}

[atulcthakur]

/ok to test 6a34c19

[atulcthakur]

/ok to test 63c6161

[atulcthakur]

/ok to test 369a08b