From 6ce049d30056d82da7148e1edbab2a6459385f82 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 05:17:50 +0800 Subject: [PATCH] fix(tf): validate GPU tabulation sizes Reject zero and oversized tabulation layer widths before launching any TensorFlow CUDA or ROCm wrapper. Cover all forward, gradient, and grad-gradient variants with GPU-only regression cases. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/op/tf/tabulate_multi_device.cc | 39 +++-- .../tf/test_tabulate_gpu_size_validation.py | 162 ++++++++++++++++++ 2 files changed, 185 insertions(+), 16 deletions(-) create mode 100644 source/tests/tf/test_tabulate_gpu_size_validation.py diff --git a/source/op/tf/tabulate_multi_device.cc b/source/op/tf/tabulate_multi_device.cc index 174faf5213..ec2d036353 100644 --- a/source/op/tf/tabulate_multi_device.cc +++ b/source/op/tf/tabulate_multi_device.cc @@ -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 class TabulateFusionSeAOp : public OpKernel { public: @@ -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); @@ -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, @@ -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, @@ -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, @@ -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, @@ -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, @@ -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); @@ -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, @@ -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, @@ -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); @@ -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); @@ -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); diff --git a/source/tests/tf/test_tabulate_gpu_size_validation.py b/source/tests/tf/test_tabulate_gpu_size_validation.py new file mode 100644 index 0000000000..e1c55dee44 --- /dev/null +++ b/source/tests/tf/test_tabulate_gpu_size_validation.py @@ -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()