Diff Coverage

    Returns:
        dict: Optimizer state dict with FQN-based keys.
    """
    return platform.get_optim_state_dict(model, optimizer, options=options)


def set_optim_state_dict(
    model: Any,

        optim_state_dict: The optimizer state dict to load.
        options: Optional configuration (full_state_dict, cpu_offload,
            strict, broadcast_from_rank0, etc.).
    """
    platform.set_optim_state_dict(model, optimizer, optim_state_dict, options=options)

        )

    @staticmethod
    def get_optim_state_dict(model, optimizer, *, options=None):
        raise NotImplementedError(
            "issue240 optimizer state_dict adapter is currently supported only on Torch backend"
        )

    @staticmethod

        )

    @staticmethod
    def set_optim_state_dict(model, optimizer, optim_state_dict, *, options=None):
        raise NotImplementedError(
            "issue240 optimizer state_dict adapter is currently supported only on Torch backend"
        )

    @staticmethod

        Returns:
            dict: The optimizer state dictionary with FQN-based keys.
        """
        raise NotImplementedError(
            "Platform subclasses must implement get_optim_state_dict"
        )

    @staticmethod

            optimizer: The optimizer instance.
            optim_state_dict: The optimizer state dict to load.
            options: Optional configuration for state dict loading.
        """
        raise NotImplementedError(
            "Platform subclasses must implement set_optim_state_dict"
        )

    @staticmethod

This module does **not** call PyTorch DCP's
``get_optimizer_state_dict``/``set_optimizer_state_dict``.
"""
from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union

import torch
import torch.distributed as dist
from torch import nn

from hyper_parallel.core.dtensor.dtensor import DTensor
from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
from hyper_parallel.core.dtensor.placement_types import Replicate, Shard

logger = logging.getLogger(__name__)

_SCALAR_STATE_KEYS = {"step"}
_TENSOR_STATE_KEYS_ADAM = {"exp_avg", "exp_avg_sq", "max_exp_avg_sq"}
_TENSOR_STATE_KEYS_SGD = {"momentum_buffer"}


class UnsupportedConfigurationError(RuntimeError):
    """Raised when the requested configuration is not supported."""


def _is_dtensor_param(param: Any) -> bool:
    return isinstance(param, DTensor)


def _is_replicate_group_root(
    mesh: DeviceMesh,
    placements: Sequence,
) -> bool:
    """Return True if current rank is the root (coordinate-0) in every

    its global rank is not 0 (e.g. PP stage-1 ranks in a 3-D mesh).

    If there are no Replicate dimensions, falls back to ``dist.get_rank() == 0``.
    """
    coord = mesh.get_coordinate()
    if coord is None:
        return False
    replicate_dims = [i for i, p in enumerate(placements) if isinstance(p, Replicate)]
    if not replicate_dims:
        return dist.get_rank() == 0
    return all(coord[d] == 0 for d in replicate_dims)


def _param_to_fqn(model: nn.Module) -> Dict[nn.Parameter, str]:
    """Build a mapping from parameter object to its fully-qualified name."""
    param_to_name: Dict[nn.Parameter, str] = {}
    for name, param in model.named_parameters():
        if param in param_to_name:
            raise ValueError(
                f"Parameter {name} shares the same object as "
                f"{param_to_name[param]}. Duplicate parameter objects are not supported."
            )
        param_to_name[param] = name
    return param_to_name


def _build_id_to_fqn(
    optimizer: torch.optim.Optimizer,
    model: nn.Module,
) -> Tuple[Dict[int, str], Dict[str, int]]:
    """Build mappings between optimizer saved IDs and FQNs.

    Returns:
        Tuple of (saved_id_to_fqn, fqn_to_saved_id).
    """
    param_to_name = _param_to_fqn(model)
    raw_sd = optimizer.state_dict()
    saved_id_to_fqn: Dict[int, str] = {}
    saved_id_to_param: Dict[int, nn.Parameter] = {}

    for runtime_group, saved_group in zip(optimizer.param_groups, raw_sd["param_groups"]):
        for parameter, saved_id in zip(runtime_group["params"], saved_group["params"]):
            if saved_id in saved_id_to_param and saved_id_to_param[saved_id] is not parameter:
                raise ValueError(
                    f"saved_id {saved_id} maps to different parameters; "
                    f"this is not supported."
                )
            saved_id_to_param[saved_id] = parameter

    for saved_id, param in saved_id_to_param.items():
        if param not in param_to_name:
            raise ValueError(
                f"Parameter with saved_id={saved_id} not found in "
                f"model.named_parameters(). Ensure the model matches the optimizer."
            )
        saved_id_to_fqn[saved_id] = param_to_name[param]

    fqn_to_saved_id = {v: k for k, v in saved_id_to_fqn.items()}
    return saved_id_to_fqn, fqn_to_saved_id


def _get_param_dtensor_info(
    param: nn.Parameter,
) -> Optional[Tuple[DeviceMesh, Sequence]]:
    """Return (device_mesh, placements) if param is DTensor, else None."""
    if _is_dtensor_param(param):
        return param.device_mesh, param.placements
    return None


def _convert_state_tensor(
    tensor: torch.Tensor,
    param: nn.Parameter,
    full_state_dict: bool,
    cpu_offload: bool,

    Returns:
        Converted tensor, or None if this rank should not keep the result.
    """
    dtensor_info = _get_param_dtensor_info(param)
    is_dtensor_value = isinstance(tensor, DTensor)

    if full_state_dict and dtensor_info is not None:
        if is_dtensor_value:
            full_t = tensor.full_tensor()
        else:
            mesh, placements = dtensor_info
            full_t = DTensor.from_local(tensor, mesh, placements).full_tensor()
        if cpu_offload:
            if not is_rank0:
                return None
            return full_t.cpu()
        return full_t

    if is_dtensor_value:
        local_t = tensor.to_local()
        if cpu_offload:
            return local_t.cpu()
        return local_t

    if cpu_offload:
        return tensor.cpu()

    return tensor


