def allreduce_grads(all_grads, average, mode="nccl"):
    """
    All-reduce average the gradients among K devices. Results are broadcasted to all devices.

    Args:
        all_grads (K x N): List of list of gradients. N is the number of variables.
        average (bool): average gradients or not.
        mode (str): "nccl", "collective"

    Returns:
        K x N: same as input, but each grad is replaced by the average over K devices.
    """
    assert mode in ["nccl", "collective"], mode

    nr_tower = len(all_grads)
    if nr_tower == 1:
        return all_grads
    new_all_grads = []  # N x K
    for grads in zip(*all_grads):
        # k grads
        if mode == "nccl":
            if get_tf_version_tuple() <= (1, 12):
                from tensorflow.contrib import nccl  # deprecated
            else:
                from tensorflow.python.ops import nccl_ops as nccl
            summed = nccl.all_sum(grads)
        else:
            from tensorflow.python.ops import collective_ops
            summed = []
            shared_cnt = _get_shared_cnt()
            for t in grads:
                with tf.device(t.device):
                    t = collective_ops.all_reduce(
                        t, len(grads),
                        42,   # group key is any fixed integer for a fixed group of devices
                        shared_cnt + 100,
                        'Add', 'Id', communication_hint='nccl')
                    summed.append(t)

        grads_for_devices = []  # K
        for g in summed:
            with tf.device(g.device):
                # tensorflow/benchmarks didn't average gradients
                if average:
                    g = tf.multiply(g, 1.0 / nr_tower)
            grads_for_devices.append(g)
        new_all_grads.append(grads_for_devices)

    # transpose to K x N
    ret = list(zip(*new_all_grads))
    return ret

    All-reduce average the gradients among K devices. Results are broadcasted to all devices.

    Args:
        all_grads (K x N): List of list of gradients. N is the number of variables.
        average (bool): average gradients or not.
        mode (str): "nccl", "collective"

    Returns:
        K x N: same as input, but each grad is replaced by the average over K devices.
    
tensorpack/graph_builder/utils.py