






import importlib.util
import json
import os
import warnings
from dataclasses import dataclass, field

import torch

from ..training_args import TrainingArguments
from ..utils import cached_property, is_sagemaker_dp_enabled, logging


lowercase    				=   logging.get_logger(__name__)





def 		__UpperCAmelCase     (	):
							# Get the sagemaker specific mp parameters from smp_options variable.
							snake_case_           =				os.getenv('SM_HP_MP_PARAMETERS'				,	'{}')
							try:
														# Parse it and check the field "partitions" is included, it is required for model parallel.
														snake_case_           =				json.loads(a_)
														if "partitions" not in smp_options:
																					return False
							except json.JSONDecodeError:
														return False

							# Get the sagemaker specific framework parameters from mpi_options variable.
							snake_case_           =				os.getenv('SM_FRAMEWORK_PARAMS'				,	'{}')
							try:
														# Parse it and check the field "sagemaker_distributed_dataparallel_enabled".
														snake_case_           =				json.loads(a_)
														if not mpi_options.get('sagemaker_mpi_enabled'				,	a_):
																					return False
							except json.JSONDecodeError:
														return False
							# Lastly, check if the `smdistributed` module is present.
							return importlib.util.find_spec('smdistributed') is not None


if is_sagemaker_model_parallel_available():
	import smdistributed.modelparallel.torch as smp

	smp.init()




@dataclass
class 					UpperCamelCase_ (     snake_case_							):



	'''simple docstring'''




	lowerCAmelCase						=					field(
	    default=''''''			,   metadata={'''help''': '''Used by the SageMaker launcher to send mp-specific args. Ignored in SageMakerTrainer'''}			,   )




	def 				_UpperCamelCase     (			self   )			->       List[Any]:
								super().__post_init__()
								warnings.warn(
								    '`SageMakerTrainingArguments` is deprecated and will be removed in v5 of Transformers. You can use '
								    '`TrainingArguments` instead.'			,     a			,     )




	@cached_property
	def 				_UpperCamelCase     (			self   )			->       "torch.device":
								logger.info('PyTorch: setting up devices'   )
								if torch.distributed.is_available() and torch.distributed.is_initialized() and self.local_rank == -1:
															logger.warning(
															    'torch.distributed process group is initialized, but local_rank == -1. '
															    'In order to use Torch DDP, launch your script with `python -m torch.distributed.launch'   )
								if self.no_cuda:
															snake_case_           =				torch.device('cpu'   )
															snake_case_           =				0
								elif is_sagemaker_model_parallel_available():
															snake_case_           =				smp.local_rank()
															snake_case_           =				torch.device('cuda'			,     a   )
															snake_case_           =				1
								elif is_sagemaker_dp_enabled():
															import smdistributed.dataparallel.torch.torch_smddp  # noqa: F401

															torch.distributed.init_process_group(backend='smddp'			,     timeout=self.ddp_timeout_delta   )
															snake_case_           =				int(os.getenv('SMDATAPARALLEL_LOCAL_RANK'   )   )
															snake_case_           =				torch.device('cuda'			,     self.local_rank   )
															snake_case_           =				1
								elif self.local_rank == -1:
															# if n_gpu is > 1 we'll use nn.DataParallel.
															# If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
															# Explicitly set CUDA to the first (index 0) CUDA device, otherwise `set_device` will
															# trigger an error that a device index is missing. Index 0 takes into account the
															# GPUs available in the environment, so `CUDA_VISIBLE_DEVICES=1,2` with `cuda:0`
															# will use the first GPU in that env, i.e. GPU#1
															snake_case_           =				torch.device('cuda:0' if torch.cuda.is_available() else 'cpu'   )
															# Sometimes the line in the postinit has not been run before we end up here, so just checking we're not at
															# the default value.
															snake_case_           =				torch.cuda.device_count()
								else:
															# Here, we'll use torch.distributed.
															# Initializes the distributed backend which will take care of synchronizing nodes/GPUs
															if not torch.distributed.is_initialized():
																						torch.distributed.init_process_group(backend='nccl'			,     timeout=self.ddp_timeout_delta   )
															snake_case_           =				torch.device('cuda'			,     self.local_rank   )
															snake_case_           =				1

								if device.type == "cuda":
															torch.cuda.set_device(a   )

								return device




	@property
	def 				_UpperCamelCase     (			self   )			->       Any:
								if is_sagemaker_model_parallel_available():
															return smp.dp_size()

								return super().world_size




	@property
	def 				_UpperCamelCase     (			self   )			->       List[str]:
								return not is_sagemaker_model_parallel_available()






	@property
	def 				_UpperCamelCase     (			self   )			->       str:
								return False













