def apply_gradients(self, apply_grads_func, grads_and_vars, **kwargs):
gradients, variables = list(zip(*grads_and_vars))
if self._nccl:
# FIXME: We have a limitation that KungFu schedules NCCL operations
# in the order of the given gradients. This order is sub-optimal
# to the topological sorting order of dataflow. We get around of this issue by
# fusing all gradients. We need to figure out H ow to get the optimal topological s
# sortting order from TensorFlow.
if self._nccl_fusion:
fused_grad = fuse(gradients)
summed_fused_gradients = group_nccl_all_reduce([fused_grad])
summed_gradients = defuse(summed_fused_gradients[0],
[g.shape for g in gradients])
else:
summed_gradients = group_nccl_all_reduce(gradients)
else:
summed_gradients = group_all_reduce(gradients)
reduced_grads = map_maybe(lambda g: g / self._num_workers,
summed_gradients)
# We need to re-zip gradients and variables as grads_and_vars can be only unzipped once.
reduced_grads_and_vars = zip(reduced_grads, variables)
return apply_grads_func(reduced_grads_and_vars, **kwargs)
def build_fake_train_op(use_nccl):
xs = [tf.Variable(tf.ones((2, 5)))]
if use_nccl:
from kungfu.tensorflow.ops import group_nccl_all_reduce
ys = group_nccl_all_reduce(xs)
else:
from kungfu.tensorflow.ops import group_all_reduce
ys = group_all_reduce(xs)
return ys