Skip to content

fix(tf): validate GPU tabulation sizes #5846

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

Open
njzjz-bot wants to merge 1 commit into
base: master
Choose a base branch
from
Open

fix(tf): validate GPU tabulation sizes #5846

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
39 changes: 23 additions & 16 deletions source/op/tf/tabulate_multi_device.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,17 @@ REGISTER_OP("TabulateFusionSeRGradGrad")
.Input("descriptor: T")
.Output("dz_dy: T");

static deepmd::tf_compat::Status validate_gpu_last_layer_size(
const int last_layer_size) {
// GPU tabulation kernels use this dimension either as the block size or to
// size dynamic shared memory, so reject invalid models before any launch.
if (last_layer_size <= 0 || last_layer_size > 1024) {
return deepmd::tf_compat::InvalidArgument(
"last_layer_size must be between 1 and 1024 for GPU tabulation");
}
return deepmd::tf_compat::Status();
}

template <typename Device, typename FPTYPE>
class TabulateFusionSeAOp : public OpKernel {
public:
Expand Down Expand Up @@ -210,6 +221,7 @@ class TabulateFusionSeAOp : public OpKernel {
const int nnei = em_tensor.shape().dim_size(1);

if (device == "GPU") {
OP_REQUIRES_OK(context, validate_gpu_last_layer_size(last_layer_size));
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_a_gpu(descriptor, table, table_info, em_x, em,
two_embed, nloc, nnei, last_layer_size);
Expand Down Expand Up @@ -276,6 +288,7 @@ class TabulateFusionSeAGradOp : public OpKernel {
const int last_layer_size = descriptor_tensor.shape().dim_size(2);

if (device == "GPU") {
OP_REQUIRES_OK(context, validate_gpu_last_layer_size(last_layer_size));
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_a_grad_gpu(dy_dem_x, dy_dem, dy_dtwo, table,
table_info, em_x, em, two_embed, dy,
Expand Down Expand Up @@ -336,15 +349,12 @@ class TabulateFusionSeAGradGradOp : public OpKernel {
const int last_layer_size = descriptor_tensor.shape().dim_size(2);

if (device == "GPU") {
OP_REQUIRES_OK(context, validate_gpu_last_layer_size(last_layer_size));
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_a_grad_grad_gpu(
dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted);
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
OP_REQUIRES(context, (last_layer_size <= 1024),
deepmd::tf_compat::InvalidArgument(
"In the process of model compression, the size of the "
"last layer of embedding net must be less than 1024!"));
} else if (device == "CPU") {
deepmd::tabulate_fusion_se_a_grad_grad_cpu(
dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
Expand Down Expand Up @@ -409,6 +419,7 @@ class TabulateFusionSeAttenOp : public OpKernel {
const int nnei = em_tensor.shape().dim_size(1);

if (device == "GPU") {
OP_REQUIRES_OK(context, validate_gpu_last_layer_size(last_layer_size));
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_a_gpu(descriptor, table, table_info, em_x, em,
two_embed, nloc, nnei, last_layer_size,
Expand Down Expand Up @@ -485,6 +496,7 @@ class TabulateFusionSeAttenGradOp : public OpKernel {
const int last_layer_size = descriptor_tensor.shape().dim_size(2);

if (device == "GPU") {
OP_REQUIRES_OK(context, validate_gpu_last_layer_size(last_layer_size));
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_a_grad_gpu(
dy_dem_x, dy_dem, dy_dtwo, table, table_info, em_x, em, two_embed, dy,
Expand Down Expand Up @@ -553,15 +565,12 @@ class TabulateFusionSeAttenGradGradOp : public OpKernel {
const int last_layer_size = descriptor_tensor.shape().dim_size(2);

if (device == "GPU") {
OP_REQUIRES_OK(context, validate_gpu_last_layer_size(last_layer_size));
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_a_grad_grad_gpu(
dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted);
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
OP_REQUIRES(context, (last_layer_size <= 1024),
deepmd::tf_compat::InvalidArgument(
"In the process of model compression, the size of the "
"last layer of embedding net must be less than 1024!"));
} else if (device == "CPU") {
deepmd::tabulate_fusion_se_a_grad_grad_cpu(
dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
Expand Down Expand Up @@ -623,6 +632,7 @@ class TabulateFusionSeTOp : public OpKernel {
const int nnei_j = em_tensor.shape().dim_size(2);

if (device == "GPU") {
OP_REQUIRES_OK(context, validate_gpu_last_layer_size(last_layer_size));
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_t_gpu(descriptor, table, table_info, em_x, em,
nloc, nnei_i, nnei_j, last_layer_size);
Expand Down Expand Up @@ -687,6 +697,7 @@ class TabulateFusionSeTGradOp : public OpKernel {
const int last_layer_size = descriptor_tensor.shape().dim_size(1);

if (device == "GPU") {
OP_REQUIRES_OK(context, validate_gpu_last_layer_size(last_layer_size));
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_t_grad_gpu(dy_dem_x, dy_dem, table, table_info,
em_x, em, dy, nloc, nnei_i, nnei_j,
Expand Down Expand Up @@ -744,15 +755,12 @@ class TabulateFusionSeTGradGradOp : public OpKernel {
const int last_layer_size = descriptor_tensor.shape().dim_size(1);

if (device == "GPU") {
OP_REQUIRES_OK(context, validate_gpu_last_layer_size(last_layer_size));
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_t_grad_grad_gpu(
dz_dy, table, table_info, em_x, em, dz_dy_dem_x, dz_dy_dem, nloc,
nnei_i, nnei_j, last_layer_size);
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
OP_REQUIRES(context, (last_layer_size <= 1024),
deepmd::tf_compat::InvalidArgument(
"In the process of model compression, the size of the "
"last layer of embedding net must be less than 1024!"));
} else if (device == "CPU") {
deepmd::tabulate_fusion_se_t_grad_grad_cpu(
dz_dy, table, table_info, em_x, em, dz_dy_dem_x, dz_dy_dem, nloc,
Expand Down Expand Up @@ -806,6 +814,7 @@ class TabulateFusionSeROp : public OpKernel {
const int nnei = em_tensor.shape().dim_size(1);

if (device == "GPU") {
OP_REQUIRES_OK(context, validate_gpu_last_layer_size(last_layer_size));
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_r_gpu(descriptor, table, table_info, em, nloc,
nnei, last_layer_size);
Expand Down Expand Up @@ -861,6 +870,7 @@ class TabulateFusionSeRGradOp : public OpKernel {
const int last_layer_size = descriptor_tensor.shape().dim_size(2);

if (device == "GPU") {
OP_REQUIRES_OK(context, validate_gpu_last_layer_size(last_layer_size));
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_r_grad_gpu(dy_dem, table, table_info, em, dy,
nloc, nnei, last_layer_size);
Expand Down Expand Up @@ -909,14 +919,11 @@ class TabulateFusionSeRGradGradOp : public OpKernel {
const int last_layer_size = descriptor_tensor.shape().dim_size(2);

if (device == "GPU") {
OP_REQUIRES_OK(context, validate_gpu_last_layer_size(last_layer_size));
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_r_grad_grad_gpu(
dz_dy, table, table_info, em, dz_dy_dem, nloc, nnei, last_layer_size);
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
OP_REQUIRES(context, (last_layer_size <= 1024),
deepmd::tf_compat::InvalidArgument(
"In the process of model compression, the size of the "
"last layer of embedding net must be less than 1024!"));
} else if (device == "CPU") {
deepmd::tabulate_fusion_se_r_grad_grad_cpu(
dz_dy, table, table_info, em, dz_dy_dem, nloc, nnei, last_layer_size);
Expand Down
162 changes: 162 additions & 0 deletions source/tests/tf/test_tabulate_gpu_size_validation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,162 @@
# SPDX-License-Identifier: LGPL-3.0-or-later
import unittest

from deepmd.tf.env import (
op_module,
tf,
)


@unittest.skipUnless(
tf.test.is_gpu_available(), reason="GPU tabulation validation requires a GPU"
)
class TestTabulateGpuSizeValidation(unittest.TestCase):
"""Ensure invalid tabulation widths are rejected before GPU launch."""

error_message = "last_layer_size must be between 1 and 1024"

@staticmethod
def _build_ops(last_layer_size: int) -> dict[str, object]:
"""Build all TensorFlow tabulation variants with a common width."""
dtype = tf.float64
table_width = max(1, 6 * last_layer_size)
with tf.device("/CPU:0"):
table_info = tf.constant([0.0, 1.0, 1.0, 1.0, 1.0, 1.0], dtype=dtype)

with tf.device("/GPU:0"):
table = tf.zeros([1, table_width], dtype=dtype)
em_x = tf.zeros([1, 1], dtype=dtype)

em_a = tf.zeros([1, 1, 4], dtype=dtype)
descriptor_a = tf.zeros([1, 4, last_layer_size], dtype=dtype)
dy_a = tf.zeros_like(descriptor_a)
dz_dem_x_a = tf.zeros_like(em_x)
dz_dem_a = tf.zeros_like(em_a)
two_embed = tf.zeros([1, last_layer_size], dtype=dtype)
dz_dtwo = tf.zeros_like(two_embed)

em_t = tf.zeros([1, 1, 1], dtype=dtype)
descriptor_t = tf.zeros([1, last_layer_size], dtype=dtype)
dy_t = tf.zeros_like(descriptor_t)
dz_dem_x_t = tf.zeros_like(em_x)
dz_dem_t = tf.zeros_like(em_t)

em_r = tf.zeros([1, 1], dtype=dtype)
descriptor_r = tf.zeros([1, 1, last_layer_size], dtype=dtype)
dy_r = tf.zeros_like(descriptor_r)
dz_dem_r = tf.zeros_like(em_r)

return {
"se_a_forward": op_module.tabulate_fusion_se_a(
table,
table_info,
em_x,
em_a,
last_layer_size=last_layer_size,
),
"se_a_grad": op_module.tabulate_fusion_se_a_grad(
table, table_info, em_x, em_a, dy_a, descriptor_a
),
"se_a_grad_grad": op_module.tabulate_fusion_se_a_grad_grad(
table,
table_info,
em_x,
em_a,
dz_dem_x_a,
dz_dem_a,
descriptor_a,
),
"se_atten_forward": op_module.tabulate_fusion_se_atten(
table,
table_info,
em_x,
em_a,
two_embed,
last_layer_size=last_layer_size,
),
"se_atten_grad": op_module.tabulate_fusion_se_atten_grad(
table,
table_info,
em_x,
em_a,
two_embed,
dy_a,
descriptor_a,
),
"se_atten_grad_grad": (
op_module.tabulate_fusion_se_atten_grad_grad(
table,
table_info,
em_x,
em_a,
two_embed,
dz_dem_x_a,
dz_dem_a,
dz_dtwo,
descriptor_a,
)
),
"se_t_forward": op_module.tabulate_fusion_se_t(
table,
table_info,
em_x,
em_t,
last_layer_size=last_layer_size,
),
"se_t_grad": op_module.tabulate_fusion_se_t_grad(
table, table_info, em_x, em_t, dy_t, descriptor_t
),
"se_t_grad_grad": op_module.tabulate_fusion_se_t_grad_grad(
table,
table_info,
em_x,
em_t,
dz_dem_x_t,
dz_dem_t,
descriptor_t,
),
"se_r_forward": op_module.tabulate_fusion_se_r(
table,
table_info,
em_r,
last_layer_size=last_layer_size,
),
"se_r_grad": op_module.tabulate_fusion_se_r_grad(
table, table_info, em_r, dy_r, descriptor_r
),
"se_r_grad_grad": op_module.tabulate_fusion_se_r_grad_grad(
table, table_info, em_r, dz_dem_r, descriptor_r
),
}

def _assert_invalid(self, last_layer_size: int, names: tuple[str, ...]) -> None:
graph = tf.Graph()
with graph.as_default():
ops = self._build_ops(last_layer_size)

config = tf.ConfigProto(allow_soft_placement=False)
config.gpu_options.allow_growth = True
with tf.Session(graph=graph, config=config) as sess:
for name in names:
with self.subTest(op=name, last_layer_size=last_layer_size):
with self.assertRaisesRegex(
tf.errors.InvalidArgumentError, self.error_message
):
sess.run(ops[name])

def test_rejects_oversized_width_for_all_gpu_paths(self) -> None:
"""Cover every forward, first-gradient, and grad-grad GPU wrapper."""
names = tuple(
f"{descriptor}_{stage}"
for descriptor in ("se_a", "se_atten", "se_t", "se_r")
for stage in ("forward", "grad", "grad_grad")
)
self._assert_invalid(1025, names)

def test_rejects_zero_width(self) -> None:
"""Cover attribute-derived and descriptor-derived zero widths."""
self._assert_invalid(0, ("se_a_forward", "se_t_grad"))


if __name__ == "__main__":
unittest.main()
Loading