def _is_scalar_state(key: str) -> bool:
    return key in _SCALAR_STATE_KEYS


def _convert_state_scalar(
    tensor: torch.Tensor,
    cpu_offload: bool,
) -> torch.Tensor:
    if isinstance(tensor, DTensor):
        tensor = tensor.to_local()
    if cpu_offload:
        return tensor.cpu()
    return tensor


def _determine_is_root(
    full_state_dict: bool,
    cpu_offload: bool,
    dtensor_info: Optional[Tuple[Any, Any]],
) -> bool:

    cpu_offload: bool,
    dtensor_info: Optional[Tuple[Any, Any]],
) -> bool:
    """Determine if current rank is the replicate-group root."""
    if full_state_dict and cpu_offload and dtensor_info is not None:
        return _is_replicate_group_root(dtensor_info[0], dtensor_info[1])
    if full_state_dict and cpu_offload:
        return (not dist.is_initialized()) or (dist.get_rank() == 0)
    return True


def _convert_state_entries(
    state: Dict[str, Any],
    param: nn.Parameter,
    full_state_dict: bool,
    cpu_offload: bool,

    full_state_dict: bool,
    cpu_offload: bool,
) -> Dict[str, Any]:
    """Convert optimizer state entries for a single parameter."""
    dtensor_info = _get_param_dtensor_info(param)
    is_root = _determine_is_root(full_state_dict, cpu_offload, dtensor_info)
    skip_non_root = full_state_dict and cpu_offload and not is_root

    converted: Dict[str, Any] = {}
    for key, value in state.items():
        if isinstance(value, torch.Tensor):
            if _is_scalar_state(key):
                result = _convert_state_scalar(value, cpu_offload)
                if not skip_non_root:
                    converted[key] = result
            else:
                result = _convert_state_tensor(
                    value, param, full_state_dict, cpu_offload, is_root,
                )
                if result is not None:
                    converted[key] = result
        else:
            if not skip_non_root:
                converted[key] = value
    return converted


def _build_result_param_groups(
    optimizer: torch.optim.Optimizer,
    raw_sd: Dict[str, Any],
    saved_id_to_fqn: Dict[int, str],
) -> List[Dict[str, Any]]:

    raw_sd: Dict[str, Any],
    saved_id_to_fqn: Dict[int, str],
) -> List[Dict[str, Any]]:
    """Build result param_groups with FQN keys from saved id-based groups."""
    result_param_groups: List[Dict[str, Any]] = []
    for saved_group in raw_sd["param_groups"]:
        pg: Dict[str, Any] = {}
        for k, v in saved_group.items():
            if k == "params":
                pg["params"] = [saved_id_to_fqn[sid] for sid in v]
            else:
                pg[k] = v
        result_param_groups.append(pg)
    return result_param_groups


