Grover 模型的多 GPU 训练

Question

我正在尝试使用 tensorflow==1.13.1 在多个 GPU 上训练Grover 的鉴别器模型。该模型的默认配置假设在多个 TPU 上进行训练，并被证明适用于具有多个 TPU 的系统。然而，由于 TPU 不可用，我试图通过使用 tensorflow 分布策略更改模型的运行配置来将训练分布在多个 GPU 上，如下所示：

dist_strategy = tf.contrib.distribute.MirroredStrategy(
        num_gpus=8,
        cross_device_ops=AllReduceCrossDeviceOps('nccl', num_packs=8),
        # cross_device_ops=AllReduceCrossDeviceOps('hierarchical_copy')
)

然后我尝试更改 GPU 的运行配置，如下所示：

 run_config = RunConfig(
        train_distribute=dist_strategy,
        eval_distribute=dist_strategy,
        log_step_count_steps=log_every_n_steps,
        model_dir=FLAGS.output_dir,
        save_checkpoints_steps=FLAGS.iterations_per_loop)

它似乎正确分配数据（我可以通过 GPU 使用情况和其他指标来判断），但是在收集梯度时它失败并出现以下错误：

ValueError: You must specify an aggregation method to update a MirroredVariable in Replica Context.

但是，我在分发策略中提供了跨设备操作。对我来说，最终的解决方案似乎是使用分布策略代码修改损失函数，但我认为它应该可以与 TPU 一起使用。任何帮助表示高度赞赏！