chore: import upstream snapshot with attribution

test/distributed_passes/CMakeLists.txt

@@ -0,0 +1,28 @@
+file(
+  GLOB TEST_OPS
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+if((NOT WITH_GPU) AND (NOT WITH_XPU))
+  list(REMOVE_ITEM TEST_OPS "test_dist_fuse_all_reduce_pass")
+endif()
+
+if(NOT ((WITH_GPU) AND (CUDA_VERSION GREATER_EQUAL 11.6)))
+  list(REMOVE_ITEM TEST_OPS test_dist_fuse_gemm_epilogue_pass)
+endif()
+
+if(NOT WITH_CUDNN_FRONTEND)
+  list(REMOVE_ITEM TEST_OPS "test_dist_fuse_resunit_pass")
+endif()
+
+foreach(TEST_OP ${TEST_OPS})
+  py_test_modules(${TEST_OP} MODULES ${TEST_OP} ENVS "NVIDIA_TF32_OVERRIDE=0")
+  list(APPEND DIST_TEST_OPS ${TEST_OP})
+  set_tests_properties(${TEST_OP} PROPERTIES TIMEOUT 250)
+  set_tests_properties(${TEST_OP} PROPERTIES LABELS "RUN_TYPE=DIST")
+endforeach()
+
+if(WITH_CUDNN_FRONTEND)
+  set_tests_properties(test_dist_fuse_resunit_pass PROPERTIES TIMEOUT 350)
+endif()

test/distributed_passes/check_pass_conflict_example.py

@@ -0,0 +1,47 @@
+# 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 unittest
+
+from dist_pass_test_base import PassConflictChecker
+from model_zoo import resnet_model
+
+from paddle.distributed.passes import new_pass
+
+
+class CheckPassConflictTest1(PassConflictChecker):
+    def pass_config(self):
+        return [
+            new_pass("fuse_all_reduce", {"max_memory_size": 1024 * 1024}),
+            new_pass("fuse_elewise_add_act"),
+        ]
+
+    def test_resnet(self):
+        self.check_main(resnet_model, batch_size=32)
+
+
+class CheckPassConflictTest2(PassConflictChecker):
+    def pass_config(self):
+        return [
+            new_pass("fuse_elewise_add_act"),
+            new_pass("fuse_all_reduce", {"max_memory_size": 1024 * 1024}),
+        ]
+
+    def test_resnet(self):
+        with self.assertRaises(Exception):  # noqa: B017
+            self.check_main(resnet_model, batch_size=32)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/distributed_passes/dist_pass_test_base.py

@@ -0,0 +1,294 @@
+# 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 inspect
+import os
+import pickle
+import shlex
+import shutil
+import sys
+import unittest
+from collections import OrderedDict
+
+import numpy as np
+
+import paddle
+from paddle.distributed.fleet.launch_utils import run_with_coverage
+from paddle.distributed.passes.pass_base import PassBase, PassManager
+
+
+def prepare_python_path_and_return_module(path):
+    dirname, filename = os.path.split(path)
+    py_suffix = ".py"
+    assert filename.endswith(py_suffix), filename
+
+    env_name = 'PYTHONPATH'
+    python_path = os.environ.get(env_name, '')
+    if python_path:
+        paths = [p for p in python_path.split(":") if p]
+        if dirname not in paths:
+            paths.append(dirname)
+        python_path = ":".join(paths)
+    else:
+        python_path = dirname
+    os.environ[env_name] = python_path
+    print('GLOG_v=', os.environ.get('GLOG_v', None), flush=1)
+    return filename[: -len(py_suffix)]
+
+
+def remove_path_if_exists(path):
+    if not os.path.exists(path):
+        return
+
+    if os.path.isfile(path):
+        os.remove(path)
+    else:
+        shutil.rmtree(path)
+
+
+# NOTE: only support GPU now
+class DistPassTestBase(unittest.TestCase):
+    def setUp(self):
+        paddle.enable_static()
+        if paddle.is_compiled_with_cuda():
+            paddle.set_flags({'FLAGS_cudnn_deterministic': 1})
+
+        seed = int(os.environ.get('SEED', -1))
+        if seed <= 0:
+            seed = np.random.randint(low=1, high=1000000, size=[1])[0]
+            os.environ['SEED'] = str(seed)
+        self.seed = seed
+        paddle.seed(self.seed)
+
+        self.rtol = 1e-5
+        self.atol = 1e-8
+        self.equal_nan = False
+
+        self.init()
+
+    def init(self):
+        pass
+
+    def get_model(self, place, **kwargs):
+        raise NotImplementedError
+
+    def apply_passes(self, main_prog, startup_prog):
+        raise NotImplementedError
+
+    def check_main(self, model=None, gpus=None, **kwargs):
+        pass_rets = self._distributed_launch(
+            model=model, apply_pass=True, gpus=gpus, **kwargs
+        )
+        no_pass_rets = self._distributed_launch(
+            model=model, apply_pass=False, gpus=gpus, **kwargs
+        )
+        self.check_results(no_pass_rets, pass_rets)
+
+    def check_results(self, no_pass_rets, pass_rets):
+        self.assertEqual(len(no_pass_rets), len(pass_rets))
+        for no_pass_ret, pass_ret in zip(no_pass_rets, pass_rets):
+            self.assertEqual(len(no_pass_ret), len(pass_ret))
+            for i, (out_var_no_pass, out_var_pass) in enumerate(
+                zip(no_pass_ret, pass_ret)
+            ):
+                if out_var_no_pass is None:
+                    self.assertIsNone(out_var_pass)
+                else:
+                    self.assertEqual(len(out_var_pass), len(out_var_no_pass))
+                    for i in range(0, len(out_var_pass)):
+                        np.testing.assert_allclose(
+                            out_var_no_pass[i],
+                            out_var_pass[i],
+                            rtol=self.rtol,
+                            atol=self.atol,
+                            equal_nan=self.equal_nan,
+                        )
+
+    @classmethod
+    def _to_var_names(cls, names_or_vars):
+        if not isinstance(names_or_vars, (list, tuple)):
+            names_or_vars = [names_or_vars]
+        ret_var_names = []
+        for name_or_var in names_or_vars:
+            if isinstance(name_or_var, str):
+                ret_var_names.append(name_or_var)
+            else:
+                ret_var_names.append(name_or_var.name)
+        return ret_var_names
+
+    def _run_gpu_main(self, model, apply_pass, dump_file, **kwargs):
+        gpu_id = int(os.environ.get('FLAGS_selected_gpus', 0))
+        place = paddle.CUDAPlace(gpu_id)
+        scope = paddle.static.Scope()
+        if model is None:
+            model = self.get_model
+        with (
+            paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ),
+            paddle.static.scope_guard(scope),
+            paddle.base.unique_name.guard(),
+        ):
+            main_prog, startup_prog, inputs, outputs, reader = model(
+                place, **kwargs
+            )
+            inputs = self._to_var_names(inputs)
+            outputs = self._to_var_names(outputs)
+            if apply_pass:
+                self.apply_passes(main_prog, startup_prog)
+
+        all_fetch_values = []
+        exe = paddle.static.Executor(place)
+        with paddle.static.scope_guard(scope):
+            exe.run(startup_prog)
+            for batch_id, input_data in enumerate(reader()):
+                assert len(input_data) == len(inputs), (
+                    f"{len(input_data)} vs {len(inputs)}"
+                )
+                feed = dict(zip(inputs, input_data))
+                fetch_values = exe.run(main_prog, feed=feed, fetch_list=outputs)
+                if paddle.distributed.get_rank() == 0:
+                    output_dict = OrderedDict(zip(outputs, fetch_values))
+                    print(f'batch {batch_id}, outputs {output_dict}')
+                all_fetch_values.append(fetch_values)
+        with open(dump_file, "wb") as f:
+            pickle.dump(all_fetch_values, f)
+
+    @classmethod
+    def _get_default_gpu_lists(cls):
+        visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
+        if visible_devices is None:
+            visible_devices = os.getenv("FLAGS_selected_gpus")
+
+        if visible_devices is None:
+            num_gpus = paddle.device.cuda.device_count()
+            return list(range(num_gpus))
+        else:
+            return [
+                int(s.strip()) for s in visible_devices.split(",") if s.strip()
+            ]
+
+    def _distributed_launch(self, model, apply_pass, gpus=None, **kwargs):
+        if gpus is None:
+            gpus = self._get_default_gpu_lists()
+
+        num_gpus = len(gpus)
+        gpus = ','.join([str(gpu_id) for gpu_id in gpus])
+
+        pid = os.getpid()
+        if apply_pass:
+            output_dir = f"test_with_pass_{pid}"
+        else:
+            output_dir = f"test_without_pass_{pid}"
+        remove_path_if_exists(output_dir)
+        os.makedirs(output_dir, mode=0o777)
+
+        input_dump_file = os.path.join(output_dir, 'inputs.bin')
+        model_dump_file = os.path.join(output_dir, 'model.bin')
+
+        if os.environ.get("WITH_COVERAGE", "OFF") == "ON":
+            run_with_coverage(True)
+            coverage_args = ["-m", "coverage", "run", "--branch", "-p"]
+        else:
+            coverage_args = []
+
+        file_dir = os.path.dirname(os.path.abspath(__file__))
+
+        try:
+            with open(input_dump_file, 'wb') as f:
+                pickle.dump(kwargs, f)
+
+            if model is not None:
+                with open(model_dump_file, 'wb') as f:
+                    pickle.dump(model, f)
+
+            cmd = [
+                sys.executable,
+                "-u",
+                *coverage_args,
+                "-m",
+                "launch",
+                "--log_dir",
+                output_dir,
+                "--gpus",
+                gpus,
+                os.path.join(file_dir, "pass_run_main.py"),
+                "--file_path",
+                inspect.getfile(type(self)),
+                "--class_name",
+                type(self).__name__,
+                "--input_file",
+                input_dump_file,
+                "--output_dir",
+                output_dir,
+            ]
+            if apply_pass:
+                cmd += ["--apply_pass"]
+            if model is not None:
+                cmd += ["--model_file", model_dump_file]
+            cmd = [shlex.quote(c) for c in cmd]
+            prepare_python_path_and_return_module(__file__)
+            exitcode = os.system(' '.join(cmd))
+            self.assertEqual(
+                exitcode,
+                0,
+                f"Pass test failed with apply_pass = {apply_pass}, please view log in {output_dir}",
+            )
+
+            results = []
+            for i in range(num_gpus):
+                dump_file = f'{output_dir}/{i}.bin'
+                self.assertTrue(
+                    os.path.exists(dump_file),
+                    f"Pass test failed with apply_pass = {apply_pass}, please view log in {output_dir}",
+                )
+                with open(dump_file, "rb") as f:
+                    results.append(pickle.load(f))
+            return results
+        finally:
+            if int(os.environ.get("DEBUG", 0)) == 0:
+                remove_path_if_exists(output_dir)
+
+
+class PassConflictChecker(DistPassTestBase):
+    def setUp(self):
+        os.environ['DEBUG'] = '0'
+        super().setUp()
+
+    def pass_config(self):
+        raise NotImplementedError
+
+    def apply_passes(self, main_prog, startup_prog):
+        passes = self.pass_config()
+        if not isinstance(passes, (list, tuple)):
+            passes = [passes]
+        for p in passes:
+            self.assertTrue(isinstance(p, PassBase))
+
+        auto_pass_manager = PassManager(passes, auto_solve_conflict=True)
+        new_passes = auto_pass_manager.passes
+        self.assertEqual(
+            len(passes),
+            len(new_passes),
+            f"After solving conflicts, the left passes are: {auto_pass_manager.names}",
+        )
+
+        for i, (p1, p2) in enumerate(zip(passes, new_passes)):
+            self.assertEqual(
+                id(p1),
+                id(p2),
+                f"After solving conflicts, the {i}-th pass is different: {p1.name} vs {p2.name}",
+            )
+
+        auto_pass_manager.apply([main_prog], [startup_prog])

test/distributed_passes/launch.py

@@ -0,0 +1,23 @@
+# 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 os
+
+from paddle.distributed.fleet import launch
+from paddle.distributed.fleet.launch_utils import run_with_coverage
+
+if __name__ == "__main__":
+    if os.environ.get("WITH_COVERAGE", "OFF") == "ON":
+        run_with_coverage(True)
+    launch.launch()

test/distributed_passes/model_zoo.py

@@ -0,0 +1,109 @@
+# 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 os
+
+import numpy as np
+
+import paddle
+from paddle import nn
+from paddle.distributed import fleet
+from paddle.vision.models import resnet50 as resnet
+
+__all__ = [
+    'resnet_model',
+]
+
+
+def get_seed_from_env():
+    return int(os.environ.get("SEED", 0))
+
+
+def resnet_model(
+    place, batch_size, image_shape=[3, 224, 224], num_classes=1000
+):
+    image = paddle.static.data(
+        shape=[batch_size, *image_shape], dtype='float32', name='image'
+    )
+    label = paddle.static.data(
+        shape=[batch_size, 1], dtype='int64', name='label'
+    )
+    model = resnet(pretrained=False)
+    loss_fn = nn.loss.CrossEntropyLoss()
+    pred_out = model(image)
+    loss = loss_fn(pred_out, label)
+    optimizer = paddle.optimizer.Adam(learning_rate=1e-3)
+
+    dist_strategy = fleet.DistributedStrategy()
+    dist_strategy.fuse_all_reduce_ops = False
+    dist_strategy.without_graph_optimization = True
+    fleet.init(is_collective=True, strategy=dist_strategy)
+    optimizer = fleet.distributed_optimizer(optimizer)
+    optimizer.minimize(loss)
+
+    rank = paddle.distributed.get_rank()
+
+    def reader():
+        seed = get_seed_from_env()
+        np.random.seed(seed + rank)
+        for _ in range(10):
+            image_np = np.random.random(size=image.shape).astype('float32')
+            label_np = np.random.randint(
+                low=0, high=num_classes, size=label.shape
+            ).astype('int64')
+            yield image_np, label_np
+
+    main_program = paddle.static.default_main_program()
+    startup_program = paddle.static.default_startup_program()
+    return main_program, startup_program, [image, label], [loss], reader
+
+
+def simple_net(place, batch_size, image_shape=[784], num_classes=10):
+    image = paddle.static.data(
+        shape=[batch_size, *image_shape], dtype='float32', name='image'
+    )
+    label = paddle.static.data(
+        shape=[batch_size, 1], dtype='int64', name='label'
+    )
+    linears = [nn.Linear(784, 784) for _ in range(3)]
+    hidden = image
+    for linear in linears:
+        hidden = linear(hidden)
+        hidden = nn.ReLU()(hidden)
+    loss_fn = nn.loss.CrossEntropyLoss()
+    loss = loss_fn(hidden, label)
+    optimizer = paddle.optimizer.Adam(learning_rate=1e-3)
+
+    dist_strategy = fleet.DistributedStrategy()
+    dist_strategy.fuse_all_reduce_ops = False
+    dist_strategy.without_graph_optimization = True
+    fleet.init(is_collective=True, strategy=dist_strategy)
+    optimizer = fleet.distributed_optimizer(optimizer)
+    optimizer.minimize(loss)
+
+    rank = paddle.distributed.get_rank()
+
+    def reader():
+        seed = get_seed_from_env()
+        np.random.seed(seed + rank)
+        for _ in range(10):
+            image_np = np.random.random(size=image.shape).astype('float32')
+            label_np = np.random.randint(
+                low=0, high=num_classes, size=label.shape
+            ).astype('int64')
+            yield image_np, label_np
+
+    main_program = paddle.static.default_main_program()
+    startup_program = paddle.static.default_startup_program()
+    return main_program, startup_program, [image, label], [loss], reader

