# Source code for tensorplay.distributed.fsdp.api

Source: https://www.tensorplay.cn/docs/_modules/tensorplay/distributed/fsdp/api.html

```
"""Configuration objects for fully sharded module execution."""

from collections.abc import Sequence
from dataclasses import dataclass
from enum import Enum, auto
from typing import Any

from tensorplay.nn.modules.batchnorm import _BatchNorm

__all__ = [
    "ShardingStrategy",
    "BackwardPrefetch",
    "MixedPrecision",
    "CPUOffload",
    "StateDictType",
    "StateDictConfig",
    "FullStateDictConfig",
    "LocalStateDictConfig",
    "ShardedStateDictConfig",
    "OptimStateDictConfig",
    "FullOptimStateDictConfig",
    "LocalOptimStateDictConfig",
    "ShardedOptimStateDictConfig",
    "StateDictSettings",
]

[docs]
class ShardingStrategy(Enum):
    FULL_SHARD = auto()
    SHARD_GRAD_OP = auto()
    NO_SHARD = auto()
    HYBRID_SHARD = auto()
    _HYBRID_SHARD_ZERO2 = auto()

[docs]
class BackwardPrefetch(Enum):
    BACKWARD_PRE = auto()
    BACKWARD_POST = auto()

[docs]
@dataclass
class MixedPrecision:
    param_dtype: Any = None
    reduce_dtype: Any = None
    buffer_dtype: Any = None
    keep_low_precision_grads: bool = False
    cast_forward_inputs: bool = False
    cast_root_forward_inputs: bool = True
    _module_classes_to_ignore: Sequence[type] = (_BatchNorm,)

[docs]
@dataclass
class CPUOffload:
    offload_params: bool = False

[docs]
class StateDictType(Enum):
    FULL_STATE_DICT = auto()
    LOCAL_STATE_DICT = auto()
    SHARDED_STATE_DICT = auto()

[docs]
@dataclass
class StateDictConfig:
    offload_to_cpu: bool = False

[docs]
@dataclass
class FullStateDictConfig(StateDictConfig):
    rank0_only: bool = False

[docs]
@dataclass
class LocalStateDictConfig(StateDictConfig):
    pass

[docs]
@dataclass
class ShardedStateDictConfig(StateDictConfig):
    _use_dtensor: bool = False

[docs]
@dataclass
class OptimStateDictConfig:
    offload_to_cpu: bool = True

[docs]
@dataclass
class FullOptimStateDictConfig(OptimStateDictConfig):
    rank0_only: bool = False

[docs]
@dataclass
class LocalOptimStateDictConfig(OptimStateDictConfig):
    offload_to_cpu: bool = False

[docs]
@dataclass
class ShardedOptimStateDictConfig(OptimStateDictConfig):
    _use_dtensor: bool = False

[docs]
@dataclass
class StateDictSettings:
    state_dict_type: StateDictType
    state_dict_config: StateDictConfig
    optim_state_dict_config: OptimStateDictConfig
```