



"""simple docstring"""





import operator as op


__lowerCAmelCase							:       Any          =       'scaler.pt'
__lowerCAmelCase							:       List[Any]          =       'pytorch_model'
__lowerCAmelCase							:       Tuple          =       'random_states'
__lowerCAmelCase							:       Dict          =       'optimizer'
__lowerCAmelCase							:       List[str]          =       'scheduler'
__lowerCAmelCase							:       Optional[Any]          =       'pytorch_model.bin'
__lowerCAmelCase							:       List[str]          =       'pytorch_model.bin.index.json'
__lowerCAmelCase							:       int          =       'model.safetensors'
__lowerCAmelCase							:       Union[str, Any]          =       'model.safetensors.index.json'
__lowerCAmelCase							:       Any          =       '1.10.2'
__lowerCAmelCase							:       Optional[int]          =       'py38'
__lowerCAmelCase							:       List[Any]          =       '4.17.0'
__lowerCAmelCase							:       List[str]          =       ['ml.p3.16xlarge', 'ml.p3dn.24xlarge', 'ml.p4dn.24xlarge']
__lowerCAmelCase							:       str          =       ['FULL_SHARD', 'SHARD_GRAD_OP', 'NO_SHARD', 'HYBRID_SHARD', 'HYBRID_SHARD_ZERO2']
__lowerCAmelCase							:       Dict          =       ['TRANSFORMER_BASED_WRAP', 'SIZE_BASED_WRAP', 'NO_WRAP']
__lowerCAmelCase							:       List[Any]          =       ['BACKWARD_PRE', 'BACKWARD_POST', 'NO_PREFETCH']
__lowerCAmelCase							:       Any          =       ['FULL_STATE_DICT', 'LOCAL_STATE_DICT', 'SHARDED_STATE_DICT']
__lowerCAmelCase							:       Dict          =       '2.0.1'
__lowerCAmelCase							:       List[str]          =       ['pdsh', 'standard', 'openmpi', 'mvapich']
__lowerCAmelCase							:       Tuple          =       ['default', 'reduce-overhead', 'max-autotune']

__lowerCAmelCase							:       int          =       {'>': op.gt, '>=': op.ge, '==': op.eq, '!=': op.ne, '<=': op.le, '<': op.lt}

# These are the args for `torch.distributed.launch` for pytorch < 1.9
__lowerCAmelCase							:       int          =       [
    'nnodes',
    'nproc_per_node',
    'rdzv_backend',
    'rdzv_endpoint',
    'rdzv_id',
    'rdzv_conf',
    'standalone',
    'max_restarts',
    'monitor_interval',
    'start_method',
    'role',
    'module',
    'm',
    'no_python',
    'run_path',
    'log_dir',
    'r',
    'redirects',
    't',
    'tee',
    'node_rank',
    'master_addr',
    'master_port',
]

__lowerCAmelCase							:       Tuple          =       ['DEEPSPEED', 'MULTI_GPU', 'FSDP', 'MEGATRON_LM']
__lowerCAmelCase							:       Tuple          =       ['DEEPSPEED', 'MULTI_XPU', 'FSDP']








a__: str							   =						'0.21.0'

from .accelerator import Accelerator
from .big_modeling import (
    cpu_offload,
    cpu_offload_with_hook,
    disk_offload,
    dispatch_model,
    init_empty_weights,
    init_on_device,
    load_checkpoint_and_dispatch,
)
from .data_loader import skip_first_batches
from .launchers import debug_launcher, notebook_launcher
from .state import PartialState
from .utils import (
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    FullyShardedDataParallelPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    find_executable_batch_size,
    infer_auto_device_map,
    is_rich_available,
    load_checkpoint_in_model,
    synchronize_rng_states,
)


if is_rich_available():
       from .utils import rich





