From 6ce049d30056d82da7148e1edbab2a6459385f82 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 05:17:50 +0800 Subject: [PATCH 1/4] 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() From 03d65018e1166ca3260f4138af660ffb0bb0c370 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Mon, 27 Jul 2026 11:51:21 +0800 Subject: [PATCH 2/4] fix(pt): validate GPU tabulation sizes before launch The PyTorch tabulation ops share the CUDA kernels with the TensorFlow ops, so the same last_layer_size > 1024 block-dim overflow reaches users through the PT backend. The grad-grad TORCH_CHECK ran after the launch, and the forward / first-gradient ops had no check at all. --- source/op/pt/tabulate_multi_device.cc | 35 +++++--- .../pt/test_tabulate_gpu_size_validation.py | 89 +++++++++++++++++++ 2 files changed, 112 insertions(+), 12 deletions(-) create mode 100644 source/tests/pt/test_tabulate_gpu_size_validation.py diff --git a/source/op/pt/tabulate_multi_device.cc b/source/op/pt/tabulate_multi_device.cc index cede1d03d9..84f2bc9499 100644 --- a/source/op/pt/tabulate_multi_device.cc +++ b/source/op/pt/tabulate_multi_device.cc @@ -18,6 +18,17 @@ void GetTensorDevice(const torch::Tensor& t, std::string& str) { } } +static void ValidateGpuLastLayerSize(const int64_t last_layer_size) { + // The GPU tabulation kernels use this dimension either as the block size or + // to size dynamic shared memory, so reject invalid models before any launch: + // an out-of-range value overflows the block dimension (or divides by zero) + // inside the launch configuration. + TORCH_CHECK(last_layer_size > 0 && last_layer_size <= 1024, + "last_layer_size must be between 1 and 1024 for GPU tabulation, " + "but got ", + last_layer_size); +} + template void TabulateFusionSeAForward(const torch::Tensor& table_tensor, const torch::Tensor& table_info_tensor, @@ -58,6 +69,7 @@ void TabulateFusionSeAForward(const torch::Tensor& table_tensor, const int64_t nnei = em_tensor.size(1); // compute if (device == "GPU") { + ValidateGpuLastLayerSize(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); @@ -110,6 +122,7 @@ void TabulateFusionSeAGradForward(const torch::Tensor& table_tensor, const int64_t last_layer_size = descriptor_tensor.size(2); // compute if (device == "GPU") { + ValidateGpuLastLayerSize(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, @@ -170,6 +183,7 @@ void TabulateFusionSeAGradGradForward(const torch::Tensor& table_tensor, const int64_t last_layer_size = descriptor_tensor.size(2); // compute if (device == "GPU") { + ValidateGpuLastLayerSize(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, @@ -179,9 +193,6 @@ void TabulateFusionSeAGradGradForward(const torch::Tensor& table_tensor, "The input tensor is on the GPU, but the GPU support for the " "customized OP library is not enabled."); #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM - TORCH_CHECK(last_layer_size <= 1024, - "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, @@ -221,6 +232,7 @@ void TabulateFusionSeTForward(const torch::Tensor& table_tensor, const int64_t nnei_j = em_tensor.size(2); // compute if (device == "GPU") { + ValidateGpuLastLayerSize(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); @@ -266,6 +278,7 @@ void TabulateFusionSeTGradForward(const torch::Tensor& table_tensor, const int64_t last_layer_size = descriptor_tensor.size(1); // compute if (device == "GPU") { + ValidateGpuLastLayerSize(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, @@ -316,6 +329,7 @@ void TabulateFusionSeTGradGradForward(const torch::Tensor& table_tensor, const int64_t last_layer_size = descriptor_tensor.size(1); // compute if (device == "GPU") { + ValidateGpuLastLayerSize(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, @@ -325,9 +339,6 @@ void TabulateFusionSeTGradGradForward(const torch::Tensor& table_tensor, "The input tensor is on the GPU, but the GPU support for the " "customized OP library is not enabled."); #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM - TORCH_CHECK(last_layer_size <= 1024, - "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, @@ -368,6 +379,7 @@ void TabulateFusionSeTTebdForward(const torch::Tensor& table_tensor, const int64_t nnei_j = em_tensor.size(2); // compute if (device == "GPU") { + ValidateGpuLastLayerSize(last_layer_size); #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM deepmd::tabulate_fusion_se_t_tebd_gpu(descriptor, table, table_info, em_x, em, nloc, nnei_i, nnei_j, @@ -414,6 +426,7 @@ void TabulateFusionSeTTebdGradForward(const torch::Tensor& table_tensor, // compute if (device == "GPU") { + ValidateGpuLastLayerSize(last_layer_size); #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM deepmd::tabulate_fusion_se_t_tebd_grad_gpu(dy_dem_x, table, table_info, em_x, em, dy, nloc, nnei_i, @@ -460,6 +473,7 @@ void TabulateFusionSeTTebdGradGradForward( const int64_t last_layer_size = descriptor_tensor.size(3); // compute if (device == "GPU") { + ValidateGpuLastLayerSize(last_layer_size); #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM deepmd::tabulate_fusion_se_t_tebd_grad_grad_gpu( dz_dy, table, table_info, em_x, em, dz_dy_dem_x, nloc, nnei_i, nnei_j, @@ -469,9 +483,6 @@ void TabulateFusionSeTTebdGradGradForward( "The input tensor is on the GPU, but the GPU support for the " "customized OP library is not enabled."); #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM - TORCH_CHECK(last_layer_size <= 1024, - "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_tebd_grad_grad_cpu( dz_dy, table, table_info, em_x, em, dz_dy_dem_x, nloc, nnei_i, nnei_j, @@ -505,6 +516,7 @@ void TabulateFusionSeRForward(const torch::Tensor& table_tensor, const int64_t nnei = em_tensor.size(1); // compute if (device == "GPU") { + ValidateGpuLastLayerSize(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); @@ -545,6 +557,7 @@ void TabulateFusionSeRGradForward(const torch::Tensor& table_tensor, const int64_t last_layer_size = descriptor_tensor.size(2); // compute if (device == "GPU") { + ValidateGpuLastLayerSize(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); @@ -585,6 +598,7 @@ void TabulateFusionSeRGradGradForward(const torch::Tensor& table_tensor, const int64_t last_layer_size = descriptor_tensor.size(2); // compute if (device == "GPU") { + ValidateGpuLastLayerSize(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); @@ -593,9 +607,6 @@ void TabulateFusionSeRGradGradForward(const torch::Tensor& table_tensor, "The input tensor is on the GPU, but the GPU support for the " "customized OP library is not enabled."); #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM - TORCH_CHECK(last_layer_size <= 1024, - "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/pt/test_tabulate_gpu_size_validation.py b/source/tests/pt/test_tabulate_gpu_size_validation.py new file mode 100644 index 0000000000..1cb0ed950b --- /dev/null +++ b/source/tests/pt/test_tabulate_gpu_size_validation.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Reject invalid tabulation widths before any PyTorch GPU kernel launch.""" + +import unittest + +import torch + +from deepmd.pt.cxx_op import ( + ENABLE_CUSTOMIZED_OP, +) + +ERROR_MESSAGE = "last_layer_size must be between 1 and 1024" + + +@unittest.skipIf(not ENABLE_CUSTOMIZED_OP, "PyTorch customized OPs are not built") +@unittest.skipUnless( + torch.cuda.is_available(), "GPU tabulation validation requires a GPU" +) +class TestTabulateGpuSizeValidation(unittest.TestCase): + """The GPU wrappers must fail on the host, not inside a launch config.""" + + dtype = torch.float64 + + def _zeros(self, *shape: int) -> torch.Tensor: + return torch.zeros(shape, dtype=self.dtype, device="cuda") + + def _table(self, last_layer_size: int) -> torch.Tensor: + return self._zeros(1, max(1, 6 * last_layer_size)) + + def _table_info(self) -> torch.Tensor: + # table_info stays on the CPU for every backend. + return torch.tensor([0.0, 1.0, 1.0, 1.0, 1.0, 1.0], dtype=self.dtype) + + def _forward_calls(self, last_layer_size: int) -> dict[str, callable]: + table = self._table(last_layer_size) + table_info = self._table_info() + em_x = self._zeros(1, 1) + em_a = self._zeros(1, 1, 4) + em_t = self._zeros(1, 1, 1) + em_r = self._zeros(1, 1) + two_embed = self._zeros(1, max(1, last_layer_size)) + return { + "se_a": lambda: torch.ops.deepmd.tabulate_fusion_se_a( + table, table_info, em_x, em_a, last_layer_size + ), + "se_atten": lambda: torch.ops.deepmd.tabulate_fusion_se_atten( + table, table_info, em_x, em_a, two_embed, last_layer_size, True + ), + "se_t": lambda: torch.ops.deepmd.tabulate_fusion_se_t( + table, table_info, em_x, em_t, last_layer_size + ), + "se_t_tebd": lambda: torch.ops.deepmd.tabulate_fusion_se_t_tebd( + table, table_info, em_x, em_t, last_layer_size + ), + "se_r": lambda: torch.ops.deepmd.tabulate_fusion_se_r( + table, table_info, em_r, last_layer_size + ), + } + + def _assert_rejected(self, last_layer_size: int) -> None: + for name, call in self._forward_calls(last_layer_size).items(): + with self.subTest(op=name, last_layer_size=last_layer_size): + with self.assertRaisesRegex(RuntimeError, ERROR_MESSAGE): + call() + + def test_rejects_oversized_width(self) -> None: + """A width past the maximum block dimension must be refused.""" + self._assert_rejected(1025) + + def test_rejects_zero_width(self) -> None: + """A zero width would divide by zero while sizing the launch.""" + self._assert_rejected(0) + + def test_gradient_paths_reject_oversized_width(self) -> None: + """The autograd wrappers derive the width from the descriptor.""" + last_layer_size = 1025 + table = self._table(last_layer_size) + table_info = self._table_info() + em_x = self._zeros(1, 1).requires_grad_(True) + em_a = self._zeros(1, 1, 4).requires_grad_(True) + with self.assertRaisesRegex(RuntimeError, ERROR_MESSAGE): + descriptor = torch.ops.deepmd.tabulate_fusion_se_a( + table, table_info, em_x, em_a, last_layer_size + )[0] + descriptor.sum().backward() + + +if __name__ == "__main__": + unittest.main() From 81838508b0e473df5435ac5456a658b2dc3be09a Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 30 Jul 2026 11:35:18 +0800 Subject: [PATCH 3/4] test(pt): cover tabulation gradient size guards Address the outstanding requested-change review comments. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- .../pt/test_tabulate_gpu_size_validation.py | 133 ++++++++++++++++-- 1 file changed, 121 insertions(+), 12 deletions(-) diff --git a/source/tests/pt/test_tabulate_gpu_size_validation.py b/source/tests/pt/test_tabulate_gpu_size_validation.py index 1cb0ed950b..47e6f000b3 100644 --- a/source/tests/pt/test_tabulate_gpu_size_validation.py +++ b/source/tests/pt/test_tabulate_gpu_size_validation.py @@ -2,6 +2,9 @@ """Reject invalid tabulation widths before any PyTorch GPU kernel launch.""" import unittest +from collections.abc import ( + Callable, +) import torch @@ -31,7 +34,7 @@ def _table_info(self) -> torch.Tensor: # table_info stays on the CPU for every backend. return torch.tensor([0.0, 1.0, 1.0, 1.0, 1.0, 1.0], dtype=self.dtype) - def _forward_calls(self, last_layer_size: int) -> dict[str, callable]: + def _forward_calls(self, last_layer_size: int) -> dict[str, Callable[[], object]]: table = self._table(last_layer_size) table_info = self._table_info() em_x = self._zeros(1, 1) @@ -67,22 +70,128 @@ def test_rejects_oversized_width(self) -> None: """A width past the maximum block dimension must be refused.""" self._assert_rejected(1025) - def test_rejects_zero_width(self) -> None: - """A zero width would divide by zero while sizing the launch.""" + def test_rejects_nonpositive_width(self) -> None: + """Zero and negative widths are invalid launch dimensions.""" self._assert_rejected(0) - - def test_gradient_paths_reject_oversized_width(self) -> None: - """The autograd wrappers derive the width from the descriptor.""" - last_layer_size = 1025 + self._assert_rejected(-1) + + def _gradient_cases( + self, + ) -> dict[ + str, + tuple[ + Callable[[], torch.Tensor], + tuple[torch.Tensor, ...], + tuple[int, ...], + ], + ]: + """Build valid forwards whose saved descriptors can test grad guards.""" + last_layer_size = 2 table = self._table(last_layer_size) table_info = self._table_info() em_x = self._zeros(1, 1).requires_grad_(True) em_a = self._zeros(1, 1, 4).requires_grad_(True) - with self.assertRaisesRegex(RuntimeError, ERROR_MESSAGE): - descriptor = torch.ops.deepmd.tabulate_fusion_se_a( - table, table_info, em_x, em_a, last_layer_size - )[0] - descriptor.sum().backward() + two_embed = self._zeros(1, last_layer_size).requires_grad_(True) + em_t = self._zeros(1, 1, 1).requires_grad_(True) + em_t_tebd_x = self._zeros(1, 1).requires_grad_(True) + em_t_tebd = self._zeros(1, 1, 1).requires_grad_(True) + em_r = self._zeros(1, 1).requires_grad_(True) + return { + "se_a": ( + lambda: torch.ops.deepmd.tabulate_fusion_se_a( + table, table_info, em_x, em_a, last_layer_size + )[0], + (em_x, em_a), + (1, 4, last_layer_size), + ), + "se_atten": ( + lambda: torch.ops.deepmd.tabulate_fusion_se_atten( + table, + table_info, + em_x, + em_a, + two_embed, + last_layer_size, + True, + )[0], + (em_x, em_a, two_embed), + (1, 4, last_layer_size), + ), + "se_t": ( + lambda: torch.ops.deepmd.tabulate_fusion_se_t( + table, table_info, em_x, em_t, last_layer_size + )[0], + (em_x, em_t), + (1, last_layer_size), + ), + "se_t_tebd": ( + lambda: torch.ops.deepmd.tabulate_fusion_se_t_tebd( + table, + table_info, + em_t_tebd_x, + em_t_tebd, + last_layer_size, + )[0], + (em_t_tebd_x,), + (1, 1, 1, last_layer_size), + ), + "se_r": ( + lambda: torch.ops.deepmd.tabulate_fusion_se_r( + table, table_info, em_r, last_layer_size + )[0], + (em_r,), + (1, 1, last_layer_size), + ), + } + + def _oversized_saved_descriptor( + self, descriptor_shape: tuple[int, ...] + ) -> Callable[[torch.Tensor], torch.Tensor]: + """Replace only a saved descriptor with an oversized-width sentinel.""" + + def pack(tensor: torch.Tensor) -> torch.Tensor: + if tuple(tensor.shape) == descriptor_shape: + return self._zeros(*descriptor_shape[:-1], 1025) + return tensor + + return pack + + @staticmethod + def _unpack_saved_tensor(tensor: torch.Tensor) -> torch.Tensor: + """Return saved tensors unchanged when their autograd node reloads.""" + return tensor + + def test_first_gradient_wrappers_reject_oversized_width(self) -> None: + """Every first-gradient wrapper validates its saved descriptor width.""" + for name, (forward, inputs, descriptor_shape) in self._gradient_cases().items(): + with self.subTest(op=name): + pack = self._oversized_saved_descriptor(descriptor_shape) + with torch.autograd.graph.saved_tensors_hooks( + pack, self._unpack_saved_tensor + ): + descriptor = forward() + with self.assertRaisesRegex(RuntimeError, ERROR_MESSAGE): + torch.autograd.grad(descriptor.sum(), inputs) + + def test_second_gradient_wrappers_reject_oversized_width(self) -> None: + """Every grad-grad wrapper validates the descriptor before launching.""" + for name, (forward, inputs, descriptor_shape) in self._gradient_cases().items(): + with self.subTest(op=name): + descriptor = forward() + pack = self._oversized_saved_descriptor(descriptor_shape) + with torch.autograd.graph.saved_tensors_hooks( + pack, self._unpack_saved_tensor + ): + first_gradients = torch.autograd.grad( + descriptor.sum(), inputs, create_graph=True + ) + differentiable_sum = sum( + gradient.sum() + for gradient in first_gradients + if gradient.requires_grad + ) + with self.assertRaisesRegex(RuntimeError, ERROR_MESSAGE): + torch.autograd.grad(differentiable_sum, inputs) if __name__ == "__main__": From b88dd447f263a15d3ec4880eb5293e2c3b30999e Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 31 Jul 2026 13:33:30 +0800 Subject: [PATCH 4/4] fix(tabulate): validate GPU widths before allocation Reject invalid TensorFlow widths before output allocation on every GPU kernel specialization, and reject negative PyTorch widths before descriptor tensors are created. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/op/pt/tabulate_multi_device.cc | 14 +++++++ source/op/tf/tabulate_multi_device.cc | 58 +++++++++++++++++---------- 2 files changed, 51 insertions(+), 21 deletions(-) diff --git a/source/op/pt/tabulate_multi_device.cc b/source/op/pt/tabulate_multi_device.cc index 84f2bc9499..49efc7a2ab 100644 --- a/source/op/pt/tabulate_multi_device.cc +++ b/source/op/pt/tabulate_multi_device.cc @@ -29,6 +29,15 @@ static void ValidateGpuLastLayerSize(const int64_t last_layer_size) { last_layer_size); } +static void ValidateGpuLastLayerSizeBeforeAllocation( + const torch::Tensor& table_tensor, const int64_t last_layer_size) { + // Forward wrappers allocate descriptor tensors using last_layer_size, so a + // negative GPU width must be rejected before torch::empty sees the shape. + if (table_tensor.device().is_cuda()) { + ValidateGpuLastLayerSize(last_layer_size); + } +} + template void TabulateFusionSeAForward(const torch::Tensor& table_tensor, const torch::Tensor& table_info_tensor, @@ -787,6 +796,7 @@ class TabulateFusionSeAOp const torch::Tensor& em_x_tensor, const torch::Tensor& em_tensor, int64_t last_layer_size) { + ValidateGpuLastLayerSizeBeforeAllocation(table_tensor, last_layer_size); // allocate output tensors auto options = torch::TensorOptions() .dtype(table_tensor.dtype()) @@ -973,6 +983,7 @@ class TabulateFusionSeAttenOp const torch::Tensor& two_embed_tensor, int64_t last_layer_size, bool is_sorted) { + ValidateGpuLastLayerSizeBeforeAllocation(table_tensor, last_layer_size); // allocate output tensors auto options = torch::TensorOptions() .dtype(table_tensor.dtype()) @@ -1141,6 +1152,7 @@ class TabulateFusionSeTOp const torch::Tensor& em_x_tensor, const torch::Tensor& em_tensor, int64_t last_layer_size) { + ValidateGpuLastLayerSizeBeforeAllocation(table_tensor, last_layer_size); // allocate output tensors auto options = torch::TensorOptions() .dtype(table_tensor.dtype()) @@ -1294,6 +1306,7 @@ class TabulateFusionSeROp const torch::Tensor& table_info_tensor, const torch::Tensor& em_tensor, int64_t last_layer_size) { + ValidateGpuLastLayerSizeBeforeAllocation(table_tensor, last_layer_size); // allocate output tensors auto options = torch::TensorOptions() .dtype(table_tensor.dtype()) @@ -1452,6 +1465,7 @@ class TabulateFusionSeTTebdOp const torch::Tensor& em_x_tensor, const torch::Tensor& em_tensor, int64_t last_layer_size) { + ValidateGpuLastLayerSizeBeforeAllocation(table_tensor, last_layer_size); // allocate output tensors auto options = torch::TensorOptions() .dtype(table_tensor.dtype()) diff --git a/source/op/tf/tabulate_multi_device.cc b/source/op/tf/tabulate_multi_device.cc index ec2d036353..365b63efb7 100644 --- a/source/op/tf/tabulate_multi_device.cc +++ b/source/op/tf/tabulate_multi_device.cc @@ -1,5 +1,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later +#include + #include "custom_op.h" #include "tabulate.h" @@ -162,11 +164,13 @@ REGISTER_OP("TabulateFusionSeRGradGrad") .Input("descriptor: T") .Output("dz_dy: T"); +template 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) { + if (std::is_same::value && + (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"); } @@ -200,6 +204,8 @@ class TabulateFusionSeAOp : public OpKernel { deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (em_tensor.shape().dims() == 3), deepmd::tf_compat::InvalidArgument("Dim of input should be 3")); + OP_REQUIRES_OK(context, + validate_gpu_last_layer_size(last_layer_size)); TensorShape descriptor_shape; descriptor_shape.AddDim(em_tensor.shape().dim_size(0)); descriptor_shape.AddDim(4); // be careful here; @@ -221,7 +227,6 @@ 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); @@ -260,6 +265,9 @@ class TabulateFusionSeAGradOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (dy_tensor.shape().dims() == 3), deepmd::tf_compat::InvalidArgument("Dim of table should be 3")); + const int last_layer_size = descriptor_tensor.shape().dim_size(2); + OP_REQUIRES_OK(context, + validate_gpu_last_layer_size(last_layer_size)); int context_output_index = 0; Tensor* dy_dem_x_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, @@ -285,10 +293,8 @@ class TabulateFusionSeAGradOp : public OpKernel { const FPTYPE* dy = dy_tensor.flat().data(); const int nloc = em_tensor.shape().dim_size(0); const int nnei = em_tensor.shape().dim_size(1); - 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, @@ -327,6 +333,9 @@ class TabulateFusionSeAGradGradOp : public OpKernel { deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (dz_dy_dem_tensor.shape().dims() == 3), deepmd::tf_compat::InvalidArgument("Dim of input should be 3")); + const int last_layer_size = descriptor_tensor.shape().dim_size(2); + OP_REQUIRES_OK(context, + validate_gpu_last_layer_size(last_layer_size)); int context_output_index = 0; Tensor* dz_dy_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, @@ -346,10 +355,8 @@ class TabulateFusionSeAGradGradOp : public OpKernel { const FPTYPE* dz_dy_dtwo = nullptr; const int nloc = em_tensor.shape().dim_size(0); const int nnei = em_tensor.shape().dim_size(1); - 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, @@ -398,6 +405,8 @@ class TabulateFusionSeAttenOp : public OpKernel { deepmd::tf_compat::InvalidArgument("Dim of input should be 3")); OP_REQUIRES(context, (two_embed_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); + OP_REQUIRES_OK(context, + validate_gpu_last_layer_size(last_layer_size)); TensorShape descriptor_shape; descriptor_shape.AddDim(em_tensor.shape().dim_size(0)); descriptor_shape.AddDim(4); // be careful here; @@ -419,7 +428,6 @@ 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, @@ -464,6 +472,9 @@ class TabulateFusionSeAttenGradOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (dy_tensor.shape().dims() == 3), deepmd::tf_compat::InvalidArgument("Dim of table should be 3")); + const int last_layer_size = descriptor_tensor.shape().dim_size(2); + OP_REQUIRES_OK(context, + validate_gpu_last_layer_size(last_layer_size)); int context_output_index = 0; Tensor* dy_dem_x_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, @@ -493,10 +504,8 @@ class TabulateFusionSeAttenGradOp : public OpKernel { const FPTYPE* dy = dy_tensor.flat().data(); const int nloc = em_tensor.shape().dim_size(0); const int nnei = em_tensor.shape().dim_size(1); - 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, @@ -543,6 +552,9 @@ class TabulateFusionSeAttenGradGradOp : public OpKernel { deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (dz_dy_dem_tensor.shape().dims() == 3), deepmd::tf_compat::InvalidArgument("Dim of input should be 3")); + const int last_layer_size = descriptor_tensor.shape().dim_size(2); + OP_REQUIRES_OK(context, + validate_gpu_last_layer_size(last_layer_size)); int context_output_index = 0; Tensor* dz_dy_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, @@ -562,10 +574,8 @@ class TabulateFusionSeAttenGradGradOp : public OpKernel { const FPTYPE* dz_dy_dtwo = dz_dy_dtwo_tensor.flat().data(); const int nloc = em_tensor.shape().dim_size(0); const int nnei = em_tensor.shape().dim_size(1); - 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, @@ -612,6 +622,8 @@ class TabulateFusionSeTOp : public OpKernel { OP_REQUIRES( context, (em_tensor.shape().dims() == 3), deepmd::tf_compat::InvalidArgument("Dim of em_tensor should be 3")); + OP_REQUIRES_OK(context, + validate_gpu_last_layer_size(last_layer_size)); TensorShape descriptor_shape; descriptor_shape.AddDim(em_tensor.shape().dim_size(0)); descriptor_shape.AddDim(last_layer_size); @@ -632,7 +644,6 @@ 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); @@ -671,6 +682,9 @@ class TabulateFusionSeTGradOp : public OpKernel { OP_REQUIRES( context, (dy_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of dy_tensor should be 2")); + const int last_layer_size = descriptor_tensor.shape().dim_size(1); + OP_REQUIRES_OK(context, + validate_gpu_last_layer_size(last_layer_size)); int context_output_index = 0; Tensor* dy_dem_x_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, @@ -694,10 +708,8 @@ class TabulateFusionSeTGradOp : public OpKernel { const int nloc = em_tensor.shape().dim_size(0); const int nnei_i = em_tensor.shape().dim_size(1); const int nnei_j = em_tensor.shape().dim_size(2); - 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, @@ -734,6 +746,9 @@ class TabulateFusionSeTGradGradOp : public OpKernel { deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (dz_dy_dem_tensor.shape().dims() == 3), deepmd::tf_compat::InvalidArgument("Dim of input should be 3")); + const int last_layer_size = descriptor_tensor.shape().dim_size(1); + OP_REQUIRES_OK(context, + validate_gpu_last_layer_size(last_layer_size)); int context_output_index = 0; Tensor* dz_dy_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, @@ -752,10 +767,8 @@ class TabulateFusionSeTGradGradOp : public OpKernel { const int nloc = em_tensor.shape().dim_size(0); const int nnei_i = em_tensor.shape().dim_size(1); const int nnei_j = em_tensor.shape().dim_size(2); - 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, @@ -795,6 +808,8 @@ class TabulateFusionSeROp : public OpKernel { deepmd::tf_compat::InvalidArgument("Dim of table should be 2")); OP_REQUIRES(context, (em_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); + OP_REQUIRES_OK(context, + validate_gpu_last_layer_size(last_layer_size)); TensorShape descriptor_shape; descriptor_shape.AddDim(em_tensor.shape().dim_size(0)); descriptor_shape.AddDim(em_tensor.shape().dim_size(1)); // be careful here; @@ -814,7 +829,6 @@ 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); @@ -851,6 +865,9 @@ class TabulateFusionSeRGradOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (dy_tensor.shape().dims() == 3), deepmd::tf_compat::InvalidArgument("Dim of table should be 3")); + const int last_layer_size = descriptor_tensor.shape().dim_size(2); + OP_REQUIRES_OK(context, + validate_gpu_last_layer_size(last_layer_size)); int context_output_index = 0; Tensor* dy_dem_tensor = NULL; OP_REQUIRES_OK(context, @@ -867,10 +884,8 @@ class TabulateFusionSeRGradOp : public OpKernel { const FPTYPE* dy = dy_tensor.flat().data(); const int nloc = em_tensor.shape().dim_size(0); const int nnei = em_tensor.shape().dim_size(1); - 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); @@ -901,6 +916,9 @@ class TabulateFusionSeRGradGradOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (dz_dy_dem_tensor.shape().dims() == 2), deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); + const int last_layer_size = descriptor_tensor.shape().dim_size(2); + OP_REQUIRES_OK(context, + validate_gpu_last_layer_size(last_layer_size)); int context_output_index = 0; Tensor* dz_dy_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, @@ -916,10 +934,8 @@ class TabulateFusionSeRGradGradOp : public OpKernel { const FPTYPE* dz_dy_dem = dz_dy_dem_tensor.flat().data(); const int nloc = em_tensor.shape().dim_size(0); const int nnei = em_tensor.shape().dim_size(1); - 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);