Diff Coverage

def _pad_dim0_for_communication(tensor: ms.Tensor, padded_dim0: int) -> ms.Tensor:
    """Pad a dim-0 shard with mint operators for fixed-size collectives."""
    actual_dim0 = tensor.shape[0]
    if actual_dim0 == padded_dim0:
        return tensor
    padding_shape = (padded_dim0 - actual_dim0, *tensor.shape[1:])
    padding = ms.mint.zeros(padding_shape, dtype=tensor.dtype)
    if normalize_runtime_device(padding.device) != normalize_runtime_device(tensor.device):
        padding = padding.to(normalize_runtime_device(tensor.device))
    return ms.mint.cat((tensor, padding), dim=0)


def make_contiguous_strides_for(shape, row_major=True):

        self.offload_to_cpu: bool = isinstance(offload_policy, CPUOffloadPolicy)
        self.pin_memory = (
            self.offload_to_cpu and cast(CPUOffloadPolicy, offload_policy).pin_memory
        )
        self._parameter_hook_migrator = ParameterHookMigrator()
        if isinstance(param, DTensor) != (
            source_shard_info is not None and source_shard_info.origin_is_dtensor
        ):
            raise ValueError(
                "source_shard_info.origin_is_dtensor must be True exactly for native DTensor parameters, "

        )
        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._parameter_hook_migrator._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._grad = None
        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()

        source_param: Parameter,
        source_local_tensor: ms.Tensor,
    ) -> Layout:
        """Build the final layout after data and model parallel sharding."""
        logical_global_stride = None
        if isinstance(source_param, DTensor):
            logical_global_size = source_param.shape
            logical_global_stride = source_param.layout.tensor_stride
        elif self.source_shard_info is not None:
            source_sharding_spec = Layout.from_device_mesh(self.source_shard_info.mesh)
            source_sharding_spec.set_placements(self.source_shard_info.placements)
            source_sharding_spec.placement_to_tensor_map(source_local_tensor.ndim)
            logical_global_size = source_sharding_spec.get_global_shape(source_local_tensor.shape)
        else:
            logical_global_size = source_local_tensor.shape

        if logical_global_stride is None:
            logical_global_stride = make_contiguous_strides_for(logical_global_size)

        sharding_spec = Layout.from_device_mesh(self._spmd_mesh)
        sharding_spec.set_placements(self._spmd_placements)
        sharding_spec.placement_to_tensor_map(source_local_tensor.ndim)
        sharding_spec.set_tensor_meta(
            logical_global_size,
            logical_global_stride,
            source_local_tensor.dtype,
        )
        return sharding_spec

    @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:
        """Store reduce-scatter results accumulated before the final micro-step."""
        self._reduce_partial_output = value

    def reduce_comm_dtype(self, grad: Optional[ms.Tensor] = None) -> Optional[ms.Type]:
        """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 reduce_scatter_output(self) -> Optional[ms.Tensor]:
        """Return cached reduce-scatter output after waiting asynchronous work."""
        if self.reduce_scatter_comm_ctx.reduce_scatter_handle is not None:

    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
        elif 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: ms.Tensor) -> bool:

            Whether the caller must synchronize after a CPU offload.
        """
        if self.mp_policy.apply_grad_on_fp32_main_grad:
            if not hasattr(self.sharded_param, "main_grad"):
                self.sharded_param.main_grad = None
            sharded_grad = self.sharded_param.main_grad
        else:
            sharded_grad = self.sharded_param.grad

            else:
                self.sharded_param.grad = self.to_sharded_dtensor(reduced_grad)
        else:
            if self.mp_policy.apply_grad_on_fp32_main_grad:
                self.sharded_param.main_grad._local_tensor.add_(reduced_grad)
                self.sharded_param.grad = None
            else:
                self.sharded_param.grad._local_tensor.add_(reduced_grad)

        self.clear_unsharded_source_grad()

        if not isinstance(hsdp_placement, Shard):
            raise AssertionError(
                f"Expected Shard, got {type(hsdp_placement)}: {hsdp_placement}"
            )
        return hsdp_placement

    def _init_shard_metadata(
        self,
        param: Parameter,

        """Initialize parameter shape and mesh metadata used by sharding."""

        self.hsdp_placement = hsdp_placement
        base_placements = list(self._get_base_spmd_placements())
        param_data = param.to_local() if self._orig_param_is_dtensor else param
        shard_dim = hsdp_placement.dim
        if param_data.ndim == 0:
            raise ValueError("fully_shard does not support scalar parameters")
        if shard_dim < 0 or shard_dim >= param_data.ndim:
            raise ValueError(
                f"Invalid fully_shard dim {shard_dim} for parameter "
                f"{self._module_info.param_name} with shape {tuple(param_data.shape)}"
            )
        self._orig_size = param_data.shape

            isinstance(self.mesh_info, DDPMeshInfo)
            and not isinstance(self.mesh_info, HSDPMeshInfo)
        )
        self.is_sharded = self.shard_world_size > 1
        dim_shard_size = (param_data.shape[shard_dim] + self.shard_world_size - 1) // self.shard_world_size
        return base_placements, param_data, shard_dim, dim_shard_size

    def _init_shard_placements(
        self,
        param_data: ms.Tensor,

            self._apply_data_parallel_placements(base_placements, self.hsdp_placement)
        )
        if param_data.shape[shard_dim] % self.shard_world_size != 0:
            if self._spmd_shard_mesh_dim is None:
                raise AssertionError("Uneven FSDP sharding requires a shard mesh dimension")
            fsdp_placement = spmd_placements[self._spmd_shard_mesh_dim]
            if isinstance(fsdp_placement, StridedShard):
                fsdp_placement = StridedShard(
                    fsdp_placement.dim,

        padded_sharded_size = list(param_data.shape)
        padded_sharded_size[shard_dim] = dim_shard_size
        self.padded_sharded_param_size = tuple(padded_sharded_size)
        if self.offload_to_cpu and not sharded_param.is_meta:
            sharded_param = sharded_param.to("cpu")
            if self.pin_memory:
                sharded_param = sharded_param.pin_memory()

        if self.sharded_size == self.padded_sharded_param_size:
            padded_sharded_param = sharded_param
        else:
            padded_sharded_param = _pad_dim0_for_communication(
                sharded_param,
                self.padded_sharded_param_size[0],

                sharded_param,
                self.padded_sharded_param_size[0],
            )
            if self.pin_memory and not padded_sharded_param.is_meta:
                padded_sharded_param = padded_sharded_param.pin_memory()
        self._sharded_param_data = padded_sharded_param.reshape(-1)
        # MindSpore optimizers must update the independent logical shard, not a
        # narrow view into padded communication storage.
        return sharded_param

        param: Parameter,
        shard_placement_fn: Optional[Callable],
    ) -> None:
        """Initialize the persistent sharded parameter and communication storage."""
        hsdp_placement = self._resolve_hsdp_placement(param, shard_placement_fn)
        base_placements, param_data, shard_dim, dim_shard_size = self._init_shard_metadata(
            param,
            hsdp_placement,
        )
        self._init_shard_placements(param_data, shard_dim, base_placements)
        sharded_param = self._build_sharded_param_data(
            param_data,
            shard_dim,
            dim_shard_size,
        )

            shard_dim,
            dim_shard_size,
        )

        self._sharding_spec = self._build_sharding_spec(param, param_data)

        shard_dtensor = DTensor.from_local(
            sharded_param,
            self._spmd_mesh,
            self._spmd_placements,
            shape=self._sharding_spec.tensor_shape,

    def init_unsharded_param(self) -> None:
        """Initialize the logical full parameter from its all-gather storage."""
        if len(self.unsharded_param_buffers) != 1:
            raise AssertionError(
                f"Expected 1 unsharded_param_buffer, got {len(self.unsharded_param_buffers)}"
            )

        all_gather_output = self.allgather_comm_ctx.allgather_output

        unsharded_numel = math.prod(self._orig_size)
        unsharded_param = self.unsharded_param_buffers[0].narrow(0, 0, unsharded_numel)
        unsharded_param = unsharded_param.reshape(self._orig_size)
        if self._orig_param_is_dtensor:
            unsharded_param = DTensor.from_local(
                unsharded_param,
                self._orig_dtensor_mesh,
                self._orig_dtensor_placements,
            )

            self._module_info.module._params[self._module_info.param_name] = param
        else:
            # slow path
            setattr(self._module_info.module, self._module_info.param_name, param)
        self._parameter_hook_migrator._save_backward_hooks(self.sharded_param)
        self._parameter_hook_migrator._migrate_backward_hooks(param)

        # Iterate through all modules that share this parameter to prevent pointer desync.
        for shared_module, shared_param_name in zip(
            self._module_info.shared_modules, self._module_info.shared_param_names

            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.add_(unsharded_grad)

    def accumulate_unsharded_grad_if_needed(self) -> None:
        if (
            self.unsharded_accumulated_grad is not None

            and self.unsharded_param.grad is not None
        ):
            grad = self._to_local_unsharded_grad(self.unsharded_param.grad)
            if self.reduce_dtype is not None and grad.dtype != self.reduce_dtype:
                grad = grad.to(self.reduce_dtype)
            self.unsharded_param.grad = None
            self.unsharded_accumulated_grad.add_(grad)

    def alloc_unsharded_param_buffers(self) -> None:

            )

    def _resolve_reset_param(self):
        """Resolve the possibly replaced module parameter before resetting storage."""
        module_info = self._module_info
        new_param = getattr(module_info.module, module_info.param_name)
        if new_param is self.sharded_param:
            return new_param
        if isinstance(new_param, DTensor):
            self.sharded_param = new_param
            if not getattr(self.sharded_param, "_hsdp_param_initialized", None):
                self.sharded_param._hsdp_param_initialized = True
        return new_param

    def _is_same_sharded_local_tensor(self, local_tensor: ms.Tensor) -> bool:
        """Whether the cached flat shard view already points to the ``local_tensor`` storage."""
        if not isinstance(self._sharded_param_data, ms.Tensor):

        if local_tensor.shape != self.sharded_size:
            raise AssertionError(
                f"Expected sharded_size to be {self.sharded_size}, got {local_tensor.shape}"
            )
        return local_tensor

    def _pin_reset_local_tensor_if_needed(self, local_tensor: ms.Tensor) -> Tuple[ms.Tensor, bool]:
        """Pin the local tensor memory when CPU offload requires it."""
        if self.pin_memory and not local_tensor.is_pinned():

        local_tensor: ms.Tensor,
    ) -> None:
        """Rebuild padded communication storage and refresh the DTensor local view."""
        if self.sharded_size == self.padded_sharded_param_size:
            padded_local_tensor = local_tensor
        else:
            padded_local_tensor = _pad_dim0_for_communication(
                local_tensor,
                self.padded_sharded_param_size[0],

                local_tensor,
                self.padded_sharded_param_size[0],
            )
            if self.pin_memory and not padded_local_tensor.is_meta:
                padded_local_tensor = padded_local_tensor.pin_memory()
        self._sharded_param_data = padded_local_tensor.reshape(-1)
        set_requires_grad_if_needed(self.sharded_param, local_tensor)
        self.sharded_param._local_tensor = local_tensor
        if not self.sharded_param._local_tensor.is_contiguous():

            )

    def reset_sharded_param(self) -> None:
        """Reset the sharded param after ``load_state_dict``."""
        new_param = self._resolve_reset_param()
        local_tensor = new_param._local_tensor if isinstance(new_param, DTensor) else new_param
        if local_tensor.is_meta:
            return
        # local_tensor can be padded twice

        # this makes it possible for trainer to call `sd = model.state_dict()` before the training loop
        # and use `sd` without calling .state_dict() per iteration
        same_local_tensor = self._is_same_sharded_local_tensor(local_tensor)
        if not same_local_tensor:
            local_tensor = self._validate_reset_local_tensor(local_tensor)
        local_tensor, pinned_local_tensor = self._pin_reset_local_tensor_if_needed(local_tensor)
        if not isinstance(self.sharded_param, DTensor):
            raise AssertionError(f"Expected DTensor, got {type(self.sharded_param)}")
        if not same_local_tensor or pinned_local_tensor:
            self._refresh_sharded_local_tensor(local_tensor)
        self._sharding_spec.set_tensor_meta(
            self._sharding_spec.tensor_shape,
            self._sharding_spec.tensor_stride,
            local_tensor.dtype,
        )
        self.sharded_param._layout = self._sharding_spec
        self.sharded_param._placements = tuple(self._sharding_spec.placements)
        self._setattr_on_modules(self.sharded_param)

    @_no_grad()
    def _get_unsharded_param_data(
        self,

        # unshard cycle.
        self.reset_sharded_param()
        all_gather_input = self.all_gather_inputs[0]

        shard_group = self.mesh_info.shard_process_group if isinstance(self.mesh_info, FSDPMeshInfo) else None
        if not self.is_sharded or shard_group is None or self.shard_world_size <= 1:
            if not self.unsharded_param_buffers:
                self.unsharded_param_buffers = [all_gather_input]
            elif self.unsharded_param_buffers[0] is not all_gather_input:
                self.alloc_unsharded_param_buffers()
                copy_without_bumping_version(self.unsharded_param_buffers[0], all_gather_input)
            self.allgather_comm_ctx.allgather_output = None
            self.allgather_comm_ctx.allgather_handle = None
            return

        self.init_unsharded_param_buffers(
            all_gather_input_numels=[all_gather_input.numel()],
            all_gather_input_dtypes=[all_gather_input.dtype],

            device=self._sharded_param_data.device.split(":")[0],
        )
        self.alloc_unsharded_param_buffers()

        self.allgather_comm_ctx.allgather_output = self.unsharded_param_buffers[0]
        if self.hsdp_placement.dim != 0:
            # Non-dim-0 shards require chunk + cat after the collective. The
            # stable full-parameter buffer must not be overwritten beforehand.
            self.allgather_comm_ctx.allgather_output = ms.mint.empty_like(
                self.unsharded_param_buffers[0]
            )
        self.allgather_comm_ctx.allgather_handle = dist.all_gather_into_tensor(
            self.allgather_comm_ctx.allgather_output,
            all_gather_input,
            group=shard_group,
            async_op=async_op,

            all_gather_input,
            group=shard_group,
            async_op=async_op,
        )
        if self.allgather_comm_ctx.allgather_output is self.unsharded_param_buffers[0]:
            self.allgather_comm_ctx.allgather_output = None

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

        self._get_unsharded_param_data(async_op=async_op)

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

            grad = self.unsharded_grad_data
        self._grad = grad.to(self.reduce_comm_dtype(grad))
        shard_dim = self.hsdp_placement.dim
        if self.shard_world_size <= 1:
            self._grad = self._grad.reshape(-1)
        elif shard_dim != 0:
            grad_chunks = self._grad.chunk(self.shard_world_size, dim=shard_dim)
            self._grad = ms.mint.cat(grad_chunks, dim=0).reshape(-1)
        else:
            padded_unsharded_dim0 = self.padded_sharded_param_size[0] * self.shard_world_size
            if self._grad.shape[0] != padded_unsharded_dim0:
                self._grad = _pad_dim0_for_communication(self._grad, padded_unsharded_dim0)

        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:
            if output_buffer is not None:
                copy_without_bumping_version(output_buffer, self._grad)
                self.reduce_scatter_comm_ctx.reduce_scatter_output = output_buffer
            else:
                self.reduce_scatter_comm_ctx.reduce_scatter_output = self._grad
            self.reduce_scatter_comm_ctx.reduce_scatter_handle = None
            return

        # Calculate output size

            if output_buffer.numel() != output_numel:
                raise ValueError(
                    f"output_buffer size mismatch: expected {output_numel}, got {output_buffer.numel()}"
                )
            if output_buffer.dtype != self._grad.dtype:
                raise ValueError(
                    f"output_buffer dtype mismatch: expected {self._grad.dtype}, got {output_buffer.dtype}"
                )
            self.reduce_scatter_comm_ctx.reduce_scatter_output = output_buffer

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

            return

        # Ascend HCCL accepts contiguous views but rejects non-contiguous input.
        if not grad.is_contiguous():
            grad = ms.mint.cat((grad,), dim=0)

        self.all_reduce_comm_ctx.all_reduce_output = grad
        self.all_reduce_comm_ctx.all_reduce_handle = dist.all_reduce(
            grad,

        if result is None or result.all_gather_output is None:
            return
        if result.handle is not None:
            result.handle.wait()
            result.handle = None
        gathered_rows = result.all_gather_output.reshape(self.shard_world_size, -1)
        for input_numels, input_dtypes, hsdp_param in zip(
            self.metadata.param_input_numels,
            self.metadata.param_input_dtypes,

            self.metadata.param_input_numels,
            self.hsdp_params,
        ):
            if hsdp_param.hsdp_placement.dim != 0 and len(input_numels) != 1:
                raise NotImplementedError(
                    "Fused non-dim-0 all-gather expects one local shard tensor per parameter."
                )
            for input_numel, output_buffer in zip(
                input_numels,

                    source,
                )
                column_offset += input_numel
        if column_offset != gathered_rows.shape[1]:
            raise AssertionError(
                "Fused all-gather copy-out consumed an unexpected number of elements: "
                f"{column_offset} != {gathered_rows.shape[1]}"
            )
        result.all_gather_input = None

                    *unsharded_grad.shape[1:],
                )
                padding = ms.mint.zeros(padding_shape, dtype=unsharded_grad.dtype)
                if _normalize_device(padding.device) != _normalize_device(unsharded_grad.device):
                    padding = padding.to(_normalize_device(unsharded_grad.device))
                unsharded_grad = ms.mint.cat((unsharded_grad, padding), dim=0)
            packed_grad = unsharded_grad.reshape(world_size, -1)
        else:
            grad_chunks = unsharded_grad.chunk(world_size, dim=shard_dim)

            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.div_(bucket.shard_world_size)

    @staticmethod
    def _issue_all_reduce_buckets(
        all_reduce_buckets: list[AllReduceBucket],

                partial_output = self.reduce_partial_outputs.get(bucket.bucket_key)
                if partial_output is None:
                    self.reduce_partial_outputs[bucket.bucket_key] = current_output
                else:
                    partial_output.add_(current_output)
            self.reduce_scatter_buckets = []
            self.all_reduce_buckets = []
            return

        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.add_(partial_output)
        all_reduce_by_source = {
            id(bucket.source_reduce_scatter_bucket): bucket
            for bucket in self.all_reduce_buckets
        }

            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_reducescatter_output()
                continue
            reduce_scatter_output = rs_bucket.move_reducescatter_output()
            for hsdp_param, param_numel, param_offset in zip(
                rs_bucket.hsdp_params,
                rs_bucket.layout.param_numels,
                rs_bucket.param_offsets,

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

        )

    def get_param_grad_view(self, index: int, target_shape: tuple[int, ...]) -> ms.Tensor:
        """Return one parameter's actual-shard gradient view."""
        return self.get_param_buffer_view(index).narrow(
            0,
            0,
            _shape_numel(target_shape),
        ).reshape(target_shape)

    def accumulate_reduce_partial_outputs(self) -> None:
        """Merge no-all-reduce micro-step outputs into the current buffer."""
        if self.fused_buffer is None:
            return
        for index, hsdp_param in enumerate(self.hsdp_params):
            if hsdp_param.reduce_partial_output is None:
                continue
            partial_output = hsdp_param.reduce_partial_output
            if partial_output.dtype != self.reduce_dtype:
                partial_output = partial_output.to(self.reduce_dtype)
            self.get_param_buffer_view(index).add_(partial_output.reshape(-1))
            hsdp_param.reduce_partial_output = None

    def issue_async_allreduce(self) -> None:
        """Launch SUM all-reduce; AVG division is applied when splitting."""
        if self.fused_buffer is None:

                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_source_replicate_grad_inplace(
                    reduced_grad,
                    hsdp_state.reduce_op_type,
                )
                need_synchronize = hsdp_param.apply_reduced_grad(reduced_grad) or need_synchronize

    @staticmethod
    def _build_param_source_shard_info(param):
        """Build normalized source-layout metadata for a native DTensor parameter."""
        if not isinstance(param, DTensor):
            return None
        return SourceShardMetaInfo(
            mesh=param.device_mesh,
            placements=tuple(param.placements),

                previous_group.hsdp_params,
            )
            for hsdp_param in previous_group.hsdp_params:
                hsdp_param.reduce_scatter_output()
                hsdp_param.clear_reduce_scatter_output()
                if hsdp_param.unsharded_accumulated_grad_data is not None:
                    hsdp_param.unsharded_accumulated_grad = None
                elif hsdp_param.unsharded_param.grad is not None:
                    hsdp_param.unsharded_param.grad = None
        return previous_groups

    def _issue_prev_fused_all_reduce(self, previous_groups: List[AllReduceParamGroup]) -> None:
        """Launch the previous module's fused all-reduce asynchronously."""

            if not self.requires_all_reduce:
                if hsdp_param.reduce_partial_output is None:
                    hsdp_param.reduce_partial_output = reduced_grad
                else:
                    hsdp_param.reduce_partial_output.add_(reduced_grad)
                hsdp_param.clear_reduce_scatter_output()
            elif hsdp_param.reduce_partial_output is not None:
                reduced_grad.add_(hsdp_param.reduce_partial_output)
                hsdp_param.reduce_partial_output = None

            if hsdp_param.unsharded_accumulated_grad_data is not None:
                hsdp_param.unsharded_accumulated_grad = None
            elif hsdp_param.unsharded_param.grad is not None:
                hsdp_param.unsharded_param.grad = None

    def wait_and_split_all_reduce_work_groups(self) -> None: