chore: import upstream snapshot with attribution

python/paddle/distributed/utils/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+
+__all__ = []

python/paddle/distributed/utils/launch_utils.py

@@ -0,0 +1,589 @@
+# 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 copy
+import os
+import platform
+import signal
+import socket
+import subprocess
+import sys
+import time
+from collections.abc import Sequence
+from contextlib import closing
+
+from paddle.distributed.fleet.launch_utils import get_backend_by_compile_flag
+from paddle.utils import strtobool
+
+from ..utils.log_utils import get_logger
+
+logger = get_logger("INFO", "root")
+
+
+def get_cluster_from_args(args, selected_gpus):
+    node_ips = [x.strip() for x in args.cluster_node_ips.split(',')]
+    node_ip = args.node_ip
+    node_rank = node_ips.index(node_ip)
+
+    logger.debug(
+        f"parsed from args:node_ips:{node_ips} node_ip:{node_ip} node_rank:{node_rank}"
+    )
+
+    free_ports = None
+    if (
+        not args.use_paddlecloud
+        and len(node_ips) <= 1
+        and args.started_port is None
+    ):
+        free_ports = find_free_ports(len(selected_gpus))
+        if free_ports is not None:
+            free_ports = list(free_ports)
+    else:
+        started_port = 6070
+        if args.started_port is not None:
+            started_port = args.started_port
+
+        free_ports = list(
+            range(started_port, started_port + len(selected_gpus))
+        )
+
+    trainer_endpoints = []
+    for ip in node_ips:
+        trainer_endpoints.append([f"{ip}:{port}" for port in free_ports])
+    return get_cluster(node_ips, node_ip, trainer_endpoints, selected_gpus)
+
+
+def get_gpus(selected_gpus):
+    if selected_gpus is None:
+        from paddle.framework import core
+
+        gpus_num = core.get_cuda_device_count()
+        gpus = [str(x) for x in range(0, gpus_num)]
+    else:
+        cuda_visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
+        if cuda_visible_devices is None or cuda_visible_devices == "":
+            gpus = [x.strip() for x in selected_gpus.split(',')]
+        else:
+            # change selected_gpus into relative values
+            # e.g. CUDA_VISIBLE_DEVICES=4,5,6,7; args.selected_gpus=4,5,6,7;
+            # therefore selected_gpus=0,1,2,3
+            cuda_visible_devices_list = cuda_visible_devices.split(',')
+            for x in selected_gpus.split(','):
+                assert x in cuda_visible_devices_list, (
+                    "Can't find "
+                    f"your selected_gpus {x} in CUDA_VISIBLE_DEVICES[{cuda_visible_devices}]."
+                )
+            gpus = [
+                cuda_visible_devices_list.index(x.strip())
+                for x in selected_gpus.split(',')
+            ]
+            logger.info(
+                f"Change selected_gpus into relative values. --ips:{selected_gpus} "
+                f"will change into relative_ips:{gpus} according to your "
+                f"CUDA_VISIBLE_DEVICES:{cuda_visible_devices_list}"
+            )
+
+    return gpus
+
+
+class Hdfs:
+    def __init__(self):
+        self.hdfs_ugi = None
+        self.hdfs_name = None
+        self.hdfs_path = None
+
+    def is_valid(self):
+        return (
+            self.hdfs_ugi is not None
+            and self.hdfs_name is not None
+            and self.hdfs_path is not None
+        )
+
+    def __str__(self):
+        return f"hdfs_ugi:{self.hdfs_ugi} hdfs_name:{self.hdfs_name} hdfs_path{self.hdfs_path}"
+
+    def __eq__(self, n):
+        return (
+            self.hdfs_ugi == n.hdfs_ugi
+            and self.hdfs_name == n.hdfs_name
+            and self.hdfs_path == n.hdfs_path
+        )
+
+    def __ne__(self, n):
+        return not self == n
+
+
+class Cluster:
+    def __init__(self, hdfs):
+        self.job_server = None
+        self.pods = []
+        self.hdfs = None
+        self.job_stage_flag = None
+
+    def __str__(self):
+        return f"job_server:{self.job_server} pods:{[str(pod) for pod in self.pods]} job_stage_flag:{self.job_stage_flag} hdfs:{self.hdfs}"
+
+    def __eq__(self, cluster):
+        if len(self.pods) != len(cluster.pods):
+            return False
+
+        for a, b in zip(self.pods, cluster.pods):
+            if a != b:
+                return False
+
+        if self.job_stage_flag != cluster.job_stage_flag:
+            return False
+
+        return True
+
+    def __ne__(self, cluster):
+        return not self.__eq__(cluster)
+
+    def update_pods(self, cluster):
+        self.pods = copy.copy(cluster.pods)
+
+    def trainers_nranks(self):
+        return len(self.trainers_endpoints())
+
+    def pods_nranks(self):
+        return len(self.pods)
+
+    def trainers_endpoints(self):
+        r = []
+        for pod in self.pods:
+            for t in pod.trainers:
+                r.append(t.endpoint)
+        return r
+
+    def pods_endpoints(self):
+        r = []
+        for pod in self.pods:
+            ep = f"{pod.addr}:{pod.port}"
+            assert pod.port is not None and pod.addr is not None, (
+                f"{ep} not a valid endpoint"
+            )
+            r.append(ep)
+
+        return r
+
+    def get_pod_by_id(self, pod_id):
+        for pod in self.pods:
+            if str(pod_id) == str(pod.id):
+                return pod
+
+        return None
+
+
+class JobServer:
+    def __init__(self):
+        self.endpoint = None
+
+    def __str__(self):
+        return f"{self.endpoint}"
+
+    def __eq__(self, j):
+        return self.endpoint == j.endpoint
+
+    def __ne__(self, j):
+        return not self == j
+
+
+class Trainer:
+    def __init__(self):
+        self.gpus = []
+        self.endpoint = None
+        self.rank = None
+
+    def __str__(self):
+        return f"gpu:{self.gpus} endpoint:{self.endpoint} rank:{self.rank}"
+
+    def __eq__(self, t):
+        if len(self.gpus) != len(t.gpus):
+            return False
+
+        if self.endpoint != t.endpoint or self.rank != t.rank:
+            return False
+
+        for a, b in zip(self.gpus, t.gpus):
+            if a != b:
+                return False
+
+        return True
+
+    def __ne__(self, t):
+        return not self == t
+
+    def get_rank(self):
+        return self.rank
+
+
+class Pod:
+    def __init__(self):
+        self.rank = None
+        self.id = None
+        self.addr = None
+        self.port = None
+        self.trainers = []
+        self.gpus = []
+
+    def __str__(self):
+        return f"rank:{self.rank} id:{self.id} addr:{self.addr} port:{self.port} visible_gpu:{self.gpus} trainers:{[str(t) for t in self.trainers]}"
+
+    def __eq__(self, pod):
+        if (
+            self.rank != pod.rank
+            or self.id != pod.id
+            or self.addr != pod.addr
+            or self.port != pod.port
+        ):
+            logger.debug(f"pod {self} != {pod}")
+            return False
+
+        if len(self.trainers) != len(pod.trainers):
+            logger.debug(f"trainers {self.trainers} != {pod.trainers}")
+            return False
+
+        for i in range(len(self.trainers)):
+            if self.trainers[i] != pod.trainers[i]:
+                logger.debug(f"trainer {self.trainers[i]} != {pod.trainers[i]}")
+                return False
+
+        return True
+
+    def __ne__(self, pod):
+        return not self == pod
+
+    def parse_response(self, res_pods):
+        pass
+
+    def get_visible_gpus(self):
+        r = ""
+        for g in self.gpus:
+            r += f"{g},"
+
+        assert r != "", f"this pod {self} can't see any gpus"
+
+        r = r[:-1]
+        return r
+
+
+def get_cluster(node_ips, node_ip, trainer_endpoints, selected_gpus):
+    assert type(trainer_endpoints) is list, "trainer_endpoints must be list"
+    cluster = Cluster(hdfs=None)
+    trainer_rank = 0
+    for node_rank, ip in enumerate(node_ips):
+        pod = Pod()
+        pod.rank = node_rank
+        pod.addr = ip
+        cur_node_endpoints = trainer_endpoints[node_rank]
+        # when use paddlecloud, endpoints may > selected_gpus(user_defined)
+        assert len(cur_node_endpoints) >= len(selected_gpus), (
+            "current trainer_endpoints size should be greater equal than selected_gpus size."
+        )
+        for i in range(len(selected_gpus)):
+            trainer = Trainer()
+            trainer.gpus.append(selected_gpus[i])
+            trainer.endpoint = f"{cur_node_endpoints[i]}"
+            trainer.rank = trainer_rank
+            trainer_rank += 1
+
+            pod.trainers.append(trainer)
+        cluster.pods.append(pod)
+
+    pod_rank = node_ips.index(node_ip)
+    return cluster, cluster.pods[pod_rank]
+
+
+def terminate_local_procs(procs):
+    for p in procs:
+        if p.proc.poll() is None:
+            p.proc.terminate()
+            if p.log_fn:
+                p.log_fn.close()
+            logger.debug(f"terminate process id:{p.proc.pid}")
+
+    # wait all process terminated
+    time.sleep(3)
+    for step in range(0, 50):
+        alive = False
+        for p in procs:
+            if p.proc.poll() is None:  # not terminate
+                os.kill(p.proc.pid, signal.SIGKILL)
+                alive = True
+
+        if not alive:
+            logger.info("terminate all the procs")
+            return
+
+        time.sleep(3)
+
+    logger.fatal("can't kill all process and exit")
+    sys.exit(1)
+
+
+def get_host_name_ip():
+    try:
+        host_name = socket.gethostname()
+        host_ip = socket.gethostbyname(host_name)
+        return host_name, host_ip
+    except:
+        return None
+
+
+def add_arguments(argname, type, default, help, argparser, **kwargs):
+    """Add argparse's argument.
+    Examples:
+        .. code-block:: pycon
+
+            >>> import argparse
+            >>> from paddle.distributed.utils import launch_utils
+            >>> parser = argparse.ArgumentParser()
+            >>> launch_utils.add_arguments("name", str, "Jonh", "User name.", parser)
+            >>> args = parser.parse_args()
+
+    """
+    type = strtobool if type == bool else type
+    argparser.add_argument(
+        "--" + argname,
+        default=default,
+        type=type,
+        help=help + ' Default: %(default)s.',
+        **kwargs,
+    )
+
+
+def find_free_ports(num):
+    def __free_port():
+        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+            s.bind(('', 0))
+            return s.getsockname()[1]
+
+    port_set = set()
+    step = 0
+    while True:
+        port = __free_port()
+        if port not in port_set:
+            port_set.add(port)
+
+        if len(port_set) >= num:
+            return port_set
+
+        step += 1
+        if step > 100:
+            print(
+                "can't find available port and use the specified static port now!"
+            )
+            return None
+
+    return None
+
+
+def _prepare_trainer_env(cluster, trainer, backend=None):
+    if backend is None:
+        backend = get_backend_by_compile_flag()  # for compatibility
+    if backend == 'bkcl':
+        proc_env = {
+            "FLAGS_selected_xpus": "{}".format(
+                ",".join([str(g) for g in trainer.gpus])
+            ),
+            "PADDLE_TRAINER_ID": str(trainer.rank),
+            "PADDLE_CURRENT_ENDPOINT": str(trainer.endpoint),
+            "PADDLE_TRAINERS_NUM": str(cluster.trainers_nranks()),
+            "PADDLE_TRAINER_ENDPOINTS": ",".join(cluster.trainers_endpoints()),
+        }
+    elif backend == 'nccl':
+        proc_env = {
+            "FLAGS_selected_gpus": "{}".format(
+                ",".join([str(g) for g in trainer.gpus])
+            ),
+            "PADDLE_TRAINER_ID": str(trainer.rank),
+            "PADDLE_CURRENT_ENDPOINT": str(trainer.endpoint),
+            "PADDLE_TRAINERS_NUM": str(cluster.trainers_nranks()),
+            "PADDLE_TRAINER_ENDPOINTS": ",".join(cluster.trainers_endpoints()),
+        }
+    elif backend == 'gloo':
+        # NOTE (xiongkun) default fall back into cpu only
+        proc_env = {
+            "PADDLE_TRAINER_ID": str(trainer.rank),
+            "PADDLE_CURRENT_ENDPOINT": str(trainer.endpoint),
+            "PADDLE_TRAINERS_NUM": str(cluster.trainers_nranks()),
+            "PADDLE_TRAINER_ENDPOINTS": ",".join(cluster.trainers_endpoints()),
+            "PADDLE_DISTRI_BACKEND": backend,  # only add here, other will be auto
+        }
+    elif backend == 'xccl':
+        from paddle.framework import core
+
+        custom_device_name = core.get_all_custom_device_type()[0]
+        proc_env = {
+            f"FLAGS_selected_{custom_device_name}s": "{}".format(
+                ",".join([str(g) for g in trainer.gpus])
+            ),
+            "PADDLE_TRAINER_ID": str(trainer.rank),
+            "PADDLE_CURRENT_ENDPOINT": str(trainer.endpoint),
+            "PADDLE_TRAINERS_NUM": str(cluster.trainers_nranks()),
+            "PADDLE_TRAINER_ENDPOINTS": ",".join(cluster.trainers_endpoints()),
+        }
+    else:
+        raise ValueError("backend must be one of 'gloo, nccl, bkcl'")
+
+    return proc_env
+
+
+class TrainerProc:
+    def __init__(self):
+        self.proc = None
+        self.log_fn = None
+        self.log_offset = None
+        self.rank = None
+        self.local_rank = None
+        self.cmd = None
+
+
+def start_local_trainers(
+    cluster, pod, training_script, training_script_args, log_dir=None
+):
+    current_env = copy.copy(os.environ.copy())
+    # paddle broadcast ncclUniqueId use socket, and
+    # proxy maybe make trainers unreachable, so delete them.
+    # if we set them to "", grpc will log error message "bad uri"
+    # so just delete them.
+    current_env.pop("http_proxy", None)
+    current_env.pop("https_proxy", None)
+
+    procs = []
+    for idx, t in enumerate(pod.trainers):
+        proc_env = _prepare_trainer_env(cluster, t)
+        current_env.update(proc_env)
+
+        logger.debug(f"trainer proc env:{current_env}")
+
+        cmd = [sys.executable, "-u", training_script, *training_script_args]
+
+        logger.info(f"start trainer proc:{cmd} env:{proc_env}")
+
+        fn = None
+        if log_dir is not None:
+            os.makedirs(log_dir, exist_ok=True)
+            fn = open(f"{log_dir}/workerlog.{idx}", "a")
+            proc = subprocess.Popen(cmd, env=current_env, stdout=fn, stderr=fn)
+        else:
+            proc = subprocess.Popen(cmd, env=current_env)
+
+        tp = TrainerProc()
+        tp.proc = proc
+        tp.rank = t.rank
+        tp.local_rank = idx
+        tp.log_fn = fn
+        tp.log_offset = fn.tell() if fn else None
+        tp.cmd = cmd
+
+        procs.append(tp)
+
+    return procs
+
+
+def pull_worker_log(tp):
+    if tp.log_fn:
+        with open(tp.log_fn.name, 'r') as fin:
+            fin.seek(tp.log_offset, 0)
+            for line in fin:
+                try:
+                    sys.stdout.write(line)
+                except UnicodeEncodeError:
+                    sys.stdout.write(
+                        'UnicodeEncodeError occurs at this line. '
+                        f'Please refer to the original log file "{tp.log_fn.name}"\n'
+                    )
+            tp.log_offset = fin.tell()
+
+
+def watch_local_trainers(procs, nranks):
+    try:
+        error = False
+        error_rank = []
+        # wait all process finish or one error
+        alive = False
+        for p in procs:
+            if p.log_fn and p.local_rank == 0:
+                pull_worker_log(p)
+
+            ret = p.proc.poll()
+            if ret is None:
+                alive = True
+            elif ret != 0:
+                error = True
+                error_rank.append(p.rank)
+
+        if error:
+            terminate_local_procs(procs)
+            sys.exit(1)
+
+    except KeyboardInterrupt:
+        logger.warning("KeyboardInterrupt, exit")
+        terminate_local_procs(procs)
+        raise
+    except SystemExit:
+        logger.error(
+            f"ABORT!!! Out of all {nranks} trainers, the trainer process with rank={error_rank} was aborted. Please check its log."
+        )
+        terminate_local_procs(procs)
+        raise
+    except:
+        logger.error(
+            f"ABORT!!! Out of all {nranks} trainers, the trainer process with rank={error_rank} was aborted. Please check its log."
+        )
+        terminate_local_procs(procs)
+        raise
+
+    return alive
+
+
+def _print_arguments(args):
+    print("-----------  Configuration Arguments -----------")
+    for arg, value in sorted(vars(args).items()):
+        print(f"{arg}: {value}")
+    print("------------------------------------------------")
+
+
+def filter_pids(processes: Sequence[str], self_pid: int) -> list[int]:
+    """Filter valid PIDs from a list of strings, excluding the current self_pid."""
+    pids_to_kill = []
+    for process in processes:
+        pid_str = process.strip()
+        if not pid_str.isdigit():
+            continue
+        pid_int = int(pid_str)
+        if pid_int == self_pid:
+            continue
+        pids_to_kill.append(pid_int)
+    return pids_to_kill
+
+
+def terminate_processes(processes: Sequence[int]) -> bool:
+    """
+    Terminate a list of processes by their PIDs.
+    Returns True if all processes were successfully terminated (or already dead).
+    Returns False if any process failed to terminate due to permissions.
+    """
+    sig = signal.SIGKILL if platform.system() != "Windows" else signal.SIGTERM
+    success = True
+    for pid in processes:
+        try:
+            os.kill(pid, sig)
+        except ProcessLookupError:
+            # Target already exited.
+            pass
+        except PermissionError:
+            success = False
+    return success

