def _shard_orig_param_state(
    fsdp_param_info: FSDPParamInfo,
    fqn: str,
    optim_state: Dict[str, Any],
) -> Dict[str, Any]:
    """
    Shard the optimizer state for the original parameter with the name ``fqn``.
    This API should only be used when ``use_orig_params`` is True.
    """
    if not optim_state:
        return {}
    fsdp_state = fsdp_param_info.state
    flat_param = fsdp_param_info.flat_param
    param_idx = fsdp_param_info.param_indices[fqn]

    optim_state = _gather_state_dict(optim_state, fsdp_state.process_group)
    start, end = flat_param._shard_indices  # type: ignore[attr-defined]
    if not (start <= param_idx <= end and flat_param._shard_param_offsets):  # type: ignore[attr-defined]
        return {}
    param_start, param_end = flat_param._shard_param_offsets[param_idx - start]  # type: ignore[attr-defined]

    # Flatten and shard the state.
    new_optim_state: Dict[str, Any] = {}
    for state_name, value in optim_state.items():
        if torch.is_tensor(value) and value.dim() > 0:
            value = value.flatten()[param_start : param_end + 1]
        new_optim_state[state_name] = value
    return new_optim_state
Shard the optimizer state for the original parameter with the name ``fqn``.
This API should only be used when ``use_orig_params`` is True.
_optim_utils.py