Add strategy config for MP allreduce matmul_grad overlapping (#58304)

From00 · web-flow · commit 9cb5a38cdc7f · 2023-10-23T15:04:13.000+08:00
* Add strategy config for MP allreduce matmul_grad overlapping

* set RUN_TYPE=EXCLUSIVE for test_mp_allreduce_matmul_grad_overlapping

* Set timeout
diff --git a/python/paddle/distributed/auto_parallel/constants.py b/python/paddle/distributed/auto_parallel/constants.py
@@ -157,3 +157,11 @@ def set_field_default_config(category, field, default_value):
 set_field_default_config(DP_OPTIMIZATION, "fuse_all_reduce_ops", True)
 set_field_default_config(DP_OPTIMIZATION, "fuse_grad_size_in_MB", 32)
 set_field_default_config(DP_OPTIMIZATION, "overlap_comm_cacl", True)
+
+#########################################
+# model parallel configuration
+#########################################
+MP_OPTIMIZATION = "mp_optimization"
+set_field_default_config(
+    MP_OPTIMIZATION, "allreduce_matmul_grad_overlapping", False
+)
diff --git a/python/paddle/distributed/auto_parallel/static/parallelizer_v2.py b/python/paddle/distributed/auto_parallel/static/parallelizer_v2.py
@@ -14,7 +14,6 @@
 
 import copy
 import logging
-import os
 import time
 
 from paddle.distributed.passes import PassManager, new_pass
@@ -355,14 +354,11 @@ def _apply_post_optimization(
             )
             params_grads = self._pass_context.get_attr("params_grads")
 
-        mp_async_allreduce_in_backward = os.getenv(
-            "FLAGS_mp_async_allreduce_in_backward"
-        ) in [1, "1", True, "True"]
-        if mp_async_allreduce_in_backward:
-            column_parallel_linear_backward_overlapping_pass = new_pass(
-                "column_parallel_linear_backward_overlapping", {}
+        if self._strategy.mp_optimization.allreduce_matmul_grad_overlapping:
+            allreduce_matmul_grad_overlapping_pass = new_pass(
+                "allreduce_matmul_grad_overlapping", {}
             )
-            column_parallel_linear_backward_overlapping_pass.apply(
+            allreduce_matmul_grad_overlapping_pass.apply(
                 [main_program], [startup_program], self._pass_context
             )
 
diff --git a/python/paddle/distributed/auto_parallel/strategy.py b/python/paddle/distributed/auto_parallel/strategy.py
@@ -136,6 +136,12 @@ def __init__(self, config_dict=None):
         super().__init__(category, config_dict)
 
 
+class MPOptimizationConfig(BaseConfig):
+    def __init__(self, config_dict=None):
+        category = constants.MP_OPTIMIZATION
+        super().__init__(category, config_dict)
+
+
 class Strategy(BaseConfig):
     """
     The `Strategy` object is used to configure the parallelization and optimization behaviors.
@@ -214,3 +220,6 @@ def __init__(self, config=None):
 
         config_dict = self._config_dict.get(constants.DP_OPTIMIZATION, None)
         self.dp_optimization = DPOptimizationConfig(config_dict)
+
+        config_dict = self._config_dict.get(constants.MP_OPTIMIZATION, None)
+        self.mp_optimization = MPOptimizationConfig(config_dict)
diff --git a/python/paddle/distributed/passes/__init__.py b/python/paddle/distributed/passes/__init__.py
@@ -24,7 +24,7 @@
 from .auto_parallel_grad_clip import *  # noqa: F403
 from .auto_parallel_supplement_explicit_dependencies import *  # noqa: F403
 from .auto_parallel_pipeline import *  # noqa: F403
-from .column_parallel_linear_backward_overlapping import *  # noqa: F403
+from .allreduce_matmul_grad_overlapping import *  # noqa: F403
 from .cpp_pass import *  # noqa: F403
 from .fuse_all_reduce import *  # noqa: F403
 from .pipeline_scheduler_pass import *  # noqa: F403
diff --git a/python/paddle/distributed/passes/allreduce_matmul_grad_overlapping.py b/python/paddle/distributed/passes/allreduce_matmul_grad_overlapping.py
@@ -27,8 +27,8 @@
 #   dY = matmul(X^T, dOut)
 #
 # Then the c_allreduce_sum can overlap with the compute of dY.
-@register_pass("column_parallel_linear_backward_overlapping")
-class ColumnParallelLinearBackwardOverlappingPass(PassBase):
+@register_pass("allreduce_matmul_grad_overlapping")
+class AllreduceMatmulGradOverlappingPass(PassBase):
     def __init__(self):
         super().__init__()
         self.set_attr("allreduce_stream", None)
diff --git a/test/auto_parallel/CMakeLists.txt b/test/auto_parallel/CMakeLists.txt
@@ -10,6 +10,10 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
                   test_auto_parallel_relaunch)
   set_tests_properties(test_auto_parallel_relaunch
                        PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 120)
+  py_test_modules(test_mp_allreduce_matmul_grad_overlapping MODULES
+                  test_mp_allreduce_matmul_grad_overlapping)
+  set_tests_properties(test_mp_allreduce_matmul_grad_overlapping
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 120)
   py_test_modules(test_relaunch_with_planner MODULES test_relaunch_with_planner)
   set_tests_properties(test_relaunch_with_planner
                        PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 120)
diff --git a/test/auto_parallel/mp_allreduce_matmul_grad_overlapping_unittest.py b/test/auto_parallel/mp_allreduce_matmul_grad_overlapping_unittest.py
@@ -0,0 +1,93 @@
+# Copyright (c) 2023 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
+import unittest
+
+import numpy as np
+from get_gpt_model import FakeDataset, generate_model
+
+import paddle
+from paddle.distributed.fleet import auto
+
+paddle.enable_static()
+
+
+def reset_prog():
+    paddle.base.framework.switch_main_program(paddle.static.Program())
+    paddle.base.framework.switch_startup_program(paddle.static.Program())
+
+
+class TestMPAllreduceMatmulGradOverlapping(unittest.TestCase):
+    def setUp(self):
+        self.rtol = 1e-5
+        self.atol = 1e-8
+        self.batch_size = 1
+        self.batch_num = 10
+        self.clip_norm = 0.2
+        self.dataset = FakeDataset(self.batch_size * self.batch_num)
+
+    def init(self, engine):
+        paddle.seed(2023)
+        np.random.seed(2023)
+        random.seed(2023)
+        place = paddle.base.CUDAPlace(paddle.distributed.ParallelEnv().dev_id)
+        engine._executor = paddle.static.Executor(place)
+
+    def get_mp_engine(self, allreduce_matmul_grad_overlapping):
+        reset_prog()
+
+        strategy = auto.Strategy()
+        strategy.auto_mode = "semi"
+        strategy.reinit = True
+        strategy.mp_optimization.allreduce_matmul_grad_overlapping = (
+            allreduce_matmul_grad_overlapping
+        )
+
+        clip = paddle.nn.ClipGradByGlobalNorm(self.clip_norm)
+        opt = paddle.optimizer.AdamW(learning_rate=0.00001, grad_clip=clip)
+        model, loss = generate_model("mp")
+
+        engine = auto.Engine(model, loss, opt, strategy=strategy)
+        self.init(engine)
+        return engine
+
+    def run_mp(self, allreduce_matmul_grad_overlapping):
+        mp_engine = self.get_mp_engine(allreduce_matmul_grad_overlapping)
+        history = mp_engine.fit(self.dataset, 3, batch_size=self.batch_size)
+        return np.array(history.history["loss"])
+
+    def check_results(self, ref_losses, check_losses, rtol=None, atol=None):
+        np.testing.assert_allclose(
+            ref_losses,
+            check_losses,
+            rtol=rtol or self.rtol,
+            atol=atol or self.atol,
+            err_msg='pass {} has wrong results!, \nu={}\nv={}\ndiff={}'.format(
+                __class__, ref_losses, check_losses, ref_losses - check_losses
+            ),
+        )
+
+    def test_mp_allreduce_matmul_grad_overlapping(self):
+        losses_with_allreduce_matmul_grad_overlapping = self.run_mp(True)
+        losses_without_allreduce_matmul_grad_overlapping = self.run_mp(False)
+
+        np.testing.assert_equal(
+            losses_with_allreduce_matmul_grad_overlapping,
+            losses_without_allreduce_matmul_grad_overlapping,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/auto_parallel/test_mp_allreduce_matmul_grad_overlapping.py b/test/auto_parallel/test_mp_allreduce_matmul_grad_overlapping.py
@@ -0,0 +1,57 @@
+# Copyright (c) 2023 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 os
+import subprocess
+import sys
+import tempfile
+import unittest
+
+
+class TestMPAllreduceMatmulGradOverlapping(unittest.TestCase):
+    def test_mp_allreduce_matmul_grad_overlapping(self):
+        file_dir = os.path.dirname(os.path.abspath(__file__))
+        launch_model_path = os.path.join(
+            file_dir, "mp_allreduce_matmul_grad_overlapping_unittest.py"
+        )
+
+        if os.environ.get("WITH_COVERAGE", "OFF") == "ON":
+            coverage_args = ["-m", "coverage", "run", "--branch", "-p"]
+        else:
+            coverage_args = []
+
+        tmp_dir = tempfile.TemporaryDirectory()
+        cmd = (
+            [sys.executable, "-u"]
+            + coverage_args
+            + [
+                "-m",
+                "paddle.distributed.launch",
+                "--devices",
+                "0,1",
+                "--log_dir",
+                tmp_dir.name,
+                launch_model_path,
+            ]
+        )
+
+        process = subprocess.Popen(cmd)
+        process.wait()
+        self.assertEqual(process.returncode, 0)
+
+        tmp_dir.cleanup()
+
+
+if __name__ == "__main__":
+    unittest.main()
