[BUG] _TensorFlow Hybrid Descriptor hangs during `compute_input_stats()` on RTX 5090 (Blackwell)

[fmatusalem]

Bug summary

TensorFlow Hybrid Descriptor hangs during compute_input_stats() on RTX 5090 (Blackwell)

Summary

The TensorFlow backend of DeePMD-kit 3.1.3 hangs indefinitely during compute_input_stats() when using the Hybrid Descriptor on an NVIDIA RTX 5090 (Blackwell, SM 12.0).

The issue is reproducible using the official water example included with DeePMD-kit.

Individual descriptors work correctly, but any Hybrid descriptor hangs before training starts.

The exact same environment works correctly on an RTX 4090.

---

Environment

GPU

• NVIDIA GeForce RTX 5090
• Compute Capability: SM 12.0
• Driver: 595.71.05

Software

• DeePMD-kit 3.1.3
• Installed using the official offline installer:
  • deepmd-kit-3.1.3-cuda129-Linux-x86_64.sh
• TensorFlow 2.19.1
• CUDA Runtime 12.9
• cuDNN 9.10.2

Python reports:

import tensorflow as tf
import deepmd

print(deepmd.__version__)
print(tf.__version__)
print(tf.config.list_physical_devices("GPU"))

Output:

DeepMD: 3.1.3
TensorFlow: 2.19.1
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

TensorFlow correctly detects the GPU.

---

Tested descriptors

Using the official water example:

Descriptor	RTX 4090	RTX 5090
se_e2_a	✅ Works	✅ Works
se_e2_r	✅ Works	✅ Works
se_e3	✅ Works	✅ Works
hybrid(se_e2_a, se_e2_a)	✅ Works	❌ Hangs
hybrid(se_e2_a, se_e2_r)	✅ Works	❌ Hangs

Exactly the same behavior is observed with my production dataset.

---

Expected behavior

Training should continue after

DEEPMD INFO    data stating... (this step may take long time)

and proceed to model training.

---

Actual behavior

Training never progresses beyond

DEEPMD INFO    data stating... (this step may take long time)

Even after several hours.

GPU utilization remains at 100%.

---

GPU utilization

While hanging:

• GPU utilization: 100%
• GPU memory: ~1.7 GB
• GPU power: ~120 W

Example:

GPU Util: 100%
Memory: ~1700 MiB
Power: ~122 W

---

py-spy output

py-spy dump shows the main thread blocked inside compute_input_stats():

compute_input_stats()
    deepmd/tf/descriptor/hybrid.py
        compute_input_stats()
            deepmd/tf/descriptor/se_a.py
                _compute_dstats_sys_smth()
                    tensorflow Session.run()

More specifically:

_call_tf_sessionrun
_run_fn
_do_call
_do_run
_run
run
run_sess
_compute_dstats_sys_smth
compute_input_stats
compute_input_stats (hybrid.py)
_compute_input_stat
data_stat
build
train

---

Additional observations

The issue is independent of the dataset.

It is reproducible using the official DeePMD water example.

The problem only appears when using the TensorFlow Hybrid Descriptor.

Individual descriptors execute normally.

---

Additional information

I also tested:

• DeePMD 3.1.3 from conda-forge
• DeePMD 3.1.3 from the official offline installer

Both exhibit exactly the same behavior.

Therefore, the issue does not appear to be related to packaging.

---

Comparison

The same installation, same TensorFlow version, same CUDA runtime, same example, and same input files:

• RTX 4090 → Works correctly
• RTX 5090 → Hangs indefinitely

The only hardware difference is the GPU architecture (Ada vs Blackwell).

---

Question

Could this be a TensorFlow backend issue affecting the Hybrid Descriptor on Blackwell (SM 12.0)?

I'd be happy to test patches or provide additional debugging information if needed.

DeePMD-kit Version

DeePMD 3.1.3

Backend and its version

TensorFlow: 2.19.1

How did you download the software?

Offline packages

Input Files, Running Commands, Error Log, etc.

The issue is reproducible using the official water example included with DeePMD-kit.

Steps to Reproduce

The issue is reproducible using the official water example included with DeePMD-kit. Just use the water hybrid descriptor example.

Further Information, Files, and Links

No response

[dosubot]

Hi @fmatusalem! I'm Dosu and I'm helping the DeepModeling team.

This is a confirmed bug that has already been reproduced and root-caused by the maintainers. It affects the TensorFlow backend on RTX 5090/5090D (Blackwell, SM 12.0) and is tracked in issue #5117 [1].

Root cause: A CUDA synchronization bug in the parallel_prefix_scan kernel in source/lib/src/gpu/neighbor_list.cu. The kernel uses a per-thread loop around block-wide collectives (cub::BlockScan and __syncthreads()). When nall > THREADS_PER_BLOCK and the final segment is partial, only some threads enter the last iteration — this is divergent use of a block-wide collective, which is undefined behavior in CUDA and deadlocks on Blackwell. On older architectures (Ada, Ampere, Volta) the same invalid pattern happened to "work" by accident due to differences in code generation and barrier scheduling [1].

Fix: PR #5575 has been opened with a verified fix. It makes every thread in the block enter every BlockScan/__syncthreads() iteration and masks out inactive tail lanes. The fix was validated on RTX 5090 hardware — training completes successfully [1].

Workarounds until the fix is released:

1. Use the PyTorch backend — switch to dp --pt train instead of dp --tf train. The PyTorch backend does not exhibit this problem.
2. Run TF training on a non-Blackwell GPU (V100/A100/H100/RTX 4090 all work).
3. Build from source with the patch from PR fix(gpu): avoid neighbor-list prefix scan deadlock #5575 applied, using CMAKE_CUDA_ARCHITECTURES=120.

Your observation that individual descriptors work but hybrid hangs is consistent with this — the hybrid descriptor calls compute_input_stats() for each sub-descriptor with PBC-expanded coordinates that push nall past the 256-thread block boundary, triggering the deadlock condition.

To reply, just mention @dosu.

---

Share context across your team and agents. Try Dosu.