
'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


__lowerCAmelCase   =       {
    """configuration_bloom""": ["""BLOOM_PRETRAINED_CONFIG_ARCHIVE_MAP""", """BloomConfig""", """BloomOnnxConfig"""],
}
try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __lowerCAmelCase   =       ["""BloomTokenizerFast"""]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __lowerCAmelCase   =       [
      """BLOOM_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """BloomForCausalLM""",
      """BloomModel""",
      """BloomPreTrainedModel""",
      """BloomForSequenceClassification""",
      """BloomForTokenClassification""",
      """BloomForQuestionAnswering""",
  ]

if TYPE_CHECKING:
  from .configuration_bloom import BLOOM_PRETRAINED_CONFIG_ARCHIVE_MAP, BloomConfig, BloomOnnxConfig

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_bloom_fast import BloomTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_bloom import (
        BLOOM_PRETRAINED_MODEL_ARCHIVE_LIST,
        BloomForCausalLM,
        BloomForQuestionAnswering,
        BloomForSequenceClassification,
        BloomForTokenClassification,
        BloomModel,
        BloomPreTrainedModel,
    )

else:
  import sys

  __lowerCAmelCase   =       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)










"""simple docstring"""





import os
import sys
import tempfile

import torch

from .state import AcceleratorState
from .utils import PrecisionType, PrepareForLaunch, is_mps_available, patch_environment



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Any     ,    SCREAMING_SNAKE_CASE							:				List[str]=()     ,    SCREAMING_SNAKE_CASE							:				Union[str, Any]=None     ,    SCREAMING_SNAKE_CASE							:				Dict="no"     ,    SCREAMING_SNAKE_CASE							:				int="29500" ):


							'''simple docstring'''







							__lowerCamelCase   :   Optional[int]       =False
							__lowerCamelCase   :   Tuple       =False
							if any(key.startswith('''KAGGLE''' ) for key in os.environ.keys() ):
														__lowerCamelCase   :   Optional[int]       =True
							elif "IPython" in sys.modules:
														__lowerCamelCase   :   int       ='''google.colab''' in str(sys.modules['''IPython'''].get_ipython() )

							try:
														__lowerCamelCase   :   str       =PrecisionType(mixed_precision.lower() )
							except ValueError:
														raise ValueError(
														    F'Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}.' )

							if (in_colab or in_kaggle) and (os.environ.get('''TPU_NAME'''     ,    SCREAMING_SNAKE_CASE ) is not None):
														# TPU launch
														import torch_xla.distributed.xla_multiprocessing as xmp

														if len(AcceleratorState._shared_state ) > 0:
																					raise ValueError(
																					    '''To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside '''
																					    '''your training function. Restart your notebook and make sure no cells initializes an '''
																					    '''`Accelerator`.''' )
														if num_processes is None:
																					__lowerCamelCase   :   int       =8

														__lowerCamelCase   :   Optional[Any]       =PrepareForLaunch(SCREAMING_SNAKE_CASE     ,    distributed_type='''TPU''' )
														print(F'Launching a training on {num_processes} TPU cores.' )
														xmp.spawn(SCREAMING_SNAKE_CASE     ,    args=SCREAMING_SNAKE_CASE     ,    nprocs=SCREAMING_SNAKE_CASE     ,    start_method='''fork''' )
							elif in_colab:
														# No need for a distributed launch otherwise as it's either CPU or one GPU.
														if torch.cuda.is_available():
																					print('''Launching training on one GPU.''' )
														else:
																					print('''Launching training on one CPU.''' )
														function(*SCREAMING_SNAKE_CASE )
							else:
														if num_processes is None:
																					raise ValueError(
																					    '''You have to specify the number of GPUs you would like to use, add `num_processes=...` to your call.''' )

														if num_processes > 1:
																					# Multi-GPU launch
																					from torch.multiprocessing import start_processes
																					from torch.multiprocessing.spawn import ProcessRaisedException

																					if len(AcceleratorState._shared_state ) > 0:
																												raise ValueError(
																												    '''To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized '''
																												    '''inside your training function. Restart your notebook and make sure no cells initializes an '''
																												    '''`Accelerator`.''' )

																					if torch.cuda.is_initialized():
																												raise ValueError(
																												    '''To launch a multi-GPU training from your notebook, you need to avoid running any instruction '''
																												    '''using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA '''
																												    '''function.''' )

																					# torch.distributed will expect a few environment variable to be here. We set the ones common to each
																					# process here (the other ones will be set be the launcher).
																					with patch_environment(
																					    world_size=SCREAMING_SNAKE_CASE     ,    master_addr='''127.0.01'''     ,    master_port=SCREAMING_SNAKE_CASE     ,    mixed_precision=SCREAMING_SNAKE_CASE ):
																												__lowerCamelCase   :   List[str]       =PrepareForLaunch(SCREAMING_SNAKE_CASE     ,    distributed_type='''MULTI_GPU''' )
																												print(F'Launching training on {num_processes} GPUs.' )
																												try:
																																			start_processes(SCREAMING_SNAKE_CASE     ,    args=SCREAMING_SNAKE_CASE     ,    nprocs=SCREAMING_SNAKE_CASE     ,    start_method='''fork''' )
																												except ProcessRaisedException as e:
																																			if "Cannot re-initialize CUDA in forked subprocess" in e.args[0]:
																																										raise RuntimeError(
																																										    '''CUDA has been initialized before the `notebook_launcher` could create a forked subprocess. '''
																																										    '''This likely stems from an outside import causing issues once the `notebook_launcher()` is called. '''
																																										    '''Please review your imports and test them when running the `notebook_launcher()` to identify '''
																																										    '''which one is problematic.''' ) from e

														else:
																					# No need for a distributed launch otherwise as it's either CPU, GPU or MPS.
																					if is_mps_available():
																												__lowerCamelCase   :   Optional[Any]       ='''1'''
																												print('''Launching training on MPS.''' )
																					elif torch.cuda.is_available():
																												print('''Launching training on one GPU.''' )
																					else:
																												print('''Launching training on CPU.''' )
																					function(*SCREAMING_SNAKE_CASE )



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				str     ,    SCREAMING_SNAKE_CASE							:				Optional[int]=()     ,    SCREAMING_SNAKE_CASE							:				Optional[Any]=2 ):


							'''simple docstring'''







							from torch.multiprocessing import start_processes

							with tempfile.NamedTemporaryFile() as tmp_file:
														# torch.distributed will expect a few environment variable to be here. We set the ones common to each
														# process here (the other ones will be set be the launcher).
														with patch_environment(
														    world_size=SCREAMING_SNAKE_CASE     ,    master_addr='''127.0.01'''     ,    master_port='''29500'''     ,    accelerate_mixed_precision='''no'''     ,    accelerate_debug_rdv_file=tmp_file.name     ,    accelerate_use_cpu='''yes'''     ,    ):
																					__lowerCamelCase   :   Tuple       =PrepareForLaunch(SCREAMING_SNAKE_CASE     ,    debug=SCREAMING_SNAKE_CASE )
																					start_processes(SCREAMING_SNAKE_CASE     ,    args=SCREAMING_SNAKE_CASE     ,    nprocs=SCREAMING_SNAKE_CASE     ,    start_method='''fork''' )