test/distributed_passes/pass_run_main.py

@@ -0,0 +1,93 @@
+# 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 argparse
+import importlib
+import os
+import pickle
+
+from dist_pass_test_base import (
+    DistPassTestBase,
+    prepare_python_path_and_return_module,
+)
+
+import paddle
+from paddle.distributed.fleet.launch_utils import run_with_coverage
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description='arguments for distributed pass tests'
+    )
+    parser.add_argument('--file_path', type=str, help='The test file path.')
+    parser.add_argument(
+        '--class_name',
+        type=str,
+        help='The test class name. It is the class name that inherits the DistPassTestBase class.',
+    )
+    parser.add_argument(
+        '--apply_pass',
+        default=False,
+        action="store_true",
+        help='Whether to apply distributed passes.',
+    )
+    parser.add_argument(
+        '--input_file',
+        type=str,
+        help='The input file which contains the dumped input arguments.',
+    )
+    parser.add_argument(
+        '--output_dir',
+        type=str,
+        help='The output directory to save the logs and output results.',
+    )
+    parser.add_argument(
+        '--model_file',
+        type=str,
+        help='The input model file which contains the dumped model function.',
+    )
+    return parser.parse_args()
+
+
+def run_main(args):
+    if os.environ.get("WITH_COVERAGE", "OFF") == "ON":
+        run_with_coverage(True)
+    module_name = prepare_python_path_and_return_module(args.file_path)
+    test_module = importlib.import_module(module_name)
+    test_class = getattr(test_module, args.class_name)
+    assert issubclass(test_class, DistPassTestBase)
+    test_obj = test_class()
+    rank = paddle.distributed.get_rank()
+    with open(args.input_file, "rb") as f:
+        kwargs = pickle.load(f)
+
+    output_file = f"{args.output_dir}/{rank}.bin"
+    if args.model_file:
+        with open(args.model_file, "rb") as f:
+            model = pickle.load(f)
+    else:
+        model = None
+
+    try:
+        test_obj.setUpClass()
+        test_obj.setUp()
+        test_obj._run_gpu_main(model, args.apply_pass, output_file, **kwargs)
+    finally:
+        test_obj.tearDown()
+        test_obj.tearDownClass()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_main(args)

