"""Module for setting optimizer states in a PyTorch model."""

import copy
import time
from collections import defaultdict
from collections.abc import Hashable, Iterable
from functools import partial
from itertools import chain, starmap
from logging import WARNING
from typing import Any, cast

import torch
from composer.utils import dist
from flwr.common import log
from torch.distributed.checkpoint.state_dict import (
    ListDictValueType,
    OptimizerStateType,
    StateDictOptions,
    _init_optim_state,  # noqa: PLC2701
    set_optimizer_state_dict,
)
from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel
from torch.optim import Optimizer
from torch.optim.optimizer import StateDict

from mt_dao.module_utils import collect_fsdp_modules, is_leaf_fsdp_module
from mt_dao.utils import match_shapes


def set_optimizers_states(
    model: torch.nn.Module,
    optimizer: Optimizer,
    step: int | None,
    opt_state_dict: dict[str, dict[str, torch.Tensor]],
) -> None:
    """Set the optimizer states for the model.

    This function sets the optimizer states for the model, handling both FSDP and
    non-FSDP models. It collects the FSDP modules from the model, initializes the
    optimizer state, and sets the optimizer state for each parameter in the model.
    If the model is wrapped in FSDP, it collects the leaf FSDP modules and sets the
    optimizer state for each parameter in those modules.

    Parameters
    ----------
    model : torch.nn.Module
        The model for which the optimizer states are being set. It can be a regular
        PyTorch model or a model wrapped in Fully Sharded Data Parallel (FSDP).
    optimizer : Optimizer
        The optimizer instance whose state is being updated. It should be compatible
        with the model's parameters.
    step : int
        The current training step, which is used to set the optimizer's step count.
    opt_state_dict : dict[str, dict[str, torch.Tensor]]
        A dictionary containing:
        - Keys as parameter names (e.g., 'linear', 'wte')
        - Keys as momentum names (e.g., 'exp_avg', 'exp_avg_sq')
        - Values as dictionaries mapping momentum types to their corresponding tensors.

    """
    # Get the model from the trainer
    # Map to store FSDP modules by their name
    fsdp_modules: dict[str, FullyShardedDataParallel] = {}
    non_fsdp_modules: dict[str, torch.nn.Module] = {}

    # Collect all FSDP modules in the model hierarchy
    if hasattr(model, "model") and isinstance(
        model.model,
        FullyShardedDataParallel,
    ):
        # First pass: collect all FSDP modules
        (fsdp_modules, non_fsdp_modules) = collect_fsdp_modules(
            "model",
            model.model,
            (fsdp_modules, non_fsdp_modules),
        )

    _init_optim_state(optimizer)
    # Handle the case where the model is wrapped in FSDP
    if fsdp_modules:
        # Get the leaf FSDP modules from the model
        leaf_fsdp_modules = {
            module_name: module
            for module_name, module in fsdp_modules.items()
            if is_leaf_fsdp_module(module_name, (fsdp_modules, non_fsdp_modules))
        }
        # Set the optimizer state for the FSDP modules. Add to the buffers the
        # information regarding the sharded states
        modules_buffer: dict[str, torch.nn.Module] = {}
        params_leftovers: dict[str, torch.nn.Parameter | torch.Tensor] = {}
        momenta_buffer: dict[str, dict[str, torch.Tensor]] = defaultdict(dict)
        # Initialize the delta for this momentum type
        for local_module_name, local_module in leaf_fsdp_modules.items():
            (c_modules_buffer, c_params_leftovers, c_momenta_buffer) = (
                set_optimizer_state_fsdp(
                    optim=optimizer,
                    local_module=local_module,
                    local_module_name=local_module_name,
                    momentum_dict=opt_state_dict,
                    step=step,
                )
            )
            modules_buffer.update(c_modules_buffer)
            params_leftovers.update(c_params_leftovers)
            for c_momentum_param_name, c_momentum_dict in c_momenta_buffer.items():
                momenta_buffer[c_momentum_param_name].update(c_momentum_dict)
        # Set the sharded states using the state dict
        set_optim_with_partial_state_dict(
            model=model,
            optimizer=optimizer,
            step_value=step,
            momenta_buffer=momenta_buffer,
        )
    else:
        for local_param_name, local_param in model.named_parameters():
            # Set the optimizer state for the current parameter
            set_optimizer_state_base(
                optimizer=optimizer,
                local_param=local_param,
                local_param_name=local_param_name,
                step=step,
                momentum_type_dict=opt_state_dict,
            )
    dist.barrier()


