mindspore 2.7.0rc1__cp310-cp310-win_amd64.whl → 2.7.1__cp310-cp310-win_amd64.whl

mindspore/mint/distributed/distributed.py

@@ -18,13 +18,16 @@ import hashlib
 import builtins
 import io
 import pickle
+from datetime import timedelta
 import numpy as np
 from mindspore import log as logger
 from mindspore.common import dtype as mstype
+from mindspore._checkparam import args_type_check
 from mindspore.ops import ReduceOp, cat
 from mindspore.common.tensor import Tensor
 from mindspore._c_expression import TensorPy as Tensor_
 from mindspore.ops.primitive import _primexpr
+from mindspore.common.api import _pynative_executor
 from mindspore.communication._comm_helper import (
     _destroy_group_helper,
     _get_rank_helper,
@@ -33,10 +36,11 @@ from mindspore.communication._comm_helper import (
     _get_group_ranks,
     _is_available,
     _is_initialized,
+    _ExistingGroup,
 )
+from mindspore.communication.management import _init_without_sched
 from mindspore.communication import (
     init,
-    release,
     get_group_size,
     get_world_rank_from_group_rank,
     create_group,
@@ -58,9 +62,11 @@ from mindspore.ops.auto_generate.gen_ops_prim import (
     dist_comm_isend_op,
     dist_comm_all_to_all_v_op,
     dist_comm_reduce_scatter_tensor_op,
+    dist_comm_reduce_scatter_tensor_uneven_op,
     dist_comm_all_to_all_v_single_op,
     dist_comm_broadcast_op,
     dist_comm_all_gather_into_tensor_op,
+    dist_comm_all_gather_into_tensor_uneven_op,
     dist_comm_irecv_op,
     dist_comm_scatter_tensor_op,
     dist_comm_gather_into_tensor_op,
@@ -70,7 +76,7 @@ from mindspore.ops.auto_generate.gen_ops_prim import (
     dist_comm_barrier_op,
     dist_comm_batch_isend_irecv_op,
 )
-from mindspore._c_expression import TCPStoreClient, GroupOptions
+from mindspore._c_expression import TCPStoreClient, GroupOptions, _finalize_collective
 
 _pickler = pickle.Pickler
 _unpickler = pickle.Unpickler
@@ -144,28 +150,26 @@ class TCPStore:
 
     Note:
         - The function is implemented by CPU and does not involve any hardware operations related to Ascend.
-        - Currently, all parameters provided by the TCPStore class constructor are not supported.
-          The master node and port number are uniformly specified by the MindSpore framework.
-          The following parameters are provided, currently not supported and settings are invalid.
-        - The current TcpStore function is limited and only supports scenarios where the key is
+        - Currently, all parameters provided by the TCPStore class constructor are not supported
+          except for `host_name`, `port`, `world_size`, `is_master`, `timeout` and `wait_for_workers`,
+          which are reserved parameters and invalid settings.
+        - The current TCPStore function is limited and only supports scenarios where the key is
           less than 4k and the value is less than 1G. Complex scenarios are to be supported.
-        - The timeout interval for message sending and receiving in the TcpStore function is controlled by
-          the `MS_RECEIVE_MSG_TIMEOUT` environment variable, in seconds, with a default value of ``15``.
-          If a timeout occurs, the user needs to increase the configuration value.
 
     Args:
-        host_name (str, invalid, optional): The hostname or IP Address the server store should run on.
-            Default is ``None``.
-        port (int, invalid, optional): The port on which the server store should listen for incoming requests.
-            Default is ``None``.
-        world_size (int, invalid, optional): The total number of store users (number of clients + 1 for the server).
-            Default is ``None`` (``None`` indicates a non-fixed number of store users).
-        is_master (bool, invalid, optional): True when initializing the server store and False for client stores.
+        host_name (str): The hostname or IP Address the server store should run on.
+            Currently only supports user input IP addresses.
+        port (int): The port on which the server store should listen for incoming requests.
+        world_size (int, optional): The total number of store users (number of clients + 1 for the server).
+            Default is ``None``, indicates a non-fixed number of store users. This parameter is
+            only valid for the server.
+        is_master (bool, optional): True when initializing the server store and False for client stores.
             Default is ``False``.
-        timeout (timedelta, invalid, optional): Timeout used by the store during initialization, Unit: seconds.
-            Default is ``300``.
-        wait_for_workers (bool, invalid, optional): Whether to wait for all the workers to connect with the server
-            store. This is only applicable when `world_size` is a fixed value. Default is ``True``.
+        timeout (timedelta, optional): Timeout used by the store during initialization. Default is
+            ``timedelta(seconds=300)``.
+        wait_for_workers (bool, optional): Whether to wait for all the workers to connect with the server
+            store. This is only applicable when `world_size` is a fixed value. Default is ``True``. This
+            parameter is only valid for the server.
         multi_tenant (bool, invalid, optional): If ``True``, all ``TCPStore`` instances in the current process with
             the same host/port will use the same underlying ``TCPServer``. Default is ``False``.
         master_listen_fd (int, invalid, optional): If specified, the underlying ``TCPServer`` will listen on this file
@@ -191,12 +195,106 @@ class TCPStore:
             for more details.
 
         >>> from mindspore.mint.distributed import TCPStore
-        >>> store = TCPStore()
+        >>> store = TCPStore("127.0.0.1", 1234)
     """
 
-    def __init__(self, host_name=None, port=None, world_size=None, is_master=False, timeout=300,
+    def __init__(self, host_name, port, world_size=None, is_master=False, timeout=timedelta(seconds=300),
                  wait_for_workers=True, multi_tenant=False, master_listen_fd=None, use_libuv=True):
-        self.instance = TCPStoreClient.get_instance()
+        if not isinstance(host_name, str):
+            raise TypeError(
+                "For 'TCPStore', the argument 'host_name' must be type of string, "
+                "but got 'host_name' type : {}.".format(type(host_name))
+            )
+        if not isinstance(port, int):
+            raise TypeError(
+                "For 'TCPStore', the argument 'port' must be type of int, "
+                "but got 'port' type : {}.".format(type(port))
+            )
+        if not isinstance(is_master, bool):
+            raise TypeError(
+                "For 'TCPStore', the argument 'is_master' must be type of bool, "
+                "but got 'is_master' type : {}.".format(type(is_master))
+            )
+        if not isinstance(timeout, timedelta):
+            raise TypeError(
+                "For 'TCPStore', the argument 'timeout' must be type of timedelta, "
+                "but got 'timeout' type : {}.".format(type(timeout))
+            )
+        if not isinstance(wait_for_workers, bool):
+            raise TypeError(
+                "For 'TCPStore', the argument 'wait_for_workers' must be type of bool, "
+                "but got 'wait_for_workers' type : {}.".format(type(wait_for_workers))
+            )
+        if world_size is None:
+            world_size = 1
+        if not isinstance(world_size, int):
+            raise TypeError(
+                "For 'TCPStore', the argument 'world_size' must be type of int, "
+                "but got 'world_size' type : {}.".format(type(world_size))
+            )
+        if port < 0 or port > 65535:
+            raise ValueError(
+                "For 'TCPStore', the argument 'port' must be legal, "
+                f"but got {port}."
+            )
+        if world_size <= 0:
+            raise ValueError(
+                "For 'TCPStore', the argument 'world_size' must be legal, "
+                f"but got {world_size}."
+            )
+        timeout_ms = int(timeout.total_seconds() * 1000)
+        self.instance = TCPStoreClient(host_name, port, is_master, timeout_ms, world_size, wait_for_workers)
+        self.host = host_name
+        self.port = port
+
+
+    def add(self, key, amount):
+        """
+        When the `add` function is called for the first time with a given key, it creates a counter in
+        the storage corresponding to that key, with the initial value set to `amount`. Subsequent calls
+        to `add` with the same key increment the counter by amount.
+
+        Args:
+            key (str): The key whose counter value will be incremented.
+            amount (int): The amount by which the counter will be incremented.
+
+        Returns:
+            int, value of counter with `key`.
+
+        Raises:
+            TypeError: If `key` is not string.
+            TypeError: If `amount` is not int.
+            RuntimeError: If the `add` and `set` pass the same `key` and the `value` passed by `set` cannot
+                be correctly converted to a numerical value, calling `add` will result in an error.
+
+        Supported Platforms:
+            ``Ascend``
+
+        Examples:
+            .. note::
+                Before running the following examples, you need to configure the communication environment variables.
+
+                For Ascend devices, it is recommended to use the msrun startup method
+                without any third-party or configuration file dependencies.
+                Please see the `msrun start up
+                <https://www.mindspore.cn/tutorials/en/master/parallel/msrun_launcher.html>`_
+                for more details.
+
+            >>> from mindspore.mint.distributed import TCPStore
+            >>> store = TCPStore("127.0.0.1", 1234)
+            >>> store.add("first_key", 1)
+        """
+        if not isinstance(key, str):
+            raise TypeError(
+                "For 'TCPStore.add', the argument 'key' must be type of string, "
+                "but got 'key' type : {}.".format(type(key))
+            )
+        if not isinstance(amount, int):
+            raise TypeError(
+                "For 'TCPStore.add', the argument 'amount' must be type of string or int, "
+                "but got 'amount' type : {}.".format(type(amount))
+            )
+        return self.instance.add(key, amount)
 
 
     def set(self, key, value):
@@ -227,7 +325,7 @@ class TCPStore:
                 for more details.
 
             >>> from mindspore.mint.distributed import TCPStore
-            >>> store = TCPStore()
+            >>> store = TCPStore("127.0.0.1", 1234)
             >>> store.set("first_key", "first_value")
         """
         if not isinstance(key, str):
@@ -245,8 +343,9 @@ class TCPStore:
 
     def get(self, key):
         """
-        Retrieves the value associated with the given `key` in the store. If `key` is not
-        present in the store, the function will return "".
+        Retrieves the value associated with the given `key` in the store. If the `key` does not exist
+        in the storage, this function will wait for the `timeout` set by the class initialization and then
+        throw an exception.
 
         Args:
             key (str): The function will return the value associated with this key.
@@ -256,6 +355,7 @@ class TCPStore:
 
         Raises:
             TypeError: If `key` is not string.
+            RuntimeError: If `get` runs out of time.
 
         Supported Platforms:
             ``Ascend``
@@ -271,7 +371,7 @@ class TCPStore:
                 for more details.
 
             >>> from mindspore.mint.distributed import TCPStore
-            >>> store = TCPStore()
+            >>> store = TCPStore("127.0.0.1", 1234)
             >>> store.set("first_key", "first_value")
             >>> data = store.get("first_key")
             >>> print(data)
@@ -299,7 +399,7 @@ class TCPStore:
             TypeError: If `key` is not string.
 
         Supported Platforms:
-            ``CPU``
+            ``Ascend``
 
         Examples:
             .. note::
@@ -312,7 +412,7 @@ class TCPStore:
                 for more details.
 
             >>> from mindspore.mint.distributed import TCPStore
-            >>> store = TCPStore()
+            >>> store = TCPStore("127.0.0.1", 1234)
             >>> store.set("first_key", "first_value")
             >>> # This should return true
             >>> store.delete_key("first_key")
@@ -387,6 +487,7 @@ def is_initialized():
     return _is_initialized()
 
 
+@args_type_check(init_method=str, timeout=timedelta, world_size=int, rank=int, store=TCPStore)
 def init_process_group(backend="hccl",
                        init_method=None,
                        timeout=None,
@@ -404,26 +505,29 @@ def init_process_group(backend="hccl",
         and the instantiation and execution of any operation and net.
 
     Args:
-        backend (str, optional): The backend to ues. default is hccl and now only support hccl.
-        init_method (str, invalid): URL specifying how to init collective communication group. Provides parameters
-            consistent with pytorch, but is not currently support, setting is invalid.
-        timeout (timedelta, invalid): Timeout for API executed. Provides parameters consistent with pytorch, but is not
-            currently support, setting is invalid.
-        world_size (int, optional): Number of the processes participating in the job.
-        rank (int, invalid): Rank of the current process. Provides parameters consistent with pytorch, but is not
-            currently support, setting is invalid.
-        store (Store, invalid): Key/Value store accessible to all workers, used to exchange connection/address
-            information. Provides parameters consistent with pytorch, but is not currently support,
-            setting is invalid.
+        backend (str, optional): The backend to ues. Default is ``"hccl"`` and now only support hccl.
+        init_method (str, optional): URL specifying how to init collective communication group. Default is ``None``.
+        timeout (timedelta, optional): Timeout for API executed. Default is ``None``. Currently, this parameter is
+            only supported for host-side cluster network configuration using `init_method` or `store`.
+        world_size (int, optional): Number of the processes participating in the job. Default is ``-1``.
+        rank (int, optional): Rank of the current process. Default is ``-1``.
+        store (Store, optional): An object that stores key/value data, facilitating the exchange of inter-process
+            communication addresses and connection information. Default is ``None``. Currently, only the
+            ``TCPStore`` type is supported.
         pg_options (ProcessGroupOptions, invalid): process group options specifying what additional options need to be
-            passed in during the construction of specific process group. Provides parameters consistent with pytorch,
-            but is not currently support, setting is invalid.
-        device_id (int, invalid): the device id to exeute. Provides parameters consistent with pytorch, but is not
-            currently support, setting is invalid.
+            passed in during the construction of specific process group. The provided parameter is a reserved
+            parameter, and the current setting does not take effect.
+        device_id (int, invalid): the device id to exeute. The provided parameter is a reserved parameter,
+            and the current setting does not take effect.
 
     Raises:
         ValueError: If `backend` is not hccl.
         ValueError: If `world_size` is not equal to -1 or process group number.
+        ValueError: If both `init_method` and `store` are set.
+        ValueError: `world_size` is not correctly set as a positive integer value, when using the initialization
+            method `init_method` or `store`.
+        ValueError: `rank` is not correctly set as a non-negative integer, when using the initialization method
+            `init_method` or `store`.
         RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails,
             or the environment variables RANK_ID/MINDSPORE_HCCL_CONFIG_PATH
             have not been exported when backend is HCCL.
@@ -447,25 +551,34 @@ def init_process_group(backend="hccl",
         >>> init_process_group()
         >>> destroy_process_group()
     """
-    if init_method is not None:
-        logger.warning("init_method is ignored, setting is invalid")
-    if timeout is not None:
-        logger.warning("timeout is ignored, setting is invalid")
-    if store is not None:
-        logger.warning("store is ignored, setting is invalid")
     if pg_options is not None:
         logger.warning("pg_options is ignored, setting is invalid")
     if device_id is not None:
         logger.warning("device_id is ignored, setting is invalid")
-    if rank != -1:
-        logger.warning("rank is ignored, setting is invalid")
     if backend != "hccl":
         raise ValueError(
             "Only support hccl now, please setting backend to hccl or using default value"
         )
 
-    # init hccl & create world group
-    init(backend)
+    if init_method is not None and store is not None:
+        raise ValueError(
+            "Only one of init_method and store is supported."
+        )
+    if init_method is not None or store is not None:
+        if world_size <= 0:
+            raise ValueError(
+                "Specified world_size must be a positive integer."
+            )
+        if rank < 0:
+            raise ValueError(
+                "Specified rank must be a non-negative integer."
+            )
+        if timeout is None:
+            timeout = timedelta(seconds=300)
+        timeout_ms = int(timeout.total_seconds() * 1000)
+        _init_without_sched(backend, init_method, timeout_ms, world_size, rank, store)
+    else:
+        init(backend)
 
     if world_size != -1 and world_size != get_group_size():
         raise ValueError(
@@ -513,7 +626,10 @@ def destroy_process_group(group=None):
     """
 
     if group == GlobalComm.WORLD_COMM_GROUP or group is None:
-        release()
+        _pynative_executor.sync()
+        _finalize_collective()
+        _ExistingGroup.ITEMS.clear()
+        _ExistingGroup.GROUP_RANKS.clear()
     elif not isinstance(group, str):
         raise TypeError(
             "For 'destroy_group', the argument 'group' must be type of string or None, "
@@ -671,6 +787,12 @@ def new_group(ranks=None,
                     hccl_config(dict)
                 }
 
+            `hccl_config` currently only supports "hccl_buffer_size" or "hccl_comm".
+
+            - hccl_buffer_size (uint32): specifies the size of the HCCL communication buffer.
+            - hccl_comm (int64): specifies an existing HcclComm pointer. If "hccl_comm" is set,
+              "hccl_buffer_size" will be ignored.
+
         use_local_synchronization (bool, invalid): Currently it is a reserved parameter.
         group_desc (str, invalid): Currently it is a reserved parameter.
 
@@ -989,6 +1111,22 @@ def _check_all_tensor_same_dtype_and_shape(*tensor_lists):
                     )
 
 
+@_primexpr
+def _check_output_shape(output, expected_shape, op_name):
+    if output.shape != expected_shape:
+        raise TypeError(
+            f"For {op_name}, the output shape should be {expected_shape}, "
+            f"but got {output.shape}.")
+
+
+@_primexpr
+def _check_output_dtype(output, expected_dtype, op_name):
+    if output.dtype != expected_dtype:
+        raise TypeError(
+            f"For {op_name}, the output dtype should be {expected_dtype}, "
+            f"but got {output.dtype}.")
+
+
 def all_reduce(tensor, op=ReduceOp.SUM, group=None, async_op=False):
     """
     Reduce tensors across all devices in such a way that all deviceswill get the same final result,
@@ -1153,6 +1291,91 @@ def all_gather_into_tensor(output_tensor, input_tensor, group=None, async_op=Fal
     return handle
 
 
+def all_gather_into_tensor_uneven(output, input, output_split_sizes=None, group=None, async_op=False):
+    r"""
+    Gathers and concatenates tensors across devices with uneven first dimensions.
+
+    Note:
+        - Input tensors must have identical shapes except for the first dimension.
+        - Output tensor's first dimension should equal to the sum of all devices' input first dimensions.
+
+    Args:
+        output (Tensor): Concatenated output tensor with shape :math:`(\sum_{i=0}^{N-1} x_{i1}, x_2, ..., x_R)`,
+            where N is the number of devices in the group.
+        input (Tensor): Local input tensor with shape :math:`(x_{k1}, x_2, ..., x_R)`, where k is current device's rank.
+        output_split_sizes (list[int], optional): Specifies first dimension sizes from each device.
+            Must match actual input dimensions when provided.
+            If ``None``, assumes equal split sizes across devices. Default: ``None``.
+        group (str, optional): The communication group to work on. If ``None``,
+            which means ``"hccl_world_group"`` in Ascend. Default: ``None``.
+        async_op (bool, optional): Whether this operator should be an async operator. Default: ``False``.
+
+    Returns:
+        CommHandle, CommHandle is an async work handle, if `async_op` is set to True.
+        CommHandle will be None, when `async_op` is False.
+
+    Raises:
+        ValueError: If the shape of `input` does not match the constraints of `output_split_sizes`.
+        RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/tutorials/en/master/parallel/msrun_launcher.html>`_
+            for more details.
+
+            This example should be run with 2 devices.
+
+        >>> import numpy as np
+        >>> import mindspore as ms
+        >>> from mindspore import ops
+        >>> from mindspore.mint.distributed import init_process_group, get_rank
+        >>> from mindspore.mint.distributed import all_gather_into_tensor_uneven
+        >>> from mindspore import Tensor
+        >>>
+        >>> ms.set_device(device_target="Ascend")
+        >>> init_process_group()
+        >>> if get_rank() == 0:
+        ...     input_tensor = Tensor(np.ones([3, 4]).astype(np.float32))
+        ... else:
+        ...     input_tensor = Tensor(np.ones([2, 4]).astype(np.float32))
+        >>> out_tensor = Tensor(np.zeros([5, 4]).astype(np.float32))
+        >>> output_split_sizes = [3, 2]
+        >>> output = all_gather_into_tensor_uneven(out_tensor, input_tensor, output_split_sizes)
+        >>> print(out_tensor)
+        [[1. 1. 1. 1.]
+         [1. 1. 1. 1.]
+         [1. 1. 1. 1.]
+         [1. 1. 1. 1.]
+         [1. 1. 1. 1.]]
+    """
+    if group is None:
+        group = GlobalComm.WORLD_COMM_GROUP
+    if not isinstance(group, str):
+        raise TypeError(
+            "The argument 'group' must be type of string, "
+            "but got 'group' type : {}.".format(type(group))
+        )
+    if not isinstance(async_op, bool):
+        raise TypeError(
+            f"The argument 'async_op' must be a bool, but got {type(async_op)}."
+        )
+    group_size = get_cache_group_size(group)
+    output_split_sizes = [] if output_split_sizes is None else output_split_sizes
+    result = dist_comm_all_gather_into_tensor_uneven_op(
+        output, input, output_split_sizes, group_size, group
+    )
+    _, handle = _deal_comm_outputs(result, async_op)
+    return handle
+
+
 def reduce_scatter_tensor(output, input, op=ReduceOp.SUM, group=None, async_op=False):
     r"""
     Reduces and scatters tensors from the specified communication group and
@@ -1243,6 +1466,101 @@ def reduce_scatter_tensor(output, input, op=ReduceOp.SUM, group=None, async_op=F
     return handle
 
 
+def reduce_scatter_tensor_uneven(output, input, input_split_sizes=None, op=ReduceOp.SUM, group=None, async_op=False):
+    r"""
+    Reduce tensors from the specified communication group and scatter to the output tensor
+    according to `input_split_sizes`.
+
+    Note:
+        - The input tensor must have identical shape and format across all processes.
+        - The first dimension of input tensor should equal to the sum of `input_split_sizes`.
+
+    Args:
+        output(Tensor): the output tensor has the same dtype as `input` with a shape of
+            :math:`(input\_split\_sizes[rank], *)`, where rank is the local rank id of the device.
+        input(Tensor): The input tensor to be reduced and scattered, Expected shape :math:`(N, *)`, where `*`
+            means any number of additional dimensions. N must equal the sum of `input_split_sizes` across ranks.
+        input_split_sizes (list[int], optional): List specifying how to split the first dimension of input tensor.
+            If ``None``, splits evenly according to group size. Default: ``None``.
+        op (str, optional): Specifies an operation used for element-wise reductions,
+            One of ReduceOp: 'SUM', 'MIN', 'MAX'. Default: ``ReduceOp.SUM``.
+        group (str, optional): The communication group to work on. If ``None``, which means ``"hccl_world_group"`` in
+            Ascend. Default: ``None``.
+        async_op (bool, optional): Whether this operator should be an async operator. Default: ``False``.
+
+    Returns:
+        CommHandle, CommHandle is an async work handle, if `async_op` is set to True.
+        CommHandle will be None, when `async_op` is False.
+
+    Raises:
+        ValueError: If the shape of `output` does not match the constraints of `input_split_sizes`.
+        RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/tutorials/en/master/parallel/msrun_launcher.html>`_
+            for more details.
+
+            This example should be run with 2 devices.
+
+        >>> import mindspore as ms
+        >>> from mindspore import Tensor
+        >>> from mindspore.mint.distributed import init_process_group, get_rank
+        >>> from mindspore.mint.distributed import reduce_scatter_tensor_uneven
+        >>> import numpy as np
+        >>>
+        >>> ms.set_device(device_target="Ascend")
+        >>> init_process_group()
+        >>> input_tensor = Tensor(np.ones([5, 8]).astype(np.float32))
+        >>> if get_rank() == 0:
+        ...     output_tensor = Tensor(np.ones([2, 8]).astype(np.float32))
+        ... else:
+        ...     output_tensor = Tensor(np.ones([3, 8]).astype(np.float32))
+        >>> input_split_sizes = [2, 3]
+        >>> output = reduce_scatter_tensor_uneven(output_tensor, input_tensor, input_split_sizes)
+        >>> print(output_tensor)
+        rank 0:
+        [[2. 2. 2. 2. 2. 2. 2. 2.]
+         [2. 2. 2. 2. 2. 2. 2. 2.]]
+        rank 1:
+        [[2. 2. 2. 2. 2. 2. 2. 2.]
+         [2. 2. 2. 2. 2. 2. 2. 2.]
+         [2. 2. 2. 2. 2. 2. 2. 2.]]
+    """
+    if not isinstance(op, str):
+        raise TypeError("For reduce_scatter_tensor_uneven, the input op type must be str")
+    if op not in ("sum", "min", "max"):
+        raise TypeError(
+            "For reduce_scatter_tensor_uneven, the input op value must be one of sum, prod, min, max"
+        )
+    if group is None:
+        group = GlobalComm.WORLD_COMM_GROUP
+    if not isinstance(group, str):
+        raise TypeError(
+            "The argument 'group' must be type of string, "
+            "but got 'group' type : {}.".format(type(group))
+        )
+    if not isinstance(async_op, bool):
+        raise TypeError(
+            f"The argument 'async_op' must be a bool, but got {type(async_op)}."
+        )
+    input_split_sizes = [] if input_split_sizes is None else input_split_sizes
+    rank_size = get_cache_group_size(group)
+    result = dist_comm_reduce_scatter_tensor_uneven_op(
+        output, input, input_split_sizes, rank_size, op, group
+    )
+    _, handle = _deal_comm_outputs(result, async_op)
+    return handle
+
+
 def reduce(tensor, dst, op=ReduceOp.SUM, group=None, async_op=False):
     """
     Reduces tensors across the processes in the specified communication group, sends the result
@@ -2386,10 +2704,7 @@ def all_to_all_single(output,
 
 def _check_tensor_list(tensor_list, tensor, group_size):
     """check all elements in tensor_list are type of Tensor or tuple or list"""
-    if not tensor_list or len(tensor_list) != group_size:
-        raise TypeError(
-            f"The argument list tensor len must be equal to group rank size, but got {len(tensor_list)}."
-        )
+    _check_group_tensor_list(tensor_list, group_size)
     if tensor.dtype != tensor_list[0].dtype:
         raise TypeError(
             f"The argument list tensor type must be equal to tensor type, but got {tensor_list[0].dtype}."
@@ -2400,13 +2715,17 @@ def _check_tensor_list(tensor_list, tensor, group_size):
         )
 
 
+def _check_group_tensor_list(tensor_list, group_size):
+    if not tensor_list or len(tensor_list) != group_size:
+        raise TypeError(
+            f"The argument list tensor len must be equal to group rank size, but got {len(tensor_list)}."
+        )
+
+
 def all_gather(tensor_list, tensor, group=None, async_op=False):
     """
     Gathers tensors from the specified communication group and returns the tensor list which is all gathered.
 
-    Note:
-        The tensors must have the same shape and format in all processes of the collection.
-
     Args:
         tensor_list (list[Tensor]): Output list.
         tensor (Tensor): The input tensor to be all gathered into tensor.
@@ -2461,7 +2780,7 @@ def all_gather(tensor_list, tensor, group=None, async_op=False):
 
     """
     _check_all_tensors(tensor_list)
-    _check_all_tensor_same_dtype_and_shape(tensor_list)
+    _check_all_tensor_same_dtype(tensor_list)
     if not isinstance(tensor, (Tensor, Tensor_)):
         raise TypeError("For all_gather_into_tensor, the input tensor must be tensor")
     if group is None:
@@ -2476,7 +2795,10 @@ def all_gather(tensor_list, tensor, group=None, async_op=False):
             f"The argument 'async_op' must be a bool, but got {type(async_op)}."
         )
     group_size = get_cache_group_size(group)
-    _check_tensor_list(tensor_list, tensor, group_size)
+    _check_group_tensor_list(tensor_list, group_size)
+    rank_id = get_group_rank_from_world_rank(get_rank(), group)
+    _check_output_shape(tensor, tensor_list[rank_id].shape, "all_gather")
+    _check_output_dtype(tensor, tensor_list[0].dtype, "all_gather")
     result = dist_comm_all_gather_op(tensor_list, tensor, group_size, group)
     _, handle = _deal_comm_outputs(result, async_op)
     return handle
@@ -2487,9 +2809,6 @@ def reduce_scatter(output, input_list, op=ReduceOp.SUM, group=None, async_op=Fal
     Reduces and scatters tensors from the specified communication group and
     returns the tensor which is reduced and scattered.
 
-    Note:
-        The tensors must have the same shape and format in all processes of the collection.
-
     Args:
         output (Tensor): the output tensor.
         input_list (list[Tensor]): List of tensors to reduce and scatter.
@@ -2543,7 +2862,7 @@ def reduce_scatter(output, input_list, op=ReduceOp.SUM, group=None, async_op=Fal
     """
 
     _check_all_tensors(input_list)
-    _check_all_tensor_same_dtype_and_shape(input_list)
+    _check_all_tensor_same_dtype(input_list)
     if not isinstance(output, (Tensor, Tensor_)):
         raise TypeError("For reduce_scatter, the output tensor must be tensor")
     if group is None:
@@ -2564,7 +2883,11 @@ def reduce_scatter(output, input_list, op=ReduceOp.SUM, group=None, async_op=Fal
             "For reduce_scatter, the input op value must be one of sum, prod, min, max"
         )
     rank_size = get_cache_group_size(group)
-    _check_tensor_list(input_list, output, rank_size)
+    _check_group_tensor_list(input_list, rank_size)
+
+    rank_id = get_group_rank_from_world_rank(get_rank(), group)
+    _check_output_shape(output, input_list[rank_id].shape, "reduce_scatter")
+    _check_output_dtype(output, input_list[0].dtype, "reduce_scatter")
     result = dist_comm_reduce_scatter_op(output, input_list, rank_size, op, group)
     _, handle = _deal_comm_outputs(result, async_op)
     return handle