diff --git a/deepmd/dpmodel/atomic_model/inter_potential.py b/deepmd/dpmodel/atomic_model/inner_potential.py similarity index 96% rename from deepmd/dpmodel/atomic_model/inter_potential.py rename to deepmd/dpmodel/atomic_model/inner_potential.py index b95aeb8da5..4e2b727688 100644 --- a/deepmd/dpmodel/atomic_model/inter_potential.py +++ b/deepmd/dpmodel/atomic_model/inner_potential.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Analytical pair potentials for Zone bridging (backend-agnostic port of -``deepmd.pt``'s ``InterPotential``). Lives in the atomic-model package: +``deepmd.pt``'s ``InnerPotential``). Lives in the atomic-model package: the atomic layer owns per-atom energy assembly, where the ZBL term is injected on the graph route. """ @@ -63,7 +63,7 @@ _A_BOHR = 0.5291772109 # Bohr radius in Å -class InterPotential(NativeOP): +class InnerPotential(NativeOP): """Analytical pair potential for Zone bridging. Supports the Ziegler-Biersack-Littmark (ZBL) screened nuclear repulsion @@ -72,7 +72,7 @@ class InterPotential(NativeOP): contributes ``V_ZBL(r_ij) / 2`` to both atom i and atom j, avoiding double-counting from the symmetric neighbor list. Backend-agnostic (array-API) port of the reference implementation in - ``deepmd.pt.model.model.sezm_model.InterPotential``. + ``deepmd.pt.model.model.sezm_model.InnerPotential``. Parameters ---------- @@ -93,7 +93,7 @@ def __init__(self, type_map: list[str], mode: str = "zbl") -> None: super().__init__() mode = str(mode).upper() if mode != "ZBL": - raise ValueError(f"Unknown InterPotential mode: {mode}") + raise ValueError(f"Unknown InnerPotential mode: {mode}") self.mode = mode self.type_map = list(type_map) self.ntypes_real = len(type_map) @@ -266,8 +266,8 @@ def call( return xp.astype(xp.reshape(atom_energy, (1, n_node, 1)), edge_vec.dtype) -@BaseAtomicModel.register("inter_potential") -class InterPotentialAtomicModel(BaseAtomicModel): +@BaseAtomicModel.register("inner_potential") +class InnerPotentialAtomicModel(BaseAtomicModel): """Analytical bridging pair potential as an ATOMIC MODEL. First-principles composition design: the analytical term maps local @@ -302,7 +302,7 @@ def __init__( **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) - self.potential = InterPotential(type_map=list(type_map), mode=mode) + self.potential = InnerPotential(type_map=list(type_map), mode=mode) self.mode = self.potential.mode self.rcut = float(rcut) self.sel = ( @@ -317,7 +317,7 @@ def change_type_map( If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. The generic base handles the public map and the stat/exclusion state; - the element lookup belongs to :class:`InterPotential`, so the update is + the element lookup belongs to :class:`InnerPotential`, so the update is delegated there rather than reimplemented here (review 3649295675 -- without it the lookup keeps the ORIGINAL elements while ``atype`` values mean new ones, and a longer new map raises ``IndexError``). @@ -403,7 +403,7 @@ def forward_atomic( ) -> dict: """Dense route unsupported: the term rides the NeighborGraph route only.""" raise NotImplementedError( - "InterPotentialAtomicModel rides the NeighborGraph route only; " + "InnerPotentialAtomicModel rides the NeighborGraph route only; " "the dense (nlist) route has no injection site for the term" ) @@ -455,7 +455,7 @@ def serialize(self) -> dict: data.update( { "@class": "Model", - "type": "inter_potential", + "type": "inner_potential", "@version": 1, "mode": self.mode, "rcut": self.rcut, @@ -465,7 +465,7 @@ def serialize(self) -> dict: return data @classmethod - def deserialize(cls, data: dict) -> "InterPotentialAtomicModel": + def deserialize(cls, data: dict) -> "InnerPotentialAtomicModel": data = data.copy() check_version_compatibility(data.pop("@version", 1), 1, 1) data.pop("@class", None) diff --git a/deepmd/dpmodel/model/dp_linear_model.py b/deepmd/dpmodel/model/dp_linear_model.py index 132dc48385..3aecf8e911 100644 --- a/deepmd/dpmodel/model/dp_linear_model.py +++ b/deepmd/dpmodel/model/dp_linear_model.py @@ -35,7 +35,7 @@ class LinearEnergyModel(DPModelCommon, DPLinearModel_): energies; on the NeighborGraph route every child consumes the same graph, so the summed energy differentiates through one shared edge backward. Used e.g. for analytical bridging compositions - (learned model + :class:`~deepmd.dpmodel.atomic_model.inter_potential.InterPotentialAtomicModel`). + (learned model + :class:`~deepmd.dpmodel.atomic_model.inner_potential.InnerPotentialAtomicModel`). """ def __init__( diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 7af56a74c3..a66f40e7f3 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -61,7 +61,7 @@ def get_standard_model(data: dict) -> BaseModel: data = copy.deepcopy(data) # Analytical bridging (e.g. ZBL): the radii feed the DESCRIPTOR's # InnerClamp/BridgingSwitch (mirrors pt's builder); the method builds the - # atomic model's InterPotential below. + # atomic model's InnerPotential below. bridging_method = str(data.get("bridging_method", "none")) bridging_enabled = bridging_method.lower() not in ("none", "") if bridging_enabled: @@ -77,8 +77,8 @@ def get_standard_model(data: dict) -> BaseModel: # Composition, not a flag (first-principles design): the analytical # bridging term is its own atomic model, summed with the learned one by the # existing linear composition machinery. - from deepmd.dpmodel.atomic_model.inter_potential import ( - InterPotentialAtomicModel, + from deepmd.dpmodel.atomic_model.inner_potential import ( + InnerPotentialAtomicModel, ) from deepmd.dpmodel.atomic_model.linear_atomic_model import ( LinearEnergyAtomicModel, @@ -87,7 +87,7 @@ def get_standard_model(data: dict) -> BaseModel: LinearEnergyModel, ) - zbl_atomic = InterPotentialAtomicModel( + zbl_atomic = InnerPotentialAtomicModel( type_map=data["type_map"], mode=bridging_method, rcut=descriptor.get_rcut(), @@ -140,7 +140,7 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel: exclusions and the analytical-bridging composition -- so ``spin`` and ``bridging_method`` combine for free: the wrapper re-classes whatever atomic model came back, be it a single learned model or a - ``LinearEnergyAtomicModel`` over ``[learned, InterPotential]`` (the + ``LinearEnergyAtomicModel`` over ``[learned, InnerPotential]`` (the analytical child accepts and ignores ``spin``; the learned child consumes it). diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index 13c8e8ae8a..d9689055da 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -232,6 +232,8 @@ def _make_wrapped_sampler( @functools.lru_cache def wrapped_sampler() -> list[dict]: sampled = sampled_func() + if not sampled: + return sampled if self.pair_excl is not None: pair_exclude_types = self.pair_excl.get_exclude_types() for sample in sampled: @@ -662,7 +664,9 @@ def compute_fitting_input_stat( """ pass - def _get_forward_wrapper_func(self) -> Callable[..., torch.Tensor]: + def _get_forward_wrapper_func( + self, + ) -> Callable[..., dict[str, torch.Tensor]]: """Get a forward wrapper of the atomic model for output bias calculation.""" def model_forward( @@ -672,7 +676,9 @@ def model_forward( fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, + spin: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: + del spin with ( torch.no_grad() ): # it's essential for pure torch forward function to use auto_batchsize diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 465ff0af19..494772d453 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -37,6 +37,9 @@ extend_input_and_build_neighbor_list, nlist_distinguish_types, ) +from deepmd.pt.utils.stat import ( + compute_output_stats, +) from deepmd.utils.path import ( DPPath, ) @@ -304,6 +307,48 @@ def get_out_bias(self) -> torch.Tensor: def set_out_bias(self, out_bias: torch.Tensor) -> None: self.atomic_model.set_out_bias(out_bias) + def predict_atomic_outputs_for_stat( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Return atomic outputs through the standard atomic-model path.""" + return self.atomic_model._get_forward_wrapper_func()( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + spin=spin, + ) + + def _change_out_bias_with_model_forward( + self, + merged: Callable[[], list[dict]] | list[dict], + model_forward: Callable[..., dict[str, torch.Tensor]], + ) -> None: + """Fit a residual output-bias shift from a complete model predictor.""" + atomic_model = self.atomic_model + delta_bias, out_std = compute_output_stats( + merged, + atomic_model.get_ntypes(), + keys=atomic_model.bias_keys, + model_forward=model_forward, + rcond=atomic_model.rcond, + preset_bias=atomic_model.preset_out_bias, + stats_distinguish_types=( + atomic_model.get_compute_stats_distinguish_types() + ), + intensive=atomic_model.get_intensive(), + ) + atomic_model._store_out_stat(delta_bias, out_std, add=True) + def change_out_bias( self, merged: Any, @@ -326,6 +371,12 @@ def change_out_bias( and do least square on the errors to obtain the target shift as bias. 'set-by-statistic' : directly use the statistic output bias in the target dataset. """ + if bias_adjust_mode == "change-by-statistic": + self._change_out_bias_with_model_forward( + merged, + self.predict_atomic_outputs_for_stat, + ) + return self.atomic_model.change_out_bias( merged, bias_adjust_mode=bias_adjust_mode, diff --git a/deepmd/pt/model/model/model.py b/deepmd/pt/model/model/model.py index b4f09b40ec..5f89ff50ad 100644 --- a/deepmd/pt/model/model/model.py +++ b/deepmd/pt/model/model/model.py @@ -49,6 +49,48 @@ def compute_or_load_stat( """ raise NotImplementedError + def predict_atomic_outputs_for_stat( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """ + Return complete atomic outputs used by residual output statistics. + + Final model classes own this prediction contract because only they know + the complete physical forward, including model-level preprocessing and + analytical contributions. Implementations must not compute derivatives + or mutate compile caches. + + Parameters + ---------- + coord + Local coordinates with shape (nf, nloc, 3). + atype + Local atom types with shape (nf, nloc). + box + Simulation cells with shape (nf, 9), or ``None``. + fparam + Optional frame parameters. + aparam + Optional atomic parameters. + charge_spin + Optional frame-level charge and spin conditions. + spin + Optional native per-atom spin vectors. + + Returns + ------- + dict[str, torch.Tensor] + Complete atomic outputs for output-statistics regression. + """ + raise NotImplementedError + @torch.jit.export def get_observed_type_list(self) -> list[str]: """Get observed types (elements) of the model during data statistics. diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index b623222a41..1cb51fe7f2 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -108,7 +108,7 @@ (``edge_vec.detach().requires_grad_(True)``), so neighbor construction, shift application and coordinate gathers live outside the autograd region (NOTE 11). -* The SeZM descriptor and the analytical ZBL term (``InterPotential``) +* The SeZM descriptor and the analytical ZBL term (``InnerPotential``) both consume that edge-vector leaf, so the energy depends on coordinates *only* through ``edge_vec``. * The fitting network predicts per-atom energy; ``apply_out_stat`` adds @@ -760,8 +760,8 @@ def __init__( self.bridging_method: str = str(bridging_method).upper() self.bridging_r_inner = float(bridging_r_inner) self.bridging_r_outer = float(bridging_r_outer) - self.inter_potential: InterPotential | None = ( - InterPotential(type_map=self.get_type_map(), mode=self.bridging_method) + self.inter_potential: InnerPotential | None = ( + InnerPotential(type_map=self.get_type_map(), mode=self.bridging_method) if self.bridging_method != "NONE" else None ) @@ -955,6 +955,7 @@ def forward_common( charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, embedding_only: bool = False, + atomic_output_only: bool = False, ) -> dict[str, torch.Tensor]: """ Return model prediction using standard neighbor list. @@ -983,6 +984,8 @@ def forward_common( clean `dens` batches may not provide corruption masks. charge_spin Frame-level charge and spin conditions with shape `(nf, 2)`. + atomic_output_only + Whether to stop after assembling complete atomic outputs. Returns ------- @@ -1007,7 +1010,7 @@ def forward_common( # === Step 2. Build geometry schema === with nvtx_range("SeZM/build_neighbor_list"): - if self.get_active_mode() == "dens": + if self.get_active_mode() == "dens" and not atomic_output_only: # extended_coord: (nf, nall, 3), extended_atype: (nf, nall) # nlist: (nf, nloc, nsel), mapping: (nf, nall) extended_coord, extended_atype, nlist, mapping = ( @@ -1017,7 +1020,7 @@ def forward_common( edge_schema = self.build_neighbor_list(cc, atype, bb) # === Step 3. Run the model compute path === - if self.get_active_mode() == "dens": + if self.get_active_mode() == "dens" and not atomic_output_only: return self.forward_common_lower_dens( extended_coord, extended_atype, @@ -1044,6 +1047,7 @@ def forward_common( spin=spin, input_prec=input_prec, embedding_only=embedding_only, + atomic_output_only=atomic_output_only, ) def forward_common_lower( @@ -1064,6 +1068,7 @@ def forward_common_lower( input_prec: torch.dtype | None = None, use_compile: bool | None = None, embedding_only: bool = False, + atomic_output_only: bool = False, ) -> dict[str, torch.Tensor]: """ Run the conservative SeZM lower interface on explicit edge vectors. @@ -1111,6 +1116,8 @@ def forward_common_lower( should_compile = ( self.should_use_compile() if use_compile is None else use_compile ) + if atomic_output_only: + should_compile = False if comm_dict is not None: if extended_atype is None: raise ValueError( @@ -1248,6 +1255,7 @@ def forward_common_lower( extended_coord_corr=extended_coord_corr, spin=spin, embedding_only=embedding_only, + atomic_output_only=atomic_output_only, ) return self._output_type_cast(model_predict, input_prec) @@ -1393,6 +1401,7 @@ def core_compute( spin: torch.Tensor | None = None, embedding_only: bool = False, conservative: bool = True, + atomic_output_only: bool = False, ) -> dict[str, torch.Tensor]: """ Compute SeZM lower outputs from the unified edge-vector schema. @@ -1446,6 +1455,10 @@ def core_compute( fitting keeps this enabled. Non-conservative property fitting disables it, so fitting outputs are reduced by their output definition without constructing edge-force gradients. + atomic_output_only + Whether computation stops after assembling complete atomic outputs, + before reductions and coordinate derivatives. Used by + output-statistics prediction. Returns ------- @@ -1470,15 +1483,29 @@ def core_compute( # scatter indices below. The embedding path produces no force, so it # keeps ``edge_vec`` detached and never allocates an autograd leaf. The # same forward-only treatment is used by non-conservative property heads. - if conservative and not embedding_only: + if conservative and not embedding_only and not atomic_output_only: edge_vec = edge_vec.detach().requires_grad_(True) # Native spin: the per-atom spin is a second autograd leaf, so the # magnetic force -dE/dspin is produced by the same backward that # scatters the edge gradient into force/virial. - if spin is not None and conservative and not embedding_only: + if ( + spin is not None + and conservative + and not embedding_only + and not atomic_output_only + ): spin = spin.detach().requires_grad_(True) + descriptor_atype = extended_atype if comm_dict is not None else atype + if descriptor_atype is None: + raise ValueError("`extended_atype` is required with `comm_dict`.") + inter_potential_edge_mask = self._make_inter_potential_edge_mask( + descriptor_atype, + edge_index, + edge_mask, + ) + # === Step 2. Descriptor forward === # ``extended_atype`` spans the extended region on the parallel path and # reduces to ``atype`` (owned atoms) on the single-domain path; the @@ -1489,7 +1516,7 @@ def core_compute( with nvtx_range("SeZM/descriptor"): descriptor, _ = descriptor_model.forward_with_edges( extended_coord=coord, - extended_atype=extended_atype if comm_dict is not None else atype, + extended_atype=descriptor_atype, edge_index=edge_index, edge_vec=edge_vec, edge_mask=edge_mask, @@ -1538,8 +1565,21 @@ def core_compute( with nvtx_range("SeZM/apply_out_stat"): fit_ret = self.atomic_model.apply_out_stat(fit_ret, atype) - # === Step 4. Apply atom mask === - for key in fit_ret.keys(): + # === Step 4. Inject analytical pair potential (edge form) === + # ZBL is evaluated from ``edge_vec`` (the autograd leaf) so its force + # and virial flow through the same edge backward as the learned energy. + if self.inter_potential is not None and "energy" in fit_ret: + fit_ret["energy"] = fit_ret["energy"] + self.inter_potential( + edge_vec=edge_vec, + edge_index=edge_index, + atype_flat=descriptor_atype.reshape(-1), + edge_mask=inter_potential_edge_mask, + n_node=nf * nloc, + real_type_count=self._get_inter_potential_real_type_count(), + ).view(nf, nloc, 1) + + # === Step 5. Apply atom mask to the complete physical output === + for key in fit_ret: out_shape = fit_ret[key].shape flat_dim = 1 for axis_size in out_shape[2:]: @@ -1550,6 +1590,9 @@ def core_compute( ).view(out_shape) fit_ret["mask"] = atom_mask + if atomic_output_only: + return fit_ret + if not conservative: return fit_output_to_model_output( fit_ret, @@ -1560,19 +1603,6 @@ def core_compute( extended_coord_corr=extended_coord_corr, ) - # === Step 5. Inject analytical pair potential (edge form) === - # ZBL is evaluated from ``edge_vec`` (the autograd leaf) so its force - # and virial flow through the same edge backward as the learned energy. - if self.inter_potential is not None and "energy" in fit_ret: - fit_ret["energy"] = fit_ret["energy"] + self.inter_potential( - edge_vec=edge_vec, - edge_index=edge_index, - atype_flat=atype.reshape(-1), - edge_mask=edge_mask, - n_node=nf * nloc, - real_type_count=self._get_inter_potential_real_type_count(), - ).view(nf, nloc, 1) - # === Step 6. Force / virial via edge-force scatter === # A single ``autograd.grad(energy, edge_vec)`` inside # ``edge_energy_deriv`` produces force, global virial and per-atom @@ -2961,6 +2991,63 @@ def canonicalize_dens_inputs( return force_input, noise_mask + # ========================================================================= + # Output Statistics + # ========================================================================= + + def predict_atomic_outputs_for_stat( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + charge_spin: torch.Tensor | None = None, + spin: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """ + Predict complete atomic outputs for residual output statistics. + + Parameters + ---------- + coord + Local coordinates with shape (nf, nloc, 3) in Å. + atype + Local atom types with shape (nf, nloc). + box + Simulation cells with shape (nf, 9), or ``None``. + fparam + Optional frame parameters. + aparam + Optional atomic parameters. + charge_spin + Optional frame-level charge and spin conditions. + spin + Optional native per-atom spin vectors. + + Returns + ------- + dict[str, torch.Tensor] + Complete atomic outputs evaluated without compilation or + coordinate derivatives. + """ + with ( + self.preserve_training_state(), + torch.no_grad(), + self.tf32_precision_ctx(), + ): + outputs = self.forward_common( + coord, + atype, + box=box, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + spin=spin, + atomic_output_only=True, + ) + return {key: value.detach() for key, value in outputs.items()} + # ========================================================================= # Output Post-Processing # ========================================================================= @@ -3117,6 +3204,54 @@ def reset_head_for_mode(self, mode: str) -> None: # Bridging Helpers # ========================================================================= + def _make_inter_potential_edge_mask( + self, + atype: torch.Tensor, + edge_index: torch.Tensor, + edge_mask: torch.Tensor, + ) -> torch.Tensor: + """Build the complete validity mask for analytical pair energies. + + Parameters + ---------- + atype + Atom types with shape (nf, nall). + edge_index + Edge source and destination indices with shape (2, E). + edge_mask + Base edge-validity mask with shape (E,). + + Returns + ------- + torch.Tensor + Boolean analytical-potential mask with shape (E,). + """ + if self.inter_potential is None: + return edge_mask + + src = edge_index[0] + dst = edge_index[1] + atype_flat = atype.reshape(-1) + keep = edge_mask + + descriptor = self.atomic_model.descriptor + if descriptor.exclude_types: + keep = keep & descriptor._edge_type_keep_mask(atype_flat, src, dst) + + atom_excl = self.atomic_model.atom_excl + if atom_excl is not None: + atom_is_present = self.atomic_model.make_atom_mask(atype) + safe_atype = torch.where(atom_is_present, atype, 0) + atom_is_included = atom_is_present & atom_excl(safe_atype).to(torch.bool) + atom_is_included = atom_is_included.reshape(-1) + keep = ( + keep + & atom_is_included.index_select(0, src) + & atom_is_included.index_select(0, dst) + ) + + return keep + def _get_inter_potential_real_type_count(self) -> int: """Return the real-type count used to mask analytical pair potentials.""" return len(self.get_type_map()) @@ -3230,6 +3365,30 @@ def deserialize(cls, data: dict[str, Any]) -> SeZMModel: # Context Managers # ========================================================================= + @contextmanager + def preserve_training_state(self) -> Generator[None, None, None]: + """ + Evaluate temporarily and restore every module training flag. + + Yields + ------ + None + Control while the complete module tree is in evaluation mode. + """ + training_states = [ + (submodule, submodule.training) for submodule in self.modules() + ] + root_training = self.training + self.eval() + try: + yield + finally: + # ``train`` invokes module-specific cache invalidation hooks before + # the exact per-module flags are restored below. + self.train(root_training) + for submodule, training in training_states: + submodule.training = training + @contextmanager def tf32_precision_ctx(self) -> Generator[None, None, None]: """Context manager to temporarily set TF32 matmul precision. @@ -3255,7 +3414,7 @@ def tf32_precision_ctx(self) -> Generator[None, None, None]: # ============================================================================= -# InterPotential: analytical pair potentials for bridging +# InnerPotential: analytical pair potentials for bridging # ============================================================================= # fmt: off @@ -3290,7 +3449,7 @@ def tf32_precision_ctx(self) -> Generator[None, None, None]: _A_BOHR = 0.5291772109 # Bohr radius in Å -class InterPotential(torch.nn.Module): +class InnerPotential(torch.nn.Module): """ Analytical pair potential module for Zone bridging. @@ -3320,7 +3479,7 @@ def __init__(self, type_map: list[str], mode: str = "zbl") -> None: super().__init__() mode = mode.upper() if mode != "ZBL": - raise ValueError(f"Unknown InterPotential mode: {mode}") + raise ValueError(f"Unknown InnerPotential mode: {mode}") self.mode = mode self.ntypes_real = len(type_map) diff --git a/deepmd/pt/model/model/sezm_property_model.py b/deepmd/pt/model/model/sezm_property_model.py index 43890d4549..e41cfddc78 100644 --- a/deepmd/pt/model/model/sezm_property_model.py +++ b/deepmd/pt/model/model/sezm_property_model.py @@ -134,6 +134,7 @@ def core_compute( extended_coord_corr: torch.Tensor | None = None, spin: torch.Tensor | None = None, embedding_only: bool = False, + atomic_output_only: bool = False, ) -> dict[str, torch.Tensor]: """Compute property outputs through the SeZM forward-only graph.""" return super().core_compute( @@ -152,6 +153,7 @@ def core_compute( spin=spin, embedding_only=embedding_only, conservative=False, + atomic_output_only=atomic_output_only, ) def _inductor_compile_options(self, *, inference: bool = False) -> dict[str, Any]: diff --git a/deepmd/pt/model/model/sezm_spin_model.py b/deepmd/pt/model/model/sezm_spin_model.py index 0be5b39c10..ace082b02b 100644 --- a/deepmd/pt/model/model/sezm_spin_model.py +++ b/deepmd/pt/model/model/sezm_spin_model.py @@ -27,7 +27,7 @@ BaseModel, ) from deepmd.pt.model.model.sezm_model import ( - InterPotential, + InnerPotential, SeZMModel, ) from deepmd.pt.model.model.spin_model import ( @@ -76,7 +76,7 @@ def __init__( real_sel: list[int], **kwargs: Any, ) -> None: - # Delay InterPotential construction until ntypes_real is available. + # Delay InnerPotential construction until ntypes_real is available. bridging_method = str(kwargs.pop("bridging_method", "none")).upper() kwargs["bridging_method"] = "none" @@ -97,7 +97,7 @@ def __init__( self.bridging_method = bridging_method self.inter_potential = ( - InterPotential(type_map=self.get_type_map(), mode=self.bridging_method) + InnerPotential(type_map=self.get_type_map(), mode=self.bridging_method) if self.bridging_method != "NONE" else None ) @@ -152,6 +152,7 @@ def forward_common( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, + atomic_output_only: bool = False, ) -> dict[str, torch.Tensor]: """Return spin-aware SeZM predictions with internal output keys.""" with nvtx_range("SeZMSpin/forward_common"): @@ -219,7 +220,10 @@ def forward_common( extended_coord_corr=extended_coord_corr[:, : nloc * 2, :].contiguous(), charge_spin=charge_spin, input_prec=input_prec, + atomic_output_only=atomic_output_only, ) + if atomic_output_only: + return model_ret return self._split_spin_common_output(model_ret, atype, nloc) def forward_lower( @@ -608,7 +612,16 @@ def _get_spin_sampled_func( @functools.lru_cache def spin_sampled_func() -> list[dict[str, Any]]: - return [_pack_spin_stat_sample(self, sys) for sys in sampled_func()] + packed_samples = [] + for sample in sampled_func(): + packed = _pack_spin_stat_sample(self, sample) + packed["model_coord"] = sample["coord"] + packed["model_atype"] = sample["atype"] + packed["model_spin"] = sample["spin"] + if "aparam" in sample: + packed["model_aparam"] = sample["aparam"] + packed_samples.append(packed) + return packed_samples return self.atomic_model._make_wrapped_sampler(spin_sampled_func) diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index ec92301b9d..334269667d 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -116,6 +116,8 @@ ) from deepmd.pt.utils.stat import ( make_stat_input, + min_pair_dist_frame_mask, + select_batch_frames, ) from deepmd.pt.utils.utils import ( to_numpy_array, @@ -425,6 +427,7 @@ def single_model_stat( _data_stat_nbatch: int, _training_data: DpLoaderSet, _stat_file_spec: StatFileSpec, + _min_pair_dist: float = 0.0, finetune_has_new_type: bool = False, preset_observed_type: list[str] | None = None, ) -> Callable[[], Any]: @@ -434,6 +437,7 @@ def get_sample() -> Any: _training_data.systems, _training_data.dataloaders, _data_stat_nbatch, + min_pair_dist=_min_pair_dist, ) return sampled @@ -541,6 +545,7 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: model_params.get("data_stat_nbatch", 10), training_data, self.stat_file_specs["Default"], + _min_pair_dist=min_pair_dist, finetune_has_new_type=self.finetune_links["Default"].get_has_new_type() if self.finetune_links is not None else False, @@ -588,7 +593,9 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: self.model[model_key] ) min_pair_dist = float( - training_params.get("training_data", {}).get("min_pair_dist", 0.0) + training_params["data_dict"][model_key] + .get("training_data", {}) + .get("min_pair_dist", 0.0) ) if min_pair_dist > 0.0: data_requirement.append( @@ -621,6 +628,7 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: model_params["model_dict"][model_key].get("data_stat_nbatch", 10), training_data[model_key], self.stat_file_specs[model_key], + _min_pair_dist=min_pair_dist, finetune_has_new_type=self.finetune_links[ model_key ].get_has_new_type() @@ -780,10 +788,41 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: self.nonfinite_grad_guard = NonFiniteGradGuard() self.lr_schedule = get_lr(config["learning_rate"]) - # Minimum pairwise distance for filtering unphysical frames during training - self.min_pair_dist = training_params.get("training_data", {}).get( - "min_pair_dist", 0.0 + # Minimum pairwise distance for filtering unphysical frames during training. + if self.multi_task: + self.min_pair_dist: float | dict[str, float] = { + model_key: float( + training_params["data_dict"][model_key] + .get("training_data", {}) + .get("min_pair_dist", 0.0) + ) + for model_key in self.model_keys + } + else: + self.min_pair_dist = float( + training_params.get("training_data", {}).get("min_pair_dist", 0.0) + ) + self.has_min_pair_filter = ( + any(value > 0.0 for value in self.min_pair_dist.values()) + if isinstance(self.min_pair_dist, dict) + else self.min_pair_dist > 0.0 ) + if self.has_min_pair_filter: + if self.multi_task: + local_training_batch_attempts = max( + max(1, len(self.training_dataloader[model_key])) + for model_key in self.model_keys + ) + else: + local_training_batch_attempts = max(1, len(self.training_dataloader)) + # Rank-specific task sampling may select loaders of different + # lengths, but validity collectives must use one shared retry count. + self._training_batch_attempts = int( + self._broadcast_value_from_rank0(local_training_batch_attempts) + ) + else: + self._training_batch_attempts = 1 + self._discarded_training_batches = 0 # JIT if JIT: @@ -1510,13 +1549,7 @@ def step(_step_id: int, task_key: str = "Default") -> None: cur_lr = self.lr_schedule.value(_step_id) pref_lr = cur_lr self.optimizer.zero_grad(set_to_none=True) - input_dict, label_dict, log_dict = self.get_data( - is_train=True, task_key=task_key - ) - # All frames filtered by min_pair_dist (single-GPU only; - # DDP path in get_data() always keeps at least one frame) - if not input_dict: - return + input_dict, label_dict, log_dict = self._next_training_batch(task_key) if SAMPLER_RECORD: print_str = f"Step {_step_id}: sample system{log_dict['sid']} frame{log_dict['fid']}\n" fout1.write(print_str) @@ -1778,16 +1811,24 @@ def log_loss_valid(_task_key: str = "Default") -> dict: input_dict, label_dict, _ = self.get_data( is_train=True, task_key=_key ) - _, loss, more_loss = self.wrapper( - **input_dict, - cur_lr=pref_lr, - label=label_dict, - task_key=_key, - ) - train_results[_key] = log_loss_train( - loss, more_loss, _task_key=_key - ) + if input_dict and not ( + self.is_distributed and self.zero_stage >= 2 + ): + _, loss, more_loss = self._get_inner_module()( + **input_dict, + cur_lr=pref_lr, + label=label_dict, + task_key=_key, + ) + train_results[_key] = log_loss_train( + loss, more_loss, _task_key=_key + ) valid_results[_key] = log_loss_valid(_task_key=_key) + if not train_results[_key] and valid_results[_key]: + train_results[_key] = dict.fromkeys( + valid_results[_key], + float("nan"), + ) if self.rank == 0: log.info( format_training_message_per_task( @@ -1964,6 +2005,7 @@ def log_loss_valid(_task_key: str = "Default") -> dict: self.t0 = time.time() self.total_train_time = 0.0 self.timed_steps = 0 + self._discarded_training_batches = 0 if self.disp_avg: # Initialize loss accumulators @@ -1980,6 +2022,13 @@ def log_loss_valid(_task_key: str = "Default") -> dict: if JIT: break + if self.rank == 0 and self._discarded_training_batches: + log.info( + "Discarded %d globally invalid batches while collecting " + "synchronized training inputs.", + self._discarded_training_batches, + ) + if ( self.change_bias_after_training and self.num_steps > self.start_step @@ -2328,6 +2377,56 @@ def save_ema_model_merged( use_ema_weights=True, ) + def _get_min_pair_dist(self, task_key: str) -> float: + """Return the minimum pair distance configured for one task.""" + if isinstance(self.min_pair_dist, dict): + return self.min_pair_dist[task_key] + return self.min_pair_dist + + def _next_training_batch( + self, + task_key: str, + ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + """Return the next batch that can produce a synchronized optimizer step. + + Parameters + ---------- + task_key + Selected training task. + + Returns + ------- + tuple[dict[str, Any], dict[str, Any], dict[str, Any]] + Model inputs, labels, and sampler metadata for a globally valid + batch. + + Raises + ------ + RuntimeError + If no globally valid batch is found within the synchronized retry + budget. + """ + max_attempts = self._training_batch_attempts + distributed_filter = ( + self.has_min_pair_filter and dist.is_available() and dist.is_initialized() + ) + + for _ in range(max_attempts): + batch = self.get_data(is_train=True, task_key=task_key) + input_dict = batch[0] + globally_valid = bool(input_dict) + if distributed_filter: + globally_valid = all_ranks_have_valid_frames(globally_valid) + if globally_valid: + return batch + self._discarded_training_batches += 1 + + raise RuntimeError( + "Unable to collect a globally valid training batch for task " + f"{task_key!r} after {max_attempts} attempts with " + f"min_pair_dist={self._get_min_pair_dist(task_key)}." + ) + def get_data( self, is_train: bool = True, task_key: str = "Default" ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: @@ -2341,26 +2440,18 @@ def get_data( return {}, {}, {} batch_data = next(iterator) # === Filter frames with atoms too close (training only) === - if is_train and self.min_pair_dist > 0.0 and "min_pair_dist" in batch_data: - min_dists = batch_data["min_pair_dist"] - if isinstance(min_dists, torch.Tensor): - valid_mask = min_dists.squeeze(-1) >= self.min_pair_dist - n_total = valid_mask.shape[0] - n_valid = int(valid_mask.sum().item()) - if n_valid == 0: - # Under distributed training (DDP/FSDP), every rank must - # participate in backward() to avoid collective communication - # deadlock. Keep one frame as a fallback instead of - # skipping the entire batch. - if dist.is_available() and dist.is_initialized(): - valid_mask[0] = True - n_valid = 1 - else: - return {}, {}, {} - if n_valid < n_total: - for key, val in batch_data.items(): - if isinstance(val, torch.Tensor) and val.shape[0] == n_total: - batch_data[key] = val[valid_mask] + min_pair_dist = self._get_min_pair_dist(task_key) + valid_mask = ( + min_pair_dist_frame_mask(batch_data, min_pair_dist) if is_train else None + ) + local_has_valid = valid_mask is None or bool(torch.any(valid_mask)) + if not local_has_valid: + return {}, {}, {} + + if valid_mask is not None: + n_valid = int(valid_mask.sum().item()) + if n_valid < valid_mask.shape[0]: + batch_data = select_batch_frames(batch_data, valid_mask) for key in batch_data.keys(): if key == "sid" or key == "fid" or key == "box" or "find_" in key: continue @@ -2474,6 +2565,29 @@ def print_on_training( fout.flush() +def all_ranks_have_valid_frames(local_has_valid: bool) -> bool: + """ + Return whether every distributed rank has a valid training frame. + + Parameters + ---------- + local_has_valid + Whether the current rank has at least one valid frame. + + Returns + ------- + bool + ``True`` only when every rank reports a valid frame. + """ + all_ranks_have_valid = torch.tensor( + int(local_has_valid), + dtype=torch.int32, + device=DEVICE, + ) + dist.all_reduce(all_ranks_have_valid, op=dist.ReduceOp.MIN) + return bool(all_ranks_have_valid.item()) + + def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: additional_data_requirement = [] if _model.get_dim_fparam() > 0: diff --git a/deepmd/pt/utils/compile_compat.py b/deepmd/pt/utils/compile_compat.py index a3cc33982f..2cf3a64f35 100644 --- a/deepmd/pt/utils/compile_compat.py +++ b/deepmd/pt/utils/compile_compat.py @@ -12,10 +12,11 @@ * helpers and workarounds common to the supported releases -- trace-shape and trace-input preparation, per-task buffer promotion, FX graph repair, the Inductor option lockdown, and the process-global configuration; and -* a workaround specific to PyTorch 2.12, which must not be applied on 2.11. +* a workaround for the Inductor symbolic-divisibility regression introduced in + PyTorch 2.12, which must not be applied on 2.11. -Only PyTorch 2.11.x and 2.12.x are permitted for compilation (see -:func:`check_compile_torch_version`). +Only the releases listed in :data:`SUPPORTED_COMPILE_TORCH` are permitted for +compilation (see :func:`check_compile_torch_version`). """ from __future__ import ( @@ -35,6 +36,7 @@ __all__ = [ "AM_PREFIX", "FIT_PREFIX", + "SUPPORTED_COMPILE_TORCH", "apply_global_compile_patches", "build_inductor_compile_options", "check_compile_torch_version", @@ -51,8 +53,33 @@ ] +#: ``(major, minor)`` releases explicitly enabled for SeZM compilation after +#: validation. This runtime allowlist is not a record of which releases happen +#: to be installed in CI. +SUPPORTED_COMPILE_TORCH = ((2, 11), (2, 12), (2, 13)) + +#: Releases carrying the Inductor symbolic-divisibility regression repaired by +#: :func:`patch_inductor_symbolic_divisibility`. PyTorch 2.11 evaluates the +#: predicate correctly and must be left alone. +_DIVISIBILITY_REGRESSION_TORCH = ((2, 12), (2, 13)) + + +def _torch_release() -> tuple[int, int]: + """Return the ``(major, minor)`` pair of the running PyTorch release. + + Returns + ------- + tuple[int, int] + The major and minor components, or ``(0, 0)`` when the version string + carries fewer than two components and therefore matches no supported + release. + """ + release = Version(torch.__version__).release + return (release[0], release[1]) if len(release) >= 2 else (0, 0) + + # ============================================================================= -# Common workarounds (PyTorch 2.11 and 2.12) +# Common workarounds (every supported release) # ============================================================================= def apply_global_compile_patches() -> None: """Apply every process-global PyTorch adjustment the compile path needs. @@ -60,8 +87,8 @@ def apply_global_compile_patches() -> None: The adjustments are mutually independent and individually idempotent. The function is intended to run exactly once, when the model module is imported, so that the global state is established before the first - compilation. The symbolic-divisibility repair is applied only on PyTorch - 2.12, where the regression exists. + compilation. The symbolic-divisibility repair is applied only on the + releases where the regression exists. """ # Silence Inductor / Triton autotune console dumps. ``torch.compile`` # reads these environment variables once, when its backend is first @@ -89,10 +116,10 @@ def apply_global_compile_patches() -> None: # supported PyTorch versions and is independent of runtime shapes. patch_inductor_force_int64_indexing() - # The symbolic-divisibility regression exists only on PyTorch 2.12; the + # The symbolic-divisibility regression was introduced in PyTorch 2.12; the # 2.11 backend evaluates the same predicate correctly and must not be # patched. - if Version(torch.__version__).release[:2] == (2, 12): + if _torch_release() in _DIVISIBILITY_REGRESSION_TORCH: patch_inductor_symbolic_divisibility() @@ -127,11 +154,20 @@ def patch_inductor_force_int64_indexing() -> None: def check_compile_torch_version() -> None: - """Fail fast when ``torch.compile`` is requested on an unsupported PyTorch.""" - version = Version(torch.__version__).release - if len(version) < 2 or (version[:2] != (2, 11) and version[:2] != (2, 12)): + """Fail fast when ``torch.compile`` is requested on an unsupported PyTorch. + + Raises + ------ + RuntimeError + If the running PyTorch release is absent from + :data:`SUPPORTED_COMPILE_TORCH`. + """ + if _torch_release() not in SUPPORTED_COMPILE_TORCH: + supported = ", ".join( + f"{major}.{minor}.x" for major, minor in SUPPORTED_COMPILE_TORCH + ) raise RuntimeError( - "deepmd `torch.compile` support requires PyTorch 2.11.x or 2.12.x; " + f"deepmd `torch.compile` support requires PyTorch {supported}; " f"found torch {torch.__version__}." ) @@ -512,10 +548,10 @@ def get_task_buffer_values( # ============================================================================= -# PyTorch 2.12-specific workarounds +# Workarounds for PyTorch 2.12 and later # ============================================================================= def patch_inductor_symbolic_divisibility() -> None: - """Repair the PyTorch 2.12 Inductor symbolic-divisibility regression. + """Repair the Inductor symbolic-divisibility regression of PyTorch 2.12+. ``SizeVarAllocator.statically_known_multiple_of`` determines whether one symbolic size is an exact multiple of another. ``SIMDKernel`` consults it @@ -526,9 +562,9 @@ def patch_inductor_symbolic_divisibility() -> None: factors polynomials, so an expression such as ``(32*s + 64) % (s + 2)`` reduces to ``0`` and the split proceeds. PyTorch 2.12 rewrote the helper and, for symbolic denominators, routes the test through Inductor's own - ``Mod`` implementation, which does not factor. ``Mod(32*s + 64, s + 2)`` - therefore stays unevaluated, the test returns ``False``, and lowering - aborts with:: + ``Mod`` implementation, which does not factor; 2.13 retains that behaviour. + ``Mod(32*s + 64, s + 2)`` therefore stays unevaluated, the test returns + ``False``, and lowering aborts with:: CantSplit: 32*s38 + 64 not divisible by s38 + 2 @@ -554,7 +590,13 @@ def patch_inductor_symbolic_divisibility() -> None: if getattr(SizeVarAllocator, "_dp_divisibility_patched", False): return - original_known_multiple_of = SizeVarAllocator.statically_known_multiple_of + original_known_multiple_of = getattr( + SizeVarAllocator, + "statically_known_multiple_of", + None, + ) + if not callable(original_known_multiple_of): + return def statically_known_multiple_of( self: Any, numerator: Any, denominator: Any diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index e0ef13de4c..5b245d187b 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -50,35 +50,131 @@ "_restore_observed_type_from_file", "_save_observed_type_to_file", "collect_observed_types", + "min_pair_dist_frame_mask", + "select_batch_frames", ] -def make_stat_input( - datasets: list[Any], dataloaders: list[Any], nbatches: int +def min_pair_dist_frame_mask( + batch: dict[str, Any], + min_pair_dist: float, +) -> torch.Tensor | None: + """ + Return the valid-frame mask for a minimum pair-distance threshold. + + Parameters + ---------- + batch + Data batch containing frame-aligned tensors. + min_pair_dist + Minimum allowed pair distance in Å. + + Returns + ------- + torch.Tensor or None + Boolean mask with shape (nframes), or ``None`` when filtering is + disabled or the distance field is unavailable. + """ + if min_pair_dist <= 0.0 or "min_pair_dist" not in batch: + return None + distances = batch["min_pair_dist"] + if not isinstance(distances, torch.Tensor): + return None + return distances.reshape(-1) >= float(min_pair_dist) + + +def select_batch_frames( + batch: dict[str, Any], + frame_mask: torch.Tensor, ) -> dict[str, Any]: + """ + Select frame-aligned tensors from one data batch. + + Parameters + ---------- + batch + Data batch containing tensors and scalar metadata. + frame_mask + Boolean selection mask with shape (nframes). + + Returns + ------- + dict[str, Any] + Batch with every frame-aligned tensor sliced by ``frame_mask``. + """ + nframes = frame_mask.shape[0] + selected: dict[str, Any] = {} + frame_keep = frame_mask.detach().cpu().tolist() + for key, value in batch.items(): + if ( + isinstance(value, torch.Tensor) + and value.ndim > 0 + and value.shape[0] == nframes + ): + selected[key] = value[frame_mask] + elif isinstance(value, list) and len(value) == nframes: + selected[key] = [ + item for item, keep in zip(value, frame_keep, strict=True) if keep + ] + else: + selected[key] = value + return selected + + +def make_stat_input( + datasets: list[Any], + dataloaders: list[Any], + nbatches: int, + min_pair_dist: float = 0.0, +) -> list[dict[str, Any]]: """Pack data for statistics. - Args: - - dataset: A list of dataset to analyze. - - nbatches: Batch count for collecting stats. + Parameters + ---------- + datasets + Data systems to analyze. + dataloaders + One data loader for each system. + nbatches + Maximum number of valid batches collected from each system. + min_pair_dist + Minimum allowed pair distance in Å. Frames below the threshold are + excluded before statistics are accumulated. Returns ------- - - a list of dicts, each of which contains data from a system + list[dict[str, Any]] + Packed statistics, one dictionary for each system that contributes at + least one valid frame. """ lst = [] - log.info(f"Packing data for statistics from {len(datasets)} systems") - for i in range(len(datasets)): + log.info("Packing data for statistics from %d systems", len(datasets)) + for system_index in range(len(datasets)): sys_stat = {} with torch.device("cpu"): - iterator = iter(dataloaders[i]) - numb_batches = min(nbatches, len(dataloaders[i])) - for _ in range(numb_batches): + dataloader = dataloaders[system_index] + dataloader_size = len(dataloader) + target_batches = min(nbatches, dataloader_size) + scan_limit = dataloader_size if min_pair_dist > 0.0 else target_batches + iterator = iter(dataloader) + accepted_batches = 0 + scanned_batches = 0 + while accepted_batches < target_batches and scanned_batches < scan_limit: try: stat_data = next(iterator) except StopIteration: - iterator = iter(dataloaders[i]) - stat_data = next(iterator) + iterator = iter(dataloader) + try: + stat_data = next(iterator) + except StopIteration: + break + scanned_batches += 1 + frame_mask = min_pair_dist_frame_mask(stat_data, min_pair_dist) + if frame_mask is not None and not torch.any(frame_mask): + continue + if frame_mask is not None: + stat_data = select_batch_frames(stat_data, frame_mask) + accepted_batches += 1 if ( "find_fparam" in stat_data and "fparam" in stat_data @@ -99,12 +195,27 @@ def make_stat_input( else: pass + if not sys_stat: + if min_pair_dist > 0.0: + log.info( + "Skipping data system %d in statistics because no frame " + "satisfies min_pair_dist=%s.", + system_index, + min_pair_dist, + ) + else: + log.info( + "Skipping data system %d in statistics because its data " + "loader produced no batch.", + system_index, + ) + continue for key in sys_stat: if isinstance(sys_stat[key], np.float32): pass elif sys_stat[key] is None or sys_stat[key][0] is None: sys_stat[key] = None - elif isinstance(stat_data[dd], torch.Tensor): + elif isinstance(sys_stat[key], list): sys_stat[key] = torch.cat(sys_stat[key], dim=0) dict_to_device(sys_stat) lst.append(sys_stat) @@ -166,20 +277,24 @@ def _post_process_stat( def _compute_model_predict( sampled: Callable[[], list[dict]] | list[dict], keys: list[str], - model_forward: Callable[..., torch.Tensor], -) -> dict[str, list[torch.Tensor]]: + model_forward: Callable[..., dict[str, torch.Tensor]], +) -> tuple[dict[str, list[np.ndarray]], list[np.ndarray]]: auto_batch_size = AutoBatchSize() model_predict = {kk: [] for kk in keys} + model_mask = [] for system in sampled: - nframes = system["coord"].shape[0] + model_coord = system.get("model_coord", system["coord"]) + model_atype = system.get("model_atype", system["atype"]) + nframes = model_coord.shape[0] coord, atype, box = ( - system["coord"], - system["atype"], - system["box"], + model_coord, + model_atype, + system.get("box"), ) fparam = system.get("fparam", None) - aparam = system.get("aparam", None) + aparam = system.get("model_aparam", system.get("aparam", None)) charge_spin = system.get("charge_spin", None) + spin = system.get("model_spin", system.get("spin", None)) def model_forward_auto_batch_size(*args: Any, **kwargs: Any) -> Any: return auto_batch_size.execute_all( @@ -190,16 +305,50 @@ def model_forward_auto_batch_size(*args: Any, **kwargs: Any) -> Any: **kwargs, ) + model_kwargs = { + "fparam": fparam, + "aparam": aparam, + "charge_spin": charge_spin, + } + if spin is not None: + model_kwargs["spin"] = spin sample_predict = model_forward_auto_batch_size( - coord, atype, box, fparam=fparam, aparam=aparam, charge_spin=charge_spin + coord, + atype, + box, + **model_kwargs, ) + sample_mask = sample_predict.get("mask", atype >= 0) + model_mask.append(to_numpy_array(sample_mask)) for kk in keys: model_predict[kk].append( to_numpy_array( sample_predict[kk] # nf x nloc x odims ) ) - return model_predict + return model_predict, model_mask + + +def _reduce_model_prediction( + prediction: np.ndarray, + mask: np.ndarray, + intensive: bool, +) -> np.ndarray: + """Reduce atomic predictions, rejecting undefined intensive means.""" + reduced = np.sum(prediction, axis=1) + if intensive: + atom_count = np.sum(mask, axis=1) + empty_frames = np.flatnonzero(atom_count == 0) + if empty_frames.size: + raise ValueError( + "Cannot reduce intensive model predictions for frames with no " + f"unmasked atoms: {empty_frames.tolist()}." + ) + atom_count = atom_count.reshape( + (atom_count.shape[0],) + (1,) * (reduced.ndim - 1) + ) + reduced = reduced / atom_count + return reduced def _make_preset_out_bias( @@ -260,10 +409,10 @@ def compute_output_stats( stat_file_path: DPPath | None = None, rcond: float | None = None, preset_bias: dict[str, list[np.ndarray | None]] | None = None, - model_forward: Callable[..., torch.Tensor] | None = None, + model_forward: Callable[..., dict[str, torch.Tensor]] | None = None, stats_distinguish_types: bool = True, intensive: bool = False, -) -> dict[str, Any]: +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. @@ -278,6 +427,8 @@ def compute_output_stats( the lazy function helps by only sampling once. ntypes : int The number of atom types. + keys : str or list[str], optional + Output labels whose per-type bias and standard deviation are computed. stat_file_path : DPPath, optional The path to the stat file. rcond : float, optional @@ -287,7 +438,7 @@ def compute_output_stats( The value is a list specifying the bias. the elements can be None or np.ndarray of output shape. For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] The `set_davg_zero` key in the descriptor should be set. - model_forward : Callable[..., torch.Tensor], optional + model_forward : Callable[..., dict[str, torch.Tensor]], optional The wrapped forward function of atomic model. If not None, the model will be utilized to generate the original energy prediction, which will be subtracted from the energy label of the data. @@ -296,6 +447,20 @@ def compute_output_stats( Whether to distinguish different element types in the statistics. intensive : bool, optional Whether the fitting target is intensive. + + Returns + ------- + tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]] + Per-output bias and standard-deviation tensors. + + Raises + ------ + ValueError + If output statistics must be computed but no sampled system contains a + valid frame. + RuntimeError + If the requested statistics cannot be computed from the available + labels. """ keys = [keys] if isinstance(keys, str) else keys assert isinstance(keys, list) @@ -308,10 +473,20 @@ def compute_output_stats( if bias_atom_e is None: # only get data once, sampled is a list of dict[str, torch.Tensor] sampled = merged() if callable(merged) else merged + if not sampled: + raise ValueError( + "Output statistics require at least one sampled system with a " + "valid frame." + ) if model_forward is not None: - model_pred = _compute_model_predict(sampled, keys, model_forward) + model_pred, model_mask = _compute_model_predict( + sampled, + keys, + model_forward, + ) else: model_pred = None + model_mask = [] # remove the keys that are not in the sample new_keys = [ @@ -338,8 +513,13 @@ def compute_output_stats( model_pred_g = ( { kk: [ - np.sum(vv[idx], axis=1) for idx in global_sampled_idx[kk] - ] # sum atomic dim + _reduce_model_prediction( + vv[idx], + model_mask[idx], + intensive, + ) + for idx in global_sampled_idx[kk] + ] for kk, vv in model_pred.items() } if model_pred diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 6e779dd23e..7bdb07de57 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -103,7 +103,7 @@ def get_sezm_model(data: dict) -> EnergyModel: return get_native_spin_model(data) # Analytical bridging (e.g. ZBL): the radii feed the DESCRIPTOR's # InnerClamp/BridgingSwitch (mirrors pt's builder); the method builds the - # atomic model's InterPotential at construction below. + # atomic model's InnerPotential at construction below. bridging_method = str(data.get("bridging_method", "none")) bridging_enabled = bridging_method.lower() not in ("none", "") if data.get("lora") is not None: @@ -167,8 +167,8 @@ def get_sezm_model(data: dict) -> EnergyModel: # Composition, not a flag (first-principles design): the analytical # bridging term is its own atomic model, summed with the learned one # by the existing linear composition machinery. - from deepmd.dpmodel.atomic_model.inter_potential import ( - InterPotentialAtomicModel, + from deepmd.dpmodel.atomic_model.inner_potential import ( + InnerPotentialAtomicModel, ) from deepmd.dpmodel.atomic_model.linear_atomic_model import ( LinearEnergyAtomicModel, @@ -177,7 +177,7 @@ def get_sezm_model(data: dict) -> EnergyModel: LinearEnergyModel, ) - zbl_atomic = InterPotentialAtomicModel( + zbl_atomic = InnerPotentialAtomicModel( type_map=data["type_map"], mode=bridging_method, rcut=descriptor.get_rcut(), diff --git a/deepmd/pt_expt/utils/__init__.py b/deepmd/pt_expt/utils/__init__.py index d30f288909..49587659cf 100644 --- a/deepmd/pt_expt/utils/__init__.py +++ b/deepmd/pt_expt/utils/__init__.py @@ -11,8 +11,8 @@ AtomExcludeMask, PairExcludeMask, ) -from .inter_potential import ( - InterPotential, +from .inner_potential import ( + InnerPotential, ) from .network import ( NetworkCollection, @@ -38,7 +38,7 @@ __all__ = [ "AtomExcludeMask", - "InterPotential", + "InnerPotential", "NetworkCollection", "PairExcludeMask", "TypeEmbedNet", diff --git a/deepmd/pt_expt/utils/inter_potential.py b/deepmd/pt_expt/utils/inner_potential.py similarity index 65% rename from deepmd/pt_expt/utils/inter_potential.py rename to deepmd/pt_expt/utils/inner_potential.py index 0647e9a73f..7e7eb9284a 100644 --- a/deepmd/pt_expt/utils/inter_potential.py +++ b/deepmd/pt_expt/utils/inner_potential.py @@ -5,8 +5,8 @@ Any, ) -from deepmd.dpmodel.atomic_model.inter_potential import ( - InterPotential as InterPotentialDP, +from deepmd.dpmodel.atomic_model.inner_potential import ( + InnerPotential as InnerPotentialDP, ) from deepmd.pt_expt.common import ( register_dpmodel_mapping, @@ -15,16 +15,16 @@ @torch_module -class InterPotential(InterPotentialDP): +class InnerPotential(InnerPotentialDP): def forward(self, *args: Any, **kwargs: Any) -> Any: return self.call(*args, **kwargs) -# InterPotential carries no trainable state (only the constant per-type +# InnerPotential carries no trainable state (only the constant per-type # atomic-number table, derived from the constructor arguments), so it # implements no serialize()/deserialize(); rebuild it fresh from # (type_map, mode). register_dpmodel_mapping( - InterPotentialDP, - lambda v: InterPotential(type_map=v.type_map, mode=v.mode), + InnerPotentialDP, + lambda v: InnerPotential(type_map=v.type_map, mode=v.mode), ) diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 897d91cec3..267add7223 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -191,7 +191,7 @@ def _needs_with_comm_artifact( atomic_model = getattr(model, "atomic_model", None) if isinstance(atomic_model, LinearEnergyAtomicModel): - # Compositions (e.g. analytical bridging: learned + InterPotential) + # Compositions (e.g. analytical bridging: learned + InnerPotential) # are single-rank on the graph route: per-edge analytical terms fold # each node's full edge set, which a single rank cannot observe for # ghost owners (pt's supports_edge_parallel()==False rationale). diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index fc8c11b635..b1d86d2ffd 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5023,9 +5023,8 @@ def training_data_args() -> list[ "Frames containing any atom pair closer than this distance are excluded " "from loss computation, as DFT labels for near-collision configurations " "are often unreliable. Set to 0 to disable (default). " - "Under distributed training (DDP/FSDP), if ALL frames in a batch are " - "filtered out on a given rank, one frame is retained to ensure every " - "rank participates in collective communication (backward all-reduce). " + "Under distributed training (DDP/FSDP), if any rank has no valid frame " + "in its current batch, every rank collectively skips that training step. " "Note: enabling this adds an O(N²) distance check per frame in the " "DataLoader workers (CPU-side), which may slow down training for large " "systems. To avoid the overhead, consider pre-cleaning the dataset instead." diff --git a/source/api_cc/tests/test_deeppot_dpa4_zbl_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa4_zbl_ptexpt.cc index 8dfe1efaaf..98949911e9 100644 --- a/source/api_cc/tests/test_deeppot_dpa4_zbl_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_dpa4_zbl_ptexpt.cc @@ -3,7 +3,7 @@ // .pt2, graph lower). // // ``bridging_method: ZBL`` builds a COMPOSITION -- LinearEnergyModel over -// [learned DPA4, InterPotentialAtomicModel] -- so this exercises the graph +// [learned DPA4, InnerPotentialAtomicModel] -- so this exercises the graph // lower of a linear composition, which no other C++ fixture covers. Before // this test ZBL bridging had NO C++ or LAMMPS coverage at all: its only // end-to-end check drove the archive through the PYTHON DeepPot, which diff --git a/source/api_cc/tests/test_deepspin_dpa4_zbl_ptexpt.cc b/source/api_cc/tests/test_deepspin_dpa4_zbl_ptexpt.cc index a60dc8863e..fa717a47ea 100644 --- a/source/api_cc/tests/test_deepspin_dpa4_zbl_ptexpt.cc +++ b/source/api_cc/tests/test_deepspin_dpa4_zbl_ptexpt.cc @@ -6,7 +6,7 @@ // test_deepspin_dpa4_graph_ptexpt.cc (native spin) and // test_deeppot_dpa4_zbl_ptexpt.cc (bridging): here ``spin`` must reach a // COMPOSITION -- LinearEnergyModel over [learned DPA4, -// InterPotentialAtomicModel] -- whose learned child consumes it and whose +// InnerPotentialAtomicModel] -- whose learned child consumes it and whose // analytical child accepts and ignores it, and the archive must still declare // ``is_spin`` so this DeepSpin path (not DeepPot) is the one that runs. // Before this test the combination had NO coverage below the pt_expt Python diff --git a/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py b/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py index 135bf04426..30de19cf8d 100644 --- a/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py +++ b/source/lmp/tests/test_lammps_dpa4_zbl_pt2.py @@ -4,7 +4,7 @@ ``source/tests/infer/gen_dpa4_zbl.py``). ``bridging_method: ZBL`` builds a COMPOSITION -- ``LinearEnergyModel`` over -``[learned DPA4, InterPotentialAtomicModel]`` with ``weights="sum"`` -- so +``[learned DPA4, InnerPotentialAtomicModel]`` with ``weights="sum"`` -- so this drives the graph lower of a linear composition through the LAMMPS pair style. The archive already had a C++ gtest (``source/api_cc/tests/test_deeppot_dpa4_zbl_ptexpt.cc``) but NO LAMMPS diff --git a/source/tests/common/dpmodel/test_inter_potential.py b/source/tests/common/dpmodel/test_inner_potential.py similarity index 86% rename from source/tests/common/dpmodel/test_inter_potential.py rename to source/tests/common/dpmodel/test_inner_potential.py index 20075900d6..87e0f6ab52 100644 --- a/source/tests/common/dpmodel/test_inter_potential.py +++ b/source/tests/common/dpmodel/test_inner_potential.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""dpmodel ``InterPotential`` (analytical ZBL bridging term) unit tests. +"""dpmodel ``InnerPotential`` (analytical ZBL bridging term) unit tests. Ports the pt reference values from -``source/tests/pt/model/test_sezm_model.py::TestInterPotential``: the exact +``source/tests/pt/model/test_sezm_model.py::TestInnerPotential``: the exact universal-ZBL formula is reproduced in-test and the half-split per-edge scatter must sum back to the full analytic pair energy. """ @@ -12,8 +12,8 @@ import numpy as np import pytest -from deepmd.dpmodel.atomic_model.inter_potential import ( - InterPotential, +from deepmd.dpmodel.atomic_model.inner_potential import ( + InnerPotential, ) _A_BOHR = 0.5291772109 @@ -48,7 +48,7 @@ def _two_atom_inputs(r: float): ) def test_zbl_known_value(type_map, atypes, zi, zj): r = 0.8 - pot = InterPotential(type_map=type_map) + pot = InnerPotential(type_map=type_map) edge_vec, edge_index, edge_mask = _two_atom_inputs(r) out = pot.call( edge_vec, @@ -67,7 +67,7 @@ def test_zbl_known_value(type_map, atypes, zi, zj): def test_virtual_types_masked(): # real_type_count=1: type 1 is a virtual/placeholder type; its edges # contribute zero, and only real-real edges survive. - pot = InterPotential(type_map=["O"]) + pot = InnerPotential(type_map=["O"]) r = 0.9 edge_vec = np.array( [[r, 0, 0], [-r, 0, 0], [0, r, 0], [0, -r, 0]], dtype=np.float64 @@ -83,7 +83,7 @@ def test_virtual_types_masked(): def test_edge_mask_zeroes_edges(): - pot = InterPotential(type_map=["O"]) + pot = InnerPotential(type_map=["O"]) edge_vec, edge_index, _ = _two_atom_inputs(0.8) out = pot.call( edge_vec, @@ -97,19 +97,19 @@ def test_edge_mask_zeroes_edges(): def test_unknown_element_raises(): with pytest.raises(ValueError, match="Unknown element symbol"): - InterPotential(type_map=["O", "Xx"]) + InnerPotential(type_map=["O", "Xx"]) def test_unknown_mode_raises(): - with pytest.raises(ValueError, match="Unknown InterPotential mode"): - InterPotential(type_map=["O"], mode="lj") + with pytest.raises(ValueError, match="Unknown InnerPotential mode"): + InnerPotential(type_map=["O"], mode="lj") def test_torch_namespace_smoke_and_gradient(): """Torch inputs match numpy at 1e-12 and edge_vec gradients exist.""" import torch - pot = InterPotential(type_map=["O", "H"]) + pot = InnerPotential(type_map=["O", "H"]) edge_vec_np, edge_index, edge_mask = _two_atom_inputs(0.8) atypes = np.array([0, 1], dtype=np.int64) ref = np.asarray(pot.call(edge_vec_np, edge_index, atypes, edge_mask, 2)) diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index 1bca3a2a9e..0db497dd84 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -2,7 +2,7 @@ """dpmodel ZBL bridging as COMPOSITION (review 3638077323, redesigned). ``bridging_method: ZBL`` builds a -``LinearEnergyModel(LinearEnergyAtomicModel([dp, InterPotentialAtomicModel], +``LinearEnergyModel(LinearEnergyAtomicModel([dp, InnerPotentialAtomicModel], weights="sum"))`` -- the analytical term is its own atomic model summed with the learned one, not a flag on it. """ @@ -12,8 +12,8 @@ import numpy as np import pytest -from deepmd.dpmodel.atomic_model.inter_potential import ( - InterPotentialAtomicModel, +from deepmd.dpmodel.atomic_model.inner_potential import ( + InnerPotentialAtomicModel, ) from deepmd.dpmodel.atomic_model.linear_atomic_model import ( LinearEnergyAtomicModel, @@ -67,8 +67,8 @@ def test_builder_composes_linear_model(): assert am.weights == "sum" kinds = [type(c).__name__ for c in am.models] assert ( - kinds == ["EnergyAtomicModel", "InterPotentialAtomicModel"] - or kinds[1] == "InterPotentialAtomicModel" + kinds == ["EnergyAtomicModel", "InnerPotentialAtomicModel"] + or kinds[1] == "InnerPotentialAtomicModel" ) # radii wired to the LEARNED child's descriptor InnerClamp dp_child = am.models[0] @@ -147,14 +147,14 @@ def test_zbl_serialize_roundtrip_energy_identical(): def test_zbl_atomic_dense_route_raises(): - zbl = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + zbl = InnerPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) with pytest.raises(NotImplementedError, match="NeighborGraph route only"): zbl.forward_atomic(None, None, None) -def test_inter_potential_supports_graph_lower(): +def test_inner_potential_supports_graph_lower(): """The analytical ZBL term is graph-capable (rides the NeighborGraph).""" - zbl = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + zbl = InnerPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) assert zbl.uses_graph_lower() is True @@ -168,10 +168,10 @@ def test_linear_graph_lower_requires_all_children(): child, which does not implement it. """ - class _GraphChild(InterPotentialAtomicModel): + class _GraphChild(InnerPotentialAtomicModel): pass # inherits uses_graph_lower() -> True - class _DenseOnlyChild(InterPotentialAtomicModel): + class _DenseOnlyChild(InnerPotentialAtomicModel): def uses_graph_lower(self) -> bool: return False @@ -202,7 +202,7 @@ def test_zbl_atomic_graph_values(): ) r = 0.8 - zbl = InterPotentialAtomicModel(type_map=["O"], rcut=4.0, sel=[8]) + zbl = InnerPotentialAtomicModel(type_map=["O"], rcut=4.0, sel=[8]) graph = NeighborGraph( n_node=np.array([2], dtype=np.int64), edge_index=np.array([[0, 1], [1, 0]], dtype=np.int64), @@ -239,22 +239,22 @@ def _pair_energy(model, natoms=2, r=1.0): return float(np.sum(out["energy"])) -class TestInterPotentialChangeTypeMap: +class TestInnerPotentialChangeTypeMap: """``change_type_map`` must rebuild the ZBL element lookup. The generic ``BaseAtomicModel.change_type_map`` only rewrites the public map and the stat/exclusion state; the nuclear-charge table belongs to - ``InterPotential`` and is rebuilt there (review 3649295675). Without it + ``InnerPotential`` and is rebuilt there (review 3649295675). Without it the lookup keeps the ORIGINAL elements while ``atype`` values already mean the new ones -- silently wrong energies, or ``IndexError`` for a longer map. """ def test_reorder_matches_a_freshly_built_model(self) -> None: - model = InterPotentialAtomicModel(type_map=["H", "O"], rcut=4.0, sel=[8]) + model = InnerPotentialAtomicModel(type_map=["H", "O"], rcut=4.0, sel=[8]) e_hh = _pair_energy(model) model.change_type_map(["O", "H"]) - fresh = InterPotentialAtomicModel(type_map=["O", "H"], rcut=4.0, sel=[8]) + fresh = InnerPotentialAtomicModel(type_map=["O", "H"], rcut=4.0, sel=[8]) e_fresh = _pair_energy(fresh) # anti-vacuity: the two element pairs must be far apart, or a stale # lookup would be indistinguishable from a rebuilt one @@ -264,7 +264,7 @@ def test_reorder_matches_a_freshly_built_model(self) -> None: assert model.potential.type_map == ["O", "H"] def test_added_element_extends_the_lookup(self) -> None: - model = InterPotentialAtomicModel(type_map=["H", "O"], rcut=4.0, sel=[8]) + model = InnerPotentialAtomicModel(type_map=["H", "O"], rcut=4.0, sel=[8]) model.change_type_map(["H", "O", "Ni"]) assert model.potential.ntypes_real == 3 assert list(model.potential.atomic_numbers) == [1.0, 8.0, 28.0] @@ -284,11 +284,11 @@ def test_added_element_extends_the_lookup(self) -> None: e_nini = float( np.sum(model.forward_common_atomic_graph(graph, graph_atype)["energy"]) ) - fresh = InterPotentialAtomicModel(type_map=["Ni"], rcut=4.0, sel=[8]) + fresh = InnerPotentialAtomicModel(type_map=["Ni"], rcut=4.0, sel=[8]) np.testing.assert_allclose(e_nini, _pair_energy(fresh), rtol=1e-12) def test_dropped_element_shrinks_the_lookup(self) -> None: - model = InterPotentialAtomicModel(type_map=["H", "O", "Ni"], rcut=4.0, sel=[8]) + model = InnerPotentialAtomicModel(type_map=["H", "O", "Ni"], rcut=4.0, sel=[8]) model.change_type_map(["Ni"]) assert model.potential.ntypes_real == 1 assert list(model.potential.atomic_numbers) == [28.0] @@ -304,7 +304,7 @@ def test_serialize_roundtrip_after_change_type_map(self) -> None: BaseAtomicModel, ) - model = InterPotentialAtomicModel(type_map=["H", "O"], rcut=4.0, sel=[8]) + model = InnerPotentialAtomicModel(type_map=["H", "O"], rcut=4.0, sel=[8]) model.change_type_map(["O", "H"]) data = model.serialize() restored = BaseAtomicModel.get_class_by_type(data["type"]).deserialize(data) @@ -324,7 +324,7 @@ class TestNativeSpinCapabilityOnAtomicModel: """ def test_analytical_term_is_not_spin_capable(self) -> None: - zbl = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + zbl = InnerPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) # inherits the concrete base default -- no descriptor, no spin input assert zbl.supports_native_spin() is False @@ -338,14 +338,14 @@ def test_composition_is_capable_when_any_child_is(self) -> None: ).atomic_model assert learned.supports_native_spin() is True kinds = [type(c).__name__ for c in learned.models] - assert kinds[1] == "InterPotentialAtomicModel", kinds + assert kinds[1] == "InnerPotentialAtomicModel", kinds # ... and the spin-free analytical child alone is not capable assert learned.models[1].supports_native_spin() is False def test_composition_without_a_spin_consumer_is_not_capable(self) -> None: """No consumer => the magnetic force would be identically zero.""" - zbl_a = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) - zbl_b = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + zbl_a = InnerPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + zbl_b = InnerPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) composed = LinearEnergyAtomicModel( [zbl_a, zbl_b], type_map=["Ni", "O"], weights="sum" ) @@ -383,7 +383,7 @@ def test_charge_spin_survives_bridging(self) -> None: assert bridged.has_chg_spin_ebd() is True # ... and the composition really is a composition assert [type(c).__name__ for c in bridged.atomic_model.models][1] == ( - "InterPotentialAtomicModel" + "InnerPotentialAtomicModel" ) def test_no_charge_spin_stays_zero(self) -> None: @@ -434,7 +434,7 @@ def test_pair_exclusion_survives_bridging(self) -> None: # ... and the bridged one really is the composition, so the value is # read off the wrapper rather than accidentally off a lone child. assert [type(c).__name__ for c in bridged.atomic_model.models][1] == ( - "InterPotentialAtomicModel" + "InnerPotentialAtomicModel" ) def test_pair_exclusion_survives_native_spin_plus_bridging(self) -> None: @@ -558,7 +558,7 @@ def test_dimension_and_default_align_with_the_learned_child(self) -> None: The analytical child consumes neither fparam nor aparam, so it is not a consumer and must not constrain either. """ - zbl = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + zbl = InnerPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) # anti-vacuity: the analytical child really is a non-consumer assert zbl.get_dim_fparam() == 0 assert zbl.get_dim_aparam() == 0 @@ -578,7 +578,7 @@ def test_matching_defaults_are_exposed(self) -> None: def test_dimension_zero_child_is_ignored(self) -> None: """Learned + ZBL must still inherit the learned default.""" - zbl = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + zbl = InnerPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) assert zbl.get_dim_fparam() == 0 # anti-vacuity: really a non-consumer m = self._compose([self._Fake(1, [0.5]), zbl]) assert m.has_default_fparam() is True @@ -605,8 +605,8 @@ def test_intensive_mixture_is_rejected_at_construction(self) -> None: composition is not physically meaningful, so it should never exist rather than exist and answer a plausible-looking default. """ - zbl_a = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) - zbl_b = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + zbl_a = InnerPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) + zbl_b = InnerPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8]) # anti-vacuity: matching children compose fine and report their value assert zbl_a.get_intensive() is False assert ( diff --git a/source/tests/infer/gen_dpa4_spin_zbl.py b/source/tests/infer/gen_dpa4_spin_zbl.py index 09d1d4121d..26812097a8 100644 --- a/source/tests/infer/gen_dpa4_spin_zbl.py +++ b/source/tests/infer/gen_dpa4_spin_zbl.py @@ -4,7 +4,7 @@ One archive, ``deeppot_dpa4_spin_zbl_graph.pt2``: a native-spin (``scheme="native"``) DPA4 that is ALSO bridged (``bridging_method: "ZBL"``), -i.e. a ``LinearEnergyModel`` over ``[learned DPA4, InterPotentialAtomicModel]`` +i.e. a ``LinearEnergyModel`` over ``[learned DPA4, InnerPotentialAtomicModel]`` with ``weights="sum"`` wrapped by the native-spin model class. Why this fixture exists @@ -144,7 +144,7 @@ ).reshape(1, _NATOMS, 3) # ZBL screening parameters, repeated here as an INDEPENDENT reference (they -# mirror deepmd/dpmodel/atomic_model/inter_potential.py, deliberately not +# mirror deepmd/dpmodel/atomic_model/inner_potential.py, deliberately not # imported from it: a reference that shares its constants with the code under # test cannot catch a wrong constant). _ZBL_A_COEFF = (0.18175, 0.50986, 0.28022, 0.028171) @@ -205,9 +205,9 @@ def _build_model_dict() -> dict: model = get_model(copy.deepcopy(SPIN_ZBL_CONFIG)) assert model.has_spin() is True kinds = [type(child).__name__ for child in model.atomic_model.models] - assert kinds[1] == "InterPotentialAtomicModel", ( + assert kinds[1] == "InnerPotentialAtomicModel", ( f"expected the bridged composition's second child to be the " - f"analytical InterPotentialAtomicModel, got {kinds!r}; without it " + f"analytical InnerPotentialAtomicModel, got {kinds!r}; without it " f"this fixture is just the plain native-spin model again." ) model_dict = model.serialize() diff --git a/source/tests/infer/gen_dpa4_zbl.py b/source/tests/infer/gen_dpa4_zbl.py index a8f5f6ade0..1b9a77d62a 100644 --- a/source/tests/infer/gen_dpa4_zbl.py +++ b/source/tests/infer/gen_dpa4_zbl.py @@ -3,7 +3,7 @@ """Generate deeppot_dpa4_zbl_graph.pt2: DPA4 with analytical ZBL bridging. ``bridging_method: ZBL`` builds a COMPOSITION -- ``LinearEnergyModel`` over -``[learned DPA4, InterPotentialAtomicModel]`` with ``weights="sum"`` -- so +``[learned DPA4, InnerPotentialAtomicModel]`` with ``weights="sum"`` -- so the frozen archive exercises a code path no other C++ fixture covers: the graph lower of a linear composition rather than a single learned model. Before this fixture, ZBL bridging had NO C++ or LAMMPS coverage at all; its @@ -49,7 +49,7 @@ # Small fp64 DPA4 + ZBL config. ``bridging_r_inner``/``r_outer`` feed the # descriptor's InnerClamp AND BridgingSwitch (they are built together from -# the same radii), and the model-level InterPotential term. +# the same radii), and the model-level InnerPotential term. ZBL_CONFIG = { "type_map": ["Ni", "O"], "descriptor": { diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index 494338a2da..125483db69 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -38,7 +38,7 @@ get_sezm_model, ) from deepmd.pt.model.model.sezm_model import ( - InterPotential, + InnerPotential, SeZMModel, ) from deepmd.pt.model.model.sezm_native_spin_model import ( @@ -53,6 +53,9 @@ from deepmd.pt.utils import ( env, ) +from deepmd.pt.utils.compile_compat import ( + SUPPORTED_COMPILE_TORCH, +) from deepmd.pt.utils.nlist import ( extend_input_and_build_neighbor_list, ) @@ -72,18 +75,18 @@ module=r"torch\._functorch\._aot_autograd\.autograd_cache", ) -# SeZM's ``torch.compile`` / AOT-export code paths are validated on torch -# 2.11.x and 2.12.x, the releases the compile pipeline supports (see -# ``deepmd.pt.utils.compile_compat``). Other torch versions can segfault or -# drift, so the compile-parity tests are skipped there. +# Keep compile-parity test gating aligned with the runtime allowlist in +# ``deepmd.pt.utils.compile_compat``. Membership in the allowlist does not imply +# that every release is installed in CI. _TORCH_VERSION = parse_version(torch.__version__) -_SKIP_OFF_COMPILE_TORCH = (_TORCH_VERSION.major, _TORCH_VERSION.minor) not in { - (2, 11), - (2, 12), -} +_SKIP_OFF_COMPILE_TORCH = ( + _TORCH_VERSION.major, + _TORCH_VERSION.minor, +) not in SUPPORTED_COMPILE_TORCH _SKIP_OFF_COMPILE_TORCH_REASON = ( - "SeZM's torch.compile path is only supported on torch 2.11.x and 2.12.x; " - f"current torch is {torch.__version__}." + "SeZM's torch.compile path is only supported on torch " + + ", ".join(f"{major}.{minor}.x" for major, minor in SUPPORTED_COMPILE_TORCH) + + f"; current torch is {torch.__version__}." ) @@ -1201,6 +1204,43 @@ def test_forward_shapes_and_reduction(self) -> None: expected = ret["atom_foo"].sum(dim=1) torch.testing.assert_close(ret["foo"], expected) + def test_change_out_bias_matches_property_reduction(self) -> None: + """Property statistics use sum or masked mean according to metadata.""" + coord, atype, box = self._make_tiny_frame() + for intensive in (False, True): + model = get_sezm_model( + self._build_model_params( + use_compile=False, + intensive=intensive, + ) + ).to(self.device) + model.eval() + label = model(coord, atype, box=box)["foo"].detach() + old_bias = model.get_out_bias().detach().clone() + sample = { + "coord": coord, + "atype": atype, + "box": box, + "foo": label, + "natoms": torch.tensor( + [[5, 5, 3, 2]], + dtype=torch.long, + device=self.device, + ), + "find_foo": torch.tensor(1.0, device=self.device), + } + + model.change_out_bias( + [sample], + bias_adjust_mode="change-by-statistic", + ) + torch.testing.assert_close( + model.get_out_bias(), + old_bias, + atol=1.0e-6, + rtol=1.0e-6, + ) + def test_property_loss_and_serialization(self) -> None: """PropertyLoss metadata and model serialization should round-trip.""" from deepmd.pt.model.model.model import ( @@ -1261,8 +1301,8 @@ def test_compile_matches_eager_and_backpropagates(self) -> None: self.assertTrue(grad_found) -class TestInterPotential(unittest.TestCase): - """Test InterPotential ZBL analytical pair potential.""" +class TestInnerPotential(unittest.TestCase): + """Test InnerPotential ZBL analytical pair potential.""" def setUp(self) -> None: self.device = env.DEVICE @@ -1285,7 +1325,7 @@ def _pair_edges( def test_zbl_known_value_OO(self) -> None: """ZBL energy for an O-O pair matches the analytic reference.""" - pot = InterPotential(type_map=["O", "H"], mode="ZBL").to(self.device) + pot = InnerPotential(type_map=["O", "H"], mode="ZBL").to(self.device) import math @@ -1308,7 +1348,7 @@ def test_zbl_known_value_OO(self) -> None: def test_zbl_known_value_OH(self) -> None: """ZBL energy for an O-H pair matches the analytic reference.""" - pot = InterPotential(type_map=["O", "H"], mode="ZBL").to(self.device) + pot = InnerPotential(type_map=["O", "H"], mode="ZBL").to(self.device) import math z_o, z_h = 8.0, 1.0 @@ -1330,7 +1370,7 @@ def test_zbl_known_value_OH(self) -> None: def test_zbl_gradient_exists(self) -> None: """ZBL produces finite gradients w.r.t. the edge vectors.""" - pot = InterPotential(type_map=["O", "H"], mode="ZBL").to(self.device) + pot = InnerPotential(type_map=["O", "H"], mode="ZBL").to(self.device) edge_vec, edge_index, atype_flat, edge_mask = self._pair_edges(1.0, [0, 1]) edge_vec = edge_vec.detach().requires_grad_(True) @@ -1340,7 +1380,7 @@ def test_zbl_gradient_exists(self) -> None: def test_virtual_spin_types_masked(self) -> None: """Edges touching a virtual spin type (>= real_type_count) contribute 0.""" - pot = InterPotential(type_map=["O", "H"], mode="ZBL").to(self.device) + pot = InnerPotential(type_map=["O", "H"], mode="ZBL").to(self.device) # Node 2 is a virtual spin atom (type 2 >= real_type_count=2). edge_vec = torch.tensor( [[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0], [0.5, 0.0, 0.0], [-0.5, 0.0, 0.0]], @@ -1371,7 +1411,7 @@ def test_virtual_spin_types_masked(self) -> None: def test_unknown_element_raises(self) -> None: """Test that unknown element raises ValueError.""" with self.assertRaises(ValueError): - InterPotential(type_map=["O", "Xx"]) + InnerPotential(type_map=["O", "Xx"]) class TestSeZMEdgeForceScatter(unittest.TestCase): @@ -1383,7 +1423,7 @@ class TestSeZMEdgeForceScatter(unittest.TestCase): float64 finite-difference checks pin the conservative-force guarantee ``F = -dE/dx`` and the PBC-correct virial ``W = -dE/deps``, and confirm the half-split per-atom virial sums back to the global virial. The ZBL - cases additionally drive ``InterPotential`` (edge form) through the + cases additionally drive ``InnerPotential`` (edge form) through the same single backward. """ @@ -1568,7 +1608,12 @@ class TestSeZMNativeSpinModel(unittest.TestCase): def setUp(self) -> None: self.device = env.DEVICE - def _build_model(self, *, use_compile: bool = False) -> SeZMNativeSpinModel: + def _build_model( + self, + *, + use_compile: bool = False, + bridging_method: str = "none", + ) -> SeZMNativeSpinModel: """Build a tiny float64 native-spin model with randomized parameters.""" params = { "type": "dpa4", @@ -1602,6 +1647,9 @@ def _build_model(self, *, use_compile: bool = False) -> SeZMNativeSpinModel: "seed": 7, }, "use_compile": use_compile, + "bridging_method": bridging_method, + "bridging_r_inner": 0.8, + "bridging_r_outer": 1.2, } model = get_model(params) # Perturb away from the near-identity initialization so the spin @@ -1644,6 +1692,37 @@ def _frame( ) return coord, atype, spin, box + def test_zbl_change_out_bias_is_invariant_for_self_labels(self) -> None: + """Native-spin statistics consume spin and the complete ZBL energy.""" + model = self._build_model(bridging_method="ZBL") + coord, atype, spin, box = self._frame() + label = model(coord, atype, spin, box=box)["energy"].detach() + old_bias = model.get_out_bias().detach().clone() + sample = { + "coord": coord, + "atype": atype, + "spin": spin, + "box": box, + "energy": label, + "natoms": torch.tensor( + [[5, 5, 3, 2]], + dtype=torch.long, + device=self.device, + ), + "find_energy": torch.tensor(1.0, device=self.device), + } + + model.change_out_bias( + [sample], + bias_adjust_mode="change-by-statistic", + ) + torch.testing.assert_close( + model.get_out_bias(), + old_bias, + atol=1.0e-12, + rtol=1.0e-12, + ) + @staticmethod def _proper_rotation(device: torch.device) -> torch.Tensor: """A deterministic proper rotation matrix (det = +1).""" @@ -1987,12 +2066,18 @@ def test_bridging_none_unchanged(self) -> None: self.assertEqual(model.bridging_method, "NONE") def test_bridging_zbl_creates_potential(self) -> None: - """Test that bridging_method='ZBL' creates InterPotential and InnerClamp.""" + """Test that bridging_method='ZBL' creates InnerPotential and InnerClamp.""" model = get_sezm_model(self._build_model_params(bridging_method="ZBL")) self.assertIsNotNone(model.inter_potential) self.assertEqual(model.bridging_method, "ZBL") self.assertIsNotNone(model.atomic_model.descriptor.inner_clamp) + def test_empty_statistics_raise(self) -> None: + """A fully filtered statistics sample cannot calibrate output bias.""" + model = get_sezm_model(self._build_model_params(bridging_method="ZBL")) + with self.assertRaisesRegex(ValueError, "at least one sampled system"): + model.compute_or_load_stat(lambda: []) + def test_zbl_adds_energy(self) -> None: """Test that ZBL bridging adds energy to the model output.""" model_plain = get_sezm_model(self._build_model_params(bridging_method="none")) @@ -2026,6 +2111,132 @@ def test_zbl_adds_energy(self) -> None: "ZBL bridging should add positive (repulsive) energy", ) + def test_change_out_bias_is_invariant_for_self_labels(self) -> None: + """Residual bias calibration uses the complete bridged energy.""" + params = self._build_model_params(bridging_method="ZBL") + params["descriptor"]["precision"] = "float64" + params["fitting_net"]["precision"] = "float64" + model = get_sezm_model(params).to(self.device).eval() + + coord = torch.tensor( + [ + [[0.0, 0.0, 0.0], [0.90, 0.0, 0.0], [1.80, 0.0, 0.0]], + [[0.0, 0.0, 0.0], [0.95, 0.0, 0.0], [1.90, 0.0, 0.0]], + ], + dtype=torch.float64, + device=self.device, + ) + atype = torch.tensor( + [[0, 0, 1], [0, 1, 1]], + dtype=torch.long, + device=self.device, + ) + box = torch.tensor( + [[10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0]] * 2, + dtype=torch.float64, + device=self.device, + ) + labels = model(coord, atype, box=box)["energy"].detach() + stat_energy = model.predict_atomic_outputs_for_stat( + coord, + atype, + box, + )["energy"].sum(dim=1) + torch.testing.assert_close(stat_energy, labels, atol=1.0e-12, rtol=1.0e-12) + + sample = { + "coord": coord, + "atype": atype, + "box": box, + "energy": labels, + "natoms": torch.tensor( + [[3, 3, 2, 1], [3, 3, 1, 2]], + dtype=torch.long, + device=self.device, + ), + "find_energy": np.float32(1.0), + } + old_bias = model.get_out_bias().detach().clone() + model.use_compile = True + model.train() + with ( + mock.patch.object( + model, + "trace_and_compile", + side_effect=AssertionError("statistics must not compile"), + ), + mock.patch( + "deepmd.pt.model.model.sezm_model.edge_energy_deriv", + side_effect=AssertionError("statistics must not compute derivatives"), + ), + ): + model.change_out_bias( + [sample], + bias_adjust_mode="change-by-statistic", + ) + + self.assertTrue(model.training) + self.assertEqual(model.compiled_core_compute_cache, {}) + self.assertTrue( + all( + module._cached_weight is None + for module in model.modules() + if isinstance(module, SO2Linear) + ) + ) + torch.testing.assert_close( + model.get_out_bias(), + old_bias, + atol=1.0e-12, + rtol=1.0e-12, + ) + self.assertTrue(all(parameter.grad is None for parameter in model.parameters())) + model.use_compile = False + model(coord, atype, box=box)["energy"].sum().backward() + self.assertTrue( + any( + parameter.grad is not None + and torch.count_nonzero(parameter.grad).item() > 0 + for parameter in model.parameters() + ) + ) + + def test_zbl_respects_exclusions(self) -> None: + """Excluded atoms and pairs contribute neither learned nor ZBL energy.""" + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [0.8, 0.0, 0.0]]], + dtype=torch.float64, + device=self.device, + ) + atype = torch.tensor([[0, 1]], dtype=torch.long, device=self.device) + box = torch.tensor( + [[10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0]], + dtype=torch.float64, + device=self.device, + ) + for exclusion_key, exclusion_value in ( + ("pair_exclude_types", [[0, 1]]), + ("atom_exclude_types", [1]), + ): + with self.subTest(exclusion_key=exclusion_key): + plain_params = self._build_model_params(bridging_method="none") + zbl_params = self._build_model_params(bridging_method="ZBL") + for params in (plain_params, zbl_params): + params[exclusion_key] = exclusion_value + params["descriptor"]["precision"] = "float64" + params["fitting_net"]["precision"] = "float64" + + model_plain = get_sezm_model(plain_params).to(self.device).eval() + model_zbl = get_sezm_model(zbl_params).to(self.device).eval() + model_zbl.load_state_dict(model_plain.state_dict(), strict=False) + + torch.testing.assert_close( + model_zbl(coord, atype, box=box)["energy"], + model_plain(coord, atype, box=box)["energy"], + atol=1.0e-12, + rtol=1.0e-12, + ) + class TestSeZMModelModes(unittest.TestCase): """Targeted regression tests for SeZM `ener` / `dens` mode routing.""" diff --git a/source/tests/pt/model/test_sezm_spin_model.py b/source/tests/pt/model/test_sezm_spin_model.py index ba48c94889..c98c80f30c 100644 --- a/source/tests/pt/model/test_sezm_spin_model.py +++ b/source/tests/pt/model/test_sezm_spin_model.py @@ -27,6 +27,9 @@ from deepmd.pt.utils import ( env, ) +from deepmd.pt.utils.compile_compat import ( + SUPPORTED_COMPILE_TORCH, +) from deepmd.pt.utils.nlist import ( extend_input_and_build_neighbor_list, ) @@ -43,16 +46,18 @@ module=r"torch\._functorch\._aot_autograd\.autograd_cache", ) -# TODO(torch-2.11): SeZM's ``torch.compile`` / AOT-export code paths are only -# stable on torch 2.11.x. CI currently pins torch 2.10, where the compiled path -# can segfault or drift, and other torch versions are similarly unstable. Skip -# the compile-parity test off 2.11 until CI standardizes on a SeZM-compatible -# torch, then drop this guard. +# Keep compile-parity test gating aligned with the runtime allowlist in +# ``deepmd.pt.utils.compile_compat``. Membership in the allowlist does not imply +# that every release is installed in CI. _TORCH_VERSION = parse_version(torch.__version__) -_SKIP_OFF_TORCH_211 = (_TORCH_VERSION.major, _TORCH_VERSION.minor) != (2, 11) -_SKIP_OFF_TORCH_211_REASON = ( - "SeZM's torch.compile path is only stable on torch 2.11.x; " - f"current torch is {torch.__version__}." +_SKIP_OFF_COMPILE_TORCH = ( + _TORCH_VERSION.major, + _TORCH_VERSION.minor, +) not in SUPPORTED_COMPILE_TORCH +_SKIP_OFF_COMPILE_TORCH_REASON = ( + "SeZM's torch.compile path is only supported on torch " + + ", ".join(f"{major}.{minor}.x" for major, minor in SUPPORTED_COMPILE_TORCH) + + f"; current torch is {torch.__version__}." ) @@ -380,7 +385,121 @@ def test_bridging_masks_virtual_pairs(self) -> None: torch.testing.assert_close(energy_with_virtual, energy_real_only) - @unittest.skipIf(_SKIP_OFF_TORCH_211, _SKIP_OFF_TORCH_211_REASON) + def test_zbl_change_out_bias_is_invariant_for_self_labels(self) -> None: + """Spin-expanded statistics include the complete bridged energy.""" + model = get_model(self._build_model_params(bridging_method="ZBL")).to( + self.device + ) + model.eval() + label = model( + self.coord, + self.atype, + spin=self.spin, + box=self.box, + )["energy"].detach() + old_bias = model.get_out_bias().detach().clone() + sample = { + "coord": self.coord, + "atype": self.atype, + "spin": self.spin, + "box": self.box, + "energy": label, + "natoms": torch.tensor( + [[3, 3, 2, 1]], + dtype=torch.long, + device=self.device, + ), + "find_energy": torch.tensor(1.0, device=self.device), + } + + model.change_out_bias( + [sample], + bias_adjust_mode="change-by-statistic", + ) + torch.testing.assert_close( + model.get_out_bias(), + old_bias, + atol=1.0e-12, + rtol=1.0e-12, + ) + + def test_zbl_statistics_use_physical_spin_neighbor_topology(self) -> None: + """Statistics preserve real-nlist expansion when `sel` is saturated.""" + params = self._build_model_params(bridging_method="ZBL") + params["descriptor"]["sel"] = [1, 1] + params["descriptor"]["precision"] = "float64" + params["fitting_net"]["precision"] = "float64" + model = get_model(params).to(self.device).eval() + coord = torch.tensor( + [ + [ + [0.0, 0.0, 0.0], + [0.8, 0.0, 0.0], + [0.0, 0.9, 0.0], + [0.9, 0.9, 0.0], + ] + ], + dtype=torch.float64, + device=self.device, + ) + atype = torch.tensor( + [[0, 1, 0, 1]], + dtype=torch.long, + device=self.device, + ) + spin = torch.tensor( + [ + [ + [0.20, 0.10, 0.00], + [0.00, 0.00, 0.00], + [0.10, 0.20, 0.10], + [0.00, 0.00, 0.00], + ] + ], + dtype=torch.float64, + device=self.device, + ) + box = torch.tensor( + [[6.0, 0.0, 0.0, 0.0, 6.0, 0.0, 0.0, 0.0, 6.0]], + dtype=torch.float64, + device=self.device, + ) + label = model(coord, atype, spin=spin, box=box)["energy"].detach() + stat_energy = model.predict_atomic_outputs_for_stat( + coord, + atype, + box, + spin=spin, + )["energy"].sum(dim=1) + torch.testing.assert_close(stat_energy, label, atol=1.0e-12, rtol=1.0e-12) + + old_bias = model.get_out_bias().detach().clone() + model.change_out_bias( + [ + { + "coord": coord, + "atype": atype, + "spin": spin, + "box": box, + "energy": label, + "natoms": torch.tensor( + [[4, 4, 2, 2]], + dtype=torch.long, + device=self.device, + ), + "find_energy": torch.tensor(1.0, device=self.device), + } + ], + bias_adjust_mode="change-by-statistic", + ) + torch.testing.assert_close( + model.get_out_bias(), + old_bias, + atol=1.0e-12, + rtol=1.0e-12, + ) + + @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) def test_compile_matches_eager(self) -> None: """Compiled SeZM spin path should match eager predictions.""" eager = get_model(self._build_model_params(use_compile=False)).to(self.device) diff --git a/source/tests/pt/test_multitask.py b/source/tests/pt/test_multitask.py index 560d89ed56..76c0ed37a1 100644 --- a/source/tests/pt/test_multitask.py +++ b/source/tests/pt/test_multitask.py @@ -252,6 +252,47 @@ def setUp(self) -> None: self.config["model"] ) + def test_filtered_tasks_share_retry_budget(self) -> None: + """All tasks share one retry budget when filtering invalid batches.""" + config = deepcopy(self.config) + data_dict = config["training"]["data_dict"] + for task_data in data_dict.values(): + task_data["training_data"]["min_pair_dist"] = 0.1 + data_dict["model_1"]["training_data"]["systems"] = [ + *data_dict["model_1"]["training_data"]["systems"], + str(Path(__file__).parent / "water/data/data_1"), + ] + config = update_deepmd_input(config, warning=False) + config = normalize(config, multi_task=True) + trainer = get_trainer(config, shared_links=self.shared_links) + + loader_lengths = { + task_key: len(trainer.training_dataloader[task_key]) + for task_key in trainer.model_keys + } + self.assertGreater(len(set(loader_lengths.values())), 1) + expected_attempts = max(loader_lengths.values()) + + for task_key in trainer.model_keys: + with self.subTest(task_key=task_key): + call_count = 0 + + def get_empty_data( + is_train: bool = True, + task_key: str = "Default", + ) -> tuple[dict, dict, dict]: + nonlocal call_count + call_count += 1 + return {}, {}, {} + + trainer.get_data = get_empty_data + with self.assertRaisesRegex( + RuntimeError, + rf"after {expected_attempts} attempts", + ): + trainer._next_training_batch(task_key) + self.assertEqual(call_count, expected_attempts) + def tearDown(self) -> None: MultiTaskTrainTest.tearDown(self) diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py index 34034585ce..a363e54d1a 100644 --- a/source/tests/pt/test_training.py +++ b/source/tests/pt/test_training.py @@ -36,6 +36,9 @@ from deepmd.pt.train.ema import ( EMA_CHECKPOINT_KEY, ) +from deepmd.pt.train.training import ( + all_ranks_have_valid_frames, +) from deepmd.pt.utils.finetune import ( get_finetune_rules, ) @@ -43,6 +46,10 @@ _cascade_top_level_defaults, preprocess_shared_params, ) +from deepmd.pt.utils.stat import ( + make_stat_input, + select_batch_frames, +) from deepmd.utils.argcheck import ( normalize, ) @@ -91,6 +98,81 @@ def raise_timeout(signum: int, frame: Any) -> None: TRAINING_TEST_TIMEOUT = _training_timeout(60) +class TestStatisticsFrameFiltering(unittest.TestCase): + """Verify statistics use the training minimum-distance frame semantics.""" + + def test_min_pair_dist_filters_every_frame_aligned_tensor(self) -> None: + batch = { + "coord": torch.arange(18, dtype=torch.float64, device="cpu").reshape( + 2, 3, 3 + ), + "atype": torch.tensor( + [[0, 0, 1], [0, 1, 1]], dtype=torch.long, device="cpu" + ), + "energy": torch.tensor( + [[1000.0], [2.0]], dtype=torch.float64, device="cpu" + ), + "natoms": torch.tensor([[3, 3, 2, 1], [3, 3, 1, 2]], device="cpu"), + "min_pair_dist": torch.tensor( + [[0.5], [1.0]], dtype=torch.float64, device="cpu" + ), + "fid": ["rejected", "accepted"], + "find_energy": np.float32(1.0), + } + filtered_batch = select_batch_frames( + batch, + torch.tensor([False, True], device="cpu"), + ) + self.assertEqual(filtered_batch["fid"], ["accepted"]) + sampled = make_stat_input( + [object()], + [[batch]], + nbatches=1, + min_pair_dist=0.8, + ) + + self.assertEqual(sampled[0]["coord"].shape[0], 1) + self.assertEqual(sampled[0]["atype"].shape[0], 1) + self.assertEqual(sampled[0]["natoms"].shape[0], 1) + torch.testing.assert_close( + sampled[0]["energy"].cpu(), + torch.tensor([[2.0]], dtype=torch.float64, device="cpu"), + ) + + def test_min_pair_dist_skips_system_without_valid_frames(self) -> None: + batch = { + "coord": torch.zeros(1, 2, 3, device="cpu"), + "atype": torch.zeros(1, 2, dtype=torch.long, device="cpu"), + "energy": torch.zeros(1, 1, device="cpu"), + "natoms": torch.tensor([[2, 2, 2]], device="cpu"), + "min_pair_dist": torch.tensor([[0.5]], device="cpu"), + "find_energy": np.float32(1.0), + } + sampled = make_stat_input( + [object()], + [[batch]], + nbatches=1, + min_pair_dist=0.8, + ) + self.assertEqual(sampled, []) + + def test_distributed_filter_skips_when_any_rank_is_empty(self) -> None: + """The global MIN decision prevents invalid fallback frames.""" + + def mark_remote_rank_empty( + valid_flag: torch.Tensor, + op: Any, + ) -> None: + self.assertEqual(op, torch.distributed.ReduceOp.MIN) + valid_flag.zero_() + + with patch( + "deepmd.pt.train.training.dist.all_reduce", + side_effect=mark_remote_rank_empty, + ): + self.assertFalse(all_ranks_have_valid_frames(local_has_valid=True)) + + class DPTrainTest: test_zbl_from_standard: bool = False @@ -1097,6 +1179,7 @@ def setUp(self) -> None: self.config["training"]["numb_steps"] = 2 self.config["training"]["save_freq"] = 2 self.config["training"]["disp_training"] = False + self.config["training"]["training_data"]["min_pair_dist"] = 0.1 self.config["validating"] = { "full_validation": False, "ema_full_validation": False, @@ -1106,7 +1189,7 @@ def tearDown(self) -> None: os.chdir(self._cwd) self._tmpdir.cleanup() - def test_skipped_batch_does_not_advance_scheduler(self) -> None: + def test_invalid_batch_is_retried_without_losing_a_step(self) -> None: trainer = get_trainer(deepcopy(self.config)) original_get_data = trainer.get_data skipped = {"done": False} @@ -1129,7 +1212,7 @@ def get_data( trainer.run() self.assertTrue(skipped["done"]) - self.assertEqual(trainer.scheduler.last_epoch, 1) + self.assertEqual(trainer.scheduler.last_epoch, 2) class TestEMATraining(unittest.TestCase): diff --git a/source/tests/pt_expt/model/test_zbl_bridging.py b/source/tests/pt_expt/model/test_zbl_bridging.py index 8151c88f7f..de266089ad 100644 --- a/source/tests/pt_expt/model/test_zbl_bridging.py +++ b/source/tests/pt_expt/model/test_zbl_bridging.py @@ -2,7 +2,7 @@ """pt_expt ZBL bridging as COMPOSITION (review 3638077323, redesigned). ``bridging_method: ZBL`` builds a linear composition -(``LinearEnergyModel`` over ``[learned, InterPotentialAtomicModel]`` with +(``LinearEnergyModel`` over ``[learned, InnerPotentialAtomicModel]`` with ``weights="sum"``); eager values still match pt's flag-architected ``SeZMModel`` bit-for-bit (identical math), pinned here as a value regression together with FD force, export/DeepEval e2e, training smoke, @@ -290,7 +290,7 @@ def test_construction_composes_and_keeps_spin(self) -> None: assert isinstance(model, NativeSpinEnergyModel) assert model.has_spin() is True kinds = [type(c).__name__ for c in model.atomic_model.models] - assert kinds[1] == "InterPotentialAtomicModel", kinds + assert kinds[1] == "InnerPotentialAtomicModel", kinds # bridging radii still reach the LEARNED child's descriptor assert float(model.atomic_model.models[0].descriptor.inner_clamp.r_inner) == 0.8 @@ -361,7 +361,7 @@ def test_native_spin_with_bridging_dpmodel() -> None: model = dp_get_model(cfg) assert isinstance(model, NativeSpinEnergyModelDP) kinds = [type(c).__name__ for c in model.atomic_model.models] - assert kinds[1] == "InterPotentialAtomicModel", kinds + assert kinds[1] == "InnerPotentialAtomicModel", kinds coord, atype, spin, box = _spin_system() out = model.call(coord.numpy(), atype.numpy(), spin.numpy(), box=box.numpy()) @@ -506,12 +506,12 @@ def test_training_smoke(self, tmp_path) -> None: os.chdir(old_cwd) -class TestInterPotentialChangeTypeMapPtExpt: +class TestInnerPotentialChangeTypeMapPtExpt: """pt_expt twin of the dpmodel ``change_type_map`` regression. Exercised through the REAL composition: ``LinearEnergyModel`` -> - ``LinearEnergyAtomicModel`` -> ``InterPotentialAtomicModel`` -> - ``InterPotential``. Inside a pt_expt module tree the element lookup is a + ``LinearEnergyAtomicModel`` -> ``InnerPotentialAtomicModel`` -> + ``InnerPotential``. Inside a pt_expt module tree the element lookup is a wrapped torch buffer, so the rebuild must land on the same device/namespace (review 3649295675) -- a numpy rebuild would desync the buffer or fail outright on CUDA.