Diff Coverage

        therefore cannot infer that this wrapper removes its internal module
        prefix. Map every direct child state name to the wrapped module so
        optimizer save and restore use the same canonical FQNs.
        """
        wrapped_module = self._swap_wrapped_module
        exposed_names = set(wrapped_module._modules)  # pylint: disable=protected-access
        exposed_names.update(wrapped_module._parameters)  # pylint: disable=protected-access
        exposed_names.update(wrapped_module._buffers)  # pylint: disable=protected-access
        if (
                getattr(wrapped_module.__class__, "get_extra_state", nn.Module.get_extra_state)
                != nn.Module.get_extra_state
        ):
            exposed_names.add(nn.modules.module._EXTRA_STATE_KEY_SUFFIX)
        return {name: _SWAP_WRAPPED_MODULE for name in exposed_names}

    @staticmethod
    def _post_state_dict_hook(
        module: nn.Module,  # pylint: disable=W0613

    Returns:
        tuple[int, ...]: Tuple of offsets for each dimension.
    """
    if dtensor_layout is None:
        # If layout is None, return all zeros (no sharding)
        return tuple(0 for _ in global_shape)

    # Validate layout attributes
    if not hasattr(dtensor_layout, 'mesh_shape') or dtensor_layout.mesh_shape is None:
        raise ValueError("Layout must have mesh_shape attribute")
    if not hasattr(dtensor_layout, 'tensor_map') or dtensor_layout.tensor_map is None:
        raise ValueError("Layout must have tensor_map attribute")
    if not hasattr(dtensor_layout, 'rank_list') or dtensor_layout.rank_list is None:
        raise ValueError("Layout must have rank_list attribute")

    if current_rank not in dtensor_layout.rank_list:
        raise ValueError(
            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,
    )

        inner_rank_id,
        global_shape,
    )
    # Extract offsets (start values) from slice_area
    return tuple(start for start, _ in slice_area)


@dataclass(frozen=True)
class CachedSaveResult:

                layout = obj.layout

                # Get chunk metadata with offsets
                if layout:
                    offsets = _compute_global_offsets(obj.shape, layout, self.rank)
                else:
                    offsets = (0,) * len(local_tensor.shape)

                sizes = local_tensor.shape

        if cpu_offload and result.device != target_device:
            result = result.to(target_device)
        return result

    mesh, _ = dtensor_info
    target_device = param.to_local().device

    scalar_placements = [Replicate()] * mesh.ndim
    if cpu_offload and tensor.device != target_device:

        """
        for param_name, param in super().named_parameters(*args, **kwargs):
            yield param_name.replace(_SWAP_PREFIX, ""), param

    def _fqn_modifiers(self) -> dict[str, str]:
        """Describe the wrapper's state-dict FQN rewrite to PyTorch DCP.

        Distributed state-dict traversal does not use ``named_modules()`` and
        therefore cannot infer that this wrapper removes its internal module

        therefore cannot infer that this wrapper removes its internal module
        prefix. Map every direct child state name to the wrapped module so
        optimizer save and restore use the same canonical FQNs.
        """
        wrapped_module = self._swap_wrapped_module
        exposed_names = set(wrapped_module._modules)  # pylint: disable=protected-access
        exposed_names.update(wrapped_module._parameters)  # pylint: disable=protected-access
        exposed_names.update(wrapped_module._buffers)  # pylint: disable=protected-access
        if (
                getattr(wrapped_module.__class__, "get_extra_state", nn.Module.get_extra_state)
                != nn.Module.get_extra_state
        ):
            exposed_names.add(nn.modules.module._EXTRA_STATE_KEY_SUFFIX)
        return {name: _SWAP_WRAPPED_MODULE for name in exposed_names}

    @staticmethod
    def _post_state_dict_hook(
        module: nn.Module,  # pylint: disable=W0613

# See the License for the specific language governing permissions and
# limitations under the License.
"""CheckpointerCallback --- save/restore policy on top of a Checkpointer."""

__all__ = ["CheckpointerCallback"]

import os
import random
from typing import TYPE_CHECKING, Any, Dict, List, Optional

    get_device_rng_state,
    set_device_rng_state,
)
from hyper_parallel.models._transformers.model_builder import apply_model_init_dtype
from .base import Callback, TrainerState


if TYPE_CHECKING:
    from hyper_parallel.trainer.base import BaseTrainer