diff --git a/README.md b/README.md index 72ebe81..2768bff 100644 --- a/README.md +++ b/README.md @@ -106,6 +106,36 @@ models include the extra communication artifact needed by LAMMPS/MPI inside the exported `.pt2` package, so the resulting file can be passed to LAMMPS in the same way as other DeePMD-kit frozen models. +#### Training MACE with `torch.compile` + +MACE training through the `pt_expt` backend can use DeePMD-kit's native +`torch.compile` path. Enable it in the `training` section: + +```json +"training": { + "enable_compile": true +} +``` + +Then run the regular exportable-backend training command: + +```sh +dp --pt-expt train input.json +``` + +On a single RTX 5090, using the water example data with `sel = "auto"`, +`hidden_irreps = "128x0e + 128x1o"`, `batch_size = 1`, and +`disp_freq = 100`, a 2000-step run was used to sample the steady-state training +speed: + +| mode | steady avg after step 200 | speedup | +| --------------- | ------------------------: | ------: | +| eager `pt_expt` | 0.2030 s/step | 1.00x | +| `torch.compile` | 0.1495 s/step | 1.36x | + +The first compiled step includes a one-time Inductor trace/compile cost +(97.17 s in this run), which is amortized over normal long training jobs. + ### Running LAMMPS + GNN models with period boundary conditions GNN models use message passing neural networks, diff --git a/deepmd_gnn/mace.py b/deepmd_gnn/mace.py index 363570b..48849ee 100644 --- a/deepmd_gnn/mace.py +++ b/deepmd_gnn/mace.py @@ -4,7 +4,7 @@ import importlib import json from copy import deepcopy -from typing import Any, Optional +from typing import Any, Optional, cast import e3nn import torch @@ -317,6 +317,50 @@ def _make_mace_network( return model +def _set_module_state_tensor( + root: torch.nn.Module, + name: str, + tensor: torch.Tensor, +) -> None: + if "." in name: + module_name, attr_name = name.rsplit(".", 1) + module = root.get_submodule(module_name) + else: + module = root + attr_name = name + setattr(module, attr_name, tensor) + + +def _link_module_state( + target: torch.nn.Module, + source: torch.nn.Module, +) -> None: + source_params = dict(source.named_parameters()) + target_params = dict(target.named_parameters()) + if set(source_params) != set(target_params): + msg = "scripted and raw MACE parameter names do not match" + raise RuntimeError(msg) + for name, param in source_params.items(): + _set_module_state_tensor(target, name, param) + + source_buffers = dict(source.named_buffers()) + target_buffers = dict(target.named_buffers()) + if set(source_buffers) != set(target_buffers): + msg = "scripted and raw MACE buffer names do not match" + raise RuntimeError(msg) + for name, buffer in source_buffers.items(): + _set_module_state_tensor(target, name, buffer) + + +@torch.jit.unused +def _is_make_fx_fake_tensor(value: Any) -> bool: # noqa: ANN401 + value_type = type(value) + return ( + value_type.__module__ == "torch._subclasses.fake_tensor" + and value_type.__name__ == "FakeTensor" + ) + + def _clear_export_guards_once(traced: torch.nn.Module) -> None: """Clear over-specialized guards from the next export of ``traced``.""" original_export = torch.export.export @@ -505,6 +549,94 @@ def descriptor(self) -> Any: # noqa: ANN401 """Provide descriptor metadata expected by deepmd-kit's pt2 export.""" return self + def get_descriptor(self) -> Any: # noqa: ANN401 + """Return descriptor metadata expected by DeePMD-kit trainers.""" + if not torch.jit.is_scripting(): + self._ensure_compile_trace_model() + return self.descriptor + + @torch.jit.unused + def _ensure_compile_trace_model(self) -> torch.nn.Module: + scripted_model_raw = self.model + if scripted_model_raw is None: + msg = "MACE network is not initialized" + raise RuntimeError(msg) + scripted_model = cast("torch.nn.Module", scripted_model_raw) + cached = cast( + "Optional[torch.nn.Module]", + self.__dict__.get("_compile_trace_model"), + ) + cached_source = self.__dict__.get("_compile_trace_model_source") + if cached is not None and cached_source is scripted_model: + cached.train(self.training) + return cached + + raw_model = _make_mace_network( + r_max=self.params["r_max"], + num_radial_basis=self.params["num_radial_basis"], + num_cutoff_basis=self.params["num_cutoff_basis"], + max_ell=self.params["max_ell"], + interaction=self.params["interaction"], + num_interactions=self.params["num_interactions"], + num_elements=self.ntypes, + hidden_irreps=self.params["hidden_irreps"], + atomic_numbers=self.atomic_numbers, + avg_num_neighbors=self.avg_num_neighbors, + pair_repulsion=self.params["pair_repulsion"], + distance_transform=self.params["distance_transform"], + correlation=self.params["correlation"], + gate=self.params["gate"], + MLP_irreps=self.params["MLP_irreps"], + std=self.params["std"], + radial_MLP=self.params["radial_MLP"], + radial_type=self.params["radial_type"], + script_model=False, + ) + raw_model.load_state_dict(scripted_model.state_dict()) + _link_module_state(raw_model, scripted_model) + raw_model.train(self.training) + self.__dict__["_compile_trace_model"] = raw_model + self.__dict__["_compile_trace_model_source"] = scripted_model + return raw_model + + @torch.jit.unused + def _forward_lower_common_compile_trace( + self, + nloc: int, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: Optional[torch.Tensor], + fparam: Optional[torch.Tensor], + aparam: Optional[torch.Tensor], + do_atomic_virial: bool, + comm_dict: Optional[dict[str, torch.Tensor]], + ) -> dict[str, torch.Tensor]: + compile_trace_model = self.__dict__.get("_compile_trace_model") + if compile_trace_model is None: + compile_trace_model = self._ensure_compile_trace_model() + + previous_model = self.model + previous_exportable_edge_index = self._use_exportable_edge_index + self.model = compile_trace_model + self._use_exportable_edge_index = True + try: + return self.forward_lower_common( + nloc, + extended_coord, + extended_atype, + nlist, + mapping, + None, + fparam, + aparam, + do_atomic_virial, + comm_dict, + ) + finally: + self.model = previous_model + self._use_exportable_edge_index = previous_exportable_edge_index + @property def observed_type(self) -> Optional[list[str]]: """Observed element types collected during statistics.""" @@ -815,18 +947,44 @@ def forward_lower( ) raise ValueError(msg) - model_ret = self.forward_lower_common( - nloc, - extended_coord, - extended_atype, - nlist, - mapping, - None, - fparam, - aparam, - do_atomic_virial, - comm_dict, - ) + if torch.jit.is_scripting(): + model_ret = self.forward_lower_common( + nloc, + extended_coord, + extended_atype, + nlist, + mapping, + None, + fparam, + aparam, + do_atomic_virial, + comm_dict, + ) + elif _is_make_fx_fake_tensor(extended_coord): + model_ret = self._forward_lower_common_compile_trace( + nloc, + extended_coord, + extended_atype, + nlist, + mapping, + fparam, + aparam, + do_atomic_virial, + comm_dict, + ) + else: + model_ret = self.forward_lower_common( + nloc, + extended_coord, + extended_atype, + nlist, + mapping, + None, + fparam, + aparam, + do_atomic_virial, + comm_dict, + ) model_predict = {} model_predict["atom_energy"] = model_ret["energy"] model_predict["energy"] = model_ret["energy_redu"]