Skip to content

build: enable PyTorch ops by default #5577

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
njzjz-bot wants to merge 1 commit into from
Closed

build: enable PyTorch ops by default #5577

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion backend/find_pytorch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ def find_pytorch() -> tuple[str | None, list[str]]:
list of str
TensorFlow requirement if not found. Empty if found.
"""
if os.environ.get("DP_ENABLE_PYTORCH", "0") == "0":
if os.environ.get("DP_ENABLE_PYTORCH", "1") == "0":
return None, []
requires = []
pt_spec = None
Expand Down
2 changes: 1 addition & 1 deletion backend/read_env.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ def get_argument_from_env() -> tuple[str, list, list, dict, str, str]:
cmake_args.append("-DENABLE_TENSORFLOW=OFF")
tf_version = None

if os.environ.get("DP_ENABLE_PYTORCH", "0") == "1":
if os.environ.get("DP_ENABLE_PYTORCH", "1") == "1":
pt_install_dir, _ = find_pytorch()
pt_version = get_pt_version(pt_install_dir)
cmake_args.extend(
Expand Down
4 changes: 2 additions & 2 deletions deepmd/pt/cxx_op.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,8 +83,8 @@ def load_library(module_name: str) -> bool:
f"incompatible and thus an error is raised when loading {module_name}. "
f"You need to install PyTorch {PT_VERSION}, or rebuild deepmd-kit "
f"against PyTorch {pt_py_version}.\nIf you are using a wheel from "
"PyPI, you may consider to install deepmd-kit execuating "
"`DP_ENABLE_PYTORCH=1 pip install deepmd-kit --no-binary deepmd-kit` "
"PyPI, you may consider to rebuild deepmd-kit from source by "
"executing `pip install deepmd-kit --no-binary deepmd-kit` "
"instead."
) from e
error_message = (
Expand Down
2 changes: 1 addition & 1 deletion doc/freeze/compress.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,6 @@ See the documentation of a specific descriptor to see whether it supports model

When compressing models in the PyTorch backend, the customized OP library for the Python interface must be installed when [freezing the model](../freeze/freeze.md).

The customized OP library for the Python interface can be installed by setting environment variable {envvar}`DP_ENABLE_PYTORCH` to `1` during [installation](../install/install-from-source.md).
The customized OP library for the Python interface is installed by default during [source-code installation](../install/install-from-source.md). Set environment variable {envvar}`DP_ENABLE_PYTORCH` to `0` during installation to disable it.

[^1]: This section is built upon Jinzhe Zeng, Duo Zhang, Denghui Lu, Pinghui Mo, Zeyu Li, Yixiao Chen, Marián Rynik, Li'ang Huang, Ziyao Li, Shaochen Shi, Yingze Wang, Haotian Ye, Ping Tuo, Jiabin Yang, Ye Ding, Yifan Li, Davide Tisi, Qiyu Zeng, Han Bao, Yu Xia, Jiameng Huang, Koki Muraoka, Yibo Wang, Junhan Chang, Fengbo Yuan, Sigbjørn Løland Bore, Chun Cai, Yinnian Lin, Bo Wang, Jiayan Xu, Jia-Xin Zhu, Chenxing Luo, Yuzhi Zhang, Rhys E. A. Goodall, Wenshuo Liang, Anurag Kumar Singh, Sikai Yao, Jingchao Zhang, Renata Wentzcovitch, Jiequn Han, Jie Liu, Weile Jia, Darrin M. York, Weinan E, Roberto Car, Linfeng Zhang, Han Wang, [J. Chem. Phys. 159, 054801 (2023)](https://doi.org/10.1063/5.0155600) licensed under a [Creative Commons Attribution (CC BY) license](http://creativecommons.org/licenses/by/4.0/).
6 changes: 3 additions & 3 deletions doc/install/install-from-source.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,7 @@ Note that TensorFlow may have specific requirements for the compiler version to

:::{tab-item} PyTorch {{ pytorch_icon }}

You can set the environment variable `export DP_ENABLE_PYTORCH=1` to enable customized C++ OPs in the PyTorch backend.
DeePMD-kit enables customized C++ OPs in the PyTorch backend by default. You can set the environment variable `export DP_ENABLE_PYTORCH=0` to disable them.
Note that PyTorch may have specific requirements for the compiler version to support the C++ standard version and [`_GLIBCXX_USE_CXX11_ABI`](https://gcc.gnu.org/onlinedocs/libstdc++/manual/using_dual_abi.html) used by PyTorch.
:::

Expand Down Expand Up @@ -196,9 +196,9 @@ The path to the ROCM toolkit directory. If `ROCM_ROOT` is not set, it will look

:::{envvar} DP_ENABLE_PYTORCH

**Choices**: `0`, `1`; **Default**: `0`
**Choices**: `0`, `1`; **Default**: `1`

{{ pytorch_icon }} Enable customized C++ OPs for the PyTorch backend. PyTorch can still run without customized C++ OPs, but features will be limited.
{{ pytorch_icon }} Enable customized C++ OPs for the PyTorch backend. Set this to `0` to disable customized C++ OPs. PyTorch can still run without customized C++ OPs, but features will be limited.
:::

:::{envvar} TENSORFLOW_ROOT
Expand Down
2 changes: 1 addition & 1 deletion doc/model/dpa2.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ DPA-2 is designed for multi-task pre-training with a shared descriptor and task-

If one wants to run the DPA-2 model on LAMMPS, the customized OP library for the Python interface must be installed when [freezing the model](../freeze/freeze.md).

The customized OP library for the Python interface can be installed by setting environment variable {envvar}`DP_ENABLE_PYTORCH` to `1` during installation.
The customized OP library for the Python interface is installed by default during source-code installation. Set environment variable {envvar}`DP_ENABLE_PYTORCH` to `0` during installation to disable it.

If one runs LAMMPS with MPI, the customized OP library for the C++ interface should be compiled against the same MPI library as the runtime MPI.
If one runs LAMMPS with MPI and CUDA devices, it is recommended to compile the customized OP library for the C++ interface with a [CUDA-Aware MPI](https://developer.nvidia.com/mpi-solutions-gpus) library and CUDA,
Expand Down
2 changes: 1 addition & 1 deletion doc/model/dpa3.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ To run the DPA3 model on LAMMPS via source code installation
the custom OP library for Python interface integration must be compiled and linked
during the [model freezing process](../freeze/freeze.md).

The customized OP library for the Python interface can be installed by setting environment variable {envvar}`DP_ENABLE_PYTORCH` to `1` during installation.
The customized OP library for the Python interface is installed by default during source-code installation. Set environment variable {envvar}`DP_ENABLE_PYTORCH` to `0` during installation to disable it.

If one runs LAMMPS with MPI, the customized OP library for the C++ interface should be compiled against the same MPI library as the runtime MPI.
If one runs LAMMPS with MPI and CUDA devices, it is recommended to compile the customized OP library for the C++ interface with a [CUDA-Aware MPI](https://developer.nvidia.com/mpi-solutions-gpus) library and CUDA,
Expand Down
Loading