官方TensorFlow 1 GPU容器中MatMul BLAS运算失败求助

使用官方tensorflow/tensorflow:1.13.1-gpu-py3镜像运行矩阵乘法代码时，MatMul BLAS运算失败。测试发现矩阵第一维度为17及以上时报错，16及以下可正常运行。GPU为RTX 3090（24GB显存），实际项目需要第一维度达到4000。

测试代码 #

import tensorflow as tf
# 创建支持GPU显存动态增长的会话配置
config = tf.ConfigProto(allow_soft_placement=True)
config.gpu_options.allow_growth = True
# 使用配置创建会话
with tf.Session(config=config) as sess:
    # 创建两个随机矩阵
    matrix_a = tf.random.normal(shape=(17, 64), dtype=tf.float32)
    matrix_b = tf.random.normal(shape=(64, 128), dtype=tf.float32)
    # 执行BLAS矩阵乘法
    result = tf.matmul(matrix_a, matrix_b)
    # 运行运算
    output = sess.run(result)
    print("BLAS运算成功！")
    # 检查结果
    print("结果：")
    print(output.shape)

报错日志 #

2024-01-01 02:04:56.923916: E tensorflow/stream_executor/cuda/cuda_blas.cc:698] failed to run cuBLAS routine cublasSgemm_v2: CUBLAS_STATUS_EXECUTION_FAILED
Traceback (most recent call last):
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py", line 1334, in _do_call
    return fn(*args)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py", line 1319, in _run_fn
    options, feed_dict, fetch_list, target_list, run_metadata)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py", line 1407, in _call_tf_sessionrun
    run_metadata)
tensorflow.python.framework.errors_impl.InternalError: Blas GEMM launch failed : a.shape=(17, 64), b.shape=(64, 128), m=17, n=128, k=64
         [[{{node MatMul}}]]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 9, in <module>
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py", line 929, in run
    run_metadata_ptr)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py", line 1152, in _run
    feed_dict_tensor, options, run_metadata)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py", line 1328, in _do_run
    run_metadata)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py", line 1348, in _do_call
    raise type(e)(node_def, op, message)
tensorflow.python.framework.errors_impl.InternalError: Blas GEMM launch failed : a.shape=(17, 64), b.shape=(64, 128), m=17, n=128, k=64
         [[node MatMul (defined at <stdin>:7) ]]

Caused by op 'MatMul', defined at:
  File "<stdin>", line 7, in <module>
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/math_ops.py", line 2455, in matmul
    a, b, transpose_a=transpose_a, transpose_b=transpose_b, name=name)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/gen_math_ops.py", line 5333, in mat_mul
    name=name)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py", line 788, in _apply_op_helper
    op_def=op_def)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py", line 3300, in create_op
    op_def=op_def)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py", line 1801, in __init__
    self._traceback = tf_stack.extract_stack()

InternalError (see above for traceback): Blas GEMM launch failed : a.shape=(17, 64), b.shape=(64, 128), m=17, n=128, k=64
         [[node MatMul (defined at <stdin>:7) ]]

核心原因 #

TensorFlow 1.13.1依赖的CUDA 10.0及对应cuBLAS版本，不支持RTX 3090的Ampere架构（sm_86）。老版本cuBLAS在处理新GPU上的小矩阵（维度超过16）运算时存在执行bug，导致报错。

解决办法 #

1. 升级TensorFlow版本（推荐） #

切换到支持Ampere架构的TensorFlow版本，最低要求TensorFlow 2.3，推荐使用2.4+版本，对应镜像可选择tensorflow/tensorflow:2.4.1-gpu-py3或更高。新版本已适配新GPU架构，能正常处理4000维度的矩阵运算。

2. 临时强制CPU运算 #

若无法升级TF，可临时指定MatMul在CPU上运行，修改代码如下：

import tensorflow as tf
config = tf.ConfigProto(allow_soft_placement=True)
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
    matrix_a = tf.random.normal(shape=(17, 64), dtype=tf.float32)
    matrix_b = tf.random.normal(shape=(64, 128), dtype=tf.float32)
    # 强制在CPU上执行矩阵乘法
    with tf.device('/cpu:0'):
        result = tf.matmul(matrix_a, matrix_b)
    output = sess.run(result)
    print("BLAS运算成功！")
    print("结果：", output.shape)

此方法会损失GPU加速性能，仅适合临时测试。

3. 自定义编译TensorFlow 1.13.1（不推荐） #

针对RTX 3090的sm_86架构重新编译TF 1.13.1，添加对应编译选项。但该方法复杂度高、耗时久，仅适合必须保留TF 1.x环境的场景。

验证示例（TF2.x兼容模式） #

升级后可直接测试大维度矩阵运算：

import tensorflow as tf
# 启用TF1.x兼容模式
tf.compat.v1.disable_eager_execution()
config = tf.compat.v1.ConfigProto(allow_soft_placement=True)
config.gpu_options.allow_growth = True
with tf.compat.v1.Session(config=config) as sess:
    matrix_a = tf.random.normal(shape=(4000, 64), dtype=tf.float32)
    matrix_b = tf.random.normal(shape=(64, 128), dtype=tf.float32)
    result = tf.matmul(matrix_a, matrix_b)
    output = sess.run(result)
    print("运算成功！结果维度：", output.shape)

内容的提问来源于stack exchange，提问作者tsadigov