Diff Coverage

                    f"Current rank {current_rank} not found in layout's rank_list {dtensor_layout.rank_list}")

            inner_rank_id = dtensor_layout.rank_list.index(current_rank)
            # Calculate slice area using infer_slice_area_by_rank
            slice_area = infer_slice_area_by_layout(
                dtensor_layout,
                inner_rank_id,
                global_shape,
            )

        layout = _build_layout(device_mesh, placements, tensor_dim)
        is_ragged = _layout_has_ragged_shard(layout)
        if is_ragged:
            if stride is not None and shape is None:
                raise ValueError("stride requires an explicit shape")
            if hasattr(local_tensor, "is_contiguous") and not local_tensor.is_contiguous():
                raise ValueError("RaggedShard local tensor must be contiguous")
            local_tensor = local_tensor.view(-1)
        elif stride is not None and shape is None:
            raise ValueError("stride requires an explicit shape")
        if run_check:
            # pylint: disable=C0415
            from hyper_parallel.core.dtensor._from_local_utils import run_from_local_checks
            run_from_local_checks(

                "device_mesh": self._device_mesh,
                "placements": self._alias_placements(),
            }
            if hasattr(self, "_global_shape"):
                constructor_kwargs["shape"] = self._global_shape
            return cls(local_tensor, **constructor_kwargs)
        layout = cp.deepcopy(self._layout)
        if layout.tensor_shape is not None:
            layout.set_tensor_meta(
                layout.tensor_shape,
                layout.tensor_stride,
                local_tensor.dtype,
            )

def _get_slice_tensor_by_layout(global_tensor, layout):
    """Transfer global tensor to local tensor by layout"""
    inner_rank_id = layout.rank_list.index(layout.mesh.rank)
    slice_area = infer_slice_area_by_layout(layout, inner_rank_id, global_tensor.shape)

    def get_slice_data(full_data, offset):
        area = ()
        for begin, end in offset:

def _infer_slice_shape_by_layout(global_shape, layout):
    """Infer slice shape from global_shape and layout"""
    inner_rank_id = layout.rank_list.index(layout.mesh.rank)
    slice_area = infer_slice_area_by_layout(layout, inner_rank_id, global_shape)
    return [end - start for start, end in slice_area]


class Layout:
    """

    def __setstate__(self, state) -> None:
        """Restore older serialized Shard values without an uneven marker."""
        self.__dict__.update(state)
        if "_uneven_shard" not in state:
            self._uneven_shard = False


class StridedShard(Shard):
    """

        if assign and state_dict and all(
            isinstance(val, DTensor) for val in state_dict.values()
        ):
            return super().load_state_dict(state_dict, strict=strict, assign=True)
        self_module = cast(ModuleClass, self)

        target_map: dict[str, TensorClass] = {}
        for name, p in platform.parameters_dict(self_module):
            target_map[name] = p
        for name, b in self_module.named_buffers():
            target_map[name] = b

        Raises:
            ValueError: If ``recursive`` is not a bool.
        """
        if not isinstance(recursive, bool):
            raise ValueError(f"recursive should be a bool, got {type(recursive)}")
        self_module = cast(ModuleClass, self)
        modules = (
            [module for _, module in platform.get_cells_and_names(self_module)]
            if recursive

    """Validate the parameter-identity metadata consumed by one fully_shard unit."""
    if tp_grad_infos is None:
        return
    if not isinstance(tp_grad_infos, Mapping):
        raise ValueError("tp_grad_infos must be a mapping from Parameter to TPShardMetaInfo")

    managed_parameters = set(managed_parameters)

    if any(isinstance(parameter, DTensor) for parameter in managed_parameters):

    platform_type = platform.platform_type
    _validate_module_for_fully_shard(module, platform_type)

    if tp_grad_infos is not None and platform_type != PlatformType.PYTORCH:
        raise NotImplementedError("tp_grad_infos is currently supported only on the Torch backend")

    if platform_type == PlatformType.MINDSPORE:
        from hyper_parallel.platform.mindspore.autograd_compat import enable_mindspore_backward_compat

        self.mesh: DeviceMesh = mesh
        self.shard_placement_fn = shard_placement_fn
        self.mp_policy = mp_policy
        self.offload_policy = offload_policy
        self.comm_fusion_policy = CommFusionPolicy(comm_fusion, comm_fusion_zero_copy)
        self.ignored_params = ignored_params
        self.replicate_params = replicate_params
        self.device = device
        self.reshard_after_forward = reshard_after_forward
        self.tp_grad_infos = tp_grad_infos
        self.scheduler_state = None
        self.forward_prefetch_cells = []
        self.backward_prefetch_cells = []
        self._backup_forward_fetch = None

        self.forward_prefetch_cells = []
        self.backward_prefetch_cells = []
        self._backup_forward_fetch = None
        # Flag to identify root module.
        self._is_root = True
        # module and its all sub-modules share one same 'HSDPSchedulerContext'
        self.scheduler_ctx = HSDPSchedulerContext()
        # When ``fully_shard`` is given multiple root modules, forward pre/post hooks coordinate
        # so unshard / PostBackward / reshard run once per forward (aligned with PyTorch FSDP2).

                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:
                        continue
                    registered_schedulers.add(id(submod_scheduler))
                    if submod_scheduler.scheduler_ctx is not tree_ctx:
                        if submod_scheduler.scheduler_ctx.root_module is not None:
                            raise ValueError(

    def _init_params_fqn(self):  # pylint: disable=W0212
        if not self._is_root or self.scheduler_ctx.root_module is None:
            return
        if self.scheduler_ctx._param_fqn_initialized:  # pylint: disable=protected-access
            return
        # Build a map from original (sharded) parameter tensor to its HSDPParam wrapper.
        param_to_hsdp_param = {}
        for submod_scheduler in self.scheduler_ctx.all_hsdp_schedulers:
            hsdp_state = submod_scheduler.hsdp_state

        """
        if self.param_group is not None:
            self.param_group.gradient_scaling_factor = factor
        else:
            for hsdp_param in self.hsdp_params:
                hsdp_param.gradient_scaling_factor = factor

    def set_requires_all_reduce(self, requires_all_reduce: bool) -> None:
        """Propagate the HSDP all-reduce switch to the active communication path."""

            root_schedulers = fsdp_schedulers[:1]
        for scheduler in root_schedulers:
            # No public API exposes root backward finalization. The hook drains
            # unconditionally, so a differentiable PP input cannot defer it.
            scheduler._root_backward_hook()  # pylint: disable=protected-access

    def _build_padded_sens(self, micro_index):
        """Build an N-length sens list aligned with the forward output structure.

        py_output = op_call(*local_args, **local_kwargs)
        op_name = platform.get_op_name(op_call)
        if op_name in _RAGGED_INPLACE_ELEMENTWISE_OPS:
            if not args or not isinstance(args[0], DTensor):
                raise ValueError(
                    f"Ragged in-place operator {op_name!r} requires a DTensor first argument"
                )
            return args[0]
        return DTensor.from_local_with_layout(

        self._module_info: ParamModuleInfo = module_info
        self.mesh_info = mesh_info
        self.mp_policy = mp_policy
        self.device = device
        self.orig_dtype = None
        self.param_dtype = None
        self.reduce_dtype = None
        self.offload_to_cpu: bool = isinstance(offload_policy, CPUOffloadPolicy)
        self.pin_memory = (
            self.offload_to_cpu and cast(CPUOffloadPolicy, offload_policy).pin_memory
        )

        )
        self._orig_param_hooks: List[Callable] = []
        self.grad_offload_event: Optional[ms.runtime.Event] = None
        dtensor_payload = unwrap_dtensor_param(param)
        if (dtensor_payload is not None) != (
            tp_grad_info is not None and tp_grad_info.origin_is_dtensor
        ):
            raise ValueError(
                "tp_grad_info.origin_is_dtensor must be True exactly for native DTensor parameters, "
                f"got parameter type {type(param).__name__} and tp_grad_info={tp_grad_info}"
            )
        self.tp_grad_info = tp_grad_info
        self._orig_param_is_dtensor = (
            tp_grad_info is not None and tp_grad_info.origin_is_dtensor
        )
        self._orig_dtensor_mesh = tp_grad_info.mesh if self._orig_param_is_dtensor else None
        self._orig_dtensor_placements = (
            tuple(tp_grad_info.placements) if self._orig_param_is_dtensor else None
        )
        self._spmd_shard_mesh_dim = self.mesh_info.shard_mesh_dim
        self._spmd_replicate_mesh_dim = self.mesh_info.replicate_mesh_dim
        self._init_sharded_param(param, shard_placement_fn)
        self._save_backward_hooks(param)
        self.unsharded_param_buffers: List[ms.Tensor] = []
        self.unsharded_accumulated_grad = None
        self._unsharded_param: Optional[Parameter] = None
        self._param_fqn: Optional[str] = None
        # Communication attributes for prefetch pattern
        self.allgather_comm_ctx = AllGatherCommCtx()
        self.reduce_scatter_comm_ctx = ReduceScatterCommCtx()
        self.all_reduce_comm_ctx = AllReduceCommCtx()
        self._accumulated_allreduced_grad = True
        self._reduce_partial_output = None
        self._post_load_hook_handle = (
            module_info.module.register_load_state_dict_post_hook(
                lambda *args, **kwargs: self.reset_sharded_param()
            )

    @property
    def reduce_partial_output(self) -> Optional[ms.Tensor]:
        """Return reduce-scatter results accumulated before the final micro-step."""
        return self._reduce_partial_output

    @reduce_partial_output.setter
    def reduce_partial_output(self, value: Optional[ms.Tensor]) -> None:
        self._reduce_partial_output = value

    def reduce_comm_dtype(self, grad: Optional[ms.Tensor] = None):
        """Resolve the communication dtype owned by this parameter."""
        if self.reduce_dtype is not None:

        if self.reduce_dtype is not None:
            return self.reduce_dtype
        if grad is not None:
            return grad.dtype
        if self.unsharded_accumulated_grad is not None:
            return self.unsharded_accumulated_grad_data.dtype
        if self.unsharded_param is not None and self.unsharded_param.grad is not None:
            return self.unsharded_grad_data.dtype
        return self.orig_dtype

    def _get_base_spmd_placements(self) -> tuple:
        """Return source-layout placements prefixed by explicit data-parallel axes."""
        if self.tp_grad_info is not None:

                [self.mesh_info.mesh, self.tp_grad_info.mesh]
            )
            dp_prefix = tuple(Replicate() for _ in range(self.mesh_info.mesh.ndim))
            return dp_prefix + tuple(self.tp_grad_info.placements)
        self._spmd_mesh = self.mesh_info.mesh
        return tuple(Replicate() for _ in range(self._spmd_mesh.ndim))

    def _apply_data_parallel_placements(
        self, placements: list, shard_placement: Shard
    ) -> tuple:

        self, placements: list, shard_placement: Shard
    ) -> tuple:
        """Apply the parameter-specific DDP/FSDP layout to source placements."""
        if len(placements) != self._spmd_mesh.ndim:
            raise AssertionError(
                f"Expected {self._spmd_mesh.ndim} unified placements, got "
                f"{len(placements)}: {placements}"
            )
        if (

        shard_dim = hsdp_placement.dim
        self._orig_size = param_data.shape
        self._contiguous_orig_stride = make_contiguous_strides_for(self._orig_size)

        if isinstance(self.mesh_info, FSDPMeshInfo):
            self.shard_rank = self.mesh_info.shard_mesh_rank
            self.shard_world_size = self.mesh_info.shard_mesh_size
        else:
            self.shard_rank = 0
            self.shard_world_size = 1
        if isinstance(self.mesh_info, DDPMeshInfo):
            self.replicate_world_size = self.mesh_info.replicate_mesh_size
        else:
            self.replicate_world_size = 1
        self.is_replicate_param = (
            isinstance(self.mesh_info, DDPMeshInfo)
            and not isinstance(self.mesh_info, HSDPMeshInfo)
        )
        self.is_sharded = self.shard_world_size > 1

        if param_data.shape[shard_dim] % self.shard_world_size != 0:
            raise NotImplementedError(
                f"Uneven sharding on dim {shard_dim} not supported: "
                f"shape={param_data.shape}, world_size={self.shard_world_size}"
            )
        chunks = ms.mint.chunk(param_data, self.shard_world_size, dim=shard_dim)
        sharded_param = chunks[self.shard_rank].clone().contiguous()
        self.sharded_size = sharded_param.shape
        self.contiguous_sharded_stride = make_contiguous_strides_for(self.sharded_size)
        self._sharded_param_data = sharded_param.view(-1)

            self._unsharded_param.requires_grad = True

    def _get_unsharded_param_from_all_gather_output(self):
        """Reconstruct the full local parameter view from the packed all-gather output."""
        if len(self.unsharded_param_buffers) != 1:
            raise AssertionError(
                f"Expected 1 unsharded_param_buffer, got {len(self.unsharded_param_buffers)}"
            )
        unsharded_tensor = self.unsharded_param_buffers[0]
        plan = build_rs_plan(
            self,
            self._sharded_local_tensor,
            self.shard_world_size if self.is_sharded else 1,

            unsharded_grad = unsharded_grad.to(self.reduce_dtype)
        if self.unsharded_accumulated_grad is None:
            self.unsharded_accumulated_grad = unsharded_grad
        else:
            self.unsharded_accumulated_grad = ms.mint.add(
                self.unsharded_accumulated_grad,
                unsharded_grad,
            )

            )
            self.unsharded_param.grad = None

    def alloc_unsharded_param_buffers(self) -> None:
        for tensor in self.unsharded_param_buffers:
            expected_size = tensor.numel() * tensor.itemsize

            storage = tensor.untyped_storage()
            if storage.size() != expected_size:

        all_gather_input = self.all_gather_inputs[0]

        # If parameter is not sharded (below threshold), no communication needed
        if not self.is_sharded:
            self.init_unsharded_param_buffers(
                all_gather_input_numels=[all_gather_input.numel()],
                all_gather_input_dtypes=[all_gather_input.dtype],
                world_size=1,
                device=all_gather_input.device.split(':')[0],

                all_gather_input_dtypes=[all_gather_input.dtype],
                world_size=1,
                device=all_gather_input.device.split(':')[0],
            )
            self.alloc_unsharded_param_buffers()
            copy_without_bumping_version(self.unsharded_param_buffers[0], all_gather_input)
            return all_gather_input, self.unsharded_param_buffers[0], None

        # Initialize output buffer
        self.init_unsharded_param_buffers(
            all_gather_input_numels=[all_gather_input.numel()],
            all_gather_input_dtypes=[all_gather_input.dtype],
            world_size=self.shard_world_size,
            device=self._sharded_param_data.device.split(':')[0],

            all_gather_input_dtypes=[all_gather_input.dtype],
            world_size=self.shard_world_size,
            device=self._sharded_param_data.device.split(':')[0],
        )
        self.alloc_unsharded_param_buffers()

        # Get communication group
        shard_group = self.mesh_info.shard_process_group if isinstance(self.mesh_info, FSDPMeshInfo) else None

        shard_group = self.mesh_info.shard_process_group if isinstance(self.mesh_info, FSDPMeshInfo) else None

        if shard_group is None or self.shard_world_size <= 1:
            # No communication needed, just copy
            copy_without_bumping_version(self.unsharded_param_buffers[0], all_gather_input)
            return all_gather_input, self.unsharded_param_buffers[0], None

        # Execute all_gather_into_tensor
        handle = dist.all_gather_into_tensor(
            self.unsharded_param_buffers[0],

            group=shard_group,
            async_op=async_op,
        )

        return all_gather_input, self.unsharded_param_buffers[0], handle

    def unshard(self, async_op: bool = False) -> None:
        if self.allgather_comm_ctx.allgather_output is not None:
            # Already triggered by HSDPState.prefetch(), so return directly.
            return  # no-op

        all_gather_input, output, handle = self._get_unsharded_param_data(async_op=async_op)
        self.allgather_comm_ctx.allgather_input = all_gather_input
        self.allgather_comm_ctx.allgather_output = output
        self.allgather_comm_ctx.allgather_handle = handle

    def wait_for_unshard(self) -> None:
        self._assert_in_states(ShardedState.SHARDED)

        # If parameter is not sharded (below threshold), no reduce-scatter needed
        if not self.is_sharded:
            if output_buffer is not None:
                copy_without_bumping_version(output_buffer, grad_flat)
                self.reduce_scatter_comm_ctx.reduce_scatter_output = output_buffer
            else:
                self.reduce_scatter_comm_ctx.reduce_scatter_output = grad_flat
            self.reduce_scatter_comm_ctx.reduce_scatter_handle = None
            return

        if shard_group is None or shard_group_size <= 1:
            if output_buffer is not None:
                copy_without_bumping_version(output_buffer, grad_flat)
                self.reduce_scatter_comm_ctx.reduce_scatter_output = output_buffer
            else:
                self.reduce_scatter_comm_ctx.reduce_scatter_output = grad_flat
            self.reduce_scatter_comm_ctx.reduce_scatter_handle = None
            return

        # Calculate output size
        output_numel = grad_flat.numel() // shard_group_size
        if output_buffer is not None:

            if output_buffer.dtype != reduce_dtype:
                raise ValueError(
                    f"output_buffer dtype mismatch: expected {reduce_dtype}, got {output_buffer.dtype}"
                )
            self.reduce_scatter_comm_ctx.reduce_scatter_output = output_buffer
        else:
            self.reduce_scatter_comm_ctx.reduce_scatter_output = ms.mint.empty(
                output_numel, dtype=reduce_dtype, device=grad.device.split(":")[0]
            )

        ``all_reduce_comm_ctx``.
        """
        grad = self.reduce_scatter_comm_ctx.reduce_scatter_output
        if grad is None:
            raise RuntimeError("all_reduce_grad requires a completed reduce-scatter output.")
        if self.reduce_dtype is not None and self.reduce_dtype != grad.dtype:
            grad = grad.to(self.reduce_dtype)
        reduce_group = (
            self.mesh_info.replicate_process_group
            if isinstance(self.mesh_info, DDPMeshInfo)
            else None

            if isinstance(self.mesh_info, DDPMeshInfo)
            else None
        )
        if reduce_group is None or self.replicate_world_size <= 1:
            self.all_reduce_comm_ctx.all_reduce_output = grad
            self.all_reduce_comm_ctx.all_reduce_handle = None
            return

        # Ascend HCCL DistCommAllReduce rejects non-contiguous tensors.
        # ``grad`` here may be a view returned by ``_to_local_unsharded_grad``
        # (DTensor.to_local() / redistribute().to_local()) or by autograd.

    def clear_unsharded_source_grad(self) -> None:
        """Release the unsharded gradient after its communication input is safe."""
        if self.unsharded_accumulated_grad is not None:
            self.unsharded_accumulated_grad = None
        if self.unsharded_param is not None and self.unsharded_param.grad is not None:
            self.unsharded_param.grad = None

    def apply_reduced_grad(self, reduced_grad):

            else:
                self.sharded_param.grad = self.to_sharded_dtensor(reduced_grad)
        else:
            if self.mp_policy.apply_grad_on_fp32_main_grad:
                accumulated_grad = ms.mint.add(
                    self.sharded_param.main_grad._local_tensor,
                    reduced_grad,
                )
                self.sharded_param.main_grad = self.to_sharded_dtensor(accumulated_grad)
                self.sharded_param.grad = None
            else:
                accumulated_grad = ms.mint.add(
                    self.sharded_param.grad._local_tensor,

        reduce_op: str,
    ) -> None:
        """All-reduce a final gradient over replicated source-layout axes."""
        if self.tp_grad_info is None or not self.tp_grad_info.placements:
            return
        source_mesh = self.tp_grad_info.mesh
        replicate_mesh_dims = tuple(
            mesh_dim
            for mesh_dim, placement in enumerate(self.tp_grad_info.placements)

            for mesh_dim, placement in enumerate(self.tp_grad_info.placements)
            if placement.is_replicate()
        )
        if not replicate_mesh_dims:
            return
        if source_mesh.mesh_dim_names is None:
            raise ValueError(
                "TP shard mesh must define mesh_dim_names to all-reduce replicated gradients."
            )
        replicate_dim_names = tuple(
            source_mesh.mesh_dim_names[mesh_dim]

            for mesh_dim in replicate_mesh_dims
        )
        replicate_mesh = source_mesh[replicate_dim_names].flatten()
        if replicate_mesh.size() <= 1:
            return
        dist.all_reduce(
            reduced_grad,
            op=reduce_op,
            group=replicate_mesh.get_group(),

    @classmethod
    def get_metadata(cls, hsdp_params, fn):
        """Retrieve or compute metadata keyed by parameter identity and version."""
        param_key = tuple((id(param), getattr(param, "version", 0)) for param in hsdp_params)
        key = hash(param_key)
        if key not in cls._cache:
            cls._cache[key] = fn(hsdp_params)
        return cls._cache[key]


@dataclass
class AllGatherBucket:

    @_no_grad()
    def copy_out(self) -> None:
        """Wait the all-gather and copy each result into stable parameter buffers."""
        result = self.all_gather_result
        if result is None or result.all_gather_output is None:
            return
        if result.handle is not None:
            result.handle.wait()
        all_gather_output = result.all_gather_output
        output_buffers = []
        for input_numels, input_dtypes, hsdp_param in zip(
            self.metadata.param_input_numels,
            self.metadata.param_input_dtypes,
            self.hsdp_params,
        ):
            hsdp_param.init_unsharded_param_buffers(
                input_numels,
                input_dtypes,
                self.shard_world_size,
                _normalize_device(all_gather_output.device),

                input_dtypes,
                self.shard_world_size,
                _normalize_device(all_gather_output.device),
            )
            hsdp_param.alloc_unsharded_param_buffers()
            output_buffers.extend(hsdp_param.unsharded_param_buffers)
        split_with_sizes_copy(
            all_gather_output.view(self.shard_world_size, -1),
            self.metadata.inp_split_sizes,
            dim=1,
            out=[tensor.view(self.shard_world_size, -1) for tensor in output_buffers],

            self.metadata.inp_split_sizes,
            dim=1,
            out=[tensor.view(self.shard_world_size, -1) for tensor in output_buffers],
        )
        self.all_gather_result = None


@dataclass
class GradientBucketLayout:

    @property
    def param_offsets(self) -> list[int]:
        """Return parameter offsets in the reduce-scatter output."""
        return self.layout.param_offsets

    @property
    def bucket_key(self) -> tuple:
        """Return the identity of this fusion class across micro-steps."""

    def move_reduce_scatter_output(self) -> ms.Tensor:
        """Transfer exclusive ownership of the completed output."""
        if self.reduce_scatter_output is None:
            raise RuntimeError("Reduce-scatter output has already been released.")
        output = self.reduce_scatter_output
        self.reduce_scatter_output = None
        return output

        param_input_numels.append(input_numels)
        inp_split_sizes.extend(input_numels)
        total_input_numel += sum(input_numels)
    if dtype is None:
        raise ValueError("Cannot build all-gather metadata for an empty parameter bucket.")
    return AllGatherMetadata(
        param_input_dtypes,
        param_input_numels,
        dtype,

    """Copy dim-1 slices from a fused all-gather into stable buffers."""
    if dim != 1:
        raise NotImplementedError("split_with_sizes_copy currently only supports dim=1")
    offset = 0
    for destination, size in zip(out, split_sizes):
        copy_without_bumping_version(
            destination,
            all_gather_output.narrow(dim, offset, size),
        )
        offset += size

        for hsdp_param in self.hsdp_params:
            if hsdp_param.shard_world_size <= 1:
                continue
            if not isinstance(hsdp_param.mesh_info, FSDPMeshInfo):
                raise ValueError(
                    f"Fused all-gather expects FSDPMeshInfo, got {type(hsdp_param.mesh_info)}"
                )
            shard_group = hsdp_param.mesh_info.shard_process_group
            communication_dtype = hsdp_param.param_dtype or hsdp_param.orig_dtype

            )

    def unshard(self, async_op: bool = False) -> None:
        """Launch fused all-gathers for every compatible bucket."""
        if self.all_gather_buckets and any(
            bucket.all_gather_result is not None for bucket in self.all_gather_buckets
        ):
            return
        self.foreach_all_gather(async_op)

    @_no_grad()
    def foreach_all_gather(self, async_op: bool = False) -> None:
        """Initialize ordered buckets and launch their all-gathers."""

        if not self.all_gather_buckets:
            self._init_all_gather_buckets()
        for hsdp_param in self.hsdp_params:
            if hsdp_param.shard_world_size <= 1:
                hsdp_param.unshard(async_op)
        for bucket in self.all_gather_buckets:
            if bucket.all_gather_result is not None:
                continue
            metadata = bucket.metadata

            )

    def wait_for_unshard(self) -> None:
        """Wait all all-gathers and install stable unsharded parameters."""
        for bucket in self.all_gather_buckets:
            bucket.copy_out()
        for hsdp_param in self.hsdp_params:
            hsdp_param.wait_for_unshard()

    @staticmethod
    def _build_gradient_bucket_layout(hsdp_params) -> GradientBucketLayout:
        """Build one compact output layout for the supplied parameter order."""

        grads_by_bucket = {}
        groups_by_bucket = {}
        for hsdp_param in self.hsdp_params:
            if not hsdp_param.sharded_param.requires_grad:
                continue
            if hsdp_param.unsharded_accumulated_grad is not None:
                unsharded_grad = hsdp_param.unsharded_accumulated_grad_data
            elif hsdp_param.unsharded_param.grad is not None:
                unsharded_grad = hsdp_param.unsharded_grad_data
            else:
                continue
            shard_group = (
                hsdp_param.mesh_info.shard_process_group
                if isinstance(hsdp_param.mesh_info, FSDPMeshInfo)
                else None

        buckets = []
        for bucket_key, hsdp_params in params_by_bucket.items():
            shard_world_size = hsdp_params[0].shard_world_size
            if any(param.shard_world_size != shard_world_size for param in hsdp_params):
                raise ValueError("A reduce-scatter bucket must use one shard world size.")
            needs_avg_div = reduce_op == "avg"
            buckets.append(
                ReduceScatterBucket(
                    layout=self._build_gradient_bucket_layout(hsdp_params),

                if any(
                    isinstance(hsdp_param.mesh_info, DDPMeshInfo)
                    for hsdp_param in rs_bucket.hsdp_params[1:]
                ):
                    raise ValueError(
                        "A reduce-scatter bucket cannot mix parameters with and without "
                        "a subsequent all-reduce."
                    )
                continue

                continue
            replicate_group = representative.mesh_info.replicate_process_group
            replicate_world_size = representative.replicate_world_size
            for hsdp_param in rs_bucket.hsdp_params[1:]:
                if (
                    not isinstance(hsdp_param.mesh_info, DDPMeshInfo)
                    or hsdp_param.mesh_info.replicate_process_group != replicate_group
                    or hsdp_param.replicate_world_size != replicate_world_size
                ):
                    raise ValueError(
                        "All parameters in a reduce-scatter bucket must share one "
                        "subsequent all-reduce group."
                    )
            buckets.append(

        self.reduce_scatter_buckets = self._build_reduce_scatter_buckets(
            reduce_scatter_reduce_op,
        )
        if not self.reduce_scatter_buckets:
            return
        self.all_reduce_buckets = self._build_all_reduce_buckets(self.reduce_scatter_buckets)
        self._issue_reduce_scatter_buckets(async_op)
        if async_op:
            self.comm_ctx.pre_param_group = self

        self._issue_reduce_scatter_buckets(async_op)
        if async_op:
            self.comm_ctx.pre_param_group = self
        else:
            self.wait_reduce_scatter_and_issue_all_reduce(async_op=False)
            self.wait_all_reduce_and_save_grad()

    def _wait_reduce_scatter_buckets(self) -> None:
        """Wait RS buckets and finish shard-dimension averaging."""
        for bucket in self.reduce_scatter_buckets:

                bucket.handle.wait()
                bucket.handle = None
            bucket.reduce_scatter_input = None
            if bucket.reduce_scatter_output is None:
                raise RuntimeError("Reduce-scatter bucket has not been prepared.")
            if bucket.needs_avg_div and bucket.shard_world_size > 1:
                bucket.reduce_scatter_output = ms.mint.div(
                    bucket.reduce_scatter_output,
                    bucket.shard_world_size,
                )

    ) -> None:
        """Launch SUM all-reduces for completed RS outputs."""
        for bucket in all_reduce_buckets:
            if bucket.all_reduce_output is None:
                raise RuntimeError("All-reduce bucket has not received its reduce-scatter output.")
            if not bucket.uses_collective:
                continue
            bucket.handle = dist.all_reduce(
                bucket.all_reduce_output,
                group=bucket.replicate_group,
                op="sum",

        for bucket in self.reduce_scatter_buckets:
            partial_output = self.reduce_partial_outputs.pop(bucket.bucket_key, None)
            if partial_output is not None:
                bucket.reduce_scatter_output = ms.mint.add(
                    bucket.reduce_scatter_output,
                    partial_output,
                )
        all_reduce_by_source = {

            all_reduce_bucket = all_reduce_by_source.get(id(rs_bucket))
            if all_reduce_bucket is not None:
                all_reduce_bucket.all_reduce_output = rs_bucket.move_reduce_scatter_output()
                continue
            reduce_scatter_output = rs_bucket.move_reduce_scatter_output()
            for hsdp_param, param_numel, param_offset in zip(
                rs_bucket.hsdp_params,
                rs_bucket.layout.param_numels,
                rs_bucket.param_offsets,
            ):
                hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_output = (
                    reduce_scatter_output.narrow(0, param_offset, param_numel)
                )
        self.reduce_scatter_buckets = []
        self._issue_all_reduce_buckets(self.all_reduce_buckets, async_op)

            if bucket.handle is not None:
                bucket.handle.wait()
                bucket.handle = None
            if bucket.all_reduce_output is None:
                raise RuntimeError("All-reduce output has already been released.")
            output = bucket.all_reduce_output
            if bucket.needs_avg_div and bucket.replicate_world_size > 1:
                output = ms.mint.div(output, bucket.replicate_world_size)
            for hsdp_param, param_numel, param_offset in zip(
                bucket.hsdp_params,
                bucket.layout.param_numels,
                bucket.layout.param_offsets,

            self.comm_ctx.all_reduce_param_group = None

    def reset_iter_state(self) -> None:
        """Drop communication references after a completed iteration."""
        for bucket in self.all_gather_buckets:
            bucket.all_gather_result = None
        for bucket in self.reduce_scatter_buckets:
            bucket.unsharded_grads = []
            bucket.reduce_scatter_input = None
            bucket.reduce_scatter_output = None
            bucket.handle = None
        for bucket in self.all_reduce_buckets:
            bucket.all_reduce_output = None
            bucket.handle = None
        self.reduce_scatter_buckets = []
        self.all_reduce_buckets = []
        self.reduce_partial_outputs.clear()
        if self.comm_ctx.pre_param_group is self:
            self.comm_ctx.pre_param_group = None
        if self.comm_ctx.all_reduce_param_group is self:
            self.comm_ctx.all_reduce_param_group = None


class AllReduceParamGroup:
    """Fuse per-parameter RS outputs for one HSDP replicate all-reduce."""

        return aligned_total_bytes // element_size

    def allocate_fused_buffer(self, device: Any) -> None:
        """Allocate and zero the fused all-reduce buffer."""
        del device
        self.fused_buffer = ms.mint.zeros(
            (self.compute_aligned_layout(),),
            dtype=self.reduce_dtype,
        )

    def get_param_buffer_view(self, index: int) -> ms.Tensor:
        """Return one parameter's communication view."""
        if self.fused_buffer is None:
            raise RuntimeError("Fused buffer not allocated. Call allocate_fused_buffer first.")
        return self.fused_buffer.narrow(
            0,
            self.param_offsets[index],
            self.param_numels[index],
        )

        param_outputs = []
        for hsdp_param in self.hsdp_params:
            reduced_output = hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_output
            if reduced_output is None:
                raise RuntimeError("All-reduce group requires one completed reduce-scatter output per parameter.")
            if reduced_output.dtype != self.reduce_dtype:
                reduced_output = reduced_output.to(self.reduce_dtype)
            partial_output = hsdp_param.reduce_partial_output
            if partial_output is not None:
                if partial_output.dtype != self.reduce_dtype:
                    partial_output = partial_output.to(self.reduce_dtype)
                reduced_output = ms.mint.add(reduced_output, partial_output)
                hsdp_param.reduce_partial_output = None
            param_outputs.append(reduced_output.reshape(-1))
            hsdp_param.clear_reduce_scatter_output()
            hsdp_param.clear_unsharded_source_grad()
        packed_output = ms.mint.cat(param_outputs, dim=0)

        if self.all_reduce_handle is not None:
            self.all_reduce_handle.wait()
            self.all_reduce_handle = None
        if self.fused_buffer is None:
            raise RuntimeError("Fused buffer has already been released.")
        output = self.fused_buffer
        if self.reduce_op == "avg" and self.replicate_world_size > 1:
            output = ms.mint.div(output, self.replicate_world_size)
        for index, hsdp_param in enumerate(self.hsdp_params):

            )
            comm_ctx.pre_param_group.wait_reduce_scatter_and_issue_all_reduce()
            comm_ctx.pre_param_group = None
        if comm_ctx.all_reduce_param_group is not None:
            comm_ctx.all_reduce_param_group.wait_all_reduce_and_save_grad()
            comm_ctx.all_reduce_param_group = None

    def _finalize_per_param_reductions(self) -> None:
        """Drain the module-tree-local comm_fusion=False communication queues."""
        previous_groups = self.hsdp_state._wait_prev_reduce_scatter()

        """Run final source-layout reductions and apply gradients for all states."""
        for hsdp_scheduler in self.scheduler_ctx.all_hsdp_schedulers:
            hsdp_state = hsdp_scheduler.hsdp_state
            if hsdp_state is None:
                continue
            need_synchronize = False
            for hsdp_param in hsdp_state.hsdp_params:
                reduced_grad = hsdp_param.all_reduce_comm_ctx.all_reduce_output
                if reduced_grad is None:
                    reduced_grad = hsdp_param.reduce_scatter_comm_ctx.reduce_scatter_output
                if reduced_grad is None:
                    continue
                hsdp_param.all_reduce_tp_replicate_grad_inplace(
                    reduced_grad,
                    hsdp_state.reduce_op_type,
                )
                need_synchronize = hsdp_param.apply_reduced_grad(reduced_grad) or need_synchronize
                hsdp_param.clear_all_reduce_output()
                hsdp_param.clear_reduce_scatter_output()
            hsdp_state._sync_current_stream_if_needed(need_synchronize)

    def reset_iter_state(self) -> None:
        """Reset MindSpore fully_shard iteration state after communication completes."""
        super().reset_iter_state()
        self.hsdp_state.reset_iter_state()
        comm_ctx = self.scheduler_ctx.param_group_comm_ctx
        comm_ctx.pre_param_group = None
        comm_ctx.all_reduce_param_group = None

    def _backward_hook(self):
        """Execute backward hook."""
        if self.scheduler_state == FSDPSchedulerState.BACKWARD:

        platform,
        scheduler_ctx,
        device=None,
    ):
        super().__init__(
            cell,
            mesh,
            shard_placement_fn,
            comm_fusion_policy,

    def _init_param_group(self) -> None:
        """Initialize fused communication for the single managed parameter list."""
        self.param_group = None
        if not self.comm_fusion_policy.enable_comm_fusion or not self.hsdp_params:
            return
        self.param_group = HSDPParamGroup(
            self.hsdp_params,
            self.device,
            self.mp_policy,

        )

    def _move_states_to_device(self) -> None:
        """Move parameters and buffers to the configured runtime device."""
        for module in self.modules:
            for param in module.get_parameters():
                if getattr(param, "_hsdp_param_initialized", False):
                    continue
                param_device = normalize_runtime_device(param.device)
                if param_device in (self.device, "meta"):
                    continue

                param_device = normalize_runtime_device(param.device)
                if param_device in (self.device, "meta"):
                    continue
                param.data = param.to(self.device)
            for buffer in module.buffers():
                if normalize_runtime_device(buffer.device) in (self.device, "meta"):
                    continue
                buffer.data = buffer.to(self.device)

    @staticmethod

    @staticmethod
    def _build_param_tp_grad_info(param):
        """Build normalized source-layout metadata for a native DTensor parameter."""
        dtensor_payload = unwrap_dtensor_param(param)
        if dtensor_payload is None:
            return None
        return TPShardMetaInfo(
            mesh=dtensor_payload.device_mesh,
            placements=tuple(dtensor_payload.placements),
            origin_is_dtensor=True,
        )

    def _init_hsdp_params(self) -> None:
        """Initialize all fully_shard-managed parameters for this module unit."""
        visited_params = set()
        filtered_params = []
        for module in self.modules:
            for _, param in module.parameters_and_names():
                if param in self.raw_ignored_params:
                    continue
                if getattr(param, "_hsdp_param_initialized", False):
                    continue
                if param in visited_params:
                    continue
                visited_params.add(param)