[Feature Request] Support PyTorch-exportable (.pt2) models for Kokkos-accelerated LAMMPS model deviation

[hcustc]

Summary

Integrate DeePMD-kit's PyTorch-exportable graph-lowering workflow into DP-GEN by exporting trained .pth models to .pt2 and using them in Kokkos-accelerated LAMMPS model-deviation tasks.

Detailed Description

Background and motivation

Recent DeePMD-kit updates provide an optimized PyTorch/Kokkos inference path for LAMMPS. This allows DeePMD models to run through pair_style deepmd/kk and benefit from Kokkos GPU acceleration, including on recent NVIDIA GPU architectures.

Unlike the regular PyTorch inference path, the Kokkos path does not directly consume the .pth model produced by PyTorch training. It uses a PyTorch-exportable, graph-lowered .pt2 model generated from the trained .pth model through the DeePMD-kit export/conversion functionality.

The required workflow is therefore:

PyTorch training
    graph.000.pth
    graph.001.pth
    graph.002.pth
    graph.003.pth
          |
          | PyTorch-exportable graph lowering
          v
    graph.000.pt2
    graph.001.pt2
    graph.002.pt2
    graph.003.pt2
          |
          | LAMMPS pair_style deepmd/kk
          v
    Kokkos-accelerated model deviation

DP-GEN currently supports:

train_backend: pytorch

and consequently produces and manages .pth models. However, it does not currently provide the intermediate .pth to .pt2 graph-lowering stage.

It also does not automatically transfer the resulting .pt2 models to model-deviation task directories or generate the LAMMPS model list using the .pt2 files. As a result, the new DeePMD-kit PyTorch/Kokkos acceleration path cannot currently be used in a fully automatic DP-GEN active-learning workflow.

Manual intervention would otherwise be required after every training iteration to export the four models, replace the model files used by the model-deviation stage, and update the generated LAMMPS input.

Requested enhancement

It would be useful to separate the model format used for training from the model format used for model-deviation deployment.

For example, DP-GEN could provide an optional configuration such as:

train_backend: pytorch
model_devi_backend: pytorch-exportable

The exact option name is only a suggestion. An equivalent configurable post-training export hook or dedicated model-deployment stage would also solve the problem.

When this option is enabled, DP-GEN would perform the following operations after training and before creating or submitting the model-deviation tasks:

1. Preserve the original four .pth models for checkpointing, restart, model reuse, and subsequent training iterations.

2. Export each trained .pth model to a PyTorch-exportable, graph-lowered .pt2 model:

   graph.000.pth -> graph.000.pt2
   graph.001.pth -> graph.001.pt2
   graph.002.pth -> graph.002.pt2
   graph.003.pth -> graph.003.pt2
   

3. Check that all expected .pt2 models were generated successfully and stop with a clear error message if any export fails.

4. Forward the four .pt2 files to each corresponding model-deviation task directory.

5. Generate the LAMMPS model list using the .pt2 files so that the tasks can run with pair_style deepmd/kk.

6. Preserve the .pth models as the canonical training outputs while treating the .pt2 files as deployment artifacts used by the model-deviation stage.

The export command, execution environment, machine, and resources should preferably be configurable. PyTorch-exportable graph lowering may need to run in an environment compatible with the target DeePMD-kit, PyTorch, CUDA, GPU architecture, and LAMMPS/Kokkos installation. This environment may be different from the machine used for model training.

A possible workflow could therefore be represented as:

00.train
    |
    | produce four `.pth` models
    v
post-training export/deployment stage
    |
    | graph-lower four `.pth` models to `.pt2`
    | validate the exported models
    v
01.model_devi
    |
    | transfer four `.pt2` models
    | generate pair_style deepmd/kk model list
    v
Kokkos-accelerated LAMMPS exploration

Backward compatibility

The proposed functionality should be optional.

The existing TensorFlow .pb and regular PyTorch .pth workflows should remain unchanged unless a PyTorch-exportable model-deviation backend or export hook is explicitly enabled.

This feature should not simply replace .pth with .pt2 globally, because the .pth files are still needed for training continuation and model reuse. Training artifacts and model-deviation deployment artifacts should be handled separately.

Problem solved

This enhancement would make the following workflow fully automatic:

PyTorch training (.pth)
    -> PyTorch-exportable graph lowering (.pt2)
    -> LAMMPS model deviation with deepmd/kk

It would eliminate the need to pause and manually modify every DP-GEN iteration and would enable the Kokkos-accelerated DeePMD inference path to be used safely and reproducibly during DP-GEN exploration.

This is particularly important for model-deviation calculations because each exploration task must load all four independently trained models. The four .pt2 models must therefore be exported, transferred, and referenced together.

Related work

PyTorch training backend selection and .pth model handling were introduced in:

• add option to select backends TF/PT #1541

A related issue concerning PyTorch model handling is:

• [BUG] Issue with DP-GEN workflow for DPA-2 model using PyTorch backend #1654

That issue concerns selecting .pth instead of .pb when the PyTorch training backend is used. The present feature request is different: it concerns adding a deployment/export stage that graph-lowers the trained .pth models to .pt2 and uses the exported models specifically for Kokkos-accelerated LAMMPS model deviation.

I can help test the proposed workflow, provide example DP-GEN configurations and execution logs, and validate four-model .pth/.pt2 model-deviation runs. I am not currently planning to implement the DP-GEN changes myself, but I would be happy to participate in testing and validation.

Further Information, Files, and Links

No response

[SchrodingersCattt]

Implemented in #1926.

The proposed DP-GEN workflow is now available through an explicit deployment backend:

"train_backend": "pytorch",
"model_devi_backend": "pytorch-exportable",
"model_format": "pt2"

DP-GEN keeps the PyTorch training checkpoint (model.ckpt.pt) for restart and reuse, freezes it through the PyTorch-exportable backend, links the resulting graph.*.pt2 files into 01.model_devi, forwards all committee models, and writes those pt2 paths into the generated LAMMPS input.

PyTorch-exportable training is also supported directly, including the DPA4C-oriented pt2/compression workflow. All new behavior is opt-in; the TensorFlow/pb default is unchanged.

DP-GEN retains model.ckpt.pt for restart and reuse, while the exported .pt2 models are used for model-deviation tasks. Kokkos/LAMMPS launch options remain part of the machine environment rather than new DP-GEN parameters.