from __future__ import annotations

from typing import Any, Optional

from paddle.distributed.collective import wait as wait
from paddle.optimizer import Optimizer

from .group_sharded_storage import GradStorage as GradStorage
from .group_sharded_storage import ParamStorage as ParamStorage
from .group_sharded_utils import GroupShardedClipGrad as GroupShardedClipGrad
from .group_sharded_utils import Type as Type
from .group_sharded_utils import device_guard as device_guard

alignment: Any
align: Any

class GroupShardedOptimizerStage2(Optimizer):
    world_size: Any = ...
    param_storages: Any = ...
    offload: Any = ...
    offload_device: str = ...
    offload_buffer_size: int = ...
    offload_param2align: Any = ...
    offload_params: Any = ...
    offload_grads: Any = ...
    dev_id: Any = ...
    def __init__(
        self,
        params: Any,
        optim: Any,
        group: Optional[Any] = ...,
        offload: bool = ...,
        device: str = ...,
        pertrain_sync_models: bool = ...,
        **kw: Any,
    ): ...
    @property
    def local_params(self): ...
    @property
    def param2rank(self): ...
    @property
    def dtype_rank_params(self): ...
    @property
    def rank_buffer_size(self): ...
    def step(self) -> None: ...
    def minimize(self) -> None: ...
    def set_state_dict(self, state_dict: Any) -> None: ...
    def state_dict(self): ...
