pair_matgl usability and performance follow-up: style registration, GPU export, ghost-row trimming, docs

Follow-up to #815/#770. Building the plugin from a clean tree per the
README exposed independent defects (full analysis, measurements and job-
level verification in the PR description):

1. The documented CMake build never registers the pair styles -- the
   ML-MATGL/KOKKOS sources are not included; add the cmake includes so a
   first build works without the CI workaround.
2. The exported TorchScript model cannot run on a GPU:
   spherical_bessel_smooth materialises torch.arange / torch.tensor on the
   default device; pin them to the input's device.
3. Ghost rows are still embedded although #815's edge folding guarantees
   they are isolated nodes (no incoming message, zero force): size the node
   buffers to nlocal. Measured 3.00x model-node work at 32 A (1,248 local
   + 2,502 ghost), 1.82x at 64 A, 1.51x at 96 A; end-to-end 1.71x / 1.46x /
   1.38x at those sizes.
4. export_matgl_checkpoint.py imports modules that no longer exist; fix to
   the public entry points.
5. python_reference.py printed a stress whose sign convention invites
   misreading against LAMMPS thermo output; label it and print both.
6. README: document the cmake includes, the single-rank constraint and the
   GPU export step.

The CPU virial sign issue we also found is deliberately NOT touched here:
it is fixed by #825 with a finite-difference CI validation, and this PR's
description records our independent corroborating measurement.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: Longyun Shen

lammps/README.md

@@ -38,17 +38,38 @@ of which you'll need for `pair_coeff`.
 
 ### 2. Build LAMMPS with the package
 
-Drop the package into a stock LAMMPS source tree and configure:
+LAMMPS builds its style tables by scanning package directories and then
+generating `style_pair.h`, and the generation happens roughly two thirds of
+the way through `cmake/CMakeLists.txt` (`GenerateStyleHeaders(...)`, line 794
+in `stable_22Jul2025_update5`). **Appending an `include()` to the END of that
+file is therefore too late**: the sources compile and libtorch links, but the
+style never reaches `style_pair.h` and LAMMPS rejects it at run time with
+
+```
+ERROR: Unrecognized pair style 'matgl' (src/force.cpp:275)
+```
+
+Register it the way LAMMPS registers its own packages instead:
 
 ```bash
-# 1) Copy or symlink the source files.
+# 1) Copy or symlink the source files into the LAMMPS src tree.
 ln -s /path/to/matgl/lammps/src/ML-MATGL <lammps>/src/ML-MATGL
 
-# 2) Tell LAMMPS' CMake about the package.
+# 2) Add ML-MATGL to the package list, so LAMMPS' own per-package loop does
+#    RegisterStyles + target_sources + include dir at the right point.
+#    In <lammps>/cmake/CMakeLists.txt, inside set(STANDARD_PACKAGES ...):
+#        ML-IAP
+#      + ML-MATGL
+#        ML-PACE
+#    (`-D PKG_ML-MATGL=ON` then works like any other package flag.)
+
+# 3) Link libtorch. This snippet only does find_package(Torch) and the link;
+#    it must NOT also add the sources, or every file compiles twice under two
+#    paths and the link fails on duplicate symbols.
 echo 'include(/path/to/matgl/lammps/cmake/ML-MATGL.cmake)' \
     >> <lammps>/cmake/CMakeLists.txt
 
-# 3) Configure + build. Match libtorch's CXX11 ABI to LAMMPS'.
+# 4) Configure + build. Match libtorch's CXX11 ABI to LAMMPS'.
 cmake -B build -S <lammps>/cmake \
     -D PKG_ML-MATGL=ON \
     -D CMAKE_PREFIX_PATH=/path/to/libtorch \
@@ -57,13 +78,26 @@ cmake -B build -S <lammps>/cmake \
 cmake --build build -j 8
 ```
 
