Diff Coverage

                cur_arg_batch_dim = 0
                if self.args_batch_dim and self.args_batch_dim[arg_idx] is not None:
                    cur_arg_batch_dim = self.args_batch_dim[arg_idx].batch_dim
                if isinstance(cur_arg, hyper_parallel.DTensor):
                    micro_arg = self.split_inputs_with_custom_shard(cur_arg, cur_arg_batch_dim, micro_idx)
                else:
                    micro_arg = self.split_inputs(cur_arg, cur_arg_batch_dim, micro_idx)
                micro_args.append(micro_arg)
            args_after_split.append(micro_args)

                cur_kwarg_batch_dim = 0
                if self.kwargs_batch_dim is not None:
                    cur_kwarg_batch_dim = self.kwargs_batch_dim[key].batch_dim
                if isinstance(cur_kwarg, hyper_parallel.DTensor):
                    micro_kwarg = self.split_inputs_with_custom_shard(cur_kwarg, cur_kwarg_batch_dim, micro_idx)
                else:
                    micro_kwarg = self.split_inputs(cur_kwarg, cur_kwarg_batch_dim, micro_idx)
                micro_kwargs[key] = micro_kwarg
            kwargs_after_split.append(micro_kwargs)

    def split_inputs_with_custom_shard(
            self, input_tensor: hyper_parallel.DTensor,
            cur_arg_batch_dim: int, micro_idx: int) -> hyper_parallel.DTensor:
        """Split a DTensor input along the batch dimension while preserving its distributed layout."""
        input_layout = input_tensor.layout
        func_wrap = hyper_parallel.custom_shard(self.split_inputs,
                                 device_mesh=input_layout.mesh,
                                 out_placements=(input_layout.placements,),
                                 in_placements=(input_layout.placements, None, None)
                                 )
        return func_wrap(input_tensor, cur_arg_batch_dim, micro_idx)

    def split_inputs(self, input_tensor: torch.Tensor, cur_arg_batch_dim: int, micro_idx: int) -> torch.Tensor:
        """
        Split the input along the specified batch_dim and micro_idx

        raise ValueError(f"pp_rank_list must contain only integer ranks, but got {pp_rank_list}.")
    if len(set(pp_rank_list)) != len(pp_rank_list):
        raise ValueError(f"pp_rank_list must not contain duplicate ranks, but got {pp_rank_list}.")
    if len(pp_rank_list) < 2:
        return []

    edge_rank_lists = {
        tuple(sorted((src_rank, dst_rank)))
        for src_rank, dst_rank in zip(pp_rank_list, pp_rank_list[1:])

    lock_key = "hyper_parallel_p2p_group_init_lock"
    lock_token = group_name
    deadline = time.monotonic() + 300
    while default_store.compare_set(lock_key, "", lock_token).decode() != lock_token:
        if time.monotonic() >= deadline:
            raise RuntimeError("Timed out waiting to initialize a batched P2P process group.")
        time.sleep(0.01)
    try:
        group_store.set(root_ready_key, "1")
        dist.barrier(group=group)
    finally:

        self._dw_cache = {}

    def clear_cache(self) -> None:
        """clear cache."""
        self.fwd_outputs_cache.clear()
        self.bwd_cache.clear()
        self._dw_cache.clear()
        self._meta_cache.clear()

    @staticmethod
    def _clear_recv_buffer(recv_info, micro_index):
        """clear fwd and bwd recv buffer."""

    def _check_pp_group(group):
        """check the type of pipeline group, if it is None, perform default initialization."""
        if group is None:
            return None
        if not isinstance(group, dist.ProcessGroup):
            raise TypeError("Argument 'group' must be type of ProcessGroup, but got type of {type(group)}.")
        return group

    @property
    def is_first_stage(self) -> bool:
        """return if is first stage."""

        else:
            if micro_index in self.args_recv_info:
                composite_args = [recv_info.buffer for recv_info in self.args_recv_info[micro_index]]
            else:
                raise RuntimeError(f"The exec order is wrong. The corresponding forward calculation \
                                    is executed before the Receive operation. micro is {micro_index}.")
        composite_kwargs = kwargs or {}
        out = self.submodule(*composite_args, **composite_kwargs)
        out_tuple = out if isinstance(out, tuple) else (out,)

        """Sens tensors for the last stage, aligned 1:1 with rg=True outputs."""
        sens_all = self.get_last_stage_sens(self.last_stage_outputs)
        if not isinstance(sens_all, list):
            return sens_all
        outputs_iter = (self.last_stage_outputs
                        if isinstance(self.last_stage_outputs, (list, tuple))
                        else [self.last_stage_outputs])
        sens = []
        for sensitivity, output in zip(sens_all, outputs_iter):
            local_output = output.to_local() if isinstance(output, hyper_parallel.DTensor) else output
            if local_output.requires_grad:
                sens.append(sensitivity)
        return sens

    def _populate_bwd_cache(self, micro_index):
        """Stash rg=True input grads so they align with peer's grad_recv_info."""
        input_grads = [recv_info.buffer.grad

        without ``grad_fn``.  ``bwd_cache`` is then populated with grads for
        rg=True inputs only, aligning 1:1 with the peer's ``grad_recv_info``.
        """
        if not self._has_backward:
            return
        recv_args = []
        if micro_index in self.grad_recv_info:
            recv_args = [recv_info.buffer for recv_info in self.grad_recv_info[micro_index]]

        fwd_output = self.fwd_cache.pop(micro_index)
        if self.is_last_stage:
            self.fwd_outputs_cache.pop(micro_index, None)

        local_output = self._filter_grad_outputs(fwd_output)

        if not local_output:
            # Nothing to backprop through (e.g. all forward outputs detached).
            self._clear_recv_buffer(self.grad_recv_info, micro_index)
            self._clear_recv_buffer(self.args_recv_info, micro_index)
            return

        grad_tensors = self._build_last_stage_sens() if self.is_last_stage else recv_args
        # MPipe owner-backward shares the tower's all-gather node across micro
        # graphs, so freeing it on the first backward breaks the later ones.

        torch.autograd.backward(local_output, grad_tensors=grad_tensors,
                                retain_graph=retain_graph)

        if not self.is_first_stage:
            self._populate_bwd_cache(micro_index)
        self._clear_recv_buffer(self.grad_recv_info, micro_index)
        self._clear_recv_buffer(self.args_recv_info, micro_index)

    def backward_input_one_chunk(self, micro_index: int) -> None:

    def backward_input_one_chunk(self, micro_index: int) -> None:
        """Compute input gradients and retain only the graph state needed by dw."""
        if not self._has_backward or self.is_first_stage:
            return

        recv_args = []
        if micro_index in self.grad_recv_info:
            recv_args = [recv_info.buffer for recv_info in self.grad_recv_info[micro_index]]

        fwd_output = self.fwd_cache.pop(micro_index)
        local_output = self._filter_grad_outputs(fwd_output)
        grad_tensors = self._build_last_stage_sens() if self.is_last_stage else recv_args
        if not isinstance(grad_tensors, (list, tuple)):
            grad_tensors = [grad_tensors]
        input_values = [
            recv_info.buffer
            for recv_info in self.args_recv_info[micro_index]
            if recv_info.requires_grad

    def backward_weight_one_chunk(self, micro_index: int) -> None:
        """Compute parameter gradients from state captured by input backward."""
        if not self._has_backward:
            return
        if self.is_first_stage:
            self.backward_one_chunk(micro_index)
            return
        if micro_index not in self._dw_cache:
            raise RuntimeError(f"stage: {self.stage_index} micro_{micro_index} dw called before dx.")

        param_groups = self._dw_cache.pop(micro_index)

        stage_backward_weight(iter(self._get_trainable_params()), param_groups)
        self._clear_recv_buffer(self.grad_recv_info, micro_index)
        self._clear_recv_buffer(self.args_recv_info, micro_index)
        if self.is_last_stage:
            self.fwd_outputs_cache.pop(micro_index, None)

    def get_last_stage_sens(self, last_stage_outputs: Any) -> Any:
        """Get last stage sens"""
        p_sens = None

    def get_last_stage_sens(self, last_stage_outputs: Any) -> Any:
        """Get last stage sens"""
        p_sens = None
        if isinstance(last_stage_outputs, (list, tuple)):
            p_sens = []
            for _, out_i in enumerate(last_stage_outputs):
                if isinstance(out_i, hyper_parallel.DTensor):
                    repeat_num = out_i.layout.repeat_num()
                    sens_i = torch.full_like(out_i.to_local(), 1.0 / repeat_num)
                else:
                    sens_i = torch.full_like(out_i, 1.0)
                p_sens.append(sens_i)
        else:
            if isinstance(last_stage_outputs, hyper_parallel.DTensor):
                repeat_num = last_stage_outputs.layout.repeat_num()
                p_sens = torch.full_like(last_stage_outputs.to_local(), 1.0 / repeat_num)
            else:
                p_sens = torch.full_like(last_stage_outputs, 1.0)

        return p_sens

            role_of = _SyncHookFunction._role_enum
            coordinator.notify_dispatched(role_of(prev_idx))
            return x

        prev_idx, next_idx = _SyncHookFunction._FWD_ROLES[hook_name]
        role_of = _SyncHookFunction._role_enum
        coordinator.notify_dispatched(role_of(prev_idx))
        coordinator.rendezvous(role_of(next_idx))
        return x

    @staticmethod
    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]:
        """Identity backward that fires a HookCoordinator rendezvous.

            result = dispatch_fn(first, *rest, **kwargs)
            if isinstance(result, tuple):
                hooked = _SyncHookFunction.apply(result[0], "B", coordinator)
                return (hooked,) + result[1:]
            return _SyncHookFunction.apply(result, "B", coordinator)

        return _wrapped

    def wrap_combine(self, combine_fn: Callable, is_last_layer: bool = False) -> Callable:

            first, rest = args[0], args[1:]
            first = _SyncHookFunction.apply(first, "C", coordinator)
            result = combine_fn(first, *rest, **kwargs)
            if isinstance(result, tuple):
                hooked = _SyncHookFunction.apply(result[0], d_hook, coordinator)
                return (hooked,) + result[1:]
            return _SyncHookFunction.apply(result, d_hook, coordinator)

        return _wrapped

        self._fwd_received.clear()
        # Snapshot the already-reduced grads from prior runs so GRAD_REDUCE
        # reduces only this run's added contribution.
        if self._owner_backward:
            self._grad_snapshot = self._snapshot_tower_grads()

    def _send_meta(self, tensor, dst) -> None:
        """Send a tensor's ``(shape, dtype)`` to ``dst``.

    def _output_arity(self):
        """Number of output tensors to communicate (cached after first use)."""
        if self._output_arity_for_comm is not None:
            return self._output_arity_for_comm
        # Fallback: derive the arity from the first retained micro.
        if self._outputs_for_stage0.get(0):
            self._output_arity_for_comm = len(self._outputs_for_stage0[0])
        elif self._outputs_for_bwd.get(0):
            self._output_arity_for_comm = len(self._outputs_for_bwd[0])
        else:
            raise ValueError("Cannot determine the arity, number of elements to be communicated")
        return self._output_arity_for_comm

    def _input_arity(self):
        """Number of input tensors to communicate (cached after first use)."""

    def _input_arity(self):
        """Number of input tensors to communicate (cached after first use)."""
        if self._input_arity_for_comm is not None:
            return self._input_arity_for_comm
        # Fallback: derive the arity from the first retained micro.
        if self._inputs_for_explicit_forward.get(0):
            self._input_arity_for_comm = len(self._inputs_for_explicit_forward[0])
        else:
            raise ValueError("Cannot determine the arity, number of elements to be communicated")
        return self._input_arity_for_comm

    def _global_rank(self, group_rank: int) -> int:
        """Global rank of group pipeline rank ``group_rank``."""

                self._output_arity_for_comm = len(out)
                self._save_grad_for_bwd(micro, out)
                self._transfer_out_for_fwd(ctx, micro, out)
            else:
                self._outputs_for_stage0[micro] = out
        # Stages built under meta-init only learn their real device here.
        if out and getattr(self._device, "type", None) in (None, "meta"):
            self._device = out[0].device

                are handed to the forward via ``_transfer_out_for_fwd`` and
                recorded via ``_save_grad_for_bwd``.
        """
        if not _RECV_BATCH:
            self._fwd_recv_one(step.micro_index, ctx)
            return
        if step.micro_index in self._fwd_received:
            return
        self._fwd_recv_all(ctx)

    def _recv_device(self):
        """Real (non-``meta``) device for receive buffers.

                             self._keep_grad.get(0)):
            if materialized:
                self._device = materialized[0].device
                return self._device
        return self._device

    def _recv_payload(self, micro):
        """Post the recv of ``micro``'s preprocess output; return ``(buffers, handles)``.

        return buffers, handles

    def _fwd_recv_one(self, micro: int, ctx: "PipelineContext") -> None:
        """Receive a single transposed micro-batch's output, waiting inline (fallback)."""
        buffers, handles = self._recv_payload(micro)
        for handle in handles:
            handle.wait()
        out = tuple(buffers)
        self._save_grad_for_bwd(micro, out)
        self._transfer_out_for_fwd(ctx, micro, out)
        self._fwd_received.add(micro)

    def _fwd_recv_all(self, ctx: "PipelineContext") -> None:
        """Post every outstanding transposed recv, then wait once so the transfers
        overlap each other and stage 0's own preprocess forward."""

            out = self._preprocess(*args, **kwargs)
        # The preprocess output may be a single tensor (text body input) or a
        # tuple (e.g. a VL visual payload: image_embeds + DeepStack levels).
        if isinstance(out, (tuple, list)):
            return tuple(t.detach() for t in out)
        return out.detach()

    def _connected_forward(self, args, kwargs):
        return self._preprocess(*args, **kwargs)

        # Backprop only the outputs that received a gradient; an absent
        # dL/dfeature is a zero contribution.
        pairs = [(out_i, g) for out_i, g in zip(out, grads) if g is not None]
        if not pairs:
            return
        outs, grad_tensors = zip(*pairs)
        torch.autograd.backward(outs, grad_tensors=grad_tensors)

    @staticmethod

        return tensor.detach().contiguous()

    def _zeros_like(self, tensor):
        # Contiguous zero grad for a feature tensor that received none.
        return torch.zeros_like(tensor).contiguous()

    def _owner_transpose_backward(self, retained_out, grads) -> None:
        # Backprop dL/dfeatures through the retained connected tower graph on this
        # rank's replica; skip non-grad-requiring outputs (else autograd raises).

        if pairs:
            outs, grad_tensors = zip(*pairs)
            torch.autograd.backward(outs, grad_tensors=grad_tensors)
        else:
            logger.debug("[mpipe] owner backward skipped: no retained output requires grad.")

    def _snapshot_tower_grads(self):
        # Clone each trainable tower grad's LOCAL shard (None where absent) so
        # _reduce_grads reduces only this run's contribution.
        return [None if p.grad is None else self._local(p.grad).detach().clone()
                for p in self._preprocess.parameters() if p.requires_grad]

    def _reduce_grads(self, group, snapshot) -> None:
        """SUM-reduce each trainable tower param's this-run contribution

def _exec_fsdp_unshard(stage):
    """Unshard every HSDPModule in the stage's submodule tree."""
    for _, module in stage.submodule.named_modules():
        if isinstance(module, HSDPModule):
            module.unshard()

def _exec_fsdp_reshard(stage):
    """Reshard every HSDPModule in the stage's submodule tree."""
    for _, module in stage.submodule.named_modules():
        if isinstance(module, HSDPModule):
            module.reshard()

                    f"No P2P multi-stream group was created for peer global rank {peer}. "
                    f"Available peers are {sorted(self._p2p_multi_stream_groups)}."
                )
        if op_type not in ("isend", "irecv"):
            raise ValueError(f"Unsupported pipeline P2P operation: {op_type!r}.")
        op = dist.isend if op_type == "isend" else dist.irecv
        return dist.P2POp(op, tensor, peer, group=group)

    def convert_stages_dict(self):

            by_peer.setdefault(item[2], []).append(item)

        for items in by_peer.values():
            ops = [self._p2p_op(op_type, tensor, peer) for op_type, tensor, peer, _ in items]
            handles = dist.batch_isend_irecv(ops)
            if not handles:
                continue
            if not self._overlap_p2p:
                self._wait_p2p(handles)
                continue
            recv_routes = [route for *_, route in items if route is not None]
            if recv_routes:
                for kind, si, mi in recv_routes:

            recv_routes = [route for *_, route in items if route is not None]
            if recv_routes:
                for kind, si, mi in recv_routes:
                    cache = self.fwd_handle_cache if kind == "fwd" else self.bwd_handle_cache
                    cache[(si, mi)] = handles
            else:
                self._send_handles.append(handles)

    def _assert_in_unshard_if_needed(self, stage, check_step):
        if not isinstance(stage.submodule, HSDPModule):
            return

                micro_batch["targets"] = targets
                for key in self.kwargs_batch_dim:
                    if key in kwarg_mbs[micro_index].keys():
                        kwarg_mbs[micro_index][key] = torch.cat(
                            [kwarg_mbs[micro_index][key], micro_batch[key]], dim=0)
                    else:
                        kwarg_mbs[micro_index][key] = micro_batch[key]
                arg_mbs[micro_index] = [input_ids]

                tensor = (arg_mbs[micro_index][0] if key == "input_ids"
                            else kwarg_mbs[micro_index][key])
                metas.append([tuple(tensor.shape), tensor.dtype])
                tensors.append(tensor)
            dist.send_object_list(metas, dst)
            handles = [dist.isend(t, dst) for t in tensors]
            self._wait_p2p(handles)

        elif step_type == MetaStepType.DATA_RECV:
            if getattr(self, "_data_src", False):

            device = self.stages[0].device
            # One recv_object_list unpacks every key's (shape, dtype) at
            # once, matching the packed DATA_SEND above.
            metas: list = [None] * len(self._DATA_KEYS)
            dist.recv_object_list(metas, src)
            handles = []
            for key, meta in zip(self._DATA_KEYS, metas):
                shape, dtype = meta
                buffer = torch.empty(shape, dtype=dtype, device=device)
                handles.append(dist.irecv(buffer, src))
                if key == "input_ids":
                    arg_mbs[micro_index] = [buffer]
                else:
                    kwarg_mbs[micro_index][key] = buffer

            shared_stage = shared_param_info.shared_stage
            group, group_ranks = self._init_shared_parameter_group(shared_stage)
            shared_param_info.group = group
            # Raw c10d collectives operate on the matching local TP/FSDP shard.
            local_param = param.to_local() if isinstance(param, DTensor) else param
            with torch.no_grad():
                dist.broadcast(local_param, group_ranks[0], group)

    def _global_rank(self, stage_index):
        real_stage_num = self.stage_num // self._virtual_chunk_num
        real_stage_index = stage_index % real_stage_num

        real_stage_index = stage_index % real_stage_num
        if self.mesh is not None:
            # mesh is a 1-D PP sub-mesh; rank_list[i] is the global rank of stage i.
            return self.mesh.rank_list[real_stage_index]
        return dist.get_global_rank(self.pp_group, real_stage_index)

    def _init_shared_parameter_group(self, shared_stage):
        """init group of shared parameter."""
        group_ranks = [self._global_rank(stage) for stage in shared_stage]

        # this avoids a world-collective ``new_group`` that would deadlock under
        # FSDP where each dp line needs its own {stage0, stage1} group.
        if self.mesh is not None and len(group_ranks) == self.mesh.size():
            return self.mesh.get_group(), group_ranks
        group = dist.new_group(group_ranks, use_local_synchronization=True)
        return group, group_ranks

    def sync_shared_parameters_grad(self) -> None:
        """sync shared parameters' grad with AllReduce."""

        """
        if self.mesh is not None:
            rank_list = self._get_layout_rank_list(layout, sender_rank)
        else:
            device_num = dist.get_world_size()
            real_stage_num = self.stage_num // self._virtual_chunk_num
            device_num_per_stage = device_num // real_stage_num
            index = self.stage_index % real_stage_num
            rank_list = tuple(range(index * device_num_per_stage, (index + 1) * device_num_per_stage))

            # layout's submesh can't be resolved by name. A PP edge shifts whole
            # stages by a fixed rank block and the within-stage tile is identical
            # across stages, so this receiver's submesh is the sender's submesh
            # (layout.rank_list) shifted by the P2P offset (me - sender_rank).
            me = dist.get_rank()
            cur_ranks = tuple(layout.rank_list or ())
            # The cached layout is reused across microbatches/steps and
            # _update_layout mutates it in place, so re-resolution must be
            # idempotent: once it holds this receiver's submesh (contains me),

        layout_dim_names = tuple(
            name for name in (layout.alias_name or ()) if name not in pp_dim_names
        )
        if not layout_dim_names:
            return (dist.get_rank(),)
        if len(layout_dim_names) == 1:
            submesh = root[layout_dim_names[0]]
        else:
            submesh = root[layout_dim_names]

        """
        requires_grad = bool(meta[-1])
        if len(meta) == 4:
            self._update_layout(meta[2], global_rank)
            buffer = DTensor.from_local(torch.empty(meta[0], dtype=meta[1], device=self.device),
                                        meta[2].mesh, meta[2].alias_placements)
        else:
            buffer = torch.empty(meta[0], dtype=meta[1], device=self.device)
        buffer.requires_grad = requires_grad
        if micro_index in self.args_recv_info:
            recv_info = self.args_recv_info[micro_index][idx]
            recv_info.buffer = buffer

    def _communicate_meta(self, global_rank, meta_send=None):
        """communicate meta."""
        if meta_send is not None:
            if self._dyn_shape or not self._meta_been_send:
                dist.send_object_list([meta_send], global_rank)
                self._meta_been_send = True
            return None

        if self._dyn_shape or not self._meta_been_recv:

            return None

        if self._dyn_shape or not self._meta_been_recv:
            obj_list = [None]
            dist.recv_object_list(obj_list, global_rank)
            self._meta_been_recv = True
            if not self._dyn_shape:
                self._meta_cache = obj_list
            return obj_list

        return specs

    def exec_fwd_recv_ops(self, micro_index):
        """Execute the forward recv operation."""
        return [dist.irecv(tensor, rank) for _, tensor, rank in self.fwd_recv_specs(micro_index)]

    def _construct_backward_recv_info(self, micro_index, idx, global_rank):
        """Reserve backward recv bookkeeping without allocating its device buffer."""
        if micro_index not in self.grad_recv_info:

        The full output meta list is also stashed in ``_fwd_output_meta``
        so backward receive buffers match the grad-requiring output slots.
        """
        return [dist.isend(tensor, rank) for _, tensor, rank in self.fwd_send_specs(micro_index)]

    def fwd_send_specs(self, micro_index):
        """Prepare forward-send tensors (+ bookkeeping) without launching.

        return specs

    def exec_bwd_recv_ops(self, micro_index):
        """Execute the backward recv operation."""
        return [dist.irecv(tensor, rank) for _, tensor, rank in self.bwd_recv_specs(micro_index)]

    def exec_bwd_send_ops(self, micro_index):
        """Execute the backward send operation.

        slots of ``args_recv_info[mi]`` (and 1:1 with the peer's
        ``grad_recv_info``).  Pairing via ``zip`` keeps send count and
        peer irecv count consistent.
        """
        return [dist.isend(tensor, rank) for _, tensor, rank in self.bwd_send_specs(micro_index)]

    def bwd_send_specs(self, micro_index):
        """Prepare backward-send (input-grad) tensors without launching.

        if micro_index not in self.args_recv_info:
            return []
        out = self.bwd_cache.pop(micro_index)
        rg_infos = [info for info in self.args_recv_info[micro_index] if info.requires_grad]
        return [
            ("isend", grad.to_local() if isinstance(grad, DTensor) else grad, info.global_rank)
            for grad, info in zip(out, rg_infos)
        ]

    def execute_reduce_grad(self) -> None:
        """Reduce nested FSDP gradients and finalize every HSDP root in the stage."""
        fsdp_schedulers = []
        seen_scheduler_ids = set()
        for _, submod in self.submodule.named_modules():
            if not isinstance(submod, HSDPModule):
                continue
            scheduler = submod.hsdp_scheduler
            scheduler_id = id(scheduler)

# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""MindSpore object-list communication helpers."""
import io
import pickle
from typing import Any, Optional


