fix(distributed): correct AIMNet2 stress reduction

[atulcthakur]

ALCHEMI Toolkit Pull Request

Description

AIMNet2 produced only rank-local stress contributions under eager domain decomposition, and compiled DD did not calculate stress.

This change:

• Sums AIMNet2 stress contributions across DD ranks.
• Enables strain-based stress for compiled DD.
• Makes the existing stress validator run without auto_fix.
• Adds assertions for the eager and compiled stress settings.

Energy and force behavior is unchanged.

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/models/aimnet2.py
• test/distributed/model/test_distributed_models.py
• test/distributed/validate/test_validate_cuda.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.

[greptile-apps]

Greptile Summary

Adds global AIMNet2 stress reduction and enables strain-derived stress for compiled domain decomposition.

• Extends the compile/recompile gate to numerically compare compiled and eager distributed stress.
• Runs the existing CUDA stress validator without automatic correction.
• Adds assertions for AIMNet2 stress reduction and compiled strain configuration.

Important Files Changed

Filename	Overview
nvalchemi/models/aimnet2.py	Declares stress as a globally reduced AIMNet2 halo output and enables strain-derived stress in compiled distributed execution.
test/distributed/model/test_aimnet2_compile_recompile_gate.py	Adds numeric compiled-versus-eager distributed stress coverage while retaining the steady-state recompilation gate.
test/distributed/model/test_distributed_models.py	Verifies that AIMNet2's distribution specification enables stress reduction and compiled strain stress.
test/distributed/validate/test_validate_cuda.py	Validates eager distributed AIMNet2 stress against the single-process reference without automatic fixes.

_{Reviews (2): Last reviewed commit: "test(distributed): cover compiled AIMNet..." | Re-trigger Greptile}

[atulcthakur]

/ok to test ca8bd9d

[atulcthakur]

/ok to test ca8bd9d