Diff Coverage

        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."""

            and self._has_trainable_preprocess
            and platform.platform_type == PlatformType.PYTORCH
        )
        if owner_backward and self._has_trainable_preprocess and not self._owner_backward:
            logger.warning(
                "[mpipe] pp_mpipe_owner_backward requested but unsupported on "
                "platform %s; falling back to the stage-0 backward.",
                platform.platform_type,
            )

            # An exception unwinds past the end-of-iteration drain, leaving
            # handles un-waited; wait them so the contract holds on that path too.
            self._drain_inflight_p2p()
            if self._swap_session is not None:
                self._swap_session.close()
                self._swap_session = None
        return losses

    def sync_shared_parameters_grad(self):

        """
        # Only active under ``run_with_dataiterator``; the legacy ``run()``
        # path pre-stages every per-micro kwarg on every rank already.
        if self.data_iterator is None:
            return
        if step_type == MetaStepType.DATA_LOAD:
            if stage_index <= 0:
                device = self.stages[0].device
                micro_batch = next(self.data_iterator)

            if stage_index <= 0:
                device = self.stages[0].device
                micro_batch = next(self.data_iterator)
                if isinstance(micro_batch, list):
                    micro_batch = micro_batch[micro_index]
                micro_batch = {
                    key: (value.to(device, non_blocking=True) if hasattr(value, "to") else value)
                    for key, value in micro_batch.items()
                }

                self.last_local_tokens += n_valid

                micro_batch["targets"] = targets
                for key in self.kwargs_batch_dim:
                    if key in kwarg_mbs[micro_index].keys():
                        kwarg_mbs[micro_index][key] = platform.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]

        elif step_type == MetaStepType.DATA_SEND:
            if getattr(self, "_data_dst", False):
                dst_stage_idx = self._data_dst[stage_index][micro_index]
            else:
                dst_stage_idx = self.stages[0].dst_stage
            dst = self.stages[0]._global_rank(dst_stage_idx)  # pylint: disable=protected-access
            # One meta round-trip per DATA_SEND instead of K; the matching
            # DATA_RECV unpacks the list in the same order.
            metas = []
            tensors = []
            for key in self._DATA_KEYS:
                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)
            platform.send_object_list(metas, dst)
            handles = [platform.isend(t, dst) for t in tensors]
            self._wait_p2p(handles)

        elif step_type == MetaStepType.DATA_RECV:
            if getattr(self, "_data_src", False):
                src_stage_idx = self._data_src[stage_index][micro_index]
            else:
                src_stage_idx = self.stages[0].src_stage
            src = self.stages[0]._global_rank(src_stage_idx)  # pylint: disable=protected-access
            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)
            platform.recv_object_list(metas, src)
            handles = []
            for key, meta in zip(self._DATA_KEYS, metas):
                shape, dtype = meta
                buffer = platform.empty(shape, dtype=dtype, device=device)
                handles.append(platform.irecv(buffer, src))
                if key == "input_ids":
                    arg_mbs[micro_index] = [buffer]
                else:
                    kwarg_mbs[micro_index][key] = buffer
            self._wait_p2p(handles)

    def _exec_pipeline_swap_step(self, cur_step, arg_mbs, kwarg_mbs):
        """Execute a pipeline activation-swap control step."""
        if self._swap_session is None:

            peer = _fwd_peer(step.stage_index)
            if peer is not None:
                new_schedule[rank].append(
                    MetaStep(step.micro_index, MetaStepType.DATA_SEND, step.stage_index))
                new_schedule[rank].append(
                    MetaStep(step.micro_index, MetaStepType.FWD_SEND, step.stage_index))
                new_schedule[peer].append(
                    MetaStep(step.micro_index, MetaStepType.DATA_RECV, step.stage_index + 1))
                new_schedule[peer].append(
                    MetaStep(step.micro_index, MetaStepType.FWD_RECV, step.stage_index + 1))
        elif step.type == MetaStepType.BWD:
            peer = _bwd_peer(step.stage_index)
            if peer is not None:

    # the RECV triggered by an overlap's second sub-step lands one slot
    # later on the receiver — matching the fact that the sender can only
    # finish emitting the second sub-step after the first completes.
    expanded = _expand_overlap_slots(aligned, real_stage_num)
    return _insert_dataload_before_fwd(
        _column_scan_insert_comms(expanded, real_stage_num, _insert_comms_for_step))


