





'''simple docstring'''

import operator as op


lowercase_					= """scaler.pt"""
lowercase_					= """pytorch_model"""
lowercase_					= """random_states"""
lowercase_					= """optimizer"""
lowercase_					= """scheduler"""
lowercase_					= """pytorch_model.bin"""
lowercase_					= """pytorch_model.bin.index.json"""
lowercase_					= """model.safetensors"""
lowercase_					= """model.safetensors.index.json"""
lowercase_					= """1.10.2"""
lowercase_					= """py38"""
lowercase_					= """4.17.0"""
lowercase_					= ["""ml.p3.16xlarge""", """ml.p3dn.24xlarge""", """ml.p4dn.24xlarge"""]
lowercase_					= ["""FULL_SHARD""", """SHARD_GRAD_OP""", """NO_SHARD""", """HYBRID_SHARD""", """HYBRID_SHARD_ZERO2"""]
lowercase_					= ["""TRANSFORMER_BASED_WRAP""", """SIZE_BASED_WRAP""", """NO_WRAP"""]
lowercase_					= ["""BACKWARD_PRE""", """BACKWARD_POST""", """NO_PREFETCH"""]
lowercase_					= ["""FULL_STATE_DICT""", """LOCAL_STATE_DICT""", """SHARDED_STATE_DICT"""]
lowercase_					= """2.0.1"""
lowercase_					= ["""pdsh""", """standard""", """openmpi""", """mvapich"""]
lowercase_					= ["""default""", """reduce-overhead""", """max-autotune"""]

lowercase_					= {""">""": op.gt, """>=""": op.ge, """==""": op.eq, """!=""": op.ne, """<=""": op.le, """<""": op.lt}

# These are the args for `torch.distributed.launch` for pytorch < 1.9
lowercase_					= [
    """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""",
]

lowercase_					= ["""DEEPSPEED""", """MULTI_GPU""", """FSDP""", """MEGATRON_LM"""]
lowercase_					= ["""DEEPSPEED""", """MULTI_XPU""", """FSDP"""]







'''simple docstring'''

from ....utils import logging


lowerCamelCase							:   Optional[Any]  						=	logging.get_logger(__name__)
class   __lowerCAmelCase						(lowercase_     ):





      '''simple docstring'''






      def __init__(self    :    Optional[int]   ,						UpperCamelCase    :    Union[str, Any]   ,						UpperCamelCase    :    Optional[Any]=None   ,						UpperCamelCase    :    int=2048 ):

             '''simple docstring'''







             lowercase__      			= config.__dict__
             lowercase__      			= modal_hidden_size
             if num_labels:
                    lowercase__      			= num_labels