def set_optimizer_state_base(
    optimizer: torch.optim.Optimizer,
    local_param: torch.nn.Parameter | torch.Tensor,
    local_param_name: str,
    step: int | None,
    momentum_type_dict: dict[str, dict[str, torch.Tensor]],
) -> None:
    """Set the optimizer state for a specific parameter.

    This function updates the optimizer state for a given parameter, setting its step
    count and momentum values. It handles name transformations to find the correct
    parameter in the provided momentum dictionaries, performs shape matching between
    remote and local momentum tensors, and updates the optimizer state accordingly.

    Parameters
    ----------
    optimizer : torch.optim.Optimizer
        The optimizer instance whose state is being updated.
    local_param : torch.nn.Parameter | torch.Tensor
        The local parameter or tensor whose optimizer state should be updated.
    local_param_name : str
        The name of the parameter in the model's state_dict.
    step : int
        The step count to set in the optimizer state.
    momentum_type_dict : dict[str, dict[str, torch.Tensor]]
        A dictionary mapping momentum types (e.g., 'exp_avg', 'exp_avg_sq') to
        dictionaries that map parameter names to their corresponding momentum values.

    Notes
    -----
    The function handles common prefixes like 'model.', 'module.', and '_fsdp_wrapped_'
    that might be present in parameter names. It also manages shape differences between
    local and remote momentum tensors using the match_shapes function.

    """
    if "step" not in optimizer.state[local_param]:
        log(WARNING, "Step not found in optimizer state for %s", local_param_name)
    if step is not None:
        # Set the step to the current server steps
        optimizer.state[local_param]["step"].data.copy_(
            torch.as_tensor(
                step,
                device=optimizer.state[local_param]["step"].device,
                dtype=optimizer.state[local_param]["step"].dtype,
            ),
        )
    # TODO(anon): Handle modules prefixes globally
    # Remove potential DDP-related prefixes and the "module." prefix
    lookup_name = (
        local_param_name.replace("model.", "")
        .replace(
            "module.",
            "",
        )
        .replace(
            "_fsdp_wrapped_",
            "",
        )
    )
    # Check if the parameter is in the momentum dictionary
    for momentum_type, momentum_dict in momentum_type_dict.items():
        # Get the new momentum value
        new_momentum = momentum_dict[lookup_name].to(
            device=local_param.device,
            dtype=local_param.dtype,
        )
        local_tensor: torch.Tensor = optimizer.state[local_param][momentum_type.lower()]
        # Reshape the remote tensor to match the local parameter's shape
        new_momentum = match_shapes(new_momentum, local_tensor)
        assert new_momentum.shape == local_tensor.shape, (
            f"Shape mismatch for {momentum_type} of {local_param_name}: "
            f"expected {local_tensor.shape}, got {new_momentum.shape}"
        )
        # Set the momentum value
        local_tensor.data.copy_(
            new_momentum.to(device=local_tensor.device, dtype=local_tensor.dtype),
        )


def set_optim_with_partial_state_dict(
    model: torch.nn.Module,
    optimizer: Optimizer,
    momenta_buffer: dict[str, dict[str, torch.Tensor]],
    step_value: int | None,
) -> float:
    """Set optimizer states using a partial state dictionary for FSDP models.

    This function creates and applies a partial optimizer state dictionary to update
    optimizer states for parameters in Fully Sharded Data Parallel (FSDP) models. It's
    designed to handle cases where regular state dictionary loading would fail due to
    complex sharding patterns. The function temporarily replaces the optimizer's
    load_state_dict method with a custom partial loading function to ensure proper
    state updates across all workers.

    Parameters
    ----------
    model : torch.nn.Module
        The model instance whose optimizer states are being updated.
    optimizer : Optimizer
        The optimizer instance whose state should be updated.
    momenta_buffer : dict[str, dict[str, torch.Tensor]]
        A nested dictionary mapping parameter names to dictionaries of momentum values.
        The inner dictionaries map momentum types (e.g., 'exp_avg', 'exp_avg_sq') to
        their corresponding tensor values.
    step_value : int
        The step count to set in the optimizer state for each parameter.

    Returns
    -------
    float
        The time in seconds taken to set the optimizer states.

    Notes
    -----
    This function works with the `set_optimizer_state_fsdp` function to handle complex
    parameter sharding in FSDP models. It uses a custom state dictionary loading
    approach by temporarily replacing the optimizer's load_state_dict method with the
    custom `load_partial_state_dict` function.

    """
    start_time_to_set = time.perf_counter()
    state: dict[str, dict[str, dict[str, torch.Tensor]] | list[dict[str, Any]]] = {}
    if dist.get_global_rank() == 0:
        my_state: dict[str, dict[str, torch.Tensor]] = defaultdict(dict)
        for param_name, momenta_dict in momenta_buffer.items():
            real_param_name = f"model.{param_name}"
            for momentum_type, momentum_tensor in momenta_dict.items():
                my_state[real_param_name] |= {
                    momentum_type.lower(): momentum_tensor,
                }
        if step_value is not None:
            for key in my_state:
                my_state[str(key)] |= {
                    "step": torch.as_tensor(step_value, dtype=torch.float32),
                }
        list_of_parameter_name = list(my_state.keys())
        local_param_groups: list[dict[str, Any]] = []
        for pg in optimizer.param_groups:
            local_param_group_dict = {k: v for k, v in pg.items() if k != "params"}
            # Since different parameter groups may exist, we need to filter
            # the parameters in the current group to only those that are in the
            # list of parameter names we have in the state
            pg_list_of_param_names = [
                param_name
                for param_name in list_of_parameter_name
                if param_name in pg["params"]
            ]
            local_param_group_dict["params"] = pg_list_of_param_names
            local_param_groups.append(local_param_group_dict)
        state["param_groups"] = cast("ListDictValueType", local_param_groups)
        state["state"] = my_state
    old_load_state_dict_fn = copy.deepcopy(optimizer.load_state_dict)
    optimizer.load_state_dict = partial(_load_partial_state_dict, optimizer)
    set_optimizer_state_dict(
        optim_state_dict=cast("OptimizerStateType", state),
        model=model,
        optimizers=optimizer,
        options=StateDictOptions(
            full_state_dict=True,
            cpu_offload=True,
            strict=False,
        ),
    )
    optimizer.load_state_dict = old_load_state_dict_fn
    return time.perf_counter() - start_time_to_set


