[auto parallel] Support 1F1B for PIR (#66810)

AndSonder · heavyrain-lzy · web-flow · commit efdd967b236d · 2024-08-08T14:40:07.000+08:00
* split program and polish the executor

* change unit test

* fix 1f1b test

* add pp in pir

* fix ut

* fit grad_merge for pir

* update test

* apply suggestions from code review

* fit for 1f1b

* fit_1f1b

* remove no related codes

* recover codes about enable_send_recv_overlap

* add test for 1f1b plan

* fix

* enable gradient_merge when pipeline pass enabled

---------

Co-authored-by: heavyrain_lzy &lt;1528794076@qq.com&gt;
diff --git a/python/paddle/distributed/auto_parallel/constants.py b/python/paddle/distributed/auto_parallel/constants.py
@@ -113,6 +113,7 @@ def set_field_default_config(category, field, default_value):
 PIPELINE = "pipeline"
 set_field_default_config(PIPELINE, "enable", False)
 set_field_default_config(PIPELINE, "schedule_mode", "1F1B")
+set_field_default_config(PIPELINE, "pp_degree", 1)
 set_field_default_config(PIPELINE, "vpp_degree", 1)
 set_field_default_config(PIPELINE, "vpp_seg_method", "")
 set_field_default_config(PIPELINE, "micro_batch_size", 1)
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -822,7 +822,6 @@ def _parallel_pir(self, mode):
         #   resolute the reshard op into special collective operation.
         #   collect the communicator created during resolution.
         apply_reshard_pass(dist_program)
-
         # print('after reshard', dist_program, flush=1)
 
         remove_other_rank_input_output_pass(dist_program)
@@ -857,6 +856,13 @@ def _parallel_pir(self, mode):
             # dist_program = apply_sharding_optimization_pass(dist_program)
             pass
 
+        if mode == "train" and self._strategy.pipeline.enable:
+            self._strategy.gradient_merge.enable = True
+            self._strategy.gradient_merge.k_steps = (
+                self._strategy.pipeline.accumulate_steps
+            )
+            self._strategy.gradient_merge.avg = True
+
         if mode == "train" and self._strategy.gradient_merge.enable:
             config = copy.deepcopy(self._strategy.gradient_merge.to_dict())
             config["gradient_sync_after_accumulate"] = (
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -21,6 +21,7 @@
     choose_reshard_func,
 )
 from .reshard_funcs.reshard_func_register import register_reshard_funcs
+from .utils import get_pp_stage_by_pp_degree
 
 register_reshard_funcs()
 
@@ -406,10 +407,15 @@ def pipeline_pass(dense_main_program, dense_starup_program, pipeline_strategy):
     pass_name = pipeline_strategy.schedule_mode
     assert pass_name in [
         "FThenB",
+        "1F1B",
     ], f"pipeline scheduler only support FThenB now, but receive {pass_name}"
 
     pass_attr = {}
     pass_attr["num_micro_batches"] = pipeline_strategy.accumulate_steps
+    pass_attr["pp_degree"] = pipeline_strategy.pp_degree
+    pass_attr["pp_stage"] = get_pp_stage_by_pp_degree(
+        pipeline_strategy.pp_degree
+    )
 
     if pass_name == "1F1B":
         # TODO(Ruibiao): Move FLAGS_1f1b_backward_forward_overlap and
diff --git a/python/paddle/distributed/auto_parallel/static/utils.py b/python/paddle/distributed/auto_parallel/static/utils.py
@@ -2341,6 +2341,14 @@ def get_pp_stage(dist_context, rank):
     return pp_idx
 
 
+def get_pp_stage_by_pp_degree(pp_degree):
+    cur_rank = paddle.distributed.get_rank()
+    word_size = paddle.distributed.get_world_size()
+    pp_group_size = word_size // pp_degree
+    pp_stage = cur_rank // pp_group_size
+    return pp_stage
+
+
 def wrap_data_for_completion(
     dist_op, input_names: list, output_names: list, attr_names: list
 ):
diff --git a/python/paddle/distributed/passes/auto_parallel_gradient_merge.py b/python/paddle/distributed/passes/auto_parallel_gradient_merge.py
@@ -21,6 +21,7 @@
     is_data_parallel_reduce_op,
     is_data_parallel_scale_op,
 )
