Skip to content

Commit 152cf8f

Browse files
OutisLiOutisLi
committed
fix(dpa1): harden higher-order moment support
1 parent 17309f4 commit 152cf8f

15 files changed

Lines changed: 508 additions & 269 deletions

File tree

‎deepmd/dpmodel/descriptor/dpa1.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -142,13 +142,13 @@ def build_dpa1_moment_basis(
142142
safe_distance = xp.sqrt(
143143
xp.where(direction_mask, distance_squared, xp.ones_like(distance_squared))
144144
)
145-
direction = diff / safe_distance
146145
basis_mask = valid_mask[..., None] & direction_mask
147146
denominator = xp.where(
148147
basis_mask,
149148
safe_distance + protection,
150149
xp.ones_like(safe_distance),
151150
)
151+
direction = diff / denominator * xp.astype(basis_mask, diff.dtype)
152152
radial = switch / denominator / radial_stddev * xp.astype(basis_mask, switch.dtype)
153153

154154
x, y, z = direction[..., 0], direction[..., 1], direction[..., 2]
@@ -1102,7 +1102,7 @@ def serialize(self) -> dict:
11021102
data = {
11031103
"@class": "Descriptor",
11041104
"type": "dpa1",
1105-
"@version": 3 if self.compress else 2,
1105+
"@version": 4 if obj.lmax != 1 else (3 if self.compress else 2),
11061106
"rcut": obj.rcut,
11071107
"rcut_smth": obj.rcut_smth,
11081108
"sel": obj.sel,
@@ -1188,7 +1188,7 @@ def serialize(self) -> dict:
11881188
def deserialize(cls, data: dict) -> "DescrptDPA1":
11891189
"""Deserialize from dict."""
11901190
data = data.copy()
1191-
check_version_compatibility(data.pop("@version"), 3, 1)
1191+
check_version_compatibility(data.pop("@version"), 4, 1)
11921192
data.pop("@class")
11931193
data.pop("type")
11941194
variables = data.pop("@variables")
@@ -1205,6 +1205,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA1":
12051205
# compat with version 1
12061206
if "use_tebd_bias" not in data:
12071207
data["use_tebd_bias"] = True
1208+
data.setdefault("lmax", 1)
12081209
obj = cls(**data)
12091210

12101211
obj.se_atten["davg"] = variables["davg"]
@@ -2332,7 +2333,7 @@ def serialize(self) -> dict:
23322333
data = {
23332334
"@class": "DescriptorBlock",
23342335
"type": "dpa1",
2335-
"@version": 1,
2336+
"@version": 2 if obj.lmax != 1 else 1,
23362337
"rcut": obj.rcut,
23372338
"rcut_smth": obj.rcut_smth,
23382339
"sel": obj.sel,
@@ -2381,7 +2382,7 @@ def serialize(self) -> dict:
23812382
def deserialize(cls, data: dict) -> "DescrptDPA1":
23822383
"""Deserialize from dict."""
23832384
data = data.copy()
2384-
check_version_compatibility(data.pop("@version"), 1, 1)
2385+
check_version_compatibility(data.pop("@version"), 2, 1)
23852386
data.pop("@class")
23862387
data.pop("type")
23872388
variables = data.pop("@variables")
@@ -2393,6 +2394,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA1":
23932394
embeddings_strip = data.pop("embeddings_strip")
23942395
else:
23952396
embeddings_strip = None
2397+
data.setdefault("lmax", 1)
23962398
obj = cls(**data)
23972399

23982400
obj["davg"] = variables["davg"]

‎deepmd/dpmodel/descriptor/se_atten_v2.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -209,7 +209,7 @@ def serialize(self) -> dict:
209209
data = {
210210
"@class": "Descriptor",
211211
"type": "se_atten_v2",
212-
"@version": 3 if self.compress else 2,
212+
"@version": 4 if obj.lmax != 1 else (3 if self.compress else 2),
213213
"rcut": obj.rcut,
214214
"rcut_smth": obj.rcut_smth,
215215
"sel": obj.sel,
@@ -292,7 +292,7 @@ def serialize(self) -> dict:
292292
def deserialize(cls, data: dict) -> "DescrptSeAttenV2":
293293
"""Deserialize from dict."""
294294
data = data.copy()
295-
check_version_compatibility(data.pop("@version"), 3, 1)
295+
check_version_compatibility(data.pop("@version"), 4, 1)
296296
data.pop("@class")
297297
data.pop("type")
298298
variables = data.pop("@variables")
@@ -305,6 +305,7 @@ def deserialize(cls, data: dict) -> "DescrptSeAttenV2":
305305
# compat with version 1
306306
if "use_tebd_bias" not in data:
307307
data["use_tebd_bias"] = True
308+
data.setdefault("lmax", 1)
308309
obj = cls(**data)
309310

310311
obj.se_atten["davg"] = variables["davg"]

‎deepmd/pt/model/descriptor/dpa1.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -512,7 +512,7 @@ def serialize(self) -> dict:
512512
data = {
513513
"@class": "Descriptor",
514514
"type": "dpa1",
515-
"@version": 2,
515+
"@version": 4 if obj.lmax != 1 else 2,
516516
"rcut": obj.rcut,
517517
"rcut_smth": obj.rcut_smth,
518518
"sel": obj.sel,
@@ -569,7 +569,7 @@ def serialize(self) -> dict:
569569
@classmethod
570570
def deserialize(cls, data: dict) -> "DescrptDPA1":
571571
data = data.copy()
572-
check_version_compatibility(data.pop("@version"), 3, 1)
572+
check_version_compatibility(data.pop("@version"), 4, 1)
573573
data.pop("@class")
574574
data.pop("type")
575575
variables = data.pop("@variables")
@@ -586,6 +586,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA1":
586586
# compat with version 1
587587
if "use_tebd_bias" not in data:
588588
data["use_tebd_bias"] = True
589+
data.setdefault("lmax", 1)
589590
obj = cls(**data)
590591

591592
def t_cvt(xx: Any) -> torch.Tensor:

‎deepmd/pt/model/descriptor/se_atten.py‎

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -90,8 +90,26 @@ def tabulate_fusion_se_atten(
9090

9191
def _safe_direction(
9292
diff: torch.Tensor,
93+
protection: float,
9394
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
94-
"""Normalize displacements with finite first and second derivatives at zero."""
95+
"""Scale displacements by the protected distance denominator.
96+
97+
Parameters
98+
----------
99+
diff
100+
Neighbor displacement vectors with shape ``(..., 3)``.
101+
protection
102+
Distance protection added to the denominator.
103+
104+
Returns
105+
-------
106+
torch.Tensor
107+
Protected displacement coordinates with shape ``(..., 3)``.
108+
torch.Tensor
109+
Unprotected distances with shape ``(..., 1)``.
110+
torch.Tensor
111+
Nonzero-distance mask with shape ``(..., 1)``.
112+
"""
95113
distance_squared = torch.sum(diff * diff, dim=-1, keepdim=True)
96114
direction_mask = distance_squared > 0.0
97115
safe_distance = torch.sqrt(
@@ -101,7 +119,12 @@ def _safe_direction(
101119
torch.ones_like(distance_squared),
102120
)
103121
)
104-
return diff / safe_distance, safe_distance, direction_mask
122+
denominator = torch.where(
123+
direction_mask,
124+
safe_distance + protection,
125+
torch.ones_like(safe_distance),
126+
)
127+
return diff / denominator, safe_distance, direction_mask
105128

106129

107130
def _compute_angular_radial(
@@ -147,7 +170,7 @@ def _build_moment_basis(
147170
rr
148171
Normalized environment matrix with shape ``(ncenter, nnei, 4)``.
149172
direction
150-
Safely normalized neighbor directions with shape
173+
Protected displacement coordinates with shape
151174
``(ncenter, nnei, 3)``.
152175
radial
153176
Zero-mean normalized radial amplitude with shape
@@ -782,7 +805,10 @@ def forward(
782805
moment_radial = rr[..., :1]
783806
direction = diff_flat
784807
if self.lmax > 1:
785-
direction, distance, direction_mask = _safe_direction(diff_flat)
808+
direction, distance, direction_mask = _safe_direction(
809+
diff_flat,
810+
self.env_protection,
811+
)
786812
radial_stddev = self.stddev[atype][..., :1].view(nfnl, nnei, 1)
787813
moment_radial = _compute_angular_radial(
788814
distance,

‎deepmd/pt/model/descriptor/se_atten_v2.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -201,7 +201,7 @@ def serialize(self) -> dict:
201201
data = {
202202
"@class": "Descriptor",
203203
"type": "se_atten_v2",
204-
"@version": 2,
204+
"@version": 4 if obj.lmax != 1 else 2,
205205
"rcut": obj.rcut,
206206
"rcut_smth": obj.rcut_smth,
207207
"sel": obj.sel,
@@ -253,7 +253,7 @@ def serialize(self) -> dict:
253253
@classmethod
254254
def deserialize(cls, data: dict) -> "DescrptSeAttenV2":
255255
data = data.copy()
256-
check_version_compatibility(data.pop("@version"), 3, 1)
256+
check_version_compatibility(data.pop("@version"), 4, 1)
257257
data.pop("@class")
258258
data.pop("type")
259259
variables = data.pop("@variables")
@@ -266,6 +266,7 @@ def deserialize(cls, data: dict) -> "DescrptSeAttenV2":
266266
# compat with version 1
267267
if "use_tebd_bias" not in data:
268268
data["use_tebd_bias"] = True
269+
data.setdefault("lmax", 1)
269270
obj = cls(**data)
270271

271272
def t_cvt(xx: Any) -> torch.Tensor:

‎source/lib/include/tabulate.h‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22
#pragma once
33

44
#include <cstdint>
5+
#include <string>
6+
7+
#include "errors.h"
58

69
namespace deepmd {
710

@@ -10,6 +13,18 @@ inline bool is_supported_se_a_basis_dimension(
1013
return ndescrpt == 4 || ndescrpt == 9 || ndescrpt == 16 || ndescrpt == 25;
1114
}
1215

16+
namespace detail {
17+
18+
inline void check_se_a_basis_dimension(const std::int64_t ndescrpt) {
19+
if (!is_supported_se_a_basis_dimension(ndescrpt)) {
20+
throw deepmd_exception(
21+
"The environment basis dimension must be 4, 9, 16, or 25, got " +
22+
std::to_string(ndescrpt));
23+
}
24+
}
25+
26+
} // namespace detail
27+
1328
template <typename FPTYPE>
1429
void tabulate_fusion_se_a_cpu(FPTYPE* out,
1530
const FPTYPE* table,

‎source/lib/src/gpu/tabulate.cu‎

Lines changed: 41 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
#include <math.h>
22

3-
#include <cassert>
3+
#if GOOGLE_CUDA
4+
#include <mutex>
5+
#include <unordered_map>
6+
#endif
47

58
#include "device.h"
69
#include "tabulate.h"
@@ -1086,6 +1089,36 @@ void launch_tabulate_fusion_se_a(FPTYPE* out,
10861089
is_sorted);
10871090
}
10881091

1092+
#if GOOGLE_CUDA
1093+
namespace {
1094+
1095+
struct CudaSharedMemoryLimits {
1096+
size_t standard;
1097+
size_t opt_in;
1098+
};
1099+
1100+
CudaSharedMemoryLimits get_cuda_shared_memory_limits(const int device) {
1101+
static std::mutex cache_mutex;
1102+
static std::unordered_map<int, CudaSharedMemoryLimits> cache;
1103+
std::lock_guard<std::mutex> lock(cache_mutex);
1104+
const auto cached = cache.find(device);
1105+
if (cached != cache.end()) {
1106+
return cached->second;
1107+
}
1108+
1109+
cudaDeviceProp properties;
1110+
DPErrcheck(cudaGetDeviceProperties(&properties, device));
1111+
const CudaSharedMemoryLimits limits{
1112+
properties.sharedMemPerBlock,
1113+
properties.sharedMemPerBlockOptin,
1114+
};
1115+
cache.emplace(device, limits);
1116+
return limits;
1117+
}
1118+
1119+
} // namespace
1120+
#endif
1121+
10891122
template <typename FPTYPE, int MTILE>
10901123
void launch_tabulate_fusion_se_a_grad(FPTYPE* dy_dem_x,
10911124
FPTYPE* dy_dem,
@@ -1103,21 +1136,18 @@ void launch_tabulate_fusion_se_a_grad(FPTYPE* dy_dem_x,
11031136
#if GOOGLE_CUDA
11041137
const size_t shared_memory = sizeof(FPTYPE) * MTILE * last_layer_size;
11051138
int device = 0;
1106-
cudaDeviceProp properties;
11071139
DPErrcheck(cudaGetDevice(&device));
1108-
DPErrcheck(cudaGetDeviceProperties(&properties, device));
1140+
const CudaSharedMemoryLimits limits = get_cuda_shared_memory_limits(device);
11091141
const size_t shared_memory_limit =
1110-
properties.sharedMemPerBlock > properties.sharedMemPerBlockOptin
1111-
? properties.sharedMemPerBlock
1112-
: properties.sharedMemPerBlockOptin;
1142+
limits.standard > limits.opt_in ? limits.standard : limits.opt_in;
11131143
if (shared_memory <= shared_memory_limit) {
11141144
auto kernel =
11151145
tabulate_fusion_se_a_grad_fifth_order_polynomial<FPTYPE, MTILE, KK,
11161146
true>;
1117-
if (shared_memory > properties.sharedMemPerBlock) {
1147+
if (shared_memory > limits.standard) {
11181148
DPErrcheck(cudaFuncSetAttribute(
11191149
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
1120-
static_cast<int>(properties.sharedMemPerBlockOptin)));
1150+
static_cast<int>(limits.opt_in)));
11211151
}
11221152
kernel<<<nloc, KK * WARP_SIZE, shared_memory>>>(
11231153
dy_dem_x, dy_dem, dy_dtwo, table, em_x, em, two_embed, dy,
@@ -1181,10 +1211,10 @@ void tabulate_fusion_se_a_gpu(FPTYPE* out,
11811211
const int last_layer_size,
11821212
const bool is_sorted,
11831213
const int ndescrpt) {
1214+
detail::check_se_a_basis_dimension(ndescrpt);
11841215
if (nloc <= 0) {
11851216
return;
11861217
}
1187-
assert(ndescrpt == 4 || ndescrpt == 9 || ndescrpt == 16 || ndescrpt == 25);
11881218
DPErrcheck(gpuGetLastError());
11891219
DPErrcheck(gpuDeviceSynchronize());
11901220
if (ndescrpt == 4) {
@@ -1223,10 +1253,10 @@ void tabulate_fusion_se_a_grad_gpu(FPTYPE* dy_dem_x,
12231253
const int last_layer_size,
12241254
const bool is_sorted,
12251255
const int ndescrpt) {
1256+
detail::check_se_a_basis_dimension(ndescrpt);
12261257
if (nloc <= 0) {
12271258
return;
12281259
}
1229-
assert(ndescrpt == 4 || ndescrpt == 9 || ndescrpt == 16 || ndescrpt == 25);
12301260
DPErrcheck(gpuGetLastError());
12311261
DPErrcheck(gpuDeviceSynchronize());
12321262
DPErrcheck(gpuMemset(dy_dem_x, 0, sizeof(FPTYPE) * nloc * nnei));
@@ -1268,10 +1298,10 @@ void tabulate_fusion_se_a_grad_grad_gpu(FPTYPE* dz_dy,
12681298
const int last_layer_size,
12691299
const bool is_sorted,
12701300
const int ndescrpt) {
1301+
detail::check_se_a_basis_dimension(ndescrpt);
12711302
if (nloc <= 0) {
12721303
return;
12731304
}
1274-
assert(ndescrpt == 4 || ndescrpt == 9 || ndescrpt == 16 || ndescrpt == 25);
12751305
DPErrcheck(gpuGetLastError());
12761306
DPErrcheck(gpuDeviceSynchronize());
12771307
DPErrcheck(

0 commit comments

Comments
 (0)