TensorFlow双GPU负载不均衡求助：3080占满4090利用率低

问题描述 #

使用TensorFlow Object Detection API搭配NVIDIA 3080（10GB）训练时出现OOM错误，购入4090（24GB）后同时使用两张显卡。但大batch size训练时，3080显存几乎被占满，4090内存使用率却偏低。希望充分利用两张显卡提升batch size，但Mirrored策略会给每张显卡分配等量数据，找不到让单张显卡承担更多负载的方法。

环境配置 #

• OS：Ubuntu 22.04
• GPUs：[0: 4090（24GB）, 1: 3080（10GB）]
• Python：3.10.9
• cudatoolkit：11.2.2（通过Anaconda安装）
• cudnn：8.1.0.77（通过Anaconda安装）

已尝试方案 #

• 切换分布式策略：从MultiWorkerMirroredStrategy改为MirroredStrategy和experimental.CentralStorageStrategy，但负载不均问题未改善，原本期望CentralStorageStrategy能优化CPU数据分配。

训练时GPU内存使用率截图 #

解决思路与方案 #

1. 自定义异构GPU的batch分配逻辑（TensorFlow原生方案） #

TensorFlow默认分布式策略会平均分配batch到各GPU，针对异构卡，可手动拆分数据集，给不同GPU分配不同大小的batch：

• 先通过tf.config.list_physical_devices('GPU')获取GPU列表，确认4090和3080的设备名称（如/GPU:0和/GPU:1）
• 自定义训练循环，按显存比例拆分总batch，比如总batch设为24，给4090分配16，3080分配8
• 训练时分别将对应batch数据放到对应GPU计算梯度，再汇总梯度更新模型

示例代码片段：

import tensorflow as tf

# 获取并配置GPU设备
gpus = tf.config.list_physical_devices('GPU')
device_4090 = tf.device('/GPU:0')
device_3080 = tf.device('/GPU:1')

# 按显存比例拆分batch：4090分配16，3080分配8
batch_split = [16, 8]

# 自定义训练步骤
def train_step(inputs):
    x1, y1 = inputs[0]  # 4090对应的batch数据
    x2, y2 = inputs[1]  # 3080对应的batch数据
    
    with tf.GradientTape(persistent=True) as tape:
        with device_4090:
            pred1 = model(x1, training=True)
            loss1 = loss_fn(y1, pred1)
        with device_3080:
            pred2 = model(x2, training=True)
            loss2 = loss_fn(y2, pred2)
        total_loss = loss1 + loss2
    
    # 计算并汇总两张显卡的梯度
    grads_4090 = tape.gradient(loss1, model.trainable_variables)
    grads_3080 = tape.gradient(loss2, model.trainable_variables)
    combined_grads = [g1 + g2 for g1, g2 in zip(grads_4090, grads_3080)]
    
    # 更新模型参数
    optimizer.apply_gradients(zip(combined_grads, model.trainable_variables))
    return total_loss

2. 使用tf.distribute.experimental.ParameterServerStrategy #

参数服务器策略可更灵活分配计算负载，可将4090设为主计算节点分配更多batch，3080作为辅助节点分配较少batch，适合精细化调度多GPU场景，不过配置相对复杂。

3. 调整显存占用与模型优化 #

• 给3080设置显存限制：tf.config.set_logical_device_configuration(gpus[1], [tf.config.LogicalDeviceConfiguration(memory_limit=9500)])，避免显存溢出，同时可进一步提高4090的batch分配比例
• 优化模型：使用TF-TRT加速，或精简模型冗余层，降低单batch显存占用，从而给4090分配更多数据

4. 版本兼容检查 #

确保TensorFlow版本与CUDA/CUDNN版本匹配（CUDA 11.2对应TensorFlow 2.5-2.8版本，建议确认当前TF版本是否符合要求）

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