nvalchemi.distributed.spec.SPEC_MPNN_HALO

nvalchemi.distributed.spec.SPEC_MPNN_HALO = MLIPSpec(distribution=DistributionSpec(policy=HaloStoragePolicy(scatter_mode='halo_correction', gather_mode='halo_read'), custom_ops=(), third_party_helpers=(), adapters=(), shard_fields=('positions', 'charges', 'atomic_numbers')), owned_only_outputs=frozenset(), all_reduce_outputs=frozenset(), output_kinds={'energy': <OutputKind.PER_GRAPH: 'per_graph'>, 'forces': <OutputKind.PER_NODE: 'per_node'>, 'stress': <OutputKind.PER_GRAPH: 'per_graph'>, 'atomic_energies': <OutputKind.PER_NODE: 'per_node'>}, system_reductions=True, node_energy_key=None, node_virial_key=None, gp_replicate_geometry=False, outputs=None, compile=None)

MACE, NequIP, Allegro, ORB. Every edge-level update is a scatter_sum into per-atom features (the halo-correction handler keeps halo rows in sync), plus a final per-graph scatter_sum on node energies (the per_system_reduce() handler drops halo rows and all-reduces across ranks), so stock model.forward produces globally-correct energy + forces with no wrapper-side post-processing.

Type:

Scatter-heavy MPNNs