@torch._disable_dynamo  # noqa: SLF001  # type: ignore[reportPrivateImportUsage]
def _load_partial_state_dict(optimizer: Optimizer, state_dict: StateDict) -> None:
    """Load a partial state dictionary into an optimizer.

    This function is a specialized version of the PyTorch optimizer's load_state_dict
    method that supports loading a state dictionary that contains only a subset of the
    optimizer's parameters. This is particularly useful with FSDP models where different
    parameters may be sharded across different processes, and the standard
    load_state_dict method might fail due to shape mismatches or missing parameters.

    The function handles parameter mapping between the saved state dictionary and the
    current optimizer state, preserving the optimizer's existing parameter structure
    while updating state values where available in the state dictionary.

    Parameters
    ----------
    optimizer : Optimizer
        The optimizer instance whose state should be partially updated.
    state_dict : StateDict
        The state dictionary containing the subset of optimizer states to load.
        Must have 'param_groups' and 'state' keys conforming to PyTorch's
        optimizer state dict format.

    Notes
    -----
    This function is used by the `set_optim_with_partial_state_dict` function to
    support updating optimizer states in complex sharding patterns when working with
    FSDP models. It temporarily replaces the optimizer's native load_state_dict method
    during state updates.

    Unlike the standard PyTorch optimizer's load_state_dict method, this function
    uses `_create_partial_id_map` to establish a mapping between parameter IDs in the
    saved state dictionary and the current optimizer parameters, allowing for partial
    updates where only a subset of parameters have state information.

    """
    # shallow copy, to be consistent with module API
    state_dict = state_dict.copy()

    for pre_hook in optimizer._optimizer_load_state_dict_pre_hooks.values():  # noqa: SLF001
        hook_result = pre_hook(optimizer, state_dict)
        if hook_result is not None:
            state_dict = hook_result

    # Validate the state_dict
    groups = optimizer.param_groups

    # Deepcopy as we write into saved_groups later to update state
    saved_groups = copy.deepcopy(state_dict["param_groups"])
    id_map = _create_partial_id_map(saved_groups, groups)

    # Copy state assigned to params (and cast tensors to appropriate types).
    # State that is not assigned to params is copied as is (needed for
    # backward compatibility).
    state: defaultdict[torch.Tensor, dict[Any, Any]] = defaultdict(dict)
    for k, v in state_dict["state"].items():
        if k in id_map:
            param = id_map[k]
            state[param] = _cast(  # type: ignore[assignment, reportArgumentType]
                param,
                v,
                param_id=k,
                param_groups=state_dict["param_groups"],
            )
        else:
            state[k] = v

    # Update parameter groups, setting their 'params' value
    def update_group(
        group: dict[str, Any],
        new_group: dict[str, Any],
    ) -> dict[str, Any]:
        new_group["params"] = group["params"]
        if "param_names" in group and "param_names" not in new_group:
            new_group["param_names"] = group["param_names"]
        # Copy all other keys since we are overwrite the parameter groups
        for key, value in group.items():
            if key not in new_group:
                new_group[key] = value
        return new_group

    param_groups = list(starmap(update_group, zip(groups, saved_groups, strict=False)))
    optimizer.__setstate__({"param_groups": param_groups})
    _recurse_set_optimizer_state_dict(optimizer, state)
    for post_hook in optimizer._optimizer_load_state_dict_post_hooks.values():  # noqa: SLF001
        post_hook(optimizer)


def _recurse_set_optimizer_state_dict(
    optimizer: Optimizer,
    state: defaultdict[torch.Tensor, dict[Any, Any]],
) -> None:
    """Recursively update the optimizer's state dictionary with new values.

    This function iterates over the optimizer's internal state and, for each parameter
    key present in both the optimizer's state and the provided state dictionary,
    updates the optimizer's state values with those from the provided state.

    Parameters
    ----------
    optimizer : Optimizer
        The optimizer whose state dictionary will be updated.
    state : defaultdict[torch.Tensor, dict[Any, Any]]
        A dictionary mapping parameters to their new state values.

    Notes
    -----
    This function is typically used as part of the optimizer state loading process,
    especially when handling partial or custom state dictionaries for distributed
    or sharded training scenarios.

    """
    for k in optimizer.__dict__["state"]:
        if k in state:
            # If the state is already set, we need to update it
            # with the new values from the state_dict
            for key, value in state[k].items():
                optimizer.__dict__["state"][k][key] = value