class ScheduleGPipe(PipelineScheduleRuntime):

        """construct_exec_order of Gpipe."""
        for stage_index in range(self.real_stage_num):
            order_list = []
            for mb_index in range(self.micro_batch_num):
                order_list.append(MetaStep(mb_index, MetaStepType.DATA_LOAD, stage_index))
                if stage_index != 0:
                    order_list.append(MetaStep(mb_index, MetaStepType.DATA_RECV, stage_index))
                    order_list.append(MetaStep(mb_index, MetaStepType.FWD_RECV, stage_index))
                order_list.append(MetaStep(mb_index, MetaStepType.FWD, stage_index))
                if stage_index != self.real_stage_num - 1:
                    order_list.append(MetaStep(mb_index, MetaStepType.DATA_SEND, stage_index))
                    order_list.append(MetaStep(mb_index, MetaStepType.FWD_SEND, stage_index))
            for mb_index in range(self.micro_batch_num):
                if stage_index != self.real_stage_num - 1:
                    order_list.append(MetaStep(mb_index, MetaStepType.BWD_RECV, stage_index))

Note:
    Exercised by the MindSpore ST gate (msrun), not the torch/CPU coverage job
    (mindspore isn't importable there).
"""
from mindspore import ops

from hyper_parallel.core.pipeline_parallel.mpipe.executor_base import MPipeTransposeExecutorBase
from hyper_parallel.platform.mindspore.pipeline_parallel.backward import forward_and_gradfn

                _raise_callable_already_wrapped(wrapped_callable)
            # A param-free module shared across siblings (one rotary embedding
            # per decoder layer) is never its own region; skip, do not flag.
            if next(submodule.parameters(recurse=True), None) is None:
                continue
            warnings.warn(
                f"Submodule '{getattr(submodule, '_swap_wrapped_module', submodule).__class__.__name__}' of "
                f"'{module.__class__.__name__}' is already wrapped. "
                "Wrapping overlapping module regions is not allowed."

            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_info, snapshot) -> None:
        """SUM-reduce each trainable tower param's this-run contribution

            seed=sampler_seed,
            drop_last=True,
        )

        self._update_dataloader_sampler_for_mpipe()

        # StatefulDataLoader supports state_dict() / load_state_dict()
        # for checkpoint resume (torchdata API, used by  + ).
        num_workers = self.args.data.num_workers

        his owned data to load
        """
        # Only for MPipe's in-schedule dataload; every other PP path expects
        # the full step batch on each rank's dataloader.
        load_mode = getattr(self.args.data, "load", None)
        pp_size = int(getattr(self.parallel_dims, "pp", 1))
        pp_schedule = str(getattr(self.args.train.accelerator, "pp_schedule", "")).lower()
        if load_mode == "single" and pp_size > 1 and pp_schedule in ("mpipe", "mpipe_transpose"):
            micro_bs = self.args.train.micro_batch_size
            dp_size = self.parallel_dims.dp_size
            pp_micro_batch_num = int(getattr(self.args.train.accelerator, "pp_micro_batch_num", 1))
            overflow_mode = str(getattr(self.args.train.accelerator,
                                        "pp_mpipe_transpose_overflow", "full"))
            pp_rank = torch.distributed.get_rank(group=self.mesh["pp"].get_group())
            owned = mpipe_owned_micros(pp_size, pp_micro_batch_num, pp_rank,
                                       mode=overflow_mode)
            if not owned:
                raise ValueError(
                    f"data.load='single': PP rank {pp_rank} owns no micros "
                    f"(M={pp_micro_batch_num} < PP={pp_size}). "
                    "Configure pp_micro_batch_num >= pp."
                )
            self.sampler = PPRankOwnedSampler(
                base_sampler=self.sampler,
                owned_micros=owned,
                micro_batch_size=micro_bs,
                pp_micro_batch_num=pp_micro_batch_num,

                pp_micro_batch_num=pp_micro_batch_num,
            )

            # Fail on an undersized dataset here rather than mid-run.
            base_len = len(self.sampler.base)
            step_chunk = self.sampler.step_size
            available_steps = base_len // step_chunk if step_chunk else 0
            max_steps = int(getattr(self.args.train, "max_steps", 0))
            if max_steps and available_steps < max_steps:
                per_step = pp_micro_batch_num * micro_bs * dp_size
                raise ValueError(
                    f"data.load='single': dataset feeds only {available_steps} "
                    f"step(s)/epoch (base sampler has {base_len} samples, "
                    f"step chunk={step_chunk}) but max_steps={max_steps}. Each "
                    f"step consumes pp_micro_batch_num*micro_batch_size*dp_size "

        Gradient clipping uses the **global** cross-stage norm
        (:meth:`_pp_clip_grad_norm`) so every stage scales by the same
        coefficient — required so the tied embed / lm_head copies stay in sync.
        """
        if getattr(self.args.data, "load", None) == "single":
            # ``last_local_tokens`` carries only this rank's micros, so PP
            # all-reduce for the cross-stage global count.
            outputs = self.pp_schedule.run_with_dataiterator(
                data_iterator,
                getattr(self, "_pp_fsdp_composed", False),
                dp_group_info=getattr(self, "_dp_group_info", None),
            )
            self.state.global_step += 1
            nt = platform.full((1,), self.pp_schedule.last_local_tokens).to(self.device)
            platform.all_reduce(nt, self._pp_group_info)
            n_valid = max(int(nt.item()), 1)
            self._last_global_tokens = n_valid
        elif getattr(self.pp_schedule, "requires_all_rank_input", False):
            outputs, n_valid = self._pp_run_all_rank_input_step(data_iterator)
        else:
            batch, targets, stop = self._pp_load_first_stage_batch(data_iterator)
            targets, attention_mask, has_attn = self._pp_prepare_broadcast_inputs(batch, targets, stop)
            self.state.global_step += 1
            self._pp_validate_rank_average_targets(targets)
            n_valid = self._pp_count_valid_tokens(targets)
            outputs = self._pp_run_schedule(batch, targets, attention_mask, has_attn)
        self._pp_post_schedule_grad_reduce()
        self._pp_average_plain_dp_grads()
        self._pp_normalize_grads(n_valid)
        grad_norm_value = self._optimizer_step_after_backward(self._pp_clip_grad_norm)

        MPipe Transpose runs the transposed preprocess on every rank, so the
        heavy (vision) inputs must be present everywhere.
        """
        micro_batches = next(data_iterator)
        batch = self._pp_concat_micro_batches(micro_batches)
        # End-of-epoch partial grad-accum group: end the epoch.
        if batch["input_ids"].shape[0] % self.pp_micro_batch_num != 0:
            raise StopIteration
        self.state.global_step += 1
        batch = {
            key: (value.to(self.device, non_blocking=True) if hasattr(value, "to") else value)
            for key, value in batch.items()
        }
        labels = batch["labels"]
        targets = torch.nn.functional.pad(labels, (0, 1), value=-100)[..., 1:].contiguous()
        self._pp_validate_rank_average_targets(targets)
        n_valid = self._pp_count_valid_tokens(targets)
        run_kwargs = {"targets": targets}
        for key in self.pp_schedule.kwargs_batch_dim:
            if key != "targets" and key in batch:
                run_kwargs[key] = batch[key]
        outputs = self.pp_schedule.run(batch["input_ids"], **run_kwargs)
        return outputs, n_valid

    def train(self):
        """Main training loop: epoch → step → micro-batch.

        Groups ``self._grad_accum`` consecutive batches into a list for
        gradient accumulation. The underlying ``StatefulDataLoader`` tracks
        iteration position, so checkpoint/resume skips consumed batches.
        """
        if getattr(self.args.data, "load", False) == "single":
            yield from self.train_dataloader
        else:
            batch_buffer = []
            for batch in self.train_dataloader:
                batch_buffer.append(batch)
                if len(batch_buffer) >= self._grad_accum:
                    yield batch_buffer
                    batch_buffer = []
            if batch_buffer:
                yield batch_buffer

    def _get_layers(self) -> list:
        """Return the repeating layers for FSDP/AC wrapping.