chore: import upstream snapshot with attribution

python/paddle/distributed/passes/fuse_all_reduce.py

@@ -0,0 +1,387 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+import paddle
+from paddle.framework import core
+from paddle.utils import unique_name
+
+from .pass_base import PassBase, PassType, register_pass
+
+
+def find_adjacent_match_sequences(
+    iterable, filter_func, adjacent_filter_func=None
+):
+    n = len(iterable)
+    match_sequences = []
+    if adjacent_filter_func is None:
+        adjacent_filter_func = lambda ref_op, new_op: True
+    i = 0
+    while True:
+        while i < n and not filter_func(iterable[i]):
+            i += 1
+        j = i + 1
+        while (
+            j < n
+            and filter_func(iterable[j])
+            and adjacent_filter_func(iterable[i], iterable[j])
+        ):
+            j += 1
+        if i < n and j <= n:
+            match_sequences.append((i, j))
+        i = j + 1
+        if i >= n:
+            break
+    return match_sequences
+
+
+def insert_fuse_all_reduce_ops(
+    block, reversed_op_indices, input_var_names, output_var_names, dtype, attrs
+):
+    fused_var = block.create_var(
+        name=unique_name.generate(f"FusedOutput_{input_var_names[0]}"),
+        dtype=dtype,
+    )
+
+    # FIXME(zengjinle): here we assume that we use
+    # c_sync_calc_stream/c_sync_comm_stream to do sync.
+    # But someone may use c_wait_compute/c_wait_comm instead.
+    if not attrs["use_calc_stream"]:
+        ring_id = attrs["ring_id"]
+        new_op_indices = list(reversed_op_indices)
+
+        for i, op_idx in enumerate(reversed_op_indices):
+            prev_op_idx = op_idx - 1
+            while (
+                prev_op_idx >= 0
+                and block.ops[prev_op_idx].type == "c_sync_calc_stream"
+            ):
+                new_op_indices.append(prev_op_idx)
+                prev_op_idx -= 1
+
+            if i > 0:
+                next_op_idx = op_idx + 1
+                n = len(block.ops)
+                while (
+                    next_op_idx < n
+                    and block.ops[next_op_idx].type == "c_sync_comm_stream"
+                ):
+                    assert block.ops[next_op_idx].attr("ring_id") == ring_id
+                    new_op_indices.append(next_op_idx)
+
+        new_op_indices = list(set(new_op_indices))
+        new_op_indices.sort(reverse=True)
+        reversed_op_indices = new_op_indices
+
+    insert_idx = reversed_op_indices[0] + 1
+    op_role_key = core.op_proto_and_checker_maker.kOpRoleAttrName()
+
+    concated_shapes = []
+    concated_ranks = []
+    for var_name in output_var_names:
+        shape = block._find_var_recursive(var_name).shape
+        concated_shapes.extend(shape)
+        concated_ranks.append(len(shape))
+
+    coalesce_tensor_op_kwargs = {
+        "type": "coalesce_tensor",
+        "inputs": {
+            "Input": input_var_names,
+        },
+        "outputs": {
+            "Output": output_var_names,
+            "FusedOutput": fused_var,
+        },
+        "attrs": {
+            "use_align": True,
+            "dtype": dtype,
+            "concated_shapes": concated_shapes,
+            "concated_ranks": concated_ranks,
+            op_role_key: attrs[op_role_key],
+        },
+    }
+
+    if not attrs["use_calc_stream"]:
+        block._insert_op_without_sync(
+            insert_idx,
+            type="c_sync_calc_stream",
+            inputs={"X": fused_var},
+            outputs={"Out": fused_var, op_role_key: attrs[op_role_key]},
+        )
+        insert_idx += 1
+
+    # all_reduce sum should insert
+    attrs["reduce_type"] = paddle.distributed.ReduceOp.SUM
+    block._insert_op_without_sync(
+        insert_idx,
+        type="all_reduce",
+        inputs={"x": fused_var},
+        outputs={"out": fused_var},
+        attrs=attrs,
+    )
+
+    for op_idx in reversed_op_indices:
+        block._remove_op(op_idx)
+
+    return coalesce_tensor_op_kwargs
+
+
+def has_same_attrs(op1, op2, attr_names):
+    for attr_name in attr_names:
+        if op1.attr(attr_name) != op2.attr(attr_name):
+            return False
+    return True
+
+
+def filter_all_collective_op_indices(block):
+    # NOTE: should add more collective ops
+    all_collective_ops = {
+        "c_broadcast",
+        "broadcast",
+        "all_gather",
+        "all_reduce",
+    }
+
+    match_op_indices = []
+    for i, op in enumerate(block.ops):
+        if op.type in all_collective_ops:
+            match_op_indices.append(i)
+    return match_op_indices
+
+
+def find_all_fuse_all_reduce_groups(block):
+    collective_op_indices = filter_all_collective_op_indices(block)
+    collective_ops = [block.ops[i] for i in collective_op_indices]
+
+    def is_valid_allreduce_op(op):
+        if op.type != "c_allreduce_sum" or op.attr("use_model_parallel"):
+            return False
+        in_var_name = op.input("X")[0]
+        out_var_name = op.output("Out")[0]
+        if in_var_name != out_var_name:
+            return False
+        in_var = block._find_var_recursive(in_var_name)
+        assert in_var is not None
+        if in_var.type != core.VarDesc.VarType.DENSE_TENSOR:
+            return False
+        shape = in_var.shape
+        if any(s <= 0 for s in shape):
+            return False
+        return True
+
+    same_attr_names = [
+        "ring_id",
+        "use_calc_stream",
+        core.op_proto_and_checker_maker.kOpRoleAttrName(),
+        core.op_proto_and_checker_maker.kOpDeviceAttrName(),
+    ]
+
+    def is_same_adjacent_op(ref_op, new_op):
+        if not has_same_attrs(ref_op, new_op, same_attr_names):
+            return False
+        ref_op_in_var = block._find_var_recursive(ref_op.input("X")[0])
+        new_op_in_var = block._find_var_recursive(new_op.input("X")[0])
+        if ref_op_in_var.dtype != new_op_in_var.dtype:
+            return False
+        return True
+
+    match_seqs = find_adjacent_match_sequences(
+        collective_ops, is_valid_allreduce_op, is_same_adjacent_op
+    )
+    new_match_seqs = []
+    for i, j in match_seqs:
+        new_match_seqs.append([collective_op_indices[k] for k in range(i, j)])
+    return new_match_seqs
+
+
+def split_fuse_all_reduce_groups_by_deps(block, groups, op_deps):
+    new_groups = []
+
+    def insert_new_group(op_indices, start_idx, end_idx):
+        if end_idx - start_idx > 1:
+            new_groups.append(op_indices[start_idx:end_idx])
+
+    for op_indices in groups:
+        n = len(op_indices)
+        assert n > 0
+        if n == 1:
+            continue
+
+        start_idx = 0
+        k = start_idx + 1
+        while k < n:
+            found_group = False
+            for prev_idx in range(start_idx, k):
+                dep = op_deps[op_indices[prev_idx]][op_indices[k]]
+                if dep == core.Node.Dep.NoDep:
+                    continue
+                # [start_idx, k) is valid groups
+                insert_new_group(op_indices, start_idx, k)
+                start_idx = k
+                break
+            k += 1
+
+        insert_new_group(op_indices, start_idx, k)
+
+    return new_groups
+
+
+def insert_coalesce_tensor_ops(block, coalesce_ops_kwargs):
+    if not coalesce_ops_kwargs:
+        return
+
+    var_infos = {}
+    for idx, op in enumerate(block.ops):
+        for var in op.input_arg_names:
+            if var not in var_infos:
+                var_infos[var] = [idx, True]
+
+        for var in op.output_arg_names:
+            if var not in var_infos:
+                var_infos[var] = [idx, False]
+
+    n = len(block.ops)
+    insert_idx_and_kwargs = []
+    for group_idx, kwargs in enumerate(coalesce_ops_kwargs):
+        all_vars = kwargs["inputs"]["Input"] + kwargs["outputs"]["Output"]
+        min_op_idx = n
+        copy_data = False
+        for var in all_vars:
+            if var not in var_infos:
+                copy_data = True
+                min_idx = 0
+                break
+            op_idx, is_input = var_infos[var]
+            if is_input:
+                copy_data = True
+            min_op_idx = min(min_op_idx, op_idx)
+        kwargs["attrs"]["copy_data"] = copy_data
+        insert_idx_and_kwargs.append((min_op_idx, kwargs))
+
+    insert_idx_and_kwargs.sort(key=lambda element: element[0], reverse=True)
+    for idx, kwargs in insert_idx_and_kwargs:
+        block._insert_op_without_sync(idx, **kwargs)
+
+
+def insert_fuse_all_reduce_by_memory_size(block, groups, max_memory_size):
+    op_role_key = core.op_proto_and_checker_maker.kOpRoleAttrName()
+    op_role_var_key = core.op_proto_and_checker_maker.kOpRoleVarAttrName()
+    op_device_key = core.op_proto_and_checker_maker.kOpDeviceAttrName()
+    coalesce_ops_kwargs = []
+    for group in reversed(groups):
+        first_op = block.ops[group[0]]
+        ring_id = first_op.attr("ring_id")
+        use_calc_stream = first_op.attr("use_calc_stream")
+        use_model_parallel = first_op.attr("use_model_parallel")
+        op_role = first_op.attr(op_role_key)
+        op_device = first_op.attr(op_device_key)
+
+        attrs = {
+            "ring_id": ring_id,
+            "use_calc_stream": use_calc_stream,
+            "use_model_parallel": use_model_parallel,
+            op_role_key: op_role,
+            op_device_key: op_device,
+        }
+        dtype = block._find_var_recursive(first_op.input("X")[0]).dtype
+        sizeof = core.size_of_dtype(dtype)
+
+        cur_mem_size = 0
+        op_role_vars = []
+        recorded_op_indices = []
+        in_var_names = []
+        out_var_names = []
+        for op_idx in reversed(group):
+            op = block.ops[op_idx]
+            in_var_name = op.input("X")[0]
+            out_var_name = op.output("Out")[0]
+            in_var = block._find_var_recursive(in_var_name)
+            mem_size = int(np.prod(in_var.shape)) * sizeof
+            if cur_mem_size + mem_size > max_memory_size:
+                if len(recorded_op_indices) > 1:
+                    attrs[op_role_var_key] = op_role_vars
+                    coalesce_op_kwargs = insert_fuse_all_reduce_ops(
+                        block,
+                        recorded_op_indices,
+                        in_var_names,
+                        out_var_names,
+                        dtype,
+                        attrs,
+                    )
+                    coalesce_ops_kwargs.append(coalesce_op_kwargs)
+
+                cur_mem_size = 0
+                op_role_vars = []
+                recorded_op_indices = []
+                in_var_names = []
+                out_var_names = []
+
+            cur_mem_size += mem_size
+            recorded_op_indices.append(op_idx)
+            in_var_names.append(in_var_name)
+            out_var_names.append(out_var_name)
+            if op.has_attr(op_role_var_key):
+                op_role_vars.extend(op.attr(op_role_var_key))
+
+        if len(recorded_op_indices) > 1:
+            attrs[op_role_var_key] = op_role_vars
+            coalesce_op_kwargs = insert_fuse_all_reduce_ops(
+                block,
+                recorded_op_indices,
+                in_var_names,
+                out_var_names,
+                dtype,
+                attrs,
+            )
+            coalesce_ops_kwargs.append(coalesce_op_kwargs)
+    block._sync_with_cpp()
+    insert_coalesce_tensor_ops(block, coalesce_ops_kwargs)
+
+
+@register_pass("fuse_all_reduce")
+class FuseAllReducePass(PassBase):
+    def __init__(self):
+        super().__init__()
+        self.set_attr("max_memory_size", -1)
+
+    def _check_self(self):
+        max_memory_size = self.get_attr("max_memory_size")
+        return max_memory_size > 0
+
+    def _check_conflict(self, other_pass):
+        return True
+
+    def _type(self):
+        return PassType.COMM_OPT
+
+    # NOTE: why FuseAllReducePass can override apply_single_impl instead of
+    # apply_impl? AllReduce is a collective operation, so the program of each
+    # rank inside the same communication group should have the same
+    # all_reduce sum operations. Therefore, FuseAllReducePass can override
+    # apply_single_impl directly.
+    def _apply_single_impl(self, main_program, startup_program, context):
+        max_memory_size = self.get_attr("max_memory_size")
+        op_deps = main_program.desc.get_op_deps()
+        num_blocks = main_program.num_blocks
+        for i in range(num_blocks):
+            block = main_program.block(i)
+            groups = find_all_fuse_all_reduce_groups(block)
+            groups = split_fuse_all_reduce_groups_by_deps(
+                block, groups, op_deps[i]
+            )
+            insert_fuse_all_reduce_by_memory_size(
+                block, groups, max_memory_size
+            )
+        main_program._sync_with_cpp()