test/distributed_passes/ps_pass_test_base.py

@@ -0,0 +1,130 @@
+# 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 os
+import shlex
+import sys
+import unittest
+
+from dist_pass_test_base import (  # noqa: F401
+    prepare_python_path_and_return_module,
+    remove_path_if_exists,
+)
+
+
+class PsPassTestBase(unittest.TestCase):
+    def init(self):
+        self.config = {}
+        self.config['ps_mode_config'] = ""
+        self.config['worker_num'] = "1"
+        self.config['server_num'] = "1"
+        self.config['run_minimize'] = "0"
+        self.config['run_single_pass'] = "0"
+        self.config['run_the_one_ps'] = '0'
+        self.config['debug_new_minimize'] = "0"
+        self.config['debug_new_pass'] = "0"
+        self.config['debug_the_one_ps'] = '0'
+        self.config['log_dir'] = ""
+        self.config['applied_pass_name'] = ""
+
+    def setUp(self):
+        print('Ps setUp...')
+
+    def tearDown(self):
+        print('Ps tearDown...')
+
+    def ps_launch(self, ps_mode="cpu-ps"):
+        if ps_mode == "cpu-ps" or ps_mode == 'heter-ps':
+            os.environ['WITH_DISTRIBUTE'] = 'ON'
+
+            cmd = [
+                sys.executable,
+                "-u",
+                "-m",
+                "launch",
+                "--log_dir",
+                self.config['log_dir'],
+                "--worker_num",
+                self.config['worker_num'],
+                "--server_num",
+                self.config['server_num'],
+            ]
+            if ps_mode == 'heter-ps':
+                os.environ['FLAGS_START_PORT'] = '12004'
+                cmd += [
+                    '--heter_worker_num',
+                    self.config['heter_worker_num'],
+                    '--heter_devices',
+                    self.config['heter_devices'],
+                ]
+
+            cmd += [
+                "../ps/ps_dnn_trainer.py",
+                "-m",
+                self.config['ps_mode_config'],
+                "--run_minimize",
+                self.config['run_minimize'],
+                "--run_single_pass",
+                self.config['run_single_pass'],
+                "--run_the_one_ps",
+                self.config['run_the_one_ps'],
+                "--debug_new_pass",
+                self.config['debug_new_pass'],
+                "--debug_new_minimize",
+                self.config['debug_new_minimize'],
+                "--applied_pass_name",
+                self.config['applied_pass_name'],
+                "--debug_the_one_ps",
+                self.config['debug_the_one_ps'],
+            ]
+        elif ps_mode == "gpu-ps":
+            os.environ['FLAGS_LAUNCH_BARRIER'] = '0'
+            os.environ['PADDLE_PSERVER_NUMS'] = '1'
+            os.environ['PADDLE_TRAINERS_NUM'] = '1'
+            os.environ['POD_IP'] = '127.0.0.1'
+            os.environ['PADDLE_PSERVERS_IP_PORT_LIST'] = '127.0.0.1:29011'
+            os.environ['PADDLE_PORT'] = '29011'
+            os.environ['FLAGS_selected_gpus'] = '0,1,2,3,4,5,6,7'
+            # pserver
+            # os.environ['TRAINING_ROLE'] = 'PSERVER'
+
+            # trainer
+            os.environ['TRAINING_ROLE'] = 'TRAINER'
+            os.environ['PADDLE_TRAINER_ID'] = '0'
+
+            cmd = [
+                sys.executable,
+                "-u",
+                "../ps/ps_dnn_trainer.py",
+                "-m",
+                self.config['ps_mode_config'],
+                "--run_minimize",
+                self.config['run_minimize'],
+                "--run_single_pass",
+                self.config['run_single_pass'],
+                "--run_the_one_ps",
+                self.config['run_the_one_ps'],
+                "--debug_new_pass",
+                self.config['debug_new_pass'],
+                "--debug_new_minimize",
+                self.config['debug_new_minimize'],
+                "--applied_pass_name",
+                self.config['applied_pass_name'],
+                "--debug_the_one_ps",
+                self.config['debug_the_one_ps'],
+            ]
+
+        cmd = [shlex.quote(c) for c in cmd]
+        prepare_python_path_and_return_module(__file__)
+        exitcode = os.system(' '.join(cmd))

test/distributed_passes/test_build_cinn_pass_resnet.py

@@ -0,0 +1,39 @@
+# Copyright (c) 2022 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 unittest
+
+from dist_pass_test_base import DistPassTestBase
+
+from paddle.distributed.passes import PassManager, new_pass
+
+
+class TestBuildCINNPass(DistPassTestBase):
+    def init(self):
+        self.atol = 0.5
+        self.rtol = 0.0
+
+    def apply_passes(self, main_prog, startup_prog):
+        pass_manager = PassManager(
+            [
+                new_pass("build_cinn"),
+                new_pass("fuse_elewise_add_act"),
+            ]
+        )
+        pass_manager.apply([main_prog], [startup_prog])
+        print(pass_manager.names)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/distributed_passes/test_dist_fuse_all_reduce_pass.py