+Check the registration before running anything:
+
+```bash
+grep matgl build/styles/style_pair.h     # expect pair_matgl.h (and _kokkos.h)
+build/lmp -h | tr ' ' '\n' | grep '^matgl'
+```
+
 ### 2b. Build the Kokkos GPU variant
 
 To get the `matgl/kk` pair style, also enable Kokkos and append the
 matching snippet to LAMMPS' CMake. CUDA example for an Ampere card
 (A100/A30):
 
 ```bash
+# Put the Kokkos sources where the KOKKOS package looks for them: its
+# RegisterStylesExt(${KOKKOS_PKG_SOURCES_DIR} kokkos ...) scans
+# <lammps>/src/KOKKOS for *_kokkos.h style headers and picks up matgl/kk
+# automatically. A separate directory is not scanned.
+cp /path/to/matgl/lammps/src/KOKKOS/pair_matgl_kokkos.* <lammps>/src/KOKKOS/
+
 echo 'include(/path/to/matgl/lammps/cmake/ML-MATGL-KOKKOS.cmake)' \
     >> <lammps>/cmake/CMakeLists.txt
 
@@ -81,9 +115,21 @@ cmake --build build -j 8
 Run with:
 
 ```bash
-mpirun -n 1 build/lmp -k on g 1 -sf kk -in in.matgl_si
+mpirun -n 1 build/lmp -k on g 1 -sf kk -pk kokkos neigh half -in in.matgl_si
 ```
 
+**`neigh half` is required, not optional.** `pair_matgl` needs `newton on`
+(it folds periodic edges back onto local rows and needs ghost contributions),
+and LAMMPS refuses `newton on` together with the Kokkos default `neigh full`:
+
+```
+ERROR: Must use 'newton off' with KOKKOS package option 'neigh full'
+(src/KOKKOS/kokkos.cpp:693)
+```
+
+Equivalently, put `package kokkos neigh half` in the input deck before
+`atom_style`.
+
 `-sf kk` makes LAMMPS prefer Kokkos pair styles, so `pair_style matgl`
 in your input deck dispatches to `matgl/kk` automatically. If you'd
 rather force it explicitly, write `pair_style matgl/kk` instead.
@@ -94,11 +140,31 @@ this explicit.
 
 Tested with:
 
-- LibTorch 2.2.x – 2.5.x (CXX11 ABI, CPU build).
+- LibTorch 2.2.x – 2.7.x (CXX11 ABI). **The Kokkos variant needs a CUDA
+  build of libtorch**, not a CPU-only one: `pair_matgl_kokkos.cpp` selects
+  `torch::Device(torch::kCUDA, gpu)` and wraps Kokkos device buffers as
+  tensors without a copy, so a CPU-only libtorch silently runs the model on
+  the host. The CPU build is what the CI job uses for the serial style.
 - LAMMPS develop branch (Aug 2024 or newer for the `add_request` /
-  `REQ_GHOST` neighbor-list API).
+  `REQ_GHOST` neighbor-list API); verified on `stable_22Jul2025_update5`.
 - C++17, MPI optional.
 
+#### Troubleshooting: CUDA 12.9 and newer toolkits
+
+libtorch's bundled Caffe2 CMake config predates two changes and each aborts
+the generate step. Both are fixed by a small file injected before
+`find_package(Torch)`, e.g. via
+`-D CMAKE_PROJECT_lammps_INCLUDE=/path/to/fixups.cmake`:
+
+- CUDA 12.9 removed the nvToolsExt shared library, so `FindCUDAToolkit` no
+  longer defines `CUDA::nvToolsExt` while `Caffe2/public/cuda.cmake` still
+  links it into `torch::nvtoolsext`. Declare it as a header-only interface
+  target (the nvtx3 headers are still shipped):
+  `add_library(CUDA::nvToolsExt INTERFACE IMPORTED GLOBAL)`.
+- On a machine without MKL, Caffe2 leaves the literal
+  `MKL_INCLUDE_DIR-NOTFOUND` inside torch's `INTERFACE_INCLUDE_DIRECTORIES`.
+  Point `MKL_INCLUDE_DIR` at any existing directory.
+
 ## LAMMPS input syntax
 
 ```lammps

lammps/cmake/ML-MATGL-KOKKOS.cmake

@@ -1,43 +1,20 @@
-# ML-MATGL Kokkos variant — drop-in CMake snippet.
+# ML-MATGL Kokkos variant -- documentation / warning fragment.
 #
-# Layered on top of ML-MATGL.cmake: include() this *after* the base snippet,
-# OR set PKG_ML-MATGL=ON and PKG_KOKKOS=ON together.
+# The Kokkos sources are picked up by LAMMPS' own KOKKOS package machinery
+# (Packages/KOKKOS.cmake scans <lammps>/src/KOKKOS for *_kokkos.* styles and
+# registers them via RegisterStylesExt), so this snippet adds NO sources:
 #
