def group_allreduce()

in blocksparse/nccl.py [0:0]

• 46 lines of code
• 14 McCabe index (conditional complexity)

def group_allreduce(
        grads, parms, search_strings=None, cast_map=None, cast_all=None, allreduce_op=allreduce, **allreduce_kwargs):
    # if no grouping specified, create one group to reduce at the end (no overlap with compute)
    if search_strings is None:
        search_strings = ["group_allreduce_all"]

    groups = [(names, list(), list()) for names in search_strings]

    last_group_idx = len(groups) - 1

    for i, (grad, param) in enumerate(zip(grads, parms)):
        for j, (names, group16, group32) in enumerate(groups):
            # each group can be a single string, or a list of strings
            # TODO: support regex's
            if isinstance(names, str):
                names = (names,)

            if j == last_group_idx or any(name in param.name for name in names):

                if cast_all is not None:
                    grad = float_cast(grad, dtype=cast_all)

                elif cast_map is not None and name in cast_map:
                    grad = float_cast(grad, dtype=cast_map[name])

                if grad.dtype.base_dtype is tf.float16:
                    group16.append((i, grad, param))
                else:
                    group32.append((i, grad, param))
                break

    for name, group16, group32 in groups:
        count = 0
        if isinstance(name, str):
            str_name = name
        else:
            str_name = "_".join(name)
        str_name = str_name.replace('/', '_')
        for group in (group16, group32):
            count += len(group)
            if len(group) > 0:
                if len(group) == 1:
                    concated = group[0][1]
                else:
                    concated = tf.concat([tf.reshape(grad, [-1]) for _, grad, _ in group], 0, name="concat_"+str_name)

                reduced = allreduce_op(concated, **allreduce_kwargs)

                if len(group) == 1:
                    grads[group[0][0]] = reduced
                else:
                    offset = 0
                    for i, grad, param in group:
                        size = param.shape.num_elements()
                        grads[i] = tf.reshape(reduced[offset: offset + size], param.shape)
                        offset += size

        if count == 0:
            print("Warning: no grads found for all_reduce group: ", name)

    # grads modified in place, but return anyway
    return grads