def _create_partial_id_map(
    remote_groups: list[dict[str, Any]],
    groups: list[dict[str, Any]],
) -> dict:
    """Create a mapping between parameter IDs in saved groups and original groups.

    This function creates a dictionary mapping between parameter IDs in remote_groups
    and a subset of parameter IDs in groups. For each group in remote_groups, it uses
    the indices in the "params" field to select which parameters to keep from the
    corresponding group in groups, then builds a mapping between the original parameter
    IDs and the selected ones.

    Parameters
    ----------
    remote_groups : list[dict[str, Any]]
        A list of dictionaries containing parameter indices in their "params" field.
        These indices indicate which parameters to keep from the original groups.
    groups : list[dict[str, Any]]
        A list of dictionaries containing the original parameter IDs in their "params"
        field.

    Returns
    -------
    dict
        A dictionary mapping parameter IDs from remote_groups to the selected parameter
        IDs from groups.

    Raises
    ------
    ValueError
        If the number of parameter groups in remote_groups does not match the number
        of parameter groups in groups. This ensures that the loaded state dictionary
        is compatible with the local optimizer's parameter structure.

    """
    if len(remote_groups) != len(groups):
        msg = (
            "The loaded state dict has %s of parameter groups,"
            " while the local optimizer has %s."
        )
        raise ValueError(msg % (len(remote_groups), len(groups)))
    selected_groups: list[dict[str, Any]] = []
    parameters_counter = 0
    for r_g, g in zip(remote_groups, groups, strict=True):
        indeces_to_keep = r_g["params"]
        # Since the indexing is global but we need to access a local list, we need
        # to translate the index to the local list of parameters
        new_params = [g["params"][idx - parameters_counter] for idx in indeces_to_keep]
        selected_groups.append({"params": new_params})
        parameters_counter += len(g["params"])
    return dict(
        zip(
            chain.from_iterable(g["params"] for g in remote_groups),
            chain.from_iterable(g["params"] for g in selected_groups),
            strict=True,
        ),
    )


def _cast(
    param: torch.Tensor,
    value: torch.Tensor | dict | Iterable,
    param_id: int,
    param_groups: list[dict[Any, Any]],
    key: Hashable = None,
) -> torch.Tensor | dict | Iterable:
    """Process optimizer state values according to parameter policies.

    This function recursively processes optimizer state values to ensure they are
    correctly cast and transformed according to the parameter's policy. It handles
    tensors, dictionaries, and other iterables by applying appropriate transformations
    to each element or nested structure.

    Parameters
    ----------
    param : torch.Tensor
        The parameter tensor associated with the optimizer state.
    value : torch.Tensor | dict | Iterable
        The value to process, which can be a tensor, dictionary, or other iterable.
    param_id : int
        The ID of the parameter in the optimizer's state.
    param_groups : list[dict[Any, Any]]
        List of parameter groups from the optimizer.
    key : Hashable, optional
        An optional key used when processing dictionary values, by default None.

    Returns
    -------
    torch.Tensor | dict | Iterable
        The processed value with the same structure as the input but potentially
        modified values according to the parameter policy.

    Notes
    -----
    This is a helper function primarily used by `load_partial_state_dict` to handle
    the casting of optimizer state values when loading a partial state dictionary.
    For tensors, it delegates to PyTorch's internal
    `_process_value_according_to_param_policy`. For dictionaries and iterables, it
    recursively processes each element.

    """
    if isinstance(value, torch.Tensor):
        return Optimizer._process_value_according_to_param_policy(  # noqa: SLF001
            param,
            value,
            param_id,
            param_groups,
            key,
        )
    if isinstance(value, dict):
        return {
            k: _cast(
                param,
                v,
                param_id=param_id,
                param_groups=param_groups,
                key=k,
            )
            for k, v in value.items()
        }
    if isinstance(value, Iterable):
        return type(value)(
            _cast(param, v, param_id=param_id, param_groups=param_groups)
            for v in value  # type: ignore[call-arg, reportCallIssue]
        )
    return None