import inspect
import os
import unittest
from dataclasses import dataclass

import torch

from accelerate import Accelerator, DistributedDataParallelKwargs, GradScalerKwargs
from accelerate.state import AcceleratorState
from accelerate.test_utils import execute_subprocess_async, require_cuda, require_multi_gpu
from accelerate.utils import KwargsHandler




@dataclass
class 					UpperCamelCase_ (     snake_case_							):



	'''simple docstring'''




	lowerCAmelCase						=					0
	lowerCAmelCase						=					False
	lowerCAmelCase						=					3.0






class 					UpperCamelCase_ (     unittest.TestCase							):



	'''simple docstring'''




	def 				_UpperCamelCase     (			self   )			->       List[Any]:
								# If no defaults are changed, `to_kwargs` returns an empty dict.
								self.assertDictEqual(MockClass().to_kwargs()			,     {}   )
								self.assertDictEqual(MockClass(a=2   ).to_kwargs()			,     {'a': 2}   )
								self.assertDictEqual(MockClass(a=2			,     b=a   ).to_kwargs()			,     {'a': 2, 'b': True}   )
								self.assertDictEqual(MockClass(a=2			,     c=2.25   ).to_kwargs()			,     {'a': 2, 'c': 2.25}   )




	@require_cuda
	def 				_UpperCamelCase     (			self   )			->       str:
								# If no defaults are changed, `to_kwargs` returns an empty dict.
								snake_case_           =				GradScalerKwargs(init_scale=10_24			,     growth_factor=2   )
								AcceleratorState._reset_state()
								snake_case_           =				Accelerator(mixed_precision='fp16'			,     kwargs_handlers=[scaler_handler]   )
								print(accelerator.use_fpaa   )
								snake_case_           =				accelerator.scaler

								# Check the kwargs have been applied
								self.assertEqual(scaler._init_scale			,     1_024.0   )
								self.assertEqual(scaler._growth_factor			,     2.0   )

								# Check the other values are at the default
								self.assertEqual(scaler._backoff_factor			,     0.5   )
								self.assertEqual(scaler._growth_interval			,     20_00   )
								self.assertEqual(scaler._enabled			,     a   )






	@require_multi_gpu
	def 				_UpperCamelCase     (			self   )			->       Optional[Any]:
								snake_case_           =				['torchrun', F'''--nproc_per_node={torch.cuda.device_count()}''', inspect.getfile(self.__class__   )]
								execute_subprocess_async(a			,     env=os.environ.copy()   )


if __name__ == "__main__":
	lowercase    				=   DistributedDataParallelKwargs(bucket_cap_mb=15, find_unused_parameters=True)
	lowercase    				=   Accelerator(kwargs_handlers=[ddp_scaler])
	lowercase    				=   torch.nn.Linear(100, 200)
	lowercase    				=   accelerator.prepare(model)

	# Check the values changed in kwargs
	lowercase    				=   ""
	lowercase    				=   model.bucket_bytes_cap // (1024 * 1024)
	if observed_bucket_cap_map != 15:
		error_msg += f"Kwargs badly passed, should have `15` but found {observed_bucket_cap_map}.\n"
	if model.find_unused_parameters is not True:
		error_msg += f"Kwargs badly passed, should have `True` but found {model.find_unused_parameters}.\n"

	# Check the values of the defaults
	if model.dim != 0:
		error_msg += f"Default value not respected, should have `0` but found {model.dim}.\n"
	if model.broadcast_buffers is not True:
		error_msg += f"Default value not respected, should have `True` but found {model.broadcast_buffers}.\n"
	if model.gradient_as_bucket_view is not False:
		error_msg += f"Default value not respected, should have `False` but found {model.gradient_as_bucket_view}.\n"

	# Raise error at the end to make sure we don't stop at the first failure.
	if len(error_msg) > 0:
		raise ValueError(error_msg)