@@ -0,0 +1,43 @@
+# 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 unittest
+
+from dist_pass_test_base import DistPassTestBase
+from model_zoo import resnet_model
+
+from paddle.distributed.passes import PassManager, new_pass
+
+
+class TestFuseAllReducePass(DistPassTestBase):
+    def init(self):
+        self.atol = 0.0
+        self.rtol = 0.0
+
+    def apply_passes(self, main_prog, startup_prog):
+        pass_manager = PassManager(
+            [
+                new_pass("fuse_elewise_add_act"),
+                new_pass("fuse_all_reduce", {"max_memory_size": 1024 * 1024}),
+            ]
+        )
+        pass_manager.apply([main_prog], [startup_prog])
+        print(pass_manager.names)
+
+    def test_bs_32(self):
+        self.check_main(resnet_model, batch_size=32)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/distributed_passes/test_dist_fuse_gemm_epilogue_pass.py

@@ -0,0 +1,210 @@
+# Copyright (c) 2022 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 unittest
+
+import numpy as np
+from dist_pass_test_base import DistPassTestBase
+
+import paddle
+from paddle import nn
+from paddle.distributed import fleet
+from paddle.distributed.passes import PassManager, new_pass
+
+paddle.enable_static()
+np.random.seed(12345)
+paddle.seed(12345)
+
+
+def verify_op_count(op_types, op_name, target_count):
+    count = 0
+    for op_type in op_types:
+        if op_type == op_name:
+            count += 1
+    return count == target_count
+
+
+class MultiFCLayer(nn.Layer):
+    def __init__(self, hidden, Activation):
+        super().__init__()
+        self.linear1 = paddle.nn.Linear(hidden, 4 * hidden)
+        self.linear2 = paddle.nn.Linear(4 * hidden, hidden)
+        self.linear3 = paddle.nn.Linear(hidden, hidden)
+
+        self.relu1 = Activation()
+        self.relu2 = Activation()
+        self.relu3 = Activation()
+
+    def forward(self, x, matmul_y, ele_y):
+        output = self.linear1(x)
+        output = self.relu1(output)
+        output = self.linear2(output)
+
+        output1 = paddle.matmul(output, matmul_y)
+        output = self.linear3(output)
+        output = self.relu2(output)
+
+        output = paddle.matmul(output, matmul_y)
+        output = paddle.add(output, ele_y)
+        output = self.relu3(output)
+        output = paddle.add(output, output1)
+        return output
+
+
+class TestFuseGemmEpiloguePassReluFP32(DistPassTestBase):
+    def init(self):
+        self.atol = 1e-3
+        self.rtol = 1e-3
+        self.activation = nn.ReLU
+        self.act_fwd_name = 'relu'
+        self.act_bwd_name = 'relu_grad'
+        self.batch = 64
+        self.seqlen = 128
+        self.hidden = 768
+        self.precision = 'FP32'  # FP32 or AMP
+
+    def get_model(self, place):
+        data = paddle.static.data(
+            name="_data", shape=[-1, self.seqlen, self.hidden], dtype='float32'
+        )
+        matmul_y = paddle.static.data(
+            name="_matmul_y",
+            shape=[1, self.hidden, self.hidden],
+            dtype='float32',
+        )
+        ele_y = paddle.static.data(
+            name="_ele_y",
+            shape=[
+                self.hidden,
+            ],
+            dtype='float32',
+        )
+
+        model = MultiFCLayer(self.hidden, self.activation)
+        out = model(data, matmul_y, ele_y)
+        loss = paddle.mean(out)
+        optimizer = paddle.optimizer.Adam(learning_rate=1e-3)
+
+        dist_strategy = fleet.DistributedStrategy()
+        dist_strategy.fuse_all_reduce_ops = False
+        dist_strategy.without_graph_optimization = True
+        if self.precision == 'AMP':
+            dist_strategy.amp = True
+            dist_strategy.amp_configs = {
+                "init_loss_scaling": 32768,
+                "use_dynamic_loss_scaling": True,
+                "custom_white_list": ['gelu'],
+            }
+        fleet.init(is_collective=True, strategy=dist_strategy)
+        optimizer = fleet.distributed_optimizer(optimizer)
+        optimizer.minimize(loss)
+
+        rank = paddle.distributed.get_rank()
+
+        def reader():
+            for _ in range(10):
+                data_arr = (
+                    np.random.random(
+                        (self.batch, self.seqlen, self.hidden)
+                    ).astype("float32")
+                    - 0.5
+                )
+                matmul_y_arr = (
+                    np.random.random((1, self.hidden, self.hidden)).astype(
+                        "float32"
+                    )
+                    - 0.5
+                )
+                ele_y_arr = (
+                    np.random.random((self.hidden,)).astype("float32") - 0.5
+                )
+                yield [data_arr, matmul_y_arr, ele_y_arr]
+
+        main_program = paddle.static.default_main_program()
+        startup_program = paddle.static.default_startup_program()
+
+        fetch_list = []
+        for p in model.parameters():
+            grad_name = p.name + '@GRAD'
+            fetch_list.append(grad_name)
+
+        fetch_list.append(loss.name)
+
+        return (
+            main_program,
+            startup_program,
+            [data, matmul_y, ele_y],
+            fetch_list,
+            reader,
+        )
+
+    def apply_passes(self, main_prog, startup_prog):
+        pass_manager = PassManager([new_pass("fuse_gemm_epilogue")])
+        pass_manager.apply([main_prog], [startup_prog])
+        print(pass_manager.names)
+
+        op_type = []
+        for op in main_prog.global_block().ops:
+            op_type.append(op.type)
+        print(op_type)
+        self.assertTrue(verify_op_count(op_type, "fused_gemm_epilogue", 3))
+        self.assertTrue(verify_op_count(op_type, "fused_gemm_epilogue_grad", 3))
+        self.assertTrue(verify_op_count(op_type, self.act_fwd_name, 1))
+        self.assertTrue(verify_op_count(op_type, self.act_bwd_name, 2))
+
+    def test_fuse_gemm_epilogue(self):
+        self.check_main()
+
+
+class TestFuseGemmEpiloguePassReluFP16(TestFuseGemmEpiloguePassReluFP32):
+    def init(self):
+        self.atol = 1e-3
+        self.rtol = 1e-3
+        self.activation = nn.ReLU
+        self.act_fwd_name = 'relu'
+        self.act_bwd_name = 'relu_grad'
+        self.batch = 64
+        self.seqlen = 128
+        self.hidden = 768
+        self.precision = 'AMP'  # FP32 or AMP
+
+
+class TestFuseGemmEpiloguePassGeluFP32(TestFuseGemmEpiloguePassReluFP32):
+    def init(self):
+        self.atol = 1e-3
+        self.rtol = 1e-3
+        self.activation = nn.GELU
+        self.act_fwd_name = 'gelu'
+        self.act_bwd_name = 'gelu_grad'
+        self.batch = 64
+        self.seqlen = 128
+        self.hidden = 768
+        self.precision = 'FP32'  # FP32 or AMP
+
+
+class TestFuseGemmEpiloguePassGeluFP16(TestFuseGemmEpiloguePassReluFP32):
+    def init(self):
+        self.atol = 5e-3
+        self.rtol = 1e-3
+        self.activation = nn.GELU
+        self.act_fwd_name = 'gelu'
+        self.act_bwd_name = 'gelu_grad'
+        self.batch = 64
+        self.seqlen = 128
+        self.hidden = 768
+        self.precision = 'AMP'  # FP32 or AMP
+
+
+if __name__ == "__main__":
+    unittest.main()