def send_object_list(obj: list[Any], dst: int = 0, group: Optional[str] = None) -> None:
    """
    Send the input Python object to dst rank.

    Args:

        dst (int, optional): Specifies the global rank that send the Python object to.
            Default: ``0``.
        group (str, optional): Communication group. Default: ``None``.
    """
    from mindspore import Tensor  # pylint: disable=C0415
    from mindspore.common import dtype as mstype  # pylint: disable=C0415
    from mindspore.communication import GlobalComm  # pylint: disable=C0415
    from mindspore.mint.distributed.distributed import _object_to_tensor, send  # pylint: disable=C0415
    if group is None:
        group = GlobalComm.WORLD_COMM_GROUP
    if not isinstance(group, str):
        raise TypeError(f"For 'send_object', the argument 'group' must be type of string, \
                          but got 'group' type : {type(group)}.")
    if not isinstance(dst, int):
        raise TypeError("For send_object, the dst must be int.")
    obj_tensor, tensor_size = _object_to_tensor(obj)
    obj_size = Tensor([tensor_size], dtype=mstype.int32)
    send(obj_size, dst, group)
    send(obj_tensor, dst, group)



def recv_object_list(recv_obj: list[Any], src: int = 0, group: Optional[str] = None) -> None:
    """
    receive Python object from src rank.

    Args:

        src (int, optional): Specifies the global rank that receive the Python object.
            Default: ``0`` .
        group (str, optional): Communication group. Default: ``None``.
    """
    from mindspore import mint  # pylint: disable=C0415
    from mindspore.common import dtype as mstype  # pylint: disable=C0415
    from mindspore.communication import GlobalComm  # pylint: disable=C0415
    from mindspore.mint.distributed.distributed import recv  # pylint: disable=C0415
    if group is None:
        group = GlobalComm.WORLD_COMM_GROUP
    if not isinstance(group, str):
        raise TypeError(f"For 'recv_object', the argument 'group' must be type of string, \
                          but got 'group' type : {type(group)}.")
    if not isinstance(src, int):
        raise TypeError("For recv_object, the src must be int.")
    obj_size = mint.zeros((1,), dtype=mstype.int32)
    recv(obj_size, src, group)
    # MindSpore PyNative ``recv`` only does a comm-stream wait; bridge to host
    # so the subsequent ``.item()`` reads the freshly-received value instead
    # of the original buffer.
    size_val = int(obj_size.item())
    obj_tensor = mint.zeros((size_val,), dtype=mstype.int8)
    recv(obj_tensor, src, group)
    buf = obj_tensor.asnumpy().tobytes()[:size_val]
    recv_obj.clear()
    recv_obj.append(pickle.Unpickler(io.BytesIO(buf)).load()[0])

    @staticmethod
    def send_object_list(obj_list, dst=None, group=None):
        # pylint: disable=C0415
        from hyper_parallel.platform.mindspore.object_collectives import send_object_list
        send_object_list(obj_list, dst, group)

    @staticmethod
    def recv_object_list(obj_list, src=None, group=None):

    @staticmethod
    def recv_object_list(obj_list, src=None, group=None):
        # pylint: disable=C0415
        from hyper_parallel.platform.mindspore.object_collectives import recv_object_list
        recv_object_list(obj_list, src, group)

    @staticmethod
    def set_tensor_requires_grad(input_tensor):