Diff Coverage

            raise ValueError(f"requires_grad_sync must be bool but got {requires_grad_sync}.")
        if not hasattr(self, "hsdp_scheduler"):
            raise ValueError("call hsdp interface first.")

        for _, module in self.named_modules():
            if isinstance(module, HSDPModule):
                module.hsdp_scheduler.set_requires_grad_sync(requires_grad_sync)

    def zero_grad(self):

        # Pass an IncompatibleKeys with the same attribute names as PyTorch
        # so external hooks can safely read .missing_keys/.unexpected_keys.
        _IK = namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"])
        incompatible_keys = _IK([], [])
        for _, module in self_module.named_modules():
            hooks = module._load_state_dict_post_hooks  # pylint: disable=protected-access
            for hook in hooks.values():
                hook(module, incompatible_keys)

                "Currently impl is equal to recurse=True, "
                "need support module_param mapping."
            )
        self_module = cast(nn.Module, self)
        for _, module in self_module.named_modules():
            if isinstance(module, HSDPModule):
                module.hsdp_scheduler.set_requires_all_reduce(requires_all_reduce)

    def set_reshard_after_forward(self, reshard_after_forward: bool, recurse: bool = True) -> None:

                "Currently impl is equal to recurse=True, "
                "need support module_param mapping."
            )
        self_module = cast(nn.Module, self)
        for _, module in self_module.named_modules():
            if isinstance(module, HSDPModule):
                module.hsdp_scheduler.set_reshard_after_forward(reshard_after_forward)

    def set_reshard_after_backward(self, reshard_after_backward: bool, recurse: bool = True) -> None:

                "Currently impl is equal to recurse=True, "
                "need support module_param mapping."
            )
        self_module = cast(nn.Module, self)
        for _, module in self_module.named_modules():
            if isinstance(module, HSDPModule):
                module.hsdp_scheduler.set_reshard_after_backward(reshard_after_backward)

    def set_reduce_op_type(self, reduce_op_type, recurse: bool = True) -> None:

        support reduce_op_type "avg" and "sum", default is "avg"
        """
        self_module = cast(nn.Module, self)
        if recurse:
            sub_modules = [m for _, m in self_module.named_modules()]
        else:
            sub_modules = [self_module]
        for module in sub_modules:
            if isinstance(module, HSDPModule):

        if self.scheduler_ctx.root_module is None:
            tree_ctx = self.scheduler_ctx
            tree_ctx.root_module = self.cell
            registered_schedulers = set()
            for module_name, module in tree_ctx.root_module.named_modules():
                from hyper_parallel.core.fully_shard.api import HSDPModule  # pylint: disable=C0415
                if isinstance(module, HSDPModule):
                    submod_scheduler = module.hsdp_scheduler
                    if submod_scheduler is None or id(submod_scheduler) in registered_schedulers:

        if dim_map == "None":
            return 0

        if isinstance(dim_map, str):
            rank = dist.get_rank()
            rank_list = layout.mesh.get_rank_list_along_axis(dim_map)
            if rank in rank_list:
                return rank_list.index(rank)
            return 0

                    f"Seq dim is sharded by multiple axes {non_none_axes}. "
                    f"Using the last axis for split_id calculation."
                )
            axis_name = non_none_axes[-1]
            rank = dist.get_rank()
            rank_list = layout.mesh.get_rank_list_along_axis(axis_name)
            if rank in rank_list:
                return rank_list.index(rank)

                    f"Seq dim is sharded by multiple axes {non_none_axes}. "
                    f"Using the last axis for split_id calculation."
                )
            axis_name = non_none_axes[-1]
            rank = dist.get_rank()
            rank_list = layout.mesh.get_rank_list_along_axis(axis_name)
            if rank in rank_list:
                return rank_list.index(rank)

def get_op_name(func):
    """Return the registry name for a Torch callable or operator overload."""
    if hasattr(func, "__name__"):
        return func.__name__
    if isinstance(func, OpOverload):
        return func.name.split("::")[-1].split(".")[0]
    if isinstance(func, OpOverloadPacket):
        return func.name.split("::")[-1]
    func_str = str(func)
    if "built-in function" in func_str:
        return func_str.split()[-1].strip(">")
    if "function" in func_str:
        return func_str.split()[1]
    return "unknown_op"


def get_cell_construct(cell):
    """Return the Torch module forward callable."""

    Returns:
        int: The rank of the current process inside ``group``.
    """
    if not torch.distributed.is_available() or not torch.distributed.is_initialized():
        return 0
    if group is None:
        return torch.distributed.get_rank()
    if hasattr(group, "rank"):
        return group.rank()
    return torch.distributed.get_group_rank(group, torch.distributed.get_rank())