test/distributed_passes/test_dist_fuse_resunit_pass.py

@@ -0,0 +1,289 @@
+# 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
+#
+#     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 os
+import unittest
+
+import numpy as np
+from dist_pass_test_base import DistPassTestBase
+
+import paddle
+from paddle import ParamAttr, nn
+from paddle.distributed import fleet
+from paddle.distributed.passes import PassManager, new_pass
+from paddle.nn import BatchNorm, Conv2D
+from paddle.nn.initializer import Constant, KaimingNormal
+
+paddle.enable_static()
+np.random.seed(12345)
+paddle.seed(12345)
+
+
+def skip_unit_test():
+    return (
+        not paddle.is_compiled_with_cuda()
+        or paddle.device.cuda.get_device_capability()[0] < 8
+        or paddle.get_cudnn_version() < 8900
+    )
+
+
+skip_msg = (
+    "only support with cuda and CUDNN 8.9 or later,"
+    " and only Ampere or later devices are supported"
+)
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        num_channels,
+        num_filters,
+        filter_size,
+        stride=1,
+        groups=1,
+        act=None,
+        lr_mult=1.0,
+        data_format="NCHW",
+        bn_weight_decay=True,
+    ):
+        super().__init__()
+        self.act = act
+        self.conv = Conv2D(
+            in_channels=num_channels,
+            out_channels=num_filters,
+            kernel_size=filter_size,
+            stride=stride,
+            padding=(filter_size - 1) // 2,
+            groups=groups,
+            weight_attr=ParamAttr(
+                learning_rate=lr_mult, initializer=KaimingNormal()
+            ),
+            bias_attr=False,
+            data_format=data_format,
+        )
+        self.bn = BatchNorm(
+            num_filters,
+            param_attr=ParamAttr(
+                learning_rate=lr_mult,
+                regularizer=(
+                    None if bn_weight_decay else paddle.regularizer.L2Decay(0.0)
+                ),
+                initializer=Constant(1.0),
+            ),
+            bias_attr=ParamAttr(
+                learning_rate=lr_mult,
+                regularizer=(
+                    None if bn_weight_decay else paddle.regularizer.L2Decay(0.0)
+                ),
+                initializer=Constant(0.0),
+            ),
+            data_layout=data_format,
+        )
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        if self.act:
+            x = self.relu(x)
+        return x
+
+
+class BottleneckBlock(nn.Layer):
+    def __init__(
+        self,
+        num_channels,
+        num_filters,
+        stride,
+        shortcut=True,
+        lr_mult=1.0,
+        data_format="NCHW",
+        bn_weight_decay=True,
+    ):
+        super().__init__()
+
+        self.conv0 = ConvBNLayer(
+            num_channels=num_channels,
+            num_filters=num_filters,
+            filter_size=1,
+            act="relu",
+            lr_mult=lr_mult,
+            data_format=data_format,
+            bn_weight_decay=bn_weight_decay,
+        )
+        self.conv1 = ConvBNLayer(
+            num_channels=num_filters,
+            num_filters=num_filters,
+            filter_size=3,
+            stride=stride,
+            act="relu",
+            lr_mult=lr_mult,
+            data_format=data_format,
+            bn_weight_decay=bn_weight_decay,
+        )
+        self.conv2 = ConvBNLayer(
+            num_channels=num_filters,
+            num_filters=num_filters * 4,
+            filter_size=1,
+            act=None,
+            lr_mult=lr_mult,
+            data_format=data_format,
+            bn_weight_decay=bn_weight_decay,
+        )
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                num_channels=num_channels,
+                num_filters=num_filters * 4,
+                filter_size=1,
+                stride=stride,
+                lr_mult=lr_mult,
+                data_format=data_format,
+                bn_weight_decay=bn_weight_decay,
+            )
+        self.relu = nn.ReLU()
+        self.shortcut = shortcut
+
+    def forward(self, x):
+        identity = x
+        x = self.conv0(x)
+        x = self.conv1(x)
+        x = self.conv2(x)
+
+        if self.shortcut:
+            short = identity
+        else:
+            short = self.short(identity)
+        x = paddle.add(x=x, y=short)
+        x = self.relu(x)
+        return x
+
+
+class ResUnitNet(nn.Layer):
+    def __init__(self, shortcut):
+        super().__init__()
+        self.shortcut = shortcut
+        self.conv1 = ConvBNLayer(
+            num_channels=3,
+            num_filters=64,
+            filter_size=3,
+            act='relu',
+            data_format='NHWC',
+        )
+
+        self.block = BottleneckBlock(
+            num_channels=64,
+            num_filters=16,
+            stride=1,
+            shortcut=self.shortcut,
+            lr_mult=1.0,
+            data_format="NHWC",
+            bn_weight_decay=True,
+        )
+
+        self.conv2 = ConvBNLayer(
+            num_channels=64,
+            num_filters=64,
+            filter_size=3,
+            act='relu',
+            data_format='NHWC',
+        )
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.block(out)
+        out = self.conv2(out)
+        out = paddle.flatten(out, 1)
+        return out
+
+
+@unittest.skipIf(skip_unit_test(), skip_msg)
+class TestFuseResUnitPass(DistPassTestBase):
+    def init(self):
+        self.atol = 1e-2
+        self.rtol = 1e-2
+        paddle.set_flags({'FLAGS_conv_workspace_size_limit': 1000})
+        self.init_attr()
+
+    def init_attr(self):
+        self.shortcut = True
+
+    def get_model(self, place, batch_size=32, image_shape=[224, 224, 3]):
+        image = paddle.static.data(
+            shape=[batch_size, *image_shape], dtype='float32', name='image'
+        )
+
+        model = ResUnitNet(self.shortcut)
+        pred_out = model(image)
+        loss = paddle.mean(pred_out)
+        optimizer = paddle.optimizer.Adam(learning_rate=1e-3)
+
+        dist_strategy = fleet.DistributedStrategy()
+        dist_strategy.fuse_all_reduce_ops = False
+        dist_strategy.without_graph_optimization = True
+        dist_strategy.amp = True
+        dist_strategy.amp_configs = {
+            "init_loss_scaling": 128.0,
+            "use_dynamic_loss_scaling": True,
+        }
+        build_strategy = paddle.static.BuildStrategy()
+        settings = {
+            "fuse_bn_act_ops": False,
+            "fuse_bn_add_act_ops": False,
+            "enable_inplace": False,
+        }
+        for k, v in settings.items():
+            setattr(build_strategy, k, v)
+        dist_strategy.build_strategy = build_strategy
+        fleet.init(is_collective=True, strategy=dist_strategy)
+        optimizer = fleet.distributed_optimizer(optimizer)
+        optimizer.minimize(loss)
+
+        rank = paddle.distributed.get_rank()
+
+        def reader():
+            seed = int(os.environ.get("SEED", 0))
+            np.random.seed(seed + rank)
+            for _ in range(10):
+                image_np = np.random.random(size=image.shape).astype('float32')
+                yield (image_np,)
+
+        main_program = paddle.static.default_main_program()
+        startup_program = paddle.static.default_startup_program()
+        return main_program, startup_program, [image], [loss], reader
+
+    def apply_passes(self, main_prog, startup_prog):
+        pass_manager = PassManager([new_pass("fuse_resunit")])
+        pass_manager.apply([main_prog], [startup_prog])
+        print(pass_manager.names)
+
+        op_type = []
+        for op in main_prog.global_block().ops:
+            op_type.append(op.type)
+        self.assertTrue("fused_scale_bias_add_relu" in op_type)
+        self.assertTrue("fused_scale_bias_relu_conv_bn" in op_type)
+        self.assertTrue("fused_dconv_drelu_dbn" in op_type)
+
+    def test_fuse_resunit(self):
+        self.check_main()
+
+
+@unittest.skipIf(skip_unit_test(), skip_msg)
+class TestFuseResUnitPassDual(TestFuseResUnitPass):
+    def init_attr(self):
+        self.shortcut = False
+
+
+if __name__ == "__main__":
+    unittest.main()