+from paddle.distributed.auto_parallel.static.pir_pass import comm_ops
 from paddle.distributed.auto_parallel.static.process_group import (
     get_world_process_group,
 )
@@ -277,6 +278,9 @@ def _pir_append_gradient_merge_backward_op(
     new_params_grads = []
 
     for param, grad in params_grads:
+        if grad is None:
+            continue
+
         assert (
             not param.is_selected_row_type()
         ), "SELECTED_ROWS is not supported in GradientMergeOptimizer for now"
@@ -286,7 +290,7 @@ def _pir_append_gradient_merge_backward_op(
         with startup_block:
             paddle.pir.set_insertion_point_to_block_end(startup_block)
             gradient_merge_var = paddle.full(
-                shape=grad.shape, fill_value=float(0), dtype=grad.dtype
+                shape=grad._local_shape, fill_value=float(0), dtype=grad.dtype
             )
             gradient_merge_var.persistable = True
 
@@ -305,7 +309,15 @@ def _pir_append_gradient_merge_backward_op(
             )
             new_gradient_merge_var.persistable = True
 
-            grad.replace_all_uses_with(new_gradient_merge_var)
+            opt_ops_use_grad = [
+                op
+                for op in grad.all_used_ops()
+                if op.op_role == int(OpRole.Optimize)
+            ]
+            grad.replace_grad_users_with(
+                new_gradient_merge_var, set(opt_ops_use_grad)
+            )
+
             new_gradient_merge_var_add = paddle._C_ops.add_(
                 new_gradient_merge_var, grad
             )
@@ -367,12 +379,7 @@ def _pir_move_reduce_to_optimizer_stage(main_program):
     main_block = main_program.global_block()
 
     for idx, op in list(enumerate(main_block.ops)):
-        if op.name() in [
-            "pd_op.c_allreduce_sum",
-            "pd_op.c_allreduce_avg",
-            "pd_op.c_reduce_sum",
-            "pd_op.c_reduce_avg",
-        ]:
+        if op.name() in comm_ops:
             op_input_names = op.get_input_names()
             # NOTE(sonder): When "@RENAME@" is in the input name, it means that the op has been renamed.
             # Such types input names are caused by shared parameter policy.
@@ -648,8 +655,8 @@ def _pir_parse_program(
         for _, new_grad in new_params_to_grads:
             scale = paddle.full([], 1.0 / k_steps)
             new_grad = paddle._C_ops.scale_(new_grad, scale, 0.0, False)
-            new_grad.get_defining_op().op_role = int(OpRole.Backward)
-            scale.get_defining_op().op_role = int(OpRole.Backward)
+            new_grad.get_defining_op().op_role = int(OpRole.Optimize)
+            scale.get_defining_op().op_role = int(OpRole.Optimize)
 
 
 @register_pass("auto_parallel_gradient_merge_pass")
diff --git a/python/paddle/distributed/passes/pipeline_scheduler_pass/pipeline_1f1b.py b/python/paddle/distributed/passes/pipeline_scheduler_pass/pipeline_1f1b.py
@@ -22,6 +22,7 @@
 from ..pass_utils import (
     AutoParallelStreamType,
     _add_event_dependency,
+    _pir_program_for_fthenb_and_1f1b,
     _program_for_fthenb_and_1f1b,
     split_program,
 )
@@ -336,6 +337,24 @@ def _partial_programs(self, program):
 
         return types, sub_programs
 
+    def _partial_pir_programs(self, program):
+        enable_send_recv_overlap = self.get_attr("enable_send_recv_overlap")
+        assert (
+            not enable_send_recv_overlap
+        ), "PIR does not support 1F1B with enable_send_recv_overlap yet."
+
+        types = [FORWARD, BACKWARD, OPT]
+        sub_program_list = _pir_program_for_fthenb_and_1f1b(
+            program, enable_send_recv_overlap
+        )
+
+        for i in range(len(types)):
+            logger.debug(
+                f"type = {types[i]}, sub_programs = {sub_program_list[i]}\n"
+            )
+        logger.debug(f"jobs_in_stable_phase = {self.jobs_in_stable_phase}")
+        return types, sub_program_list
+
     def _split_program_for_overlapping(self, job_type, program, split_points):
         assert job_type in [
             FORWARD,
diff --git a/test/auto_parallel/pir/1F1B_pass_unittest_pir.py b/test/auto_parallel/pir/1F1B_pass_unittest_pir.py
@@ -0,0 +1,159 @@
+# Copyright (c) 2024 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
+#
+#     https://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 random
+from functools import reduce
+
+import numpy as np
+from mlp_demo import PPDemoNet
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.io import BatchSampler, DataLoader
+
+BATCH_SIZE = 4
+BATCH_NUM = 4
+IMAGE_SIZE = 16
+CLASS_NUM = 8
+
+
+def apply_pass(use_1f1b=False):
+    strategy = dist.Strategy()
+    strategy.auto_mode = "semi"
+    strategy.reinit = True
+
+    if use_1f1b:
+        pipeline = strategy.pipeline
+        pipeline.enable = True
+        pipeline.schedule_mode = "1F1B"
+        pipeline.accumulate_steps = 2
+    else:
+        pipeline = strategy.pipeline
+        pipeline.enable = True
+        pipeline.schedule_mode = "FThenB"
+        pipeline.accumulate_steps = 2
+
+    return strategy
+
+
+class RandomDataset(paddle.io.Dataset):
+    def __init__(self, images, labels, num_samples, return_dict=False):
+        self.images = images
+        self.labels = labels
+        self.num_samples = num_samples
+        self.return_dict = return_dict
+
+    def __getitem__(self, idx):
+        if self.return_dict:
+            return {
+                "image": self.images[idx],
+                "label": self.labels[idx],
+            }
+        else:
+            return self.images[idx], self.labels[idx]
+
+    def __len__(self):
+        return self.num_samples
+
+
+class TestSimpleNet1F1BPass:
+    def init_dist_env(self):
+        order = ["dp", "pp", "mp"]
+        dp_degree = 1
+        mp_degree = 1
+        pp_degree = 2
+        degree = [dp_degree, pp_degree, mp_degree]
+        mesh_dims = list(filter(lambda x: x[1] > 1, list(zip(order, degree))))
+        if not mesh_dims:
+            mesh_dims = [("dp", 1)]
+        dim_names = [mesh_dim[0] for mesh_dim in mesh_dims]
+        mesh_shape = [mesh_dim[1] for mesh_dim in mesh_dims]
+        mesh_arr = np.arange(
+            0, reduce(lambda x, y: x * y, mesh_shape, 1)
+        ).reshape(mesh_shape)
+        global_mesh = dist.ProcessMesh(mesh_arr, dim_names)
+        dist.auto_parallel.set_mesh(global_mesh)
+        paddle.seed(2024)
+        np.random.seed(2024)
+        random.seed(2024)
+
+    def create_data_loader(
+        self,
+        batch_size=BATCH_SIZE,
+        batch_num=BATCH_NUM,
+        image_size=IMAGE_SIZE,
+        class_num=CLASS_NUM,
+    ):
+        nsamples = batch_size * batch_num
+        images = np.random.rand(nsamples, image_size).astype('float32')
+        labels = np.random.rand(nsamples, class_num).astype('float32')
+        dataset = RandomDataset(images, labels, nsamples)
+        data_sampler = BatchSampler(
+            dataset, batch_size=batch_size, shuffle=False, drop_last=True
+        )
+        loader = DataLoader(dataset, batch_sampler=data_sampler, num_workers=0)
+        return loader
+
+    def test_pp_1f1b(self):
+        self.init_dist_env()
+        strategy_fthenb = apply_pass(use_1f1b=False)
+        loss_fthenb = self.run_pipeline(strategy_fthenb)
+
+        self.init_dist_env()
+        strategy_1f1b = apply_pass(use_1f1b=True)
+        loss_1f1b = self.run_pipeline(strategy_1f1b)
+
+        cur_rank = paddle.distributed.get_rank()
+        if cur_rank == 1:
+            self.check_result(loss_1f1b, loss_fthenb)
+
+    def run_pipeline(self, strategy):
+        mesh1 = dist.ProcessMesh([0], dim_names=["pp"])
+        mesh2 = dist.ProcessMesh([1], dim_names=["pp"])
+        model = PPDemoNet(mesh1, mesh2)
+
+        lr_scheduler = paddle.optimizer.lr.LinearWarmup(
+            learning_rate=0.1, warmup_steps=2, start_lr=0, end_lr=0.0001
+        )
+
+        opt = paddle.optimizer.SGD(
+            learning_rate=lr_scheduler, parameters=model.parameters()
+        )
+
+        loader = self.create_data_loader()
+        dist_loader = dist.shard_dataloader(loader, meshes=[mesh1, mesh2])
+
+        loss_fn = nn.MSELoss()
+
+        dist_model = dist.to_static(model, dist_loader, loss_fn, opt, strategy)
+
+        dist_model.train()
+        loss = None
+        for _, (image, label) in enumerate(dist_loader()):
+            loss = dist_model(image, label)
+            lr_scheduler.step()
+            if int(dist.get_rank()) == 1:
+                assert loss is not None
+            else:
+                assert loss is None
+
+        return np.array(loss)
+
+    def check_result(self, loss1, loss2):
+        return np.array_equal(loss1, loss2)
+
+
+if __name__ == '__main__':
+    TestSimpleNet1F1BPass().test_pp_1f1b()
diff --git a/test/auto_parallel/pir/CMakeLists.txt b/test/auto_parallel/pir/CMakeLists.txt
@@ -35,7 +35,13 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
   py_test_modules(
     test_eliminate_transpose_pass MODULES test_eliminate_transpose_pass ENVS
     FLAGS_enable_pir_in_executor=1)
+  py_test_modules(test_pipeline_scheduler_1f1b_pir MODULES
+                  test_pipeline_scheduler_1f1b_pir ENVS FLAGS_enable_pir_api=1)
+  set_tests_properties(test_pipeline_scheduler_1f1b_pir
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 50)
   py_test_modules(test_moe_api MODULES test_moe_api ENVS FLAGS_enable_pir_api=1)
   py_test_modules(test_pir_stack_grad_spmd_rule MODULES
                   test_stack_grad_spmd_rule ENVS FLAGS_enable_pir_api=1)
 endif()
+py_test_modules(test_pir_1f1b_plan MODULES test_pir_1f1b_plan ENVS
+                FLAGS_enable_pir_api=1)
diff --git a/test/auto_parallel/pir/test_pipeline_scheduler_1f1b_pir.py b/test/auto_parallel/pir/test_pipeline_scheduler_1f1b_pir.py
@@ -0,0 +1,47 @@
+# Copyright (c) 2024 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
+#
+#     https://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 unittest
+
+import collective.test_communication_api_base as test_base
+
+
+class Test1F1BPass(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=2,
+            timeout=50,
+        )
+        self._default_envs = {
+            "FLAGS_embedding_deterministic": "1",
+            "FLAGS_cudnn_deterministic": "1",
+            "FLAGS_enable_pir_api": "1",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+        }
+
+    def test_pp(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "1F1B_pass_unittest_pir.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/auto_parallel/pir/test_pir_1f1b_plan.py b/test/auto_parallel/pir/test_pir_1f1b_plan.py
