使用以下代码(使用tensorflow.contrib.nccl.all_sum
),我希望看到字节通过NVLINK传输。
实际上,我不是。
from tensorflow.contrib.nccl import all_sum
with tf.device('/gpu:0'):
a = tf.get_variable(
"a", initializer=tf.constant(1.0, shape=(args.dim, args.dim)))
with tf.device('/gpu:1'):
b = tf.get_variable(
"b", initializer=tf.constant(2.0, shape=(args.dim, args.dim)))
with tf.device('/gpu:0'):
summed_node = all_sum([a, b])
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True,
log_device_placement=True))
init = tf.global_variables_initializer()
sess.run(init)
with tf.device('/gpu:0'):
summed = sess.run(summed_node)
我的机器是p3.8xlarge
的AWS实例。我的理解是,此配置支持NVLINK。
执行很好,但是当我使用nvidia-smi nvlink -g 0 -i 0
时,链接的Tx / Rx计数为零。