POSTED BY
2023-10-04

複数のGPUを同時に使わせるMirroredStrategyの宣言箇所で起こることがある。

import tensorflow as tf

with tf.distribute.MirroredStrategy().scope():

    ....モデル構築ブロック....

InvalidArgumentError: No OpKernel was registered to support Op 'NcclAllReduce' used by {{node Adam/NcclAllReduce}} with these attrs: [reduction="sum", shared_name="c1", T=DT_FLOAT, num_devices=2]

以下のようなcross_device_ops引数を与えると解決すると思われる。

import tensorflow as tf

with tf.distribute.MirroredStrategy(cross_device_ops=tf.distribute.HierarchicalCopyAllReduce()).scope():

    ....モデル構築ブロック....

参照
https://github.com/tensorflow/tensorflow/issues/33656