def get_device_handle(device_type: str = "npu"):
    """Return the ``torch`` device module for ``device_type``.

    Raises:
        RuntimeError: If torch exposes no module for ``device_type``.
    """
    try:
        handle = getattr(torch, device_type)
    except AttributeError as e:
        raise RuntimeError(f"Failed to resolve device handle: 'torch.{device_type}'.") from e
    return handle


def get_world_size() -> int:
    """Return the Torch distributed world size.

    Kept as a thin wrapper so single-process call sites and white-box unit
    tests share one patchable seam; ``torch.distributed.get_world_size`` itself
    raises before ``init_process_group``.
    """
    return dist.get_world_size()


# Process-group cache keyed by the string form of the sorted rank tuple, e.g.
# ``"(0, 1, 2, 3)"``. Lives here rather than in a module-specific ``utils`` so

    Under ``torch.compile`` the contiguity of a graph input is not knowable at
    trace time, so the copy is emitted unconditionally.
    """
    if torch.compiler.is_compiling():
        return x.contiguous()
    if not x.is_contiguous() or x.storage_offset() != 0:
        return x.contiguous()
    return x

# ---------------------------------------------------------------------------

def get_created_group(rank_list: Union[list[int], tuple[int, ...]]):
    """Return an existing process group by rank list, or ``None``."""
    group_key = str(tuple(sorted(rank_list)))
    if group_key in EXISTING_COMM_GROUPS:
        return EXISTING_COMM_GROUPS[group_key]
    return None


def split_group(parent_pg=None,
                split_ranks: Optional[list] = None,

    """Create split groups for every rank list in *split_ranks*.

    Returns the split process group relative to the current rank id.
    """
    del parent_pg, timeout, group_desc
    if split_ranks is None or len(split_ranks) == 0:
        raise ValueError("split_ranks cannot be None or empty")

    split_group_pg = None
    for split_rank in split_ranks:
        dist_group = get_created_group(split_rank)
        if dist_group is None:
            dist_group = dist.new_group(ranks=split_rank, pg_options=pg_options)
            EXISTING_COMM_GROUPS[str(tuple(sorted(split_rank)))] = dist_group
        if dist.get_rank() in split_rank:
            split_group_pg = dist_group

    return split_group_pg


def init_process_group(*args, **kwargs) -> None:
    """Initialize the default torch distributed process group."""
    if not dist.is_initialized():
        dist.init_process_group(*args, **kwargs)


# ---------------------------------------------------------------------------
# Differentiable collectives

    _OP_MAP['mean'] = dist.ReduceOp.AVG
else:
    # Fallback for older torch versions if necessary, though this might require manual division upstream
    # Assuming standard behavior where 'mean' implies native AVG support or upstream handling
    _OP_MAP['mean'] = dist.ReduceOp.SUM


def resolve_reduce_op(op: Union[str, Any]) -> Any:
    """Resolve a string op name (or pass through an already-resolved ``ReduceOp``)."""

    @staticmethod
    def forward(ctx, tensor: Tensor, peer_rank: int, group) -> Tensor:  # pylint: disable=arguments-differ
        """Perform symmetric bidirectional P2P exchange with ``peer_rank``."""
        ctx.peer_rank = peer_rank
        ctx.group = group
        send_buf = tensor.contiguous()
        recv_buf = torch.empty_like(send_buf)
        reqs = dist.batch_isend_irecv([
            dist.P2POp(dist.isend, send_buf, peer_rank, group),
            dist.P2POp(dist.irecv, recv_buf, peer_rank, group),
        ])
        for req in reqs:
            req.wait()
        return recv_buf

    @staticmethod
    def backward(ctx, grad_output: Tensor):
        """Perform symmetric P2P exchange for the backward gradient pass."""
        send_buf = grad_output.contiguous()
        recv_buf = torch.empty_like(send_buf)
        reqs = dist.batch_isend_irecv([
            dist.P2POp(dist.isend, send_buf, ctx.peer_rank, ctx.group),
            dist.P2POp(dist.irecv, recv_buf, ctx.peer_rank, ctx.group),
        ])
        for req in reqs:
            req.wait()
        return recv_buf, None, None


class _TorchDifferentiableVariableAllGather(torch.autograd.Function):
    """Variable dim-zero all-gather with an uneven reduce-scatter backward."""

    @staticmethod
    def forward(ctx, input_tensor, output_splits, group):  # pylint: disable=arguments-differ
        """Gather each rank's true row count without replicating inputs for A2A."""
        if input_tensor.ndim == 0:
            raise ValueError("variable all-gather input must have at least one dimension")
        splits = tuple(output_splits)
        if not splits:
            raise ValueError("output_splits must contain at least one group rank")
        if any(not isinstance(rows, int) or isinstance(rows, bool) or rows < 0 for rows in splits):
            raise ValueError(f"output_splits must contain non-negative integers, got {splits!r}")

        group_rank = dist.get_rank(group=group)
        if group_rank < 0 or group_rank >= len(splits):
            raise ValueError(f"group rank must be in [0, {len(splits)}), got {group_rank}")
        if input_tensor.shape[0] != splits[group_rank]:
            raise ValueError(
                "variable all-gather local rows must match output_splits at the group rank, "
                f"got local_rows={input_tensor.shape[0]}, group_rank={group_rank}, "
                f"output_splits={splits!r}"
            )

                f"got local_rows={input_tensor.shape[0]}, group_rank={group_rank}, "
                f"output_splits={splits!r}"
            )

        input_tensor = input_tensor.contiguous()
        feature_shape = tuple(input_tensor.shape[1:])
        if input_tensor.device.type == "npu":
            gathered = [input_tensor.new_empty((rows, *feature_shape)) for rows in splits]
            dist.all_gather(gathered, input_tensor, group=group)
        else:
            max_rows = max(splits)
            if max_rows == 0:
                gathered = [input_tensor.new_empty((0, *feature_shape)) for _ in splits]
            else:
                padded = input_tensor.new_zeros((max_rows, *feature_shape))
                if input_tensor.shape[0] > 0:
                    padded[:input_tensor.shape[0]].copy_(input_tensor)
                padded_outputs = [torch.empty_like(padded) for _ in splits]
                dist.all_gather(padded_outputs, padded, group=group)
                gathered = [
                    output[:rows].contiguous()
                    for output, rows in zip(padded_outputs, splits)
                ]

        ctx.output_splits = splits
        ctx.group = group
        ctx.group_rank = group_rank
        return torch.cat(gathered, dim=0)

    @staticmethod
    def backward(ctx, grad_output):
        """Sum replicated output gradients and return this rank's uneven shard."""
        output_rows = ctx.output_splits[ctx.group_rank]
        output = grad_output.new_empty((output_rows, *grad_output.shape[1:]))
        if sum(ctx.output_splits) == 0:
            return output, None, None

        grad_output = grad_output.contiguous()
        if grad_output.device.type == "npu":
            from torch_npu.distributed import reduce_scatter_tensor_uneven  # pylint: disable=C0415
            reduce_scatter_tensor_uneven(
                output,
                grad_output,
                input_split_sizes=list(ctx.output_splits),
                op=dist.ReduceOp.SUM,

                op=dist.ReduceOp.SUM,
                group=ctx.group,
            )
        else:
            reduced = grad_output.clone()
            dist.all_reduce(reduced, op=dist.ReduceOp.SUM, group=ctx.group)
            start = sum(ctx.output_splits[:ctx.group_rank])
            output.copy_(reduced.narrow(0, start, output_rows))
        return output, None, None


