fix(distributed): Fix UMA graph-parallel stress reduction

[atulcthakur]

ALCHEMI Toolkit Pull Request

Description

In UMA graph-partition execution, each rank calculates a local contribution
to the system stress.

Energy and forces were reduced across ranks, but stress was not. Instead,
stress entered the generic output-consolidation logic, which divided this
local contribution by the number of ranks without first summing the
contributions.

This caused each rank to return an incomplete stress tensor.

Bug example

Using uma-s-1p1 with the omat task on two H100 GPUs:

• Single-GPU stress[0, 0]: 0.01203544
• Rank 0 before the fix: 0.00298884
• Rank 1 before the fix: 0.00302888
• Maximum stress error: approximately 9.2e-03

Neither rank returned the correct system stress.

Fix

In the internal graph-partition path:

1. SUM-reduce spec-declared graph outputs, currently UMA stress, over the
   domain process group.
2. Mark the reduced outputs as already global.
3. Pass them through output consolidation without applying the usual
   autograd /world_size correction.

After the all-reduce, every rank receives the complete system stress.

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/distributed_model.py

  • SUM-reduces UMA stress and prevents later division by world_size.

• test/distributed/model/test_uma_gp_partition_multigpu.py

  • Adds single-GPU versus two-rank stress comparisons.

Testing

Extended the real UMA graph-partition multi-GPU test to:

• Calculate a single-GPU reference.
• Run the same structure with two graph-partition ranks.
• Compare energy and gathered forces.
• Compare the complete stress returned by each rank with the single-GPU
  reference.

The test passed on two H100 GPUs:

• Stress error on both ranks: 2.96e-08

• 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 2e43853

[greptile-apps]

Preserve differentiable reduction semantics

This generic loop handles every spec-declared all-reduce output with plain dist.all_reduce, bypassing the autograd-aware reduction and world-size correction used by consolidate_sharded_outputs when an output is also listed in model_config.autograd_outputs. Route overlapping outputs through the differentiable consolidation path so extending the spec does not silently produce rank-local gradients.

Knowledge Base Used: Distributed domain decomposition

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

Greptile Summary

Fixes graph-partition stress consolidation by SUM-reducing spec-declared rank partials before marking them global, and extends the multi-GPU UMA test to compare every rank’s stress against a single-GPU reference.

Important Files Changed

Filename	Overview
nvalchemi/distributed/distributed_model.py	SUM-reduces additional spec-declared graph outputs and prevents consolidation from applying the replicated-output world-size correction, though the generic loop bypasses differentiable reduction semantics for future overlapping output classifications.
test/distributed/model/test_uma_gp_partition_multigpu.py	Adds per-rank stress gathering and comparison against a nontrivial single-GPU UMA reference.

_{Reviews (1): Last reviewed commit: "fix(distributed): all-reduce graph-paral..." | Re-trigger Greptile}