nvalchemi.distributed.spec.SPEC_UMA_HALO

nvalchemi.distributed.spec.SPEC_UMA_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)

UMA (eSCN-family) via the halo storage policy, with fairchem graph parallel disabled.

Each rank holds owned + halo rows and runs a standard full forward over them. UMAWrapper.distribution_spec layers in OpAdapters for the fused Triton edge-permute kernels (edge→node aggregation gets per-layer halo correction); the per-system reductions route through per_system_reduce(), and forces/stress flow through plain autograd.