-# Usage (from a stock LAMMPS source tree):
-#   cmake -B build \
-#       -D PKG_ML-MATGL=ON -D PKG_KOKKOS=ON \
-#       -D Kokkos_ENABLE_CUDA=ON \
-#       -D Kokkos_ARCH_AMPERE80=ON \
-#       -D CMAKE_PREFIX_PATH=/path/to/libtorch \
-#       -D CMAKE_CXX_COMPILER=$LAMMPS/lib/kokkos/bin/nvcc_wrapper \
-#       <other flags>
+#   cp /path/to/matgl/lammps/src/KOKKOS/pair_matgl_kokkos.* <lammps>/src/KOKKOS/
 #
-# The `pair_matgl/kk` style is registered via the standard LAMMPS Kokkos
-# pair-style macro so users invoke it with `pair_style matgl/kk` or by
-# launching LAMMPS with `-sf kk -k on g 1`.
+# and configure with -D PKG_ML-MATGL=ON -D PKG_KOKKOS=ON (plus the usual
+# Kokkos CUDA flags). libtorch linkage comes from ML-MATGL.cmake.
 
 if(NOT PKG_ML-MATGL OR NOT PKG_KOKKOS)
     return()
 endif()
 
-if(NOT DEFINED ML_MATGL_KOKKOS_DIR)
-    get_filename_component(ML_MATGL_KOKKOS_DIR
-        "${CMAKE_CURRENT_LIST_DIR}/../src/KOKKOS" ABSOLUTE)
-endif()
-
-if(NOT EXISTS "${ML_MATGL_KOKKOS_DIR}/pair_matgl_kokkos.cpp")
-    message(FATAL_ERROR
-        "ML-MATGL-KOKKOS source not found at ${ML_MATGL_KOKKOS_DIR}. "
-        "Set -DML_MATGL_KOKKOS_DIR=<path/to/lammps/src/KOKKOS>.")
-endif()
-
-file(GLOB ML_MATGL_KOKKOS_SOURCES "${ML_MATGL_KOKKOS_DIR}/*.cpp")
-
-target_sources(lammps PRIVATE ${ML_MATGL_KOKKOS_SOURCES})
-target_include_directories(lammps PRIVATE ${ML_MATGL_KOKKOS_DIR})
-
 # Single-GPU only: warn loudly. MACE upstream issues #1294 and #322 cover
 # the multi-rank-with-libtorch breakage we inherit.
 message(STATUS
-    "ML-MATGL-KOKKOS: enabled. Single-GPU runs only — multi-rank Kokkos with "
+    "ML-MATGL-KOKKOS: enabled. Single-GPU runs only -- multi-rank Kokkos with "
     "libtorch is unreliable (see MACE issues #1294, #322).")

lammps/cmake/ML-MATGL.cmake

@@ -1,41 +1,27 @@
-# ML-MATGL package — drop-in CMake snippet for a stock LAMMPS source tree.
+# ML-MATGL package -- libtorch link fragment for a stock LAMMPS source tree.
 #
-# Usage (from a stock LAMMPS source tree):
-#   1. Copy or symlink lammps/src/ML-MATGL  →  <lammps>/src/ML-MATGL
-#   2. Append to <lammps>/cmake/CMakeLists.txt (anywhere after the `set(STANDARD_PACKAGES …)`
-#      block):
-#        include(/path/to/matgl/lammps/cmake/ML-MATGL.cmake)
-#   3. Configure with:
-#        cmake -B build \
-#            -D PKG_ML-MATGL=ON \
-#            -D CMAKE_PREFIX_PATH=/path/to/libtorch \
-#            -D CMAKE_BUILD_TYPE=Release \
-#            <other flags>
+# This snippet ONLY locates libtorch and links it into the `lammps` target.
+# It must NOT add the pair-style sources: LAMMPS registers styles by scanning
+# package directories and generating style_pair.h roughly two thirds of the
+# way through cmake/CMakeLists.txt, so sources added from an appended
+# include() are compiled but never registered -- and adding them here while
+# the package loop also adds them compiles every file twice and breaks the
+# link. See lammps/README.md, "Build LAMMPS with the package":
 #