def set_optimizer_state_fsdp(
    optim: torch.optim.Optimizer,
    local_module: FullyShardedDataParallel,
    local_module_name: str,
    momentum_dict: dict[str, dict[str, torch.Tensor]],
    step: int | None,
) -> tuple[
    dict[str, torch.nn.Module],
    dict[str, torch.nn.Parameter | torch.Tensor],
    dict[str, dict[str, torch.Tensor]],
]:
    """Set optimizer states for parameters in an FSDP module.

    This function updates optimizer states for parameters in a Fully Sharded Data
    Parallel (FSDP) module. It handles the complexity of accessing sharded parameters
    by using FSDP's context managers to temporarily gather full parameters across
    workers. For each parameter, it calls the base setter function to update its
    optimizer state. Parameters that couldn't be directly updated (typically due to
    complex sharding) are collected into buffer dictionaries for later processing
    through alternative methods.

    Parameters
    ----------
    optim : torch.optim.Optimizer
        The optimizer instance whose states are being updated.
    local_module : FullyShardedDataParallel
        The FSDP module instance containing parameters whose optimizer states should be
        updated.
    local_module_name : str
        The name of the FSDP module in the model hierarchy.
    momentum_dict : dict[str, dict[str, torch.Tensor]]
        A nested dictionary mapping parameter names to dictionaries of momentum values.
        The inner dictionaries map momentum types (e.g., 'exp_avg', 'exp_avg_sq') to
        their corresponding values as PyTorch Tensors.
    step : int
        The step count to set in the optimizer state for each parameter.

    Returns
    -------
    tuple[
        dict[str, torch.nn.Module],
        dict[str, torch.nn.Parameter | torch.Tensor],
        dict[str, dict[str, torch.Tensor]],
    ]
        A tuple containing three dictionaries:
        - A mapping of parameter names to their parent modules
        - A mapping of parameter names to parameter tensors
        - A mapping of parameter names to dictionaries of momentum tensors
        These dictionaries contain information about parameters that couldn't be
        directly updated and need special handling through other methods.

    """
    # Initialize the leftover buffer for parameters that are sharded across multiple
    # ranks not with full shards
    modules_buffer: dict[str, torch.nn.Module] = {}
    params_leftovers: dict[str, torch.nn.Parameter | torch.Tensor] = {}
    momenta_buffer: dict[str, dict[str, torch.Tensor]] = defaultdict(dict)
    # Open the FSDP context manager to summon the parameters that we need to check to
    # acknowledge the sharding structure
    with FullyShardedDataParallel.summon_full_params(
        local_module,
        writeback=True,
        recurse=False,
    ):
        for (
            internal_local_module_name,
            internal_local_module,
        ) in local_module.named_modules():
            # Iterate over the named parameters of the current internal local module
            for local_param_name, local_param in internal_local_module.named_parameters(
                recurse=False,
            ):
                # Call the base setter acting on a single parameter
                set_optimizer_state_base_fsdp(
                    optim=optim,
                    local_names=(
                        local_module_name,
                        internal_local_module_name,
                        local_param_name,
                    ),
                    buffers=(
                        modules_buffer,
                        params_leftovers,
                        momenta_buffer,
                    ),
                    local_objects=(local_module, local_param),
                    remote_objects=(momentum_dict, step),
                )
    return (
        modules_buffer,
        params_leftovers,
        momenta_buffer,
    )