python/paddle/distributed/utils/log_utils.py

@@ -0,0 +1,33 @@
+# 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 logging
+
+
+def get_logger(log_level, name="root"):
+    logger = logging.getLogger(name)
+
+    # Avoid printing multiple logs
+    logger.propagate = False
+
+    if not logger.handlers:
+        log_handler = logging.StreamHandler()
+        logger.setLevel(log_level)
+        log_format = logging.Formatter(
+            '[%(asctime)-15s] [%(levelname)8s] %(filename)s:%(lineno)s - %(message)s'
+        )
+        log_handler.setFormatter(log_format)
+        logger.addHandler(log_handler)
+
+    return logger

python/paddle/distributed/utils/moe_utils.py

@@ -0,0 +1,308 @@
+# 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.
+
+from paddle import _legacy_C_ops
+from paddle.common_ops_import import check_variable_and_dtype
+from paddle.distributed import fleet
+from paddle.framework import LayerHelper, in_dynamic_mode
+
+
+def global_scatter(
+    x, local_count, global_count, group=None, use_calc_stream=True
+):
+    """
+    The global_scatter operator distributes the data of x to n_expert * world_size experts according to local_count,
+    and then receives data according to global_count. The expert refers to a user-defined expert network,
+    n_expert refers to the number of expert networks owned by each card, and world_size refers to the number of graphics cards running the network.
+
+    As shown below, the value of the world size is 2, n_expert 2, the batch size of the x 4 and local_count is [2, 0, 2, 0].
+    The global_count of the rank 0 is [2, 0, , ], rank 1 is [2, 0, ,](Due to the limited space, only the data calculated on rank 0 is shown here).
+    In the global_scatter operator, local_count[i] represents sending local_count[i] data to the (i % n_expert)th expert of the (i // n_expert)th card,
+    global_count[i] represents receiving global_count[i] data from the (i // n_expert)th card to the (i % n_expert)th expert of this card. The rank in the
+    figure represent the rank of the current card in all cards.
+
+    The process of global_scatter sending data is as follows:
+
+    local_count[0] represents taking out 2 batches from x and sending 2 batches to the 0th expert of the 0th card;
+
+    local_count[1] represents taking out 0 batches from x and sending 0 batches to the 1st expert of the 0th card;
+
+    local_count[2] represents taking out 2 batches from x and sending 2 batches to the 0th expert of the 1st card;
+
+    local_count[3] represents taking out 0 batches from x and sending 0 batches to the 1st expert of the 1st card;
+
+    Therefore, the global_count[0] of the 0th card is equal to 2, which means that 2 batches of data are received from the 0th card to the 0th expert;
+
+    the global_count[1] of the 0th card is equal to 0, which means that 0 batches of data are received from the 0th card to the 1st expert;
+
+    the global_count[0] of the 1st card is equal to 2, which means that 2 batches of data are received from the 0th card to the 0th expert;
+
+    the global_count[1] of the 1st card is equal to 0, which means that 0 batches of data are received from the 0th card to the 1st expert.
+
+    .. image:: https://githubraw.cdn.bcebos.com/PaddlePaddle/docs/develop/docs/api/paddle/distributed/img/global_scatter_gather.png
+        :width: 800
+        :alt: global_scatter_gather
+        :align: center
+
+    Args:
+        x (Tensor): Tensor. The tensor data type should be float16, float32, float64, int32 or int64.
+        local_count (Tensor): Tensor which have n_expert * world_size elements that indicates
+            how many data needed to be sent. The tensor data type should be int64.
+        global_count (Tensor): Tensor which have n_expert * world_size elements that indicates
+            how many data needed to be received. The tensor data type should be int64.
+        group (Group, optional): The group instance return by new_group or None for global default group. Default: None.
+        use_calc_stream (bool, optional): Whether to use calculation stream (True) or communication stream. Default: True.
+
+    Returns:
+        out (Tensor): The data received from all experts.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> import paddle
+            >>> from paddle.distributed import init_parallel_env
+            >>> from paddle.distributed.utils import moe_utils
+            >>> init_parallel_env()
+            >>> n_expert = 2
+            >>> world_size = 2
+            >>> d_model = 2
+            >>> in_feat = d_model
+            >>> local_input_buf = paddle.to_tensor(
+            ...     [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]],
+            ...     dtype='float32',
+            ...     stop_gradient=False,
+            ... )
+            >>> if paddle.distributed.ParallelEnv().local_rank == 0:
+            ...     local_count = paddle.to_tensor([2, 1, 1, 1], dtype="int64")
+            ...     global_count = paddle.to_tensor([2, 1, 1, 1], dtype="int64")
+            >>> else:
+            ...     local_count = paddle.to_tensor([1, 1, 2, 1], dtype="int64")
+            ...     global_count = paddle.to_tensor([1, 1, 2, 1], dtype="int64")
+            >>> a = moe_utils.global_scatter(
+            ...     local_input_buf,
+            ...     local_count,
+            ...     global_count,
+            ... )
+            >>> a.stop_gradient = False
+            >>> print(a)
+            >>> # out for rank 0: [[1, 2], [3, 4], [1, 2], [5, 6], [3, 4]]
+            >>> # out for rank 1: [[7, 8], [5, 6], [7, 8], [9, 10], [9, 10]]
+            >>> # backward test
+            >>> c = a * a
+            >>> c.backward()
+            >>> print("local_input_buf.grad: ", local_input_buf.grad)
+            >>> # out for rank 0: [[2, 4], [6, 8], [10, 12], [14, 16], [18, 20]]
+            >>> # out for rank 1: [[2, 4], [6, 8], [10, 12], [14, 16], [18, 20]]
+
+    """
+    if group is not None and not group.is_member():
+        return
+
+    ring_id = 0 if group is None else group.id
+    if in_dynamic_mode():
+        return _legacy_C_ops.global_scatter(
+            x,
+            local_count,
+            global_count,
+            'use_calc_stream',
+            use_calc_stream,
+            'ring_id',
+            ring_id,
+        )
+    else:
+        op_type = 'global_scatter'
+        check_variable_and_dtype(
+            x,
+            'x',
+            ['float16', 'float32', 'float64', 'int32', 'int64', 'uint16'],
+            'global_scatter',
+        )
+        check_variable_and_dtype(
+            local_count, 'local_count', ['int64'], 'global_scatter'
+        )
+        check_variable_and_dtype(
+            global_count, 'global_count', ['int64'], 'global_scatter'
+        )
+
+        helper = LayerHelper(op_type, **locals())
+        out = helper.create_variable_for_type_inference(dtype=x.dtype)
+
+        helper.append_op(
+            type=op_type,
+            inputs={
+                'X': [x],
+                'local_count': [local_count],
+                'global_count': [global_count],
+            },
+            outputs={'Out': [out]},
+            attrs={'ring_id': ring_id, 'use_calc_stream': use_calc_stream},
+        )
+        return out
+
+
+def global_gather(
+    x, local_count, global_count, group=None, use_calc_stream=True
+):
+    """
+    The global_gather operator gathers the data of x into n_expert * world_size experts according to global_count, and then receives data according to local_count.
+    The expert refers to a user-defined expert network, n_expert refers to the number of expert networks owned by each card, and world_size refers to the number of graphics cards running the network.
+
+    As shown below, the value of the world size is 2, n_expert 2, the batch size of the x 4 and local_count is [2, 0, 2, 0].
+    The global_count of the rank 0 is [2, 0, , ], rank 1 is [2, 0, ,](Due to the limited space, only the data calculated on rank 0 is shown here).
+    In the global_gather operator, the meaning of the global_count and local_count is opposed to global_scatter, global_count[i] represents sending global_count[i] data to the (i % n_expert)th expert of the (i // n_expert)th card,
+    local_count[i] represents receiving local_count[i] data from the (i // n_expert)th card to the (i % n_expert)th expert of this card. The data sent will be arranged according to the experts of each card.
+    The rank in the figure represent the rank of the current card in all cards.
+
+    The process of global_gather sending data is as follows:
+
+    The global_count[0] of the 0th card represents sending 2 data to the 0th expert of the 0th card;
+
+    The global_count[1] of the 0th card represents sending 0 data to the 1st expert of the 0th card;
+
+    The global_count[0] of the 1st card represents sending 2 data to the 0th expert of the 0th card;
+
+    The global_count[1] of the 1st card represents sending 0 data to the 1st expert of the 0th card.
+
+    .. image:: https://githubraw.cdn.bcebos.com/PaddlePaddle/docs/develop/docs/api/paddle/distributed/img/global_scatter_gather.png
+        :width: 800
+        :alt: global_scatter_gather
+        :align: center
+
+
+    Args:
+        x (Tensor): Tensor. Tensor whose data type should be float16, float32, float64, int32 or int64.
+        local_count (Tensor): Tensor which have n_expert * world_size elements that indicates
+            how many data needed to be received. Tensor data type should be int64.
+        global_count (Tensor): Tensor which have n_expert * world_size elements that indicates
+            how many data needed to be sent. Tensor data type should be int64.
+        group (Group, optional): The group instance return by new_group or None for global default group. Default: None.
+        use_calc_stream (bool, optional): Whether to use calculation stream (True) or communication stream. Default: True.
+
+    Returns:
+        out (Tensor): The data received from all experts.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> import paddle
+            >>> from paddle.distributed import init_parallel_env
+            >>> from paddle.distributed.utils import moe_utils
+            >>> init_parallel_env()
+            >>> n_expert = 2
+            >>> world_size = 2
+            >>> d_model = 2
+            >>> in_feat = d_model
+            >>> local_input_buf = paddle._to_tensor(
+            ...     [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]],
+            ...     dtype='float32',
+            ...     stop_gradient=False,
+            ... )
+            >>> if paddle.distributed.ParallelEnv().local_rank == 0:
+            ...     local_count = paddle.to_tensor([2, 1, 1, 1], dtype="int64")
+            ...     global_count = paddle.to_tensor([2, 1, 1, 1], dtype="int64")
+            >>> else:
+            ...     local_count = paddle.to_tensor([1, 1, 2, 1], dtype="int64")
+            ...     global_count = paddle.to_tensor([1, 1, 2, 1], dtype="int64")
+            >>> a = moe_utils.global_gather(
+            ...     local_input_buf,
+            ...     local_count,
+            ...     global_count,
+            ... )
+            >>> print(a)
+            >>> # out for rank 0: [[1, 2], [3, 4], [7, 8], [1, 2], [7, 8]]
+            >>> # out for rank 1: [[5, 6], [9, 10], [3, 4], [5, 6], [9, 10]]
+            >>> a.stop_gradient = False
+            >>> c = a * a
+            >>> c.backward()
+            >>> print("local_input_buf.grad", local_input_buf.grad)
+            >>> # out for rank 0: [[2, 4], [6, 8], [10, 12], [14, 16], [18, 20]]
+            >>> # out for rank 1: [[2, 4], [6, 8], [10, 12], [14, 16], [18, 20]]
+
+    """
+    if group is not None and not group.is_member():
+        return
+
+    ring_id = 0 if group is None else group.id
+    if in_dynamic_mode():
+        return _legacy_C_ops.global_gather(
+            x,
+            local_count,
+            global_count,
+            'use_calc_stream',
+            use_calc_stream,
+            'ring_id',
+            ring_id,
+        )
+    else:
+        op_type = 'global_gather'
+        check_variable_and_dtype(
+            x,
+            'x',
+            ['float16', 'float32', 'float64', 'int32', 'int64', 'uint16'],
+            'global_gather',
+        )
+
+        check_variable_and_dtype(
+            local_count, 'local_count', ['int64'], 'global_gather'
+        )
+
+        check_variable_and_dtype(
+            global_count, 'global_count', ['int64'], 'global_gather'
+        )
+        helper = LayerHelper(op_type, **locals())
+        out = helper.create_variable_for_type_inference(dtype=x.dtype)
+
+        helper.append_op(
+            type=op_type,
+            inputs={
+                'X': [x],
+                'local_count': [local_count],
+                'global_count': [global_count],
+            },
+            outputs={'Out': [out]},
+            attrs={
+                'ring_id': group,
+                'use_calc_stream': use_calc_stream,
+            },
+        )
+        return out
+
+
+def get_complete_pp_mesh(mesh):
+    """
+    Get complete pp mesh with given mesh.
+
+    Args:
+        mesh (Mesh): Mesh object.
+
+    Returns:
+        Mesh: Complete mesh.
+
+    """
+    process_id = mesh.process_ids[0]
+    global_mesh = fleet.auto.get_mesh()
+
+    if global_mesh and "pp" in global_mesh.dim_names:
+        pp_degree = global_mesh.get_dim_size("pp")
+        for i in range(pp_degree):
+            pp_mesh = global_mesh.get_mesh_with_dim("pp", i)
+            if process_id in pp_mesh.process_ids:
+                return pp_mesh
+        AssertionError(
+            f"Current mesh: {mesh} not found in global mesh {global_mesh}"
+        )
+    else:
+        return mesh

