




import baseaa
import io
import json
import os
from copy import deepcopy

from ..optimizer import AcceleratedOptimizer
from ..scheduler import AcceleratedScheduler







class 				lowerCAmelCase_       :






				def __init__(			self      ,snake_case__ ):
							if isinstance(snake_case__      ,snake_case__ ):
										# Don't modify user's data should they want to reuse it (e.g. in tests), because once we
										# modified it, it will not be accepted here again, since `auto` values would have been overridden
										SCREAMING_SNAKE_CASE_   : Any					  =      deepcopy(snake_case__ )
							elif os.path.exists(snake_case__ ):
										with io.open(snake_case__      ,'r'      ,encoding='utf-8' ) as f:
													SCREAMING_SNAKE_CASE_   : Dict					  =      json.load(snake_case__ )
							else:
										try:
													SCREAMING_SNAKE_CASE_   : Dict					  =      baseaa.urlsafe_baadecode(snake_case__ ).decode('utf-8' )
													SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      json.loads(snake_case__ )
										except (UnicodeDecodeError, AttributeError, ValueError):
													raise ValueError(
													    F'Expected a string path to an existing deepspeed config, or a dictionary, or a base64 encoded string. Received: {config_file_or_dict}' )

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      config

							self.set_stage_and_offload()






				def 		snake_case	(			self ):
							# zero stage - this is done as early as possible, before model is created, to allow
							# ``is_deepspeed_zero3_enabled`` query and getting to the early deepspeed config object
							# during ``zero.Init()`` which needs to know the dtype, and some other hparams.
							SCREAMING_SNAKE_CASE_   : List[str]					  =      self.get_value('zero_optimization.stage'      ,-1 )

							# offload
							SCREAMING_SNAKE_CASE_   : int					  =      False
							if self.is_zeroa() or self.is_zeroa():
										SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      set(['cpu', 'nvme'] )
										SCREAMING_SNAKE_CASE_   : List[Any]					  =      set(
										    [
										        self.get_value('zero_optimization.offload_optimizer.device' ),
										        self.get_value('zero_optimization.offload_param.device' ),
										    ] )
										if len(offload_devices & offload_devices_valid ) > 0:
													SCREAMING_SNAKE_CASE_   : int					  =      True






				def 		snake_case	(			self      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : int					  =      self.config

							# find the config node of interest if it exists
							SCREAMING_SNAKE_CASE_   : Dict					  =      ds_key_long.split('.' )
							SCREAMING_SNAKE_CASE_   : Dict					  =      nodes.pop()
							for node in nodes:
										SCREAMING_SNAKE_CASE_   : Optional[int]					  =      config.get(snake_case__ )
										if config is None:
													return None, ds_key

							return config, ds_key






				def 		snake_case	(			self      ,snake_case__      ,snake_case__=None ):
							SCREAMING_SNAKE_CASE_  ,    SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      self.find_config_node(snake_case__ )
							if config is None:
										return default
							return config.get(snake_case__      ,snake_case__ )






				def 		snake_case	(			self      ,snake_case__      ,snake_case__=False ):
							SCREAMING_SNAKE_CASE_   : str					  =      self.config

							# find the config node of interest if it exists
							SCREAMING_SNAKE_CASE_   : int					  =      ds_key_long.split('.' )
							for node in nodes:
										SCREAMING_SNAKE_CASE_   : Any					  =      config
										SCREAMING_SNAKE_CASE_   : str					  =      config.get(snake_case__ )
										if config is None:
													if must_exist:
																raise ValueError(F'Can\'t find {ds_key_long} entry in the config: {self.config}' )
													else:
																return

        # if found remove it
							if parent_config is not None:
										parent_config.pop(snake_case__ )






				def 		snake_case	(			self      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : str					  =      self.get_value(snake_case__ )
							return False if value is None else bool(snake_case__ )






				def 		snake_case	(			self      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : int					  =      self.get_value(snake_case__ )
							return False if value is None else not bool(snake_case__ )






				def 		snake_case	(			self ):
							return self._stage == 2






				def 		snake_case	(			self ):
							return self._stage == 3






				def 		snake_case	(			self ):
							return self._offload







class 				lowerCAmelCase_       :






				def __init__(			self      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      engine






				def 		snake_case	(			self      ,snake_case__      ,**snake_case__ ):
							# runs backpropagation and handles mixed precision
							self.engine.backward(snake_case__      ,**snake_case__ )

							# Deepspeed's `engine.step` performs the following operations:
							# - gradient accumulation check
							# - gradient clipping
							# - optimizer step
							# - zero grad
							# - checking overflow
							# - lr_scheduler step (only if engine.lr_scheduler is not None)
							self.engine.step()
							# and this plugin overrides the above calls with no-ops when Accelerate runs under
							# Deepspeed, but allows normal functionality for non-Deepspeed cases thus enabling a simple
							# training loop that works transparently under many training regimes.







class 				lowerCAmelCase_       (					lowerCamelCase_							):






				def __init__(			self      ,snake_case__ ):
							super().__init__(snake_case__      ,device_placement=snake_case__      ,scaler=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Dict					  =      hasattr(self.optimizer      ,'overflow' )






				def 		snake_case	(			self      ,snake_case__=None ):
							pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed






				def 		snake_case	(			self ):
							pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed






				@property
				def 		snake_case	(			self ):
							if self.__has_overflow__:
										return self.optimizer.overflow
							return False







class 				lowerCAmelCase_       (					lowerCamelCase_							):






				def __init__(			self      ,snake_case__      ,snake_case__ ):
							super().__init__(snake_case__      ,snake_case__ )






				def 		snake_case	(			self ):
							pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed







class 				lowerCAmelCase_       :






				def __init__(			self      ,snake_case__      ,snake_case__=0.001      ,snake_case__=0      ,**snake_case__ ):
							SCREAMING_SNAKE_CASE_   : Tuple					  =      params
							SCREAMING_SNAKE_CASE_   : Dict					  =      lr
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      weight_decay
							SCREAMING_SNAKE_CASE_   : Dict					  =      kwargs







class 				lowerCAmelCase_       :






				def __init__(			self      ,snake_case__      ,snake_case__=None      ,snake_case__=0      ,**snake_case__ ):
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      optimizer
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      total_num_steps
							SCREAMING_SNAKE_CASE_   : Dict					  =      warmup_num_steps
							SCREAMING_SNAKE_CASE_   : str					  =      kwargs











# coding=utf-8
# Copyright 2020 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# this script dumps information about the environment

import os
import sys

import transformers


UpperCAmelCase__      			=						'''3'''

print('''Python version:''', sys.version)
print('''transformers version:''', transformers.__version__)

try:
			import torch

			print('''Torch version:''', torch.__version__)
			print('''Cuda available:''', torch.cuda.is_available())
			print('''Cuda version:''', torch.version.cuda)
			print('''CuDNN version:''', torch.backends.cudnn.version())
			print('''Number of GPUs available:''', torch.cuda.device_count())
			print('''NCCL version:''', torch.cuda.nccl.version())
except ImportError:
			print('''Torch version:''', None)

try:
			import deepspeed

			print('''DeepSpeed version:''', deepspeed.__version__)
except ImportError:
			print('''DeepSpeed version:''', None)

try:
			import tensorflow as tf

			print('''TensorFlow version:''', tf.__version__)
			print('''TF GPUs available:''', bool(tf.config.list_physical_devices('''GPU''')))
			print('''Number of TF GPUs available:''', len(tf.config.list_physical_devices('''GPU''')))
except ImportError:
			print('''TensorFlow version:''', None)