-# CMake variables consumed:
-#   PKG_ML-MATGL           - turn the package on/off (default OFF).
-#   CMAKE_PREFIX_PATH      - must point at a libtorch install (CXX11 ABI build).
-#   ML_MATGL_DIR           - override path to lammps/src/ML-MATGL (defaults to
-#                            ${CMAKE_CURRENT_LIST_DIR}/../src/ML-MATGL).
+#   1. Copy or symlink lammps/src/ML-MATGL  ->  <lammps>/src/ML-MATGL
+#   2. Add ML-MATGL to set(STANDARD_PACKAGES ...) in
+#      <lammps>/cmake/CMakeLists.txt, so the per-package loop does
+#      RegisterStyles + target_sources at the right point.
+#   3. Append to <lammps>/cmake/CMakeLists.txt:
+#        include(/path/to/matgl/lammps/cmake/ML-MATGL.cmake)
+#   4. Configure with -D PKG_ML-MATGL=ON and CMAKE_PREFIX_PATH at libtorch.
 
 option(PKG_ML-MATGL "Build the matgl pair_style backed by libtorch" OFF)
 
 if(NOT PKG_ML-MATGL)
     return()
 endif()
 
-# Locate the source directory.
-if(NOT DEFINED ML_MATGL_DIR)
-    get_filename_component(ML_MATGL_DIR
-        "${CMAKE_CURRENT_LIST_DIR}/../src/ML-MATGL" ABSOLUTE)
-endif()
-
-if(NOT EXISTS "${ML_MATGL_DIR}/pair_matgl.cpp")
-    message(FATAL_ERROR
-        "ML-MATGL source not found at ${ML_MATGL_DIR}. "
-        "Set -DML_MATGL_DIR=<path/to/lammps/src/ML-MATGL>.")
-endif()
-
 # Pull in libtorch.
 find_package(Torch REQUIRED)
 if(NOT TORCH_LIBRARIES)
@@ -44,24 +30,16 @@ if(NOT TORCH_LIBRARIES)
         "Did you set CMAKE_PREFIX_PATH to a libtorch install?")
 endif()
 
-# Compose the source list.
-file(GLOB ML_MATGL_SOURCES "${ML_MATGL_DIR}/*.cpp")
-
-# Hook into the LAMMPS build. This file is included from
-# <lammps>/cmake/CMakeLists.txt; the `lammps` target already exists by then.
-target_sources(lammps PRIVATE ${ML_MATGL_SOURCES})
-target_include_directories(lammps PRIVATE ${ML_MATGL_DIR})
 target_compile_features(lammps PRIVATE cxx_std_17)
 target_link_libraries(lammps PRIVATE ${TORCH_LIBRARIES})
 
 # Make sure libtorch's headers come ahead of any system Eigen/torch shims.
 target_include_directories(lammps PRIVATE ${TORCH_INCLUDE_DIRS})
 
-# LibTorch ships with -D_GLIBCXX_USE_CXX11_ABI=…; propagate it so consumers
-# (e.g. KOKKOS in Phase 3) see the same ABI.
+# LibTorch ships with -D_GLIBCXX_USE_CXX11_ABI=...; propagate it so every
+# consumer (including the KOKKOS sources) sees the same ABI.
 if(DEFINED TORCH_CXX_FLAGS)
     set_property(TARGET lammps APPEND_STRING PROPERTY COMPILE_FLAGS " ${TORCH_CXX_FLAGS}")
 endif()
 
-message(STATUS "ML-MATGL: enabled, sources from ${ML_MATGL_DIR}")
-message(STATUS "ML-MATGL: linking against TORCH_LIBRARIES=${TORCH_LIBRARIES}")
+message(STATUS "ML-MATGL: libtorch linked, TORCH_LIBRARIES=${TORCH_LIBRARIES}")

lammps/src/KOKKOS/pair_matgl_kokkos.cpp

@@ -175,11 +175,11 @@ void PairMATGLKokkos<DeviceType>::compute(int eflag, int vflag)
   const auto type_to_z = d_type_to_z_;
   Kokkos::parallel_for(
       "matgl_kk:fill_atoms",
-      Kokkos::RangePolicy<DeviceType>(0, nall),
+      Kokkos::RangePolicy<DeviceType>(0, nlocal),
       KOKKOS_LAMBDA(const int i) {
         const int t = type(i);
         d_atomic_numbers_(i) = type_to_z(t);
-        d_local_or_ghost_(i) = (i < nlocal);
+        d_local_or_ghost_(i) = true;   // only owned rows are built now
       });
 
   // local_row_of_(j) = the owned row representing the same physical atom as