def differentiable_all_gather_concat(data: Tensor, group, concat_size: int, concat_dim: int,
                                     rank_list=None) -> Tensor:

        _TorchContiguousGrad.apply(tensor)
        for tensor in dist_func.all_gather(data, group=group)
    ]
    if rank_list is not None:
        group_ranks = dist.get_process_group_ranks(group)
        if tuple(rank_list) != tuple(group_ranks):
            rank_to_idx = {int(rank): idx for idx, rank in enumerate(group_ranks)}
            output = [output[rank_to_idx[int(rank)]] for rank in rank_list]
    return torch.cat(output, dim=concat_dim)


def differentiable_all_to_all_single(input_tensor: Tensor, input_splits: Sequence[int],

def differentiable_all_to_all(input_data: Tensor, output_shape: Sequence[int], group) -> Tensor:
    """Autograd-aware all-to-all producing a tensor of ``output_shape``."""
    input_data = ensure_contiguous(input_data)
    output_tensor = torch.empty(output_shape, device=input_data.device, dtype=input_data.dtype)
    return dist_func.all_to_all_single(output_tensor, input_data, group=group)


def differentiable_all_reduce(data: Tensor, op: Union[str, Any], group) -> Tensor:
    """Autograd-aware all-reduce with string or ``ReduceOp`` *op*."""

    )

    # 'avg' maps to SUM in _OP_MAP, so the division stays manual.
    if op == 'avg':
        output_tensor = output_tensor / dev_num
    return output_tensor