def set_optimizer_state_base_fsdp(
    optim: torch.optim.Optimizer,
    local_names: tuple[str, str, str],
    buffers: tuple[
        dict[str, torch.nn.Module],
        dict[str, torch.nn.Parameter | torch.Tensor],
        dict[str, dict[str, torch.Tensor]],
    ],
    local_objects: tuple[FullyShardedDataParallel, torch.nn.Parameter | torch.Tensor],
    remote_objects: tuple[dict[str, dict[str, torch.Tensor]], int | None],
) -> None:
    """Set optimizer state for a specific parameter in an FSDP model.

    This function sets optimizer state (like momentum values) for a specific parameter
    in a Fully Sharded Data Parallel (FSDP) model. It handles the complexity of matching
    remote state values to local parameters, accounting for potential sharding across
    multiple workers. If direct application isn't possible due to shape mismatches,
    it collects parameters into buffers for later processing through alternative means.

    Parameters
    ----------
    optim : torch.optim.Optimizer
        The optimizer instance whose state is being updated.
    local_names : tuple[str, str, str]
        A tuple containing three string identifiers:
        - The local module name
        - The internal local module name
        - The local parameter name
    buffers : tuple[
        dict[str, torch.nn.Module],
        dict[str, torch.nn.Parameter | torch.Tensor],
        dict[str, dict[str, torch.Tensor]],
    ]
        A tuple of three dictionaries to collect parameters that need special handling:
        - A mapping of parameter names to their parent modules
        - A mapping of parameter names to parameter tensors
        - A mapping of parameter names to dictionaries of momentum tensors
    local_objects : tuple[FullyShardedDataParallel, torch.nn.Parameter | torch.Tensor]
        A tuple containing:
        - The local FSDP module
        - The local parameter or tensor whose optimizer state should be updated
    remote_objects : tuple[dict[str, dict[str, torch.Tensor]], int]
        A tuple containing:
        - A dictionary mapping parameter names to dictionaries of momentum values
        - The step count to set in the optimizer state

    Notes
    -----
    The function handles common FSDP name prefixes and transformations to ensure proper
    matching between remote and local parameters. When direct state updates aren't
    possible due to sharding, the function adds relevant information to the provided
    buffers for later processing by higher-level functions.

    """
    # Unpack remote objects
    opt_states, step = remote_objects
    # Unpack local objects
    local_module, local_param = local_objects
    # Unpack buffers
    modules_buffer, params_leftovers, momenta_buffer = buffers
    # Unpack local names
    local_module_name, internal_local_module_name, local_param_name = local_names
    # Construct the remote parameter name by removing the "model." and FSDP prefixes and
    # by adding the local module name, the internal local module name, and the local
    # parameter name
    # TODO(anon): Handle modules prefixes globally
    remote_param_name = (
        (
            (local_module_name + "." + internal_local_module_name)
            .replace("model.", "")
            .replace("_fsdp_wrapped_module.", "")
            .replace("._fsdp_wrapped_module", "")
            .replace("._checkpoint_wrapped_module", "")
        )
        + "."
        + local_param_name
    )
    # If the local parameter is empty, skip the setting since this shard is not
    # available in this worker
    if local_param.shape == torch.Size([0]):
        return
    # Loop over the momentum types in the remote momentum dictionary
    for momentum_type, momentum_dict in opt_states.items():
        # Get the remote parameter from the OrderedDict
        momentum_state = momentum_dict.pop(remote_param_name)
        # Get the local momentum tensor
        local_opt: torch.Tensor = optim.state[local_param][momentum_type.lower()]
        # Try to match the shapes of the two tensors
        remote_opt = match_shapes(momentum_state, local_opt)
        # If the shapes don't match it means that it is sharded across multiple
        # workers and we need to use another method to set it
        if remote_opt.shape != local_opt.shape:
            # Fill the buffers to be used by a second method later
            modules_buffer[remote_param_name] = local_module
            params_leftovers[remote_param_name] = local_param
            momenta_buffer[remote_param_name][momentum_type.lower()] = remote_opt
        else:
            # Set the momentum value
            local_opt.data.copy_(
                remote_opt.to(local_opt.device, dtype=local_opt.dtype),
            )
        # Set the step to the current server steps
        if step is not None:
            optim.state[local_param]["step"].data.copy_(
                torch.as_tensor(
                    step,
                    device=optim.state[local_param]["step"].device,
                    dtype=optim.state[local_param]["step"].dtype,
                ),
            )
