TensorFlow多GPU训练出现NaN值问题求助

大家好，我最近在尝试用多GPU训练模型，但遇到了一个棘手的问题：训练刚开始几步，损失就变成NaN了，但如果只用单GPU训练的话，一切都正常运行。我写了一个能复现问题的简易测试脚本，下面是具体的代码、TensorFlow构建信息以及GPU检测结果，希望有大佬能帮我排查下问题所在~

TensorFlow版本：2.18.0

复现代码 #

import tensorflow as tf
import numpy as np

print(tf.sysconfig.get_build_info())

# Check if GPUs are available
gpus = tf.config.list_physical_devices('GPU')
print(gpus)
if gpus:
    print(f"Number of GPUs available: {len(gpus)}")
else:
    print("No GPUs found. Training will proceed on CPU.")

# Define the strategy for multi-GPU training
strategy = tf.distribute.MirroredStrategy()

# Dummy dataset
def create_dummy_dataset(samples=50000):
    # Generate dummy input data 
    X = np.random.random((samples, 20)).astype(np.float32)
    # Generate dummy labels (binary classification)
    y = np.random.randint(0, 2, (samples, 1)).astype(np.float32)
    return tf.data.Dataset.from_tensor_slices((X, y)).shuffle(samples).batch(32)

dataset = create_dummy_dataset()

# Define the model inside the strategy scope
with strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(64, activation='relu', input_shape=(20,)),
        tf.keras.layers.Dense(32, activation='relu'),
        tf.keras.layers.Dense(1, activation='sigmoid')
    ])
    model.compile(optimizer='adam',
                  loss='binary_crossentropy',
                  metrics=['accuracy'])

# Train the model
model.fit(dataset, epochs=10)

TensorFlow构建信息输出 #

OrderedDict(
{'cpu_compiler': '/usr/lib/llvm-18/bin/clang', 'cuda_compute_capabilities': ['sm_60', 'sm_70', 'sm_80', 'sm_89', 'compute_90'], 'cuda_version': '12.5.1', 'cudnn_version': '9', 'is_cuda_build': True, 'is_rocm_build': False, 'is_tensorrt_build': False}
)

GPU检测输出 #

[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'), PhysicalDevice(name='/physical_device:GPU:1', device_type='GPU')] 
Number of GPUs available: 2 

备注：内容来源于stack exchange，提问作者D. Ramsook