GPU超算上TensorFlow脚本每行重复执行32次的问题排查

核心问题 #

你的脚本重复执行32次+GPU内存溢出，根源在Slurm启动命令误用了MPI：

• 超算节点通常默认配备32个CPU核心，xxxx-mpirun会自动按核心数启动32个独立进程，每个进程都会完整执行NN-test.py的所有代码，因此会重复打印"DFs loaded""Normalised"这类输出。
• 32个进程同时在同一个GPU上申请内存、加载模型，直接导致显存被挤占耗尽，触发CUDA_ERROR_OUT_OF_MEMORY报错。

修复步骤 #

1. 修改Slurm启动命令 #

把Slurm脚本最后一行的xxxx-mpirun python3 NN-test.py改成直接运行Python：

python3 NN-test.py

你的模型是单GPU单进程设计，完全不需要用MPI启动多进程。

2. 单GPU内存优化（可选） #

如果单GPU运行仍有内存压力，可以在Python脚本开头添加GPU内存动态增长限制，避免TensorFlow一次性占满显存：

import tensorflow as tf

# 启用GPU内存动态分配，按需使用
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

3. 多GPU训练（进阶可选） #

如果想利用超算的多GPU资源，不要用MPI启动多进程，而是用TensorFlow官方的分布式策略，比如MirroredStrategy，修改模型构建部分：

# 在模型定义前初始化分布式策略
strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    # 原有的模型定义、编译代码都放在这个作用域内
    model = keras.Sequential([
        layers.Dense(500, input_shape=(700,)),
        layers.BatchNormalization(),
        layers.Activation('relu'),
        layers.Dropout(0.2),
        layers.Dense(500),
        layers.BatchNormalization(),
        layers.Activation('relu'),
        layers.Dense(1, activation='sigmoid')
    ])
    optimiser = Adam(learning_rate = 0.0001)
    model.compile(optimizer=optimiser, loss='binary_crossentropy', metrics=['accuracy'])

同时要修改Slurm脚本的--gres=gpu:N参数（N为要使用的GPU数量），仍然用python3 NN-test.py启动，不要用MPI。

验证修改 #

修改后重新提交Slurm任务，重复输出的问题会消失，GPU内存占用恢复正常，模型即可正常训练。

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