python/paddle/distributed/utils/nccl_utils.py

@@ -0,0 +1,49 @@
+# 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.
+
+
+from paddle.base import core
+
+
+def get_nccl_version_str(ver):
+    if ver >= 10000:
+        NCCL_MAJOR_VERSION = int(ver // 10000)
+        ver = ver % 10000
+    else:
+        NCCL_MAJOR_VERSION = int(ver // 1000)
+        ver = ver % 1000
+
+    NCCL_MINOR_VERSION = int(ver // 100)
+    NCCL_PATCH_VERSION = int(ver % 100)
+
+    return f"{NCCL_MAJOR_VERSION}.{NCCL_MINOR_VERSION}.{NCCL_PATCH_VERSION}"
+
+
+def check_nccl_version_for_p2p():
+    nccl_version = core.nccl_version()
+    nccl_version_str = get_nccl_version_str(nccl_version)
+    nccl_version_baseline = 2804
+    assert nccl_version >= nccl_version_baseline, (
+        "The version of NCCL is required to be at least v2.8.4 while training with "
+        f"pipeline/MoE parallelism, but we found v{nccl_version_str}. The previous version of NCCL has "
+        "some bugs in p2p communication, and you can see more detailed description "
+        "about this issue from ReleaseNotes of NCCL v2.8.4 "
+        "(https://docs.nvidia.com/deeplearning/nccl/release-notes/rel_2-8-4.html#rel_2-8-4)."
+    )
+
+
+def check_nccl_version_for_bf16():
+    nccl_version = core.nccl_version()
+    nccl_version_baseline = 21000
+    return nccl_version >= nccl_version_baseline

python/paddle/distributed/utils/process_utils.py

@@ -0,0 +1,221 @@
+# 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 shutil
+import subprocess
+
+import paddle
+from paddle.distributed.utils.log_utils import get_logger
+
+logger = get_logger("INFO", "root")
+
+SUCCESS_CODE = 0
+FAIL_CODE = 1
+
+
+def _get_cpu_info(numa_id):
+    """
+    get cpu info from lscpu
+    """
+
+    def _process_raw_cpu_info(i):
+        processed_cpu_info = []
+        cpu_ranges = i.split(',')
+        for cpu_range in cpu_ranges:
+            start, end = (
+                int(cpu_range.split("-")[0]),
+                int(cpu_range.split("-")[1]),
+            )
+            processed_cpu_info.extend(list(range(start, end + 1)))
+        return processed_cpu_info
+
+    try:
+        cpus = None
+        cmd = ["lscpu"]
+        output = subprocess.check_output(cmd).decode("utf-8").split(os.linesep)
+        numa_key = f"node{numa_id}"
+        for line in output:
+            if line.find(numa_key) >= 0:
+                raw_cpu_info = line.strip().split()[3]
+                cpus = _process_raw_cpu_info(raw_cpu_info)
+                break
+        return cpus
+    except Exception as e:
+        logger.warning(f"_get_cpu_info failed, reason:{e}")
+        return None
+
+
+def _has_nvidia_smi():
+    """
+    check if nvidia-smi is available
+    """
+    return shutil.which("nvidia-smi")
+
+
+def _has_xpu_smi():
+    """
+    check if xpu-smi is available
+    """
+    return shutil.which("xpu-smi")
+
+
+def _get_xpu_device_from_env(str_device_list, local_rank):
+    if len(str_device_list.strip()) == 0:
+        return None
+    visible_devices = str_device_list.split(',')
+    if len(visible_devices) <= local_rank:
+        return None
+    return visible_devices[local_rank]
+
+
+def _get_xpu_device(local_rank):
+    """
+    get currently used xpu physical device id
+    """
+    # NOTE(lijin23): priority XPULINK_VISIBLE_DEVICES > XPU_VISIBLE_DEVICES >
+    # CUDA_VISIBLE_DEVICES
+    xpulink_visible_devices = os.getenv("XPULINK_VISIBLE_DEVICES")
+    if xpulink_visible_devices is not None:
+        return _get_xpu_device_from_env(xpulink_visible_devices, local_rank)
+
+    xpu_visible_devices = os.getenv("XPU_VISIBLE_DEVICES")
+    if xpu_visible_devices is not None:
+        return _get_xpu_device_from_env(xpu_visible_devices, local_rank)
+
+    cuda_visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
+    if cuda_visible_devices is not None:
+        return _get_xpu_device_from_env(cuda_visible_devices, local_rank)
+
+    return str(local_rank)
+
+
+def _get_gpu_device(local_rank):
+    """
+    get currently used gpu physical device id
+    """
+    cuda_visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
+    if cuda_visible_devices is None or cuda_visible_devices == "":
+        return str(local_rank)
+    cuda_visible_devices = cuda_visible_devices.split(',')
+    if len(cuda_visible_devices) <= local_rank:
+        return None
+    return cuda_visible_devices[local_rank]
+
+
+def _get_gpu_numa_info(gpu_id):
+    """
+    get gpu numa info from nvidia-smi
+    """
+    try:
+        cmd = ["nvidia-smi", "topo", "-C", "-i", gpu_id]
+        output = subprocess.check_output(cmd, timeout=3).decode("utf-8")
+        numa_id = output.strip().split()[-1]
+        return numa_id
+    except Exception as e:
+        logger.warning(f"_get_cpu_info failed, reason:{e}")
+        return None
+
+
+def _get_xpu_affinity_mask(xpu_id):
+    xpu_id = int(xpu_id)
+    cmd = ["xpu-smi", "topo", "-m"]
+    if os.getenv("CUDA_DEVICE_ORDER") == "OAM_ID":
+        # NOTE(lijin23): if CUDA_DEVICE_ORDER is set to OAM_ID,
+        #  we need to get the cpu affinity using OAM_ID
+        cmd = ["xpu-smi", "topo", "-mo"]
+    output = subprocess.check_output(cmd, timeout=60).decode("utf-8")
+    cpu_affinity = output.splitlines()[xpu_id + 1].split()[-2]
+    affinity_mask = []
+    for affinity_range in cpu_affinity.split(','):
+        start, end = affinity_range.split('-')
+        affinity_mask.extend(range(int(start), int(end) + 1))
+    return affinity_mask
+
+
+def set_affinity_gpu():
+    """
+    set affinity for gpu
+    """
+    if not _has_nvidia_smi():
+        logger.warning(
+            "nvidia-smi is not available, set_affinity is aborted, plz check your environment."
+        )
+        return FAIL_CODE
+    local_rank = max(int(os.getenv("PADDLE_LOCAL_RANK", "0")), 0)
+    device_id = _get_gpu_device(local_rank)
+    if device_id is None:
+        logger.warning(
+            "Failed to get device id from cuda_visible_devices, set_affinity is aborted, plz check your environment."
+        )
+        return FAIL_CODE
+    numa_id = _get_gpu_numa_info(device_id)
+    if numa_id is None:
+        logger.warning(
+            "Failed to get numa info, set_affinity is aborted, plz check your environment."
+        )
+        return FAIL_CODE
+    if numa_id == "N/A":
+        logger.warning(
+            "nvidia-smi topo return numa id as N/A, set_affinity is aborted, plz check your environment. (Notice: This is expected behavior when executed on single numa node environment)"
+        )
+        return FAIL_CODE
+    affinity_mask = _get_cpu_info(numa_id)
+    if affinity_mask is None:
+        logger.warning(
+            "Failed to get cpu info, set_affinity is aborted, plz check your environment."
+        )
+        return FAIL_CODE
+    affinity = os.sched_getaffinity(0)
+    logger.info(f"Check affinity before setting: {affinity}")
+    os.sched_setaffinity(0, affinity_mask)
+    affinity = os.sched_getaffinity(0)
+    logger.info(f"check affinity after setting: {affinity}")
+    return SUCCESS_CODE
+
+
+def set_affinity_xpu():
+    """
+    set affinity for xpu
+    """
+    if not _has_xpu_smi():
+        logger.warning(
+            "xpu-smi is not available, set_affinity is aborted, plz check your environment."
+        )
+        return FAIL_CODE
+    local_rank = max(int(os.getenv("PADDLE_LOCAL_RANK", "0")), 0)
+    device_id = _get_xpu_device(local_rank)
+    if device_id is None:
+        logger.warning(
+            "Failed to get device id, set_affinity is aborted, plz check your environment."
+        )
+        return FAIL_CODE
+    affinity_mask = _get_xpu_affinity_mask(device_id)
+    affinity = os.sched_getaffinity(0)
+    logger.info(f"Check affinity before setting: {affinity}")
+    os.sched_setaffinity(0, affinity_mask)
+    affinity = os.sched_getaffinity(0)
+    logger.info(f"Check affinity after setting: {affinity}")
+    return SUCCESS_CODE
+
+
+def set_affinity():
+    if paddle.device.is_compiled_with_cuda():
+        return set_affinity_gpu()
+    elif paddle.device.is_compiled_with_xpu():
+        return set_affinity_xpu()
+    else:
+        # TODO(@gexiao): supports other devices if needed
+        logger.warning("Currently set_affinity only supports gpu env.")
+        return FAIL_CODE

python/paddle/distributed/utils/stream_utils.py

@@ -0,0 +1,19 @@
+# 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.
+
+from enum import Enum
+
+
+class ExecutionStreamType(Enum):
+    DefaultStream = "DefaultStream"