nvalchemi.distributed.ForceStrategy#

class nvalchemi.distributed.ForceStrategy(*values)[source]#

How a model’s forces are produced under a distributed forward.

A single named choice from which the framework derives CompilePolicy.forces_via_autograd / consolidate_node_energy / energy_output — so a model declares intent once and cannot express an invalid combination.

Members#

MODEL_INTERNAL

The model computes forces inside its own forward; the framework runs the wrapper as-is and consolidates its outputs. E.g. UMA.

FRAMEWORK_FROM_NODE_ENERGY

The framework drives an energy-only forward returning per-node "atomic_energies", does the owned-only per-graph sum + cross-rank all-reduce, and takes forces = -dE/dx via autograd. The MACE pattern.

FRAMEWORK_FROM_GLOBAL_ENERGY

Same autograd force path, but the model consolidates the per-system "energy" inside its forward and the framework differentiates it as-is. The AIMNet2 pattern.