test/distributed_passes/test_fuse_allreduce_split_to_reducescatter_pass.py

@@ -0,0 +1,72 @@
+# 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
+#
+#     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 os
+import tempfile
+import unittest
+
+import requests
+from test_imperative_base import new_program_scope
+
+import paddle
+
+program_txt = '''
+{
+    (%2) = "builtin.parameter" () {parameter_name:"linear_0.w_0",persistable:[true],stop_gradient:[false]} : () -> builtin.tensor<8192x28672xbf16>
+    (%38) = "pd_op.data" () {dtype:(pd_op.DataType)bfloat16,name:"linear_0.tmp_0",persistable:[false],place:(pd_op.Place)Place(gpu:0),shape:(pd_op.IntArray)[4096,1,28672],stop_gradient:[false]} : () -> builtin.tensor<4096x1x28672xbf16>
+    (%48) = "pd_op.data" () {dtype:(pd_op.DataType)bfloat16,name:"input",persistable:[false],place:(pd_op.Place)Place(gpu:0),shape:(pd_op.IntArray)[4096,1,28672],stop_gradient:[false]} : () -> builtin.tensor<4096x1x28672xbf16>
+    (%50) = "pd_op.matmul" (%48, %2) {persistable:[false],stop_gradient:[false],transpose_x:false,transpose_y:true} : (builtin.tensor<4096x1x28672xbf16>, builtin.tensor<8192x28672xbf16>) -> builtin.tensor<4096x1x8192xbf16>
+    (%57) = "pd_op.all_reduce_" (%50) {event_to_record:"event_7989",events_to_wait:[],execution_stream:"auto_parallel_mp",force_record_event:false,persistable:[false],ring_id:(Int32)36,stop_gradient:[false],reduce_type:(Int32)0,use_model_parallel:true} : (builtin.tensor<4096x1x8192xbf16>) -> builtin.tensor<4096x1x8192xbf16>
+    (%63) = "pd_op.assign" (%57) {persistable:[false],stop_gradient:[false]} : (builtin.tensor<4096x1x8192xbf16>) -> builtin.tensor<4096x1x8192xbf16>
+    (%64) = "pd_op.full" () {dtype:(pd_op.DataType)int32,place:(pd_op.Place)Place(cpu),shape:(pd_op.IntArray)[1],stop_gradient:[true],value:(Float)0} : () -> builtin.tensor<1xi32>
+    (%65) = "pd_op.split_with_num" (%63, %64) {num:(Int32)2,persistable:[false],stop_gradient:[false]} : (builtin.tensor<4096x1x8192xbf16>, builtin.tensor<1xi32>) -> vec[builtin.tensor<2048x1x8192xbf16>,builtin.tensor<2048x1x8192xbf16>]
+    (%66) = "builtin.slice" (%65) {index:(Int32)0,persistable:[false],stop_gradient:[false]} : (vec[builtin.tensor<2048x1x8192xbf16>,builtin.tensor<2048x1x8192xbf16>]) -> builtin.tensor<2048x1x8192xbf16>
+    (%67) = "pd_op.assign" (%66) {persistable:[false],stop_gradient:[false]} : (builtin.tensor<2048x1x8192xbf16>) -> builtin.tensor<2048x1x8192xbf16>
+}'''
+
+
+class TestPass(unittest.TestCase):
+    def test_if_with_single_output(self):
+        paddle.pir.register_paddle_dialect()
+        url = "https://paddle-ci.cdn.bcebos.com/json_file/main_program.json"
+
+        response = requests.get(url)
+        response.raise_for_status()
+
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".json") as f:
+            f.write(response.content)
+            f.flush()
+            os.fsync(f.fileno())
+            temp_file_name = f.name
+
+        with paddle.pir_utils.IrGuard():
+            with new_program_scope():
+                ir_program = paddle.load(temp_file_name)
+                for op in ir_program.global_block().ops:
+                    if op.name() == 'pd_op.all_reduce_':
+                        op.set_str_attr("event_to_record", "event_7989")
+                        op.set_int_array_attr("events_to_wait", [])
+                        op.set_str_attr("execution_stream", "auto_parallel_mp")
+                        op.set_bool_attr("force_record_event", False)
+            pm = paddle.pir.PassManager()
+            pm.add_pass('fuse_allreduce_split_to_reducescatter_pass', {})
+            pm.run(ir_program)
+            self.assertEqual(ir_program.global_block().num_ops(), 6)
+            self.assertEqual(
+                ir_program.global_block().ops[-2].name(), "pd_op.reduce_scatter"
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/distributed_passes/test_ps_server_pass.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2022 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 unittest
+
+from ps_pass_test_base import PsPassTestBase
+
+
+class TestPsServerPass(PsPassTestBase):
+    def init(self):
+        pass
+
+    def setUp(self):
+        print('TestPsServerPass setUp...')
+
+    def tearDown(self):
+        print('TestPsServerPass tearDown...')
+
+    def test_add_lr_decay_table_pass(self):
+        pass
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distributed_passes/test_white_lists.py

@@ -0,0 +1,97 @@
+# 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 unittest
+
+from paddle.distributed.passes.pass_base import (
+    PassBase,
+    _make_rule_from_white_lists_dict as make_white_lists_rule,
+    new_pass,
+    register_pass,
+)
+
+
+class TestConcretePass(PassBase):
+    def __init__(self):
+        super().__init__()
+
+    def _check_self(self):
+        return True
+
+    def _check_conflict(self, other_pass):
+        return True
+
+    def _apply_single_impl(self, main_program, startup_program, context):
+        pass
+
+
+@register_pass("A")
+class A(TestConcretePass):
+    def __init__(self):
+        super().__init__()
+
+
+@register_pass("B")
+class B(TestConcretePass):
+    def __init__(self):
+        super().__init__()
+
+
+@register_pass("C")
+class C(TestConcretePass):
+    def __init__(self):
+        super().__init__()
+
+
+@register_pass("D")
+class D(TestConcretePass):
+    def __init__(self):
+        super().__init__()
+
+
+@register_pass("E")
+class E(TestConcretePass):
+    def __init__(self):
+        super().__init__()
+
+
+class TestMakeWhiteListsRule(unittest.TestCase):
+    def test_main(self):
+        before_white_lists = {"A": ["B", "C"]}
+        after_white_lists = {"D": ["C"]}
+        rule = make_white_lists_rule(before_white_lists, after_white_lists)
+
+        pass_a = new_pass("A")
+        pass_b = new_pass("B")
+        pass_c = new_pass("C")
+        pass_d = new_pass("D")
+        pass_e = new_pass("E")
+
+        self.assertTrue(rule(pass_a, pass_e))
+        self.assertTrue(rule(pass_e, pass_a))
+
+        self.assertTrue(rule(pass_a, pass_b))
+        self.assertFalse(rule(pass_b, pass_a))
+        self.assertTrue(rule(pass_a, pass_c))
+        self.assertFalse(rule(pass_c, pass_a))
+
+        self.assertFalse(rule(pass_a, pass_d))
+        self.assertFalse(rule(pass_d, pass_a))
+
+        self.assertTrue(rule(pass_c, pass_d))
+        self.assertFalse(rule(pass_d, pass_c))
+
+
+if __name__ == "__main__":
+    unittest.main()