@@ -306,8 +306,12 @@ void PairMATGLKokkos<DeviceType>::compute(int eflag, int vflag)
 
   // x is (nall,3) double already on `DeviceType`. We make a libtorch view
   // through from_blob and cast to the model's dtype if needed.
+  // Local atoms occupy rows [0, nlocal) of LAMMPS' x, so narrowing to
+  // nlocal keeps this a zero-copy view while dropping the isolated ghost
+  // rows the model would otherwise embed (see pair_matgl.cpp for why they
+  // are isolated).
   torch::Tensor positions_d = torch::from_blob(
-      x.data(), {nall, 3}, torch::TensorOptions().dtype(torch::kFloat64).device(torch_device_));
+      x.data(), {nlocal, 3}, torch::TensorOptions().dtype(torch::kFloat64).device(torch_device_));
   torch::Tensor positions = (dtype_ == torch::kFloat64)
                                 ? positions_d.clone()
                                 : positions_d.to(dtype_);
@@ -320,9 +324,9 @@ void PairMATGLKokkos<DeviceType>::compute(int eflag, int vflag)
   unit_shifts = unit_shifts.narrow(0, 0, total_edges);
 
   torch::Tensor atomic_numbers = blob_from_view(d_atomic_numbers_, torch_long_opts);
-  atomic_numbers = atomic_numbers.narrow(0, 0, nall);
+  atomic_numbers = atomic_numbers.narrow(0, 0, nlocal);
   torch::Tensor local_or_ghost = blob_from_view(d_local_or_ghost_, torch_bool_opts);
-  local_or_ghost = local_or_ghost.narrow(0, 0, nall);
+  local_or_ghost = local_or_ghost.narrow(0, 0, nlocal);
 
   // 6) Cell.
   torch::Tensor cell = torch::zeros({3, 3}, torch_real_opts);
@@ -377,11 +381,11 @@ void PairMATGLKokkos<DeviceType>::compute(int eflag, int vflag)
   // into LAMMPS' f.
   using UnmanagedF = Kokkos::View<double **, Kokkos::LayoutRight, DeviceType,
                                   Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
-  UnmanagedF d_force_in(forces_t.data_ptr<double>(), nall, 3);
+  UnmanagedF d_force_in(forces_t.data_ptr<double>(), nlocal, 3);
 
   Kokkos::parallel_for(
       "matgl_kk:add_forces",
-      Kokkos::RangePolicy<DeviceType>(0, nall),
+      Kokkos::RangePolicy<DeviceType>(0, nlocal),   // ghosts carried zero
       KOKKOS_LAMBDA(const int i) {
         f(i, 0) += d_force_in(i, 0);
         f(i, 1) += d_force_in(i, 1);

lammps/src/ML-MATGL/export_matgl_checkpoint.py

@@ -15,16 +15,14 @@
 
 import matgl
 
-# matgl/src/matgl/ext/_lammps.py imports create_line_graph_torch, which is
-# only used by the M3GNet export path and is absent from the current
-# _compute_pyg.py. Irrelevant to TensorNet exports; shim it in-memory only
-# (no matgl repo file changes) so the import below succeeds.
-import matgl.graph._compute_pyg as _cpyg
-
-if not hasattr(_cpyg, "create_line_graph_torch"):
-    _cpyg.create_line_graph_torch = _cpyg.create_line_graph
-
-from matgl.ext._lammps import export_lammps_model
+# Module layout as of matgl 4.0.3: the helpers live in matgl.graph._compute
+# and the exporter in matgl.ext.lammps. The previous names
+# (matgl.graph._compute_pyg, matgl.ext._lammps) predate that layout; the first
+# no longer exists, so the import raised ModuleNotFoundError and made the
+# "point pair_coeff at a checkpoint directory" path unusable.
+# _compute exports create_line_graph_torch directly,
+# so the in-memory shim that used to sit here is no longer needed either.
+from matgl.ext.lammps import export_lammps_model
 
 checkpoint_dir, out_path = sys.argv[1], sys.argv[2]
 potential = matgl.load_model(checkpoint_dir)