def differentiable_variable_all_gather(

def differentiable_variable_all_gather(
        input_tensor: Tensor, output_splits: Sequence[int], group) -> Tensor:
    """Gather variable dim-zero shards on HCCL or Gloo with autograd support."""
    return _TorchDifferentiableVariableAllGather.apply(
        input_tensor, tuple(output_splits), group
    )

def p2p_exchange(tensor: Tensor, peer_rank: int, group=None) -> Tensor:
    """Symmetric bidirectional P2P exchange with ``peer_rank``."""
    if peer_rank == dist.get_rank(group):
        return tensor
    return _TorchP2PExchangeFunction.apply(tensor, peer_rank, group)


def exchange_splits_via_all_to_all(input_tensor: Tensor, group) -> Tensor:
    """All-to-all a per-rank split vector, returning the received counts.

    def execute(self, tensor: Any) -> Any:
        """Execute the differentiable collective selected during lowering."""
        if self.kind == "all_gather":
            return communication.differentiable_all_gather_concat(
                tensor,
                self.group,
                self.group_size,
                self.tensor_dim,

                self.group_size,
                self.tensor_dim,
            )
        if self.kind == "all_reduce":
            return communication.differentiable_all_reduce(
                tensor,
                self.reduce_op,
                self.group,
            )

                self.reduce_op,
                self.group,
            )
        if self.kind == "reduce_scatter":
            return communication.differentiable_reduce_scatter(
                tensor,
                self.group_size,
                self.tensor_dim,
                self.reduce_op,

                self.reduce_op,
                self.group,
            )
        if self.kind == "all_reduce_shard":
            reduced = communication.differentiable_all_reduce(
                tensor,
                self.reduce_op,
                self.group,
            )

            f"multiple of 2 * cp_size ({2 * cp_mesh.size()})"
        )

    peer_rank = _head_tail_peer_rank(cp_mesh)
    query_peer = communication.p2p_exchange(
        query.narrow(2, local_q_len // 2, local_q_len // 2), peer_rank)
    global_key, global_value = flex_cp_allgather(key, value, 2, cp_mesh)
    keep_output = _run_head_tail_half(
        attention_fn, query.narrow(2, 0, local_q_len // 2),

        + list(range(scatter_dim))
        + list(range(scatter_dim + 1, split_ndim))
    )
    send = tensor.contiguous().reshape(split_shape).permute(permutation).contiguous()
    received = communication.differentiable_all_to_all(
        send, list(send.shape), cp_mesh.get_group())
    return _reconstruct_all_to_all(received, gather_dim)

    output_splits = [0] * cp_size
    if cp_rank > 0:
        output_splits[rank_to_group_index[rank_list[cp_rank - 1]]] = halo_width

    exchange_output = communication.differentiable_all_to_all_single(
        exchange_input,
        input_splits,
        output_splits,
        group=cp_group,

    a2a_input = a2a_input.contiguous()
    split_len = a2a_input.shape[0] // split_count
    input_splits = [split_len] * split_count
    output_splits = [split_len] * split_count
    output = communication.differentiable_all_to_all_single(
        a2a_input,
        input_splits,
        output_splits,
        group=device_mesh.get_group(),

    output_splits = [0] * cp_size
    if cp_rank > 0:
        output_splits[rank_to_group_index[rank_list[cp_rank - 1]]] = halo_width

    exchange_output = communication.differentiable_all_to_all_single(
        exchange_input,
        input_splits,
        output_splits,
        group=cp_group,

    all_to_all_input = all_to_all_input.flatten(0, 1)

    split_length = all_to_all_input.shape[0] // split_count
    splits = [split_length] * split_count
    output = communication.differentiable_all_to_all_single(
        all_to_all_input,
        splits,
        splits,
        group=device_mesh.get_group(),