Ubuntu Server 22.10下TensorFlow未充分利用4张Quadro RTX 6000 GPU问题

一、多GPU未并行利用的核心原因 #

• 未启用分布式训练配置：TensorFlow默认不会自动将模型分配至多GPU，即便能识别所有GPU，若未配置分布式策略（如MirroredStrategy），所有计算会集中在GPU0上，导致单卡显存溢出，其余GPU闲置。
• 模型/数据未正确分发：即便启用了分布式策略，若模型构建、数据加载未在策略作用域内执行，仍会出现仅单卡工作的情况。

二、GPU利用率为0%的潜在问题 #

• cuBLAS注册失败：cuBLAS是TensorFlow依赖的GPU线性代数核心库，注册失败会导致GPU无法执行计算任务，所有计算转至CPU处理，GPU仅加载少量模型参数（低显存占用）但无实际计算，因此利用率为0。这通常是CUDA、cuDNN与TensorFlow版本不匹配导致的。
• TensorRT库缺失：TensorRT负责GPU推理加速优化，缺失虽不直接影响训练，但反映出CUDA生态组件安装不完整，可能间接引发GPU调度相关问题。

三、显存持续上涨且崩溃后未释放的原因 #

• 显存泄漏：训练过程中未正确释放变量、数据集迭代器未关闭，或动态图模式下重复构建计算图，都会导致显存持续累积。崩溃后显存未释放是因为Python进程未被彻底终止，残留的显存占用需手动清理。
• BFC分配器显存碎片：日志中bfc_allocator的报错显示，即使总显存足够，也因显存碎片过多无法分配连续大内存块，最终导致单卡显存溢出。

四、针对性解决步骤 #

1. 配置多GPU分布式训练

   • 使用MirroredStrategy包裹模型构建与训练逻辑：

     import tensorflow as tf
     strategy = tf.distribute.MirroredStrategy()
     with strategy.scope():
         # 在此作用域内完成模型构建、编译、数据加载
         model = your_model_definition()
         model.compile(optimizer='adam', loss='your_loss_function')
     model.fit(train_dataset, epochs=your_epochs)
     

   • 确保数据集使用tf.data接口，并设置合适的batch_size，让分布式策略自动将数据分发至各GPU。

2. 修复CUDA/cuDNN/TensorFlow版本匹配问题

   • 核对TensorFlow官方版本兼容表，确保nvidia-driver、CUDA Toolkit、cuDNN版本与当前TensorFlow版本完全匹配（例如TF2.10对应CUDA11.2、cuDNN8.1）。
   • 重新安装对应版本的CUDA Toolkit（包含cuBLAS）和cuDNN，确保库文件路径被正确添加到环境变量（如LD_LIBRARY_PATH）。

3. 解决显存泄漏与残留占用问题

   • 崩溃后手动终止Python进程释放显存：

     ps aux | grep python | grep -v grep | awk '{print $2}' | xargs kill -9
     

   • 在训练代码中添加显存清理逻辑：

     import gc
     tf.keras.backend.clear_session()
     gc.collect()
     

   • 避免在循环内重复构建模型或计算图，固定计算图结构以减少显存消耗。

4. 优化显存分配策略

   • 启用显存增长模式，避免进程启动时一次性占用全部显存：

     gpus = tf.config.list_physical_devices('GPU')
     for gpu in gpus:
         tf.config.experimental.set_memory_growth(gpu, True)
     

   • 适当减小batch_size，若模型过大无法单卡容纳，可采用模型并行方案，将不同层分配至不同GPU。

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