def _check_chained_optimizer(optimizer: Any) -> None:
    if type(optimizer).__name__ == "ChainedOptimizer":
        raise ValueError(
            "ChainedOptimizer is not supported by get_optim_state_dict / "
            "set_optim_state_dict. Use ChainedOptimizer.state_dict() and "
            "ChainedOptimizer.load_state_dict() instead, which handle "
            "multi-optimizer merging internally."

            "multi-optimizer merging internally."
        )


def get_optim_state_dict(
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    *,
    options: Any = None,

    Returns:
        Optimizer state dict with FQN-based keys.
    """
    _check_chained_optimizer(optimizer)

    full_state_dict = getattr(options, "full_state_dict", False)
    cpu_offload = getattr(options, "cpu_offload", False)
    flatten = getattr(options, "flatten_optimizer_state_dict", False)

    raw_sd = optimizer.state_dict()
    saved_id_to_fqn, _ = _build_id_to_fqn(optimizer, model)

    param_by_id: Dict[int, nn.Parameter] = {}
    for runtime_group, saved_group in zip(optimizer.param_groups, raw_sd["param_groups"]):
        for parameter, saved_id in zip(runtime_group["params"], saved_group["params"]):
            param_by_id[saved_id] = parameter

    result_state: Dict[str, Dict[str, Any]] = {}
    for saved_id, state in raw_sd["state"].items():
        fqn = saved_id_to_fqn[saved_id]
        param = param_by_id[saved_id]
        converted = _convert_state_entries(
            state, param, full_state_dict, cpu_offload,
        )
        if converted:
            result_state[fqn] = converted

    result_param_groups = _build_result_param_groups(
        optimizer, raw_sd, saved_id_to_fqn,
    )

    result: Dict[str, Any] = {
        "state": result_state,
        "param_groups": result_param_groups,
    }

    if flatten:
        result = _flatten_optim_state_dict(result)

    return result


def _check_strict_fqns(
    source_fqns: set,
    target_fqns: set,
    strict: bool,
) -> None:

    target_fqns: set,
    strict: bool,
) -> None:
    """Validate FQN compatibility under strict mode."""
    if not strict:
        return
    extra_fqns = source_fqns - target_fqns
    if extra_fqns:
        raise ValueError(
            f"strict=True but checkpoint contains FQNs not in target "
            f"optimizer: {sorted(extra_fqns)}"
        )

            f"optimizer: {sorted(extra_fqns)}"
        )


def _load_state_values(
    source_state: Dict[str, Any],
    param: nn.Parameter,
    full_state_dict: bool,
    cpu_offload: bool,

    cpu_offload: bool,
    stores_dtensor: bool,
) -> Dict[str, Any]:
    """Convert and load state values for a single FQN."""
    converted: Dict[str, Any] = {}
    for key, value in source_state.items():
        if isinstance(value, torch.Tensor):
            if _is_scalar_state(key):
                converted[key] = _convert_input_scalar_to_target(
                    value, param, cpu_offload, stores_dtensor,
                )
            else:
                converted[key] = _convert_input_tensor_to_target(
                    value, param, full_state_dict, cpu_offload, stores_dtensor,
                )
        else:
            converted[key] = value
    return converted


def _merge_param_groups(
    source_param_groups: List[Dict[str, Any]],
    target_raw_sd: Dict[str, Any],
    target_fqn_to_saved_id: Dict[str, int],
) -> None:

    target_raw_sd: Dict[str, Any],
    target_fqn_to_saved_id: Dict[str, int],
) -> None:
    """Merge source param_group fields into target param_groups in-place."""
    if not source_param_groups:
        return
    target_saved_id_to_fqn = {v: k for k, v in target_fqn_to_saved_id.items()}

    target_pg_by_fqns: Dict[frozenset, Dict[str, Any]] = {}
    for saved_group in target_raw_sd["param_groups"]:
        fqns_in_group = frozenset(
            target_saved_id_to_fqn.get(sid, "")
            for sid in saved_group.get("params", [])
        )
        target_pg_by_fqns[fqns_in_group] = saved_group

    for source_pg in source_param_groups:
        source_fqn_set = frozenset(source_pg.get("params", []))
        matched_target_pg = None
        for target_fqn_set, target_pg in target_pg_by_fqns.items():
            if source_fqn_set & target_fqn_set:
                matched_target_pg = target_pg
                break
        if matched_target_pg is None:
            continue
        for k, v in source_pg.items():
            if k == "params":
                continue
            matched_target_pg[k] = v


def _build_target_fqn_mappings(
    optimizer: torch.optim.Optimizer,
    model: nn.Module,
    param_by_fqn: Dict[str, nn.Parameter],
) -> Tuple[Dict[str, int], set]:

    Returns:
        Tuple of (fqn_to_saved_id, dtensor_state_ids).
    """
    target_raw_sd = optimizer.state_dict()

    target_id_to_param: Dict[int, nn.Parameter] = {}
    for runtime_group, saved_group in zip(optimizer.param_groups, target_raw_sd["param_groups"]):
        for parameter, saved_id in zip(runtime_group["params"], saved_group["params"]):
            target_id_to_param[saved_id] = parameter

    fqn_to_saved_id: Dict[str, int] = {}
    for saved_id, param in target_id_to_param.items():
        for name, p in model.named_parameters():
            if p is param:
                fqn_to_saved_id[name] = saved_id
                break

    dtensor_state_ids: set = set()
    for saved_id, state in target_raw_sd["state"].items():
        for value in state.values():
            if isinstance(value, DTensor):
                dtensor_state_ids.add(saved_id)
                break

    return fqn_to_saved_id, dtensor_state_ids


def _load_fqn_states(
    optim_state_dict: Dict[str, Any],
    target_raw_sd: Dict[str, Any],
    target_fqn_to_saved_id: Dict[str, int],
    param_by_fqn: Dict[str, nn.Parameter],

    Returns:
        Updated state dict with saved_id keys.
    """
    new_state: Dict[int, Dict[str, Any]] = dict(target_raw_sd["state"])

    for fqn, source_state in optim_state_dict.get("state", {}).items():
        if fqn not in target_fqn_to_saved_id:
            if strict:
                raise ValueError(
                    f"strict=True but FQN '{fqn}' not found in target optimizer."
                )
            continue

        target_saved_id = target_fqn_to_saved_id[fqn]
        param = param_by_fqn[fqn]

        if target_saved_id not in new_state:
            new_state[target_saved_id] = {}

        new_state[target_saved_id].update(
            _load_state_values(
                source_state, param, full_state_dict, cpu_offload,
                target_saved_id in dtensor_state_ids,
            )

                target_saved_id in dtensor_state_ids,
            )
        )

    return new_state


def set_optim_state_dict(
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    optim_state_dict: Dict[str, Any],
    *,

                broadcast from the replicate-group root to all ranks within
                the same replicate subgroup.  In PP+HSDP, each PP stage
                broadcasts independently within its own HSDP subgroup.
    """
    full_state_dict = getattr(options, "full_state_dict", False)
    cpu_offload = getattr(options, "cpu_offload", False)
    flatten = getattr(options, "flatten_optimizer_state_dict", False)
    strict = getattr(options, "strict", True)
    broadcast_from_rank0 = getattr(options, "broadcast_from_rank0", False)

    _check_chained_optimizer(optimizer)

    if flatten:
        optim_state_dict = _unflatten_optim_state_dict(optim_state_dict, model, strict=strict)

    if full_state_dict and cpu_offload and not broadcast_from_rank0:
        has_any_state = bool(optim_state_dict.get("state", {}))
        if not has_any_state:
            raise ValueError(
                "Received empty state dict with full_state_dict=True "
                "and cpu_offload=True but broadcast_from_rank0=False. "
                "Set broadcast_from_rank0=True to allow the replicate-group "
                "root to broadcast."

                "Set broadcast_from_rank0=True to allow the replicate-group "
                "root to broadcast."
            )

    if broadcast_from_rank0:
        optim_state_dict = _broadcast_state_from_rank0(
            optim_state_dict, model, full_state_dict, cpu_offload,
        )

    param_by_fqn: Dict[str, nn.Parameter] = {}
    for name, param in model.named_parameters():
        param_by_fqn[name] = param

    target_fqn_to_saved_id, dtensor_state_ids = _build_target_fqn_mappings(
        optimizer, model, param_by_fqn,
    )

    source_fqns = set(optim_state_dict.get("state", {}).keys())
    target_fqns = set(target_fqn_to_saved_id.keys())
    _check_strict_fqns(source_fqns, target_fqns, strict)

    target_raw_sd = optimizer.state_dict()
    new_state = _load_fqn_states(
        optim_state_dict, target_raw_sd, target_fqn_to_saved_id,
        param_by_fqn, dtensor_state_ids, full_state_dict, cpu_offload, strict,
    )

    target_raw_sd["state"] = new_state

    _merge_param_groups(
        optim_state_dict.get("param_groups", []),
        target_raw_sd,
        target_fqn_to_saved_id,
    )

        target_raw_sd,
        target_fqn_to_saved_id,
    )

    optimizer.load_state_dict(target_raw_sd)


def _convert_input_tensor_to_target(
    tensor: torch.Tensor,
    param: nn.Parameter,
    full_state_dict: bool,
    cpu_offload: bool,

    Returns:
        Tensor or DTensor matching the target optimizer's expected format.
    """
    dtensor_info = _get_param_dtensor_info(param)

    if not dtensor_info:
        target_device = param.data.device
        result = tensor
        if cpu_offload and result.device != target_device:
            result = result.to(target_device)
        return result

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

    from hyper_parallel.core.dtensor.dtensor import distribute_tensor  # pylint: disable=C0415

    if isinstance(tensor, DTensor):
        if tensor.device_mesh == mesh and tensor.placements == placements:
            if stores_dtensor:
                return tensor
            return tensor.to_local()
        redistributed = tensor.redistribute(mesh, placements)
        if stores_dtensor:
            return redistributed
        return redistributed.to_local()

    if full_state_dict:
        if cpu_offload:
            tensor = tensor.to(target_device)
        dt = distribute_tensor(tensor, mesh, placements)
        if stores_dtensor:
            return dt
        return dt.to_local()

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

    if stores_dtensor:
        if tensor.shape == param.to_local().shape:
            return DTensor.from_local(tensor, mesh, placements)
        dt = distribute_tensor(tensor, mesh, placements)
        return dt

    if tensor.shape == param.to_local().shape:
        return tensor

    dt = distribute_tensor(tensor, mesh, placements)
    return dt.to_local()


def _convert_input_scalar_to_target(
    tensor: torch.Tensor,
    param: nn.Parameter,
    cpu_offload: bool,
    stores_dtensor: bool = False,

    Returns:
        Plain tensor or replicated DTensor matching the target's expected format.
    """
    dtensor_info = _get_param_dtensor_info(param)

    if not dtensor_info or not stores_dtensor:
        target_device = param.data.device if not dtensor_info else param.to_local().device
        result = tensor
        if cpu_offload and result.device != target_device:
            result = result.to(target_device)
        return result

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

    scalar_placements = [Replicate()] * mesh.ndim
    if cpu_offload and tensor.device != target_device:
        tensor = tensor.to(target_device)
    return DTensor.from_local(tensor, mesh, scalar_placements)


def _get_broadcast_groups(
    model: nn.Module,
) -> List[Dict[str, Any]]:
    """Identify broadcast groups from the model's DTensor parameters.

            pg: ProcessGroup for broadcast within the group
            src_rank: global rank of the replicate-group root
            is_root: True if current rank is the root of this group
    """
    param_by_fqn: Dict[str, nn.Parameter] = {}
    for name, param in model.named_parameters():
        param_by_fqn[name] = param

    model_fqns = list(param_by_fqn.keys())

    if not model_fqns:
        return []

    first_param = param_by_fqn[model_fqns[0]]
    dtensor_info = _get_param_dtensor_info(first_param)

    if dtensor_info is not None:
        mesh, placements = dtensor_info
        is_root = _is_replicate_group_root(mesh, placements)
        replicate_dims = [i for i, p in enumerate(placements) if isinstance(p, Replicate)]
        if replicate_dims:
            pg = mesh.get_group(replicate_dims[0])
            coord = mesh.get_coordinate()
            if coord is not None:
                root_coord = list(coord)
                for d in replicate_dims:
                    root_coord[d] = 0
                src_rank = int(mesh.mesh[tuple(root_coord)])
            else:
                src_rank = 0
        else:
            pg = dist.group.WORLD
            src_rank = 0
    else:
        is_root = dist.get_rank() == 0
        pg = dist.group.WORLD
        src_rank = 0

    return [{
        "fqns": model_fqns,
        "pg": pg,
        "src_rank": src_rank,
        "is_root": is_root,

        "param_by_fqn": param_by_fqn,
    }]


def _broadcast_fqn_list(
    is_root: bool,
    fqns: List[str],
    pg: dist.ProcessGroup,
    src_rank: int,

    pg: dist.ProcessGroup,
    src_rank: int,
) -> List[str]:
    """Broadcast FQN list from root to all ranks in the replicate group."""
    fqn_list = fqns if is_root else []
    obj = [fqn_list]
    dist.broadcast_object_list(obj, src=src_rank, group=pg)
    return obj[0]


def _build_state_schema(
    state: Dict[str, Any],
) -> Dict[str, Any]:
    """Build schema dict describing the types and shapes of state entries."""
    schema: Dict[str, Any] = {}
    for key, value in state.items():
        if isinstance(value, torch.Tensor):
            schema[key] = {
                "shape": tuple(value.shape),
                "dtype": str(value.dtype),
                "is_scalar": _is_scalar_state(key),
            }

                "dtype": str(value.dtype),
                "is_scalar": _is_scalar_state(key),
            }
        else:
            schema[key] = {"type": type(value).__name__, "value": value}
    return schema


def _broadcast_schema_per_fqn(
    fqn_list: List[str],
    optim_state_dict: Dict[str, Any],
    is_root: bool,
    pg: dist.ProcessGroup,

    pg: dist.ProcessGroup,
    src_rank: int,
) -> Dict[str, Dict[str, Any]]:
    """Broadcast schema for each FQN from root to all ranks."""
    fqn_schema: Dict[str, Dict[str, Any]] = {}
    for fqn in fqn_list:
        if is_root and fqn in optim_state_dict.get("state", {}):
            schema = _build_state_schema(optim_state_dict["state"][fqn])
        else:
            schema = {}
        schema_list = [schema] if is_root else [None]
        dist.broadcast_object_list(schema_list, src=src_rank, group=pg)
        fqn_schema[fqn] = schema_list[0]
    return fqn_schema


def _get_local_device() -> torch.device:
    """Get the device for the current local rank."""
    local_rank = dist.get_rank() if dist.is_initialized() else 0
    if torch.npu.is_available():
        return torch.device(f"npu:{local_rank}")
    return torch.device(f"cuda:{local_rank}")


def _broadcast_scalar_entry(
    key: str,
    fqn: str,
    info: Dict[str, Any],
    optim_state_dict: Dict[str, Any],

    pg: dist.ProcessGroup,
    cpu_offload: bool,
) -> torch.Tensor:
    """Broadcast a scalar state entry and return the result tensor."""
    dtype_str = info.get("dtype", "torch.float32")
    dtype = _resolve_dtype(dtype_str)
    device = _get_local_device()

    if is_root and fqn in optim_state_dict.get("state", {}):
        scalar_val = optim_state_dict["state"][fqn][key]
        t = scalar_val.clone() if isinstance(scalar_val, torch.Tensor) else torch.tensor(scalar_val)
        if t.dim() == 0:
            t = t.reshape(1)
    else:
        t = torch.zeros(1, dtype=dtype)

    if t.device.type == "cpu":
        t = t.to(device)
    dist.broadcast(t, src=src_rank, group=pg)
    return t.reshape(()).cpu() if cpu_offload else t.reshape(())


def _broadcast_tensor_entry(
    key: str,
    fqn: str,
    info: Dict[str, Any],
    optim_state_dict: Dict[str, Any],

    full_state_dict: bool,
    cpu_offload: bool,
) -> torch.Tensor:
    """Broadcast a tensor state entry and return the result tensor."""
    shape = info.get("shape", ())
    dtype_str = info.get("dtype", "torch.float32")
    dtype = _resolve_dtype(dtype_str)
    device = _get_local_device()

    if is_root and fqn in optim_state_dict.get("state", {}):
        src_tensor = optim_state_dict["state"][fqn][key]
        if src_tensor.device.type == "cpu":
            src_tensor = src_tensor.to(device)
    else:
        src_tensor = torch.zeros(shape, dtype=dtype, device=device)

    dist.broadcast(src_tensor, src=src_rank, group=pg)

    if full_state_dict and cpu_offload:
        return src_tensor.cpu() if is_root else src_tensor
    return src_tensor.cpu() if cpu_offload else src_tensor


def _broadcast_tensor_data_per_fqn(
    fqn_list: List[str],
    fqn_schema: Dict[str, Dict[str, Any]],
    optim_state_dict: Dict[str, Any],
    is_root: bool,

    full_state_dict: bool,
    cpu_offload: bool,
) -> Dict[str, Dict[str, Any]]:
    """Broadcast tensor data for each FQN from root to all ranks."""
    result_state: Dict[str, Dict[str, Any]] = {}
    for fqn in fqn_list:
        schema = fqn_schema[fqn]
        if not schema:
            continue
        result_state[fqn] = {}
        for key, info in schema.items():
            is_scalar = info.get("is_scalar", False)
            if is_scalar:
                result_state[fqn][key] = _broadcast_scalar_entry(
                    key, fqn, info, optim_state_dict,
                    is_root, src_rank, pg, cpu_offload,
                )
            else:
                result_state[fqn][key] = _broadcast_tensor_entry(
                    key, fqn, info, optim_state_dict,
                    is_root, src_rank, pg, full_state_dict, cpu_offload,
                )
    return result_state


def _broadcast_state_from_rank0(
    optim_state_dict: Dict[str, Any],
    model: nn.Module,
    full_state_dict: bool,
    cpu_offload: bool,

    subgroup), so different PP stages broadcast independently without
    cross-stage interference.  This avoids the problem where global
    rank 0's FQN list does not include other stages' parameters.
    """
    if not dist.is_initialized():
        return optim_state_dict

    groups = _get_broadcast_groups(model)

    if not groups:
        return {
            "state": {},
            "param_groups": optim_state_dict.get("param_groups", []),
        }

    group = groups[0]
    pg = group["pg"]
    src_rank = group["src_rank"]
    is_root = group["is_root"]
    fqns = group["fqns"]

    fqn_list = _broadcast_fqn_list(is_root, fqns, pg, src_rank)
    fqn_schema = _broadcast_schema_per_fqn(
        fqn_list, optim_state_dict, is_root, pg, src_rank,
    )
    result_state = _broadcast_tensor_data_per_fqn(
        fqn_list, fqn_schema, optim_state_dict,
        is_root, src_rank, pg, full_state_dict, cpu_offload,
    )

    return {
        "state": result_state,
        "param_groups": optim_state_dict.get("param_groups", []),
    }

        "param_groups": optim_state_dict.get("param_groups", []),
    }


def _flatten_optim_state_dict(
    state_dict: Dict[str, Any],
) -> Dict[str, Any]:
    """Flatten nested state dict to single-level dot-separated keys.

    optimizer state dicts using ``no_dist=True`` (see P6 test pattern),
    rather than using a coordinated DCP save with flatten format across
    all ranks.
    """
    flat: Dict[str, Any] = {}

    for fqn, state in state_dict.get("state", {}).items():
        for key, value in state.items():
            flat[f"state.{fqn}.{key}"] = value

    for group_idx, group in enumerate(state_dict.get("param_groups", [])):
        params_list = group.get("params", [])
        if not params_list:
            raise UnsupportedConfigurationError(
                f"Cannot flatten param_groups[{group_idx}]: empty 'params' list. "
                f"Empty param groups cannot be represented in flatten format "
                f"because there is no FQN to prefix the group fields. "
                f"Use non-flatten format or provide a stable group_name."

                f"Empty param groups cannot be represented in flatten format "
                f"because there is no FQN to prefix the group fields. "
                f"Use non-flatten format or provide a stable group_name."
            )
        for fqn in params_list:
            for key, value in group.items():
                if key == "params":
                    continue
                flat[f"param_group.{fqn}.{key}"] = value

    return flat


def _check_inconsistent_pg_fields(
    existing: Dict[str, Any],
    incoming: Dict[str, Any],
    fqn: str,
    strict: bool,

    fqn: str,
    strict: bool,
) -> None:
    """Check and raise/warn on inconsistent param_group fields."""
    common_keys = set(existing.keys()) & set(incoming.keys())
    for k in common_keys:
        if existing[k] != incoming[k]:
            if strict:
                raise ValueError(
                    f"strict=True but param_group field '{k}' is "
                    f"inconsistent within the same group: "
                    f"existing={existing[k]!r}, "
                    f"incoming(from {fqn})={incoming[k]!r}"

                    f"inconsistent within the same group: "
                    f"existing={existing[k]!r}, "
                    f"incoming(from {fqn})={incoming[k]!r}"
                )
            logger.warning(
                "param_group field '%s' is inconsistent within the "
                "same group: existing=%r, incoming(from %s)=%r. "
                "Keeping existing value (strict=False).",
                k, existing[k], fqn, incoming[k],

                k, existing[k], fqn, incoming[k],
            )


def _parse_unflatten_entries(
    flat_dict: Dict[str, Any],
    known_fqns_sorted: List[str],
) -> Tuple[Dict[str, Dict[str, Any]], Dict[str, Dict[str, Any]]]:
    """Parse flat dict entries into state and param_group_fields dicts.

    Returns:
        Tuple of (state, param_group_fields).
    """
    state: Dict[str, Dict[str, Any]] = {}
    param_group_fields: Dict[str, Dict[str, Any]] = {}

    for key, value in flat_dict.items():
        if key.startswith("state."):
            remainder = key[len("state."):]
            fqn = _match_fqn_from_remainder(remainder, known_fqns_sorted)
            if fqn is None:
                raise ValueError(
                    f"Cannot match FQN from flat key '{key}'. "
                    f"Known FQNs: {known_fqns_sorted}"
                )
            state_key = remainder[len(fqn) + 1:]
            state.setdefault(fqn, {})[state_key] = value

        elif key.startswith("param_group."):
            remainder = key[len("param_group."):]
            fqn = _match_fqn_from_remainder(remainder, known_fqns_sorted)
            if fqn is None:
                raise ValueError(
                    f"Cannot match FQN from flat key '{key}'. "
                    f"Known FQNs: {known_fqns_sorted}"
                )
            field_name = remainder[len(fqn) + 1:]
            param_group_fields.setdefault(fqn, {})[field_name] = value

    return state, param_group_fields


def _assemble_param_groups(
    known_fqns: List[str],
    state: Dict[str, Dict[str, Any]],
    param_group_fields: Dict[str, Dict[str, Any]],
    strict: bool,

    param_group_fields: Dict[str, Dict[str, Any]],
    strict: bool,
) -> List[Dict[str, Any]]:
    """Assemble param_groups list from FQN-based state and fields."""
    param_groups: List[Dict[str, Any]] = []
    current_group: Dict[str, Any] = {"params": []}

    for fqn in known_fqns:
        if fqn in param_group_fields or fqn in state:
            fields = param_group_fields.get(fqn, {})
            if not current_group["params"]:
                current_group["params"].append(fqn)
                current_group.update(fields)
            else:
                existing_fields = {
                    k: v for k, v in current_group.items() if k != "params"
                }
                if existing_fields == fields or not fields:
                    current_group["params"].append(fqn)
                else:
                    _check_inconsistent_pg_fields(existing_fields, fields, fqn, strict)
                    current_group["params"].append(fqn)

    if current_group["params"]:
        param_groups.append(current_group)

    if not param_groups:
        raise UnsupportedConfigurationError(
            "Cannot unflatten: empty param_group in flatten format. "
            "Provide a stable group_name or use non-flatten format."
        )

    return param_groups


def _unflatten_optim_state_dict(
    flat_dict: Dict[str, Any],
    model: nn.Module,
    strict: bool = True,
) -> Dict[str, Any]: