



'''simple docstring'''



import json
import os
import tempfile
from unittest.mock import patch

import torch
from torch.utils.data import DataLoader, TensorDataset

from accelerate import DistributedType, infer_auto_device_map, init_empty_weights
from accelerate.accelerator import Accelerator
from accelerate.state import GradientState, PartialState
from accelerate.test_utils import require_bnb, require_multi_gpu, slow
from accelerate.test_utils.testing import AccelerateTestCase, require_cuda
from accelerate.utils import patch_environment
def 					UpperCAmelCase   (						):

							'''simple docstring'''





							SCREAMING_SNAKE_CASE					:      Any     					=			torch.nn.Linear(2					,		4				)
							SCREAMING_SNAKE_CASE					:      Optional[int]     					=			torch.optim.AdamW(model.parameters()					,		lr=1.0				)
							SCREAMING_SNAKE_CASE					:      List[Any]     					=			torch.optim.lr_scheduler.OneCycleLR(A					,		max_lr=0.01					,		steps_per_epoch=2					,		epochs=1				)
							SCREAMING_SNAKE_CASE					:      Optional[int]     					=			DataLoader(TensorDataset(torch.tensor([1, 2, 3]				)				)				)
							SCREAMING_SNAKE_CASE					:      List[Any]     					=			DataLoader(TensorDataset(torch.tensor([4, 5, 6]				)				)				)

							return model, optimizer, scheduler, train_dl, valid_dl
def 					UpperCAmelCase   (						A  :				Optional[int]				):

							'''simple docstring'''





							return (model.weight.abs().sum() + model.bias.abs().sum()).item()
def 					UpperCAmelCase   (						A  :				Any				):

							'''simple docstring'''





							SCREAMING_SNAKE_CASE					:      str     					=			torch.nn.Linear(*tuple(model.weight.T.shape				)				).state_dict()
							model.load_state_dict(A				)






class 	lowerCamelCase_							(					snake_case_			):







						@require_cuda
						def 					__lowercase    (  self   :      Optional[int]				):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE					:      Tuple     					=			Accelerator()
													assert PartialState._shared_state["_cpu"] is False
													assert PartialState._shared_state["device"].type == "cuda"
													with self.assertRaises(lowerCAmelCase__				):
																				SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			Accelerator(cpu=lowerCAmelCase__				)







						def 					__lowercase    (  self   :      Any				):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE					:      Dict     					=			Accelerator()
													SCREAMING_SNAKE_CASE					:      int     					=			GradientState()
													assert state.num_steps == 1
													SCREAMING_SNAKE_CASE					:      Dict     					=			4
													assert state.num_steps == 4

													assert state.sync_gradients is True
													SCREAMING_SNAKE_CASE					:      Optional[int]     					=			False
													assert state.sync_gradients is False
													GradientState._reset_state()







						def 					__lowercase    (  self   :      Optional[int]				):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE					:      List[Any]     					=			Accelerator()
													SCREAMING_SNAKE_CASE					:      str     					=			create_components()

													(
													    SCREAMING_SNAKE_CASE
													)					:      str     					=			accelerator.prepare(lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				)

													self.assertTrue(prepared_model in accelerator._models				)
													self.assertTrue(prepared_optimizer in accelerator._optimizers				)
													self.assertTrue(prepared_scheduler in accelerator._schedulers				)
													self.assertTrue(prepared_train_dl in accelerator._dataloaders				)
													self.assertTrue(prepared_valid_dl in accelerator._dataloaders				)







						def 					__lowercase    (  self   :      str				):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			Accelerator()
													SCREAMING_SNAKE_CASE					:      Any     					=			create_components()
													accelerator.prepare(lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				)
													accelerator.free_memory()

													self.assertTrue(len(accelerator._models				) == 0				)
													self.assertTrue(len(accelerator._optimizers				) == 0				)
													self.assertTrue(len(accelerator._schedulers				) == 0				)
													self.assertTrue(len(accelerator._dataloaders				) == 0				)







						def 					__lowercase    (  self   :      str				):

													"""simple docstring"""
													PartialState._reset_state()

													# Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
													def noop(*lowerCAmelCase__   :      Any				,    **lowerCAmelCase__   :      str				):
																				pass

													with patch('''torch.cuda.set_device'''				,    lowerCAmelCase__				), patch_environment(ACCELERATE_TORCH_DEVICE='''cuda:64'''				):
																				SCREAMING_SNAKE_CASE					:      str     					=			Accelerator()
																				self.assertEqual(str(accelerator.state.device				)				,    '''cuda:64'''				)







						def 					__lowercase    (  self   :      str				):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE					:      Any     					=			Accelerator()
													SCREAMING_SNAKE_CASE					:      Any     					=			create_components()
													accelerator.prepare(lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				)

													SCREAMING_SNAKE_CASE					:      int     					=			get_signature(lowerCAmelCase__				)

													with tempfile.TemporaryDirectory() as tmpdirname:
																				accelerator.save_state(lowerCAmelCase__				)

																				# make sure random weights don't match
																				load_random_weights(lowerCAmelCase__				)
																				self.assertTrue(abs(model_signature - get_signature(lowerCAmelCase__				)				) > 1e-3				)

																				# make sure loaded weights match
																				accelerator.load_state(lowerCAmelCase__				)
																				self.assertTrue(abs(model_signature - get_signature(lowerCAmelCase__				)				) < 1e-3				)







						def 					__lowercase    (  self   :      Any				):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE					:      str     					=			Accelerator()
													SCREAMING_SNAKE_CASE					:      int     					=			create_components()
													accelerator.prepare(lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				)

													SCREAMING_SNAKE_CASE					:      str     					=			get_signature(lowerCAmelCase__				)

													# saving hook
													def save_config(lowerCAmelCase__   :      int				,    lowerCAmelCase__   :      Dict				,    lowerCAmelCase__   :      Tuple				):
																				SCREAMING_SNAKE_CASE					:      Dict     					=			{'''class_name''': models[0].__class__.__name__}

																				with open(os.path.join(lowerCAmelCase__				,    '''data.json'''				)				,    '''w'''				) as f:
																											json.dump(lowerCAmelCase__				,    lowerCAmelCase__				)

        # loading hook
													def load_config(lowerCAmelCase__   :      str				,    lowerCAmelCase__   :      Optional[Any]				):
																				with open(os.path.join(lowerCAmelCase__				,    '''data.json'''				)				,    '''r'''				) as f:
																											SCREAMING_SNAKE_CASE					:      List[Any]     					=			json.load(lowerCAmelCase__				)

																				SCREAMING_SNAKE_CASE					:      Dict     					=			config['''class_name''']

													SCREAMING_SNAKE_CASE					:      int     					=			accelerator.register_save_state_pre_hook(lowerCAmelCase__				)
													SCREAMING_SNAKE_CASE					:      int     					=			accelerator.register_load_state_pre_hook(lowerCAmelCase__				)

													with tempfile.TemporaryDirectory() as tmpdirname:
																				accelerator.save_state(lowerCAmelCase__				)

																				# make sure random weights don't match with hooks
																				load_random_weights(lowerCAmelCase__				)
																				self.assertTrue(abs(model_signature - get_signature(lowerCAmelCase__				)				) > 1e-3				)

																				# random class name to verify correct one is loaded
																				SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			'''random'''

																				# make sure loaded weights match with hooks
																				accelerator.load_state(lowerCAmelCase__				)
																				self.assertTrue(abs(model_signature - get_signature(lowerCAmelCase__				)				) < 1e-3				)

																				# mode.class_name is loaded from config
																				self.assertTrue(model.class_name == model.__class__.__name__				)

													# remove hooks
													save_hook.remove()
													load_hook.remove()

													with tempfile.TemporaryDirectory() as tmpdirname:
																				accelerator.save_state(lowerCAmelCase__				)

																				# make sure random weights don't match with hooks removed
																				load_random_weights(lowerCAmelCase__				)
																				self.assertTrue(abs(model_signature - get_signature(lowerCAmelCase__				)				) > 1e-3				)

																				# random class name to verify correct one is loaded
																				SCREAMING_SNAKE_CASE					:      Tuple     					=			'''random'''

																				# make sure loaded weights match with hooks removed
																				accelerator.load_state(lowerCAmelCase__				)
																				self.assertTrue(abs(model_signature - get_signature(lowerCAmelCase__				)				) < 1e-3				)

																				# mode.class_name is NOT loaded from config
																				self.assertTrue(model.class_name != model.__class__.__name__				)







						def 					__lowercase    (  self   :      Optional[Any]				):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE					:      Tuple     					=			Accelerator()
													SCREAMING_SNAKE_CASE					:      Dict     					=			create_components()
													SCREAMING_SNAKE_CASE					:      List[Any]     					=			None

													# This should work
													SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			accelerator.prepare(
													    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				)
													self.assertTrue(dummy_obj is None				)







						def 					__lowercase    (  self   :      Optional[int]				):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			Accelerator()
													SCREAMING_SNAKE_CASE					:      List[str]     					=			create_components()
													SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			[1, 2, 3]

													# This should work
													SCREAMING_SNAKE_CASE					:      Optional[int]     					=			accelerator.prepare(
													    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				,    lowerCAmelCase__				)
													self.assertEqual(
													    getattr(lowerCAmelCase__				,    '''_is_accelerate_prepared'''				,    lowerCAmelCase__				)				,    lowerCAmelCase__				,    '''Dummy object should have `_is_accelerate_prepared` set to `True`'''				,    )
													self.assertEqual(
													    getattr(lowerCAmelCase__				,    '''_is_accelerate_prepared'''				,    lowerCAmelCase__				)				,    lowerCAmelCase__				,    '''Model is missing `_is_accelerator_prepared` or is set to `False`'''				,    )
													self.assertEqual(
													    getattr(lowerCAmelCase__				,    '''_is_accelerate_prepared'''				,    lowerCAmelCase__				)				,    lowerCAmelCase__				,    '''Optimizer is missing `_is_accelerator_prepared` or is set to `False`'''				,    )
													self.assertEqual(
													    getattr(lowerCAmelCase__				,    '''_is_accelerate_prepared'''				,    lowerCAmelCase__				)				,    lowerCAmelCase__				,    '''Scheduler is missing `_is_accelerator_prepared` or is set to `False`'''				,    )
													self.assertEqual(
													    getattr(lowerCAmelCase__				,    '''_is_accelerate_prepared'''				,    lowerCAmelCase__				)				,    lowerCAmelCase__				,    '''Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`'''				,    )
													self.assertEqual(
													    getattr(lowerCAmelCase__				,    '''_is_accelerate_prepared'''				,    lowerCAmelCase__				)				,    lowerCAmelCase__				,    '''Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`'''				,    )







						@slow
						@require_bnb
						def 					__lowercase    (  self   :      Optional[Any]				):

													"""simple docstring"""
													from transformers import AutoModelForCausalLM

													SCREAMING_SNAKE_CASE					:      Dict     					=			AutoModelForCausalLM.from_pretrained(
													    '''EleutherAI/gpt-neo-125m'''				,    load_in_abit=lowerCAmelCase__				,    device_map={'''''': 0}				,    )
													SCREAMING_SNAKE_CASE					:      str     					=			Accelerator()

													# This should work
													SCREAMING_SNAKE_CASE					:      Dict     					=			accelerator.prepare(lowerCAmelCase__				)







						@slow
						@require_bnb
						def 					__lowercase    (  self   :      Optional[int]				):

													"""simple docstring"""
													from transformers import AutoModelForCausalLM

													SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			Accelerator()

													with init_empty_weights():
																				SCREAMING_SNAKE_CASE					:      Optional[int]     					=			AutoModelForCausalLM.from_pretrained(
																				    '''EleutherAI/gpt-neo-125m'''				,    )
																				model.tie_weights()
																				SCREAMING_SNAKE_CASE					:      Tuple     					=			infer_auto_device_map(lowerCAmelCase__				)
																				SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			'''cpu'''

													SCREAMING_SNAKE_CASE					:      str     					=			AutoModelForCausalLM.from_pretrained(
													    '''EleutherAI/gpt-neo-125m'''				,    device_map=lowerCAmelCase__				,    load_in_abit=lowerCAmelCase__				,    llm_inta_enable_fpaa_cpu_offload=lowerCAmelCase__				)

													# This should not work and get value error
													with self.assertRaises(lowerCAmelCase__				):
																				SCREAMING_SNAKE_CASE					:      Dict     					=			accelerator.prepare(lowerCAmelCase__				)







						@slow
						@require_bnb
						@require_multi_gpu
						def 					__lowercase    (  self   :      Tuple				):

													"""simple docstring"""
													from transformers import AutoModelForCausalLM

													SCREAMING_SNAKE_CASE					:      Tuple     					=			{'''distributed_type''': DistributedType.MULTI_GPU}

													with init_empty_weights():
																				SCREAMING_SNAKE_CASE					:      List[str]     					=			AutoModelForCausalLM.from_pretrained(
																				    '''EleutherAI/gpt-neo-125m'''				,    )
																				model.tie_weights()
																				SCREAMING_SNAKE_CASE					:      Tuple     					=			infer_auto_device_map(lowerCAmelCase__				)
																				SCREAMING_SNAKE_CASE					:      Any     					=			1

													SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			AutoModelForCausalLM.from_pretrained(
													    '''EleutherAI/gpt-neo-125m'''				,    load_in_abit=lowerCAmelCase__				,    device_map=lowerCAmelCase__				,    )
													SCREAMING_SNAKE_CASE					:      Any     					=			Accelerator()

													# This should not work and get value error
													with self.assertRaises(lowerCAmelCase__				):
																				SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			accelerator.prepare(lowerCAmelCase__				)

													PartialState._reset_state()







						@slow
						@require_bnb
						@require_multi_gpu
						def 					__lowercase    (  self   :      str				):

													"""simple docstring"""
													from transformers import AutoModelForCausalLM

													with init_empty_weights():
																				SCREAMING_SNAKE_CASE					:      List[Any]     					=			AutoModelForCausalLM.from_pretrained(
																				    '''EleutherAI/gpt-neo-125m'''				,    )
																				SCREAMING_SNAKE_CASE					:      List[str]     					=			infer_auto_device_map(lowerCAmelCase__				)
																				SCREAMING_SNAKE_CASE					:      List[Any]     					=			1

													SCREAMING_SNAKE_CASE					:      Optional[int]     					=			AutoModelForCausalLM.from_pretrained(
													    '''EleutherAI/gpt-neo-125m'''				,    load_in_abit=lowerCAmelCase__				,    device_map=lowerCAmelCase__				,    )
													SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			Accelerator()

													# This should work
													SCREAMING_SNAKE_CASE					:      int     					=			accelerator.prepare(lowerCAmelCase__				)







						@require_cuda
						def 					__lowercase    (  self   :      Dict				):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE					:      Tuple     					=			torch.nn.Linear(10				,    10				)
													SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			torch.optim.SGD(model.parameters()				,    lr=0.01				)
													SCREAMING_SNAKE_CASE					:      str     					=			Accelerator(cpu=lowerCAmelCase__				)
													SCREAMING_SNAKE_CASE					:      Optional[int]     					=			accelerator.prepare(lowerCAmelCase__				)









'''simple docstring'''



import os


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_doctest_list.py
lowerCAmelCase_     :       Optional[Any]								=						'.'


if __name__ == "__main__":
							lowerCAmelCase_     :       Optional[Any]								=						os.path.join(REPO_PATH, 'utils/documentation_tests.txt')
							lowerCAmelCase_     :       Dict								=						[]
							lowerCAmelCase_     :       Any								=						[]
							with open(doctest_file_path) as fp:
														for line in fp:
																					lowerCAmelCase_     :       List[str]								=						line.strip()
																					lowerCAmelCase_     :       Optional[Any]								=						os.path.join(REPO_PATH, line)
																					if not (os.path.isfile(path) or os.path.isdir(path)):
																												non_existent_paths.append(line)
																					all_paths.append(path)
							if len(non_existent_paths) > 0:
														lowerCAmelCase_     :       Dict								=						'\n'.join(non_existent_paths)
														raise ValueError(f'`utils/documentation_tests.txt` contains non-existent paths:\n{non_existent_paths}')
							if all_paths != sorted(all_paths):
														raise ValueError('Files in `utils/documentation_tests.txt` are not in alphabetical order.')




