nvalchemi.distributed.spec.SPEC_MPNN_GP

nvalchemi.distributed.spec.SPEC_MPNN_GP = MLIPSpec(distribution=DistributionSpec(policy=GraphParallelPolicy(), custom_ops=(), third_party_helpers=(), adapters=(), shard_fields=()), owned_only_outputs=frozenset({'atomic_energies', 'forces'}), 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=CompilePolicy(static_shapes=True, graph_padder=None, force_strategy=<ForceStrategy.FRAMEWORK_FROM_GLOBAL_ENERGY: 'framework_from_global_energy'>, stress_via_strain=False))

atoms split by a balanced index range (no spatial halo), each rank owning the edges into its nodes. The plain-interior promotion set (shard_fields=()) keeps the model on plain owned-row tensors; per message-passing layer the framework all-gathers the node features to a replicated tensor (reduce-scatter on the backward) so every edge sees its source, and the final per-graph node-energy sum drops to owners and all-reduces. The complement of SPEC_MPNN_HALO: index-balanced and locality-blind, it avoids the ghost overhead halo pays when the box approaches the cutoff, at the cost of a full node all-gather each layer.

Type:

Scatter-heavy MPNNs under the graph-parallel strategy