使用4块GPU训练模型时TensorFlow仍出现内存不足问题求助

CPU: Intel(R) Xeon(R) CPU E5-2698 v4 @ 2.20GHz
GPUs: 4x Tesla V100 16 GB
RAM: 256 GB DDR4

使用开源代码训练网络，数据集为约20GB的RML2018.01 Dataset，严格采用原研究训练参数（批量大小400，仅训练1个epoch用于测试）。该网络仅202,776个参数，但训练时持续出现"内存不足"错误。

为适配4GPU环境，做了两处修改：

1. 将os.environ["CUDA_VISIBLE_DEVICES"] = "0"改为os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
2. 调整模型构建代码块：

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    model = culstm.LSTMModel(weights=None,input_shape=[1024,2], classes=24)
    model.compile(loss='categorical_crossentropy', metrics=['accuracy'], optimizer='adam')
model.summary()

训练时触发OOM错误，GPU和RAM负载显示：使用os.environ["TF_GPU_ALLOCATOR"] = "cuda_malloc_async"时，TensorFlow显存占用为15GB + 3.5GB*3 = 25.5GB（总显存64GB），RAM使用率从未超过100GB，显存利用明显不合理。Docker容器内nvidia-smi可识别全部4块显卡。

• 添加os.environ["TF_GPU_ALLOCATOR"] = "cuda_malloc_async"
• 添加GPU显存增长配置：

config = ConfigProto()
config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)

• 清理GPU缓存

将数据集减半或批量大小设为8时可正常训练，但这两种方式无法满足原研究结果复现需求。原代码为单GPU训练，理论上4块GPU显存足够。怀疑TensorFlow未正确利用其余GPU显存，尽管训练开始时会在所有GPU上分配内存。

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