From 67520366cdeba1f31163dc1d660b795dcb667f1f Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 22 Jun 2026 22:29:46 +0800 Subject: [PATCH 1/3] Enable MACE pt_expt torch compile training --- README.md | 29 +++++++ deepmd_gnn/mace.py | 184 +++++++++++++++++++++++++++++++++++++++++---- 2 files changed, 200 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 72ebe81..9ef1322 100644 --- a/README.md +++ b/README.md @@ -106,6 +106,35 @@ 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 `tests/mace.json` with `batch_size = 1`, +`disp_freq = 100`, and `numb_steps = 2000`, the measured training timings were: + +| mode | total wall time | first compiled step | steady avg after step 200 | +| --------------- | --------------: | ------------------: | ------------------------: | +| eager `pt_expt` | 138.66 s | n/a | 0.0658 s/step | +| `torch.compile` | 126.64 s | 64.36 s | 0.0313 s/step | + +The compiled run was about 1.09x faster end-to-end for this 2000-step run and +about 2.1x faster after the one-time trace/compile cost. Short jobs may not +benefit because the first compiled step includes the Inductor compile time. + ### 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"] From 523d729b5dab7f57a21544632b3c8cc1881dd7f9 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 22 Jun 2026 22:43:45 +0800 Subject: [PATCH 2/3] Focus compile benchmark on steady-state speed --- README.md | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 9ef1322..b8ee9e4 100644 --- a/README.md +++ b/README.md @@ -123,17 +123,17 @@ Then run the regular exportable-backend training command: dp --pt-expt train input.json ``` -On a single RTX 5090, using `tests/mace.json` with `batch_size = 1`, -`disp_freq = 100`, and `numb_steps = 2000`, the measured training timings were: +On a single RTX 5090, using `tests/mace.json` with `batch_size = 1` and +`disp_freq = 100`, a 2000-step run was used to sample the steady-state training +speed: -| mode | total wall time | first compiled step | steady avg after step 200 | -| --------------- | --------------: | ------------------: | ------------------------: | -| eager `pt_expt` | 138.66 s | n/a | 0.0658 s/step | -| `torch.compile` | 126.64 s | 64.36 s | 0.0313 s/step | +| mode | steady avg after step 200 | speedup | +| --------------- | ------------------------: | ------: | +| eager `pt_expt` | 0.0658 s/step | 1.0x | +| `torch.compile` | 0.0313 s/step | 2.1x | -The compiled run was about 1.09x faster end-to-end for this 2000-step run and -about 2.1x faster after the one-time trace/compile cost. Short jobs may not -benefit because the first compiled step includes the Inductor compile time. +The first compiled step includes a one-time Inductor trace/compile cost +(64.36 s in this run), which is amortized over normal long training jobs. ### Running LAMMPS + GNN models with period boundary conditions From ff22c8688b1c38a17571fec0c299d4f8e17eb1e8 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 22 Jun 2026 23:04:59 +0800 Subject: [PATCH 3/3] Update compile benchmark for default MACE config --- README.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index b8ee9e4..2768bff 100644 --- a/README.md +++ b/README.md @@ -123,17 +123,18 @@ Then run the regular exportable-backend training command: dp --pt-expt train input.json ``` -On a single RTX 5090, using `tests/mace.json` with `batch_size = 1` and +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.0658 s/step | 1.0x | -| `torch.compile` | 0.0313 s/step | 2.1x | +| 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 -(64.36 s in this run), which is amortized over normal long training jobs. +(97.17 s in this run), which is amortized over normal long training jobs. ### Running LAMMPS + GNN models with period boundary conditions