






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 						lowerCAmelCase_		()		->     List[Any]:


				'''simple docstring'''




				lowerCAmelCase__         =    torch.nn.Linear(2  ,       4							)
				lowerCAmelCase__         =    torch.optim.AdamW(model.parameters()  ,       lr=1.0							)
				lowerCAmelCase__         =    torch.optim.lr_scheduler.OneCycleLR(lowercase__  ,       max_lr=0.01  ,       steps_per_epoch=2  ,       epochs=1							)
				lowerCAmelCase__         =    DataLoader(TensorDataset(torch.tensor([1, 2, 3]							)							)							)
				lowerCAmelCase__         =    DataLoader(TensorDataset(torch.tensor([4, 5, 6]							)							)							)

				return model, optimizer, scheduler, train_dl, valid_dl





def 						lowerCAmelCase_		(lowercase__    :							Dict							)		->     Tuple:


				'''simple docstring'''




				return (model.weight.abs().sum() + model.bias.abs().sum()).item()





def 						lowerCAmelCase_		(lowercase__    :							int							)		->     int:


				'''simple docstring'''




				lowerCAmelCase__         =    torch.nn.Linear(*tuple(model.weight.T.shape							)							).state_dict()
				model.load_state_dict(lowercase__							)




class 			lowerCAmelCase_							(					snake_case__       ):




							@require_cuda
							def     __snake_case       ( self  : Optional[Any]						):
											lowerCAmelCase__         =    Accelerator()
											assert PartialState._shared_state["_cpu"] is False
											assert PartialState._shared_state["device"].type == "cuda"
											with self.assertRaises(__UpperCAmelCase						):
															lowerCAmelCase__         =    Accelerator(cpu=__UpperCAmelCase						)




							def     __snake_case       ( self  : Dict						):
											lowerCAmelCase__         =    Accelerator()
											lowerCAmelCase__         =    GradientState()
											assert state.num_steps == 1
											lowerCAmelCase__         =    4
											assert state.num_steps == 4

											assert state.sync_gradients is True
											lowerCAmelCase__         =    False
											assert state.sync_gradients is False
											GradientState._reset_state()




							def     __snake_case       ( self  : Dict						):
											lowerCAmelCase__         =    Accelerator()
											lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__         =    create_components()

											(
											    (
											    lowerCAmelCase__	
											)      ,		(
											    lowerCAmelCase__	
											)      ,		(
											    lowerCAmelCase__	
											)      ,		(
											    lowerCAmelCase__	
											)      ,		(
											    lowerCAmelCase__	
											)      ,			
											)         =    accelerator.prepare(__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase						)

											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     __snake_case       ( self  : Tuple						):
											lowerCAmelCase__         =    Accelerator()
											lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__         =    create_components()
											accelerator.prepare(__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase						)
											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     __snake_case       ( self  : Optional[int]						):
											PartialState._reset_state()

											# Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
											def noop(*SCREAMING_SNAKE_CASE_  : Tuple			,				**SCREAMING_SNAKE_CASE_  : List[Any]						):
															pass

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




							def     __snake_case       ( self  : Dict						):
											lowerCAmelCase__         =    Accelerator()
											lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__         =    create_components()
											accelerator.prepare(__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase						)

											lowerCAmelCase__         =    get_signature(__UpperCAmelCase						)

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

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

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




							def     __snake_case       ( self  : int						):
											lowerCAmelCase__         =    Accelerator()
											lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__         =    create_components()
											accelerator.prepare(__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase						)

											lowerCAmelCase__         =    get_signature(__UpperCAmelCase						)

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

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

															# loading hook
											def load_config(SCREAMING_SNAKE_CASE_  : str			,				SCREAMING_SNAKE_CASE_  : str						):
															with open(os.path.join(__UpperCAmelCase			,				'''data.json'''						)			,				'''r'''						) as f:
																			lowerCAmelCase__         =    json.load(__UpperCAmelCase						)

															lowerCAmelCase__         =    config['''class_name''']

											lowerCAmelCase__         =    accelerator.register_save_state_pre_hook(__UpperCAmelCase						)
											lowerCAmelCase__         =    accelerator.register_load_state_pre_hook(__UpperCAmelCase						)

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

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

															# random class name to verify correct one is loaded
															lowerCAmelCase__         =    '''random'''

															# make sure loaded weights match with hooks
															accelerator.load_state(__UpperCAmelCase						)
															self.assertTrue(abs(model_signature - get_signature(__UpperCAmelCase						)						) < 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(__UpperCAmelCase						)

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

															# random class name to verify correct one is loaded
															lowerCAmelCase__         =    '''random'''

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

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




							def     __snake_case       ( self  : Union[str, Any]						):
											lowerCAmelCase__         =    Accelerator()
											lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__         =    create_components()
											lowerCAmelCase__         =    None

											# This should work
											lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__         =    accelerator.prepare(
											    __UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase			,				__UpperCAmelCase						)
											self.assertTrue(dummy_obj is None						)




							def     __snake_case       ( self  : Dict						):
											lowerCAmelCase__         =    Accelerator()
											lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__      ,		lowerCAmelCase__         =    create_components()
											lowerCAmelCase__         =    [1, 2, 3]

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




							@slow
							@require_bnb
							def     __snake_case       ( self  : Any						):
											from transformers import AutoModelForCausalLM

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

											# This should work
											lowerCAmelCase__         =    accelerator.prepare(__UpperCAmelCase						)




							@slow
							@require_bnb
							def     __snake_case       ( self  : str						):
											from transformers import AutoModelForCausalLM

											lowerCAmelCase__         =    Accelerator()

											with init_empty_weights():
															lowerCAmelCase__         =    AutoModelForCausalLM.from_pretrained(
															    '''EleutherAI/gpt-neo-125m'''			,				)
															model.tie_weights()
															lowerCAmelCase__         =    infer_auto_device_map(__UpperCAmelCase						)
															lowerCAmelCase__         =    '''cpu'''

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

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




							@slow
							@require_bnb
							@require_multi_gpu
							def     __snake_case       ( self  : int						):
											from transformers import AutoModelForCausalLM

											lowerCAmelCase__         =    {'''distributed_type''': DistributedType.MULTI_GPU}

											with init_empty_weights():
															lowerCAmelCase__         =    AutoModelForCausalLM.from_pretrained(
															    '''EleutherAI/gpt-neo-125m'''			,				)
															model.tie_weights()
															lowerCAmelCase__         =    infer_auto_device_map(__UpperCAmelCase						)
															lowerCAmelCase__         =    1

											lowerCAmelCase__         =    AutoModelForCausalLM.from_pretrained(
											    '''EleutherAI/gpt-neo-125m'''			,				load_in_abit=__UpperCAmelCase			,				device_map=__UpperCAmelCase			,				)
											lowerCAmelCase__         =    Accelerator()

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

											PartialState._reset_state()




							@slow
							@require_bnb
							@require_multi_gpu
							def     __snake_case       ( self  : int						):
											from transformers import AutoModelForCausalLM

											with init_empty_weights():
															lowerCAmelCase__         =    AutoModelForCausalLM.from_pretrained(
															    '''EleutherAI/gpt-neo-125m'''			,				)
															lowerCAmelCase__         =    infer_auto_device_map(__UpperCAmelCase						)
															lowerCAmelCase__         =    1

											lowerCAmelCase__         =    AutoModelForCausalLM.from_pretrained(
											    '''EleutherAI/gpt-neo-125m'''			,				load_in_abit=__UpperCAmelCase			,				device_map=__UpperCAmelCase			,				)
											lowerCAmelCase__         =    Accelerator()

											# This should work
											lowerCAmelCase__         =    accelerator.prepare(__UpperCAmelCase						)




							@require_cuda
							def     __snake_case       ( self  : str						):
											lowerCAmelCase__         =    torch.nn.Linear(10			,				10						)
											lowerCAmelCase__         =    torch.optim.SGD(model.parameters()			,				lr=0.01						)
											lowerCAmelCase__         =    Accelerator(cpu=__UpperCAmelCase						)
											lowerCAmelCase__         =    accelerator.prepare(__UpperCAmelCase						)













# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation

import warnings

from .state import AcceleratorState, GradientState


warnings.filterwarnings("ignore", category=UserWarning, module="torch.optim.lr_scheduler")






class 						lowercase_   :






      '''simple docstring'''




      def __init__( self :		Optional[int]					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		bool = True					,   __UpperCAmelCase :		bool = False    )       ->Tuple:




       """simple docstring"""


       a						    = scheduler
       a						    = optimizers if isinstance(__UpperCAmelCase					,   (list, tuple)    ) else [optimizers]
       a						    = split_batches
       a						    = step_with_optimizer
       a						    = GradientState()







      def  __lowerCAmelCase   ( self :		Tuple					,   *__UpperCAmelCase :		List[Any]					,   **__UpperCAmelCase :		Optional[Any]    )       ->int:




       """simple docstring"""


       if not self.step_with_optimizer:
        # No link between scheduler and optimizer -> just step
        self.scheduler.step(*__UpperCAmelCase					,   **__UpperCAmelCase    )
        return

       # Otherwise, first make sure the optimizer was stepped.
       if not self.gradient_state.sync_gradients:
        if self.gradient_state.adjust_scheduler:
         self.scheduler._step_count += 1
        return

       for opt in self.optimizers:
        if opt.step_was_skipped:
         return
       if self.split_batches:
        # Split batches -> the training dataloader batch size is not changed so one step per training step
        self.scheduler.step(*__UpperCAmelCase					,   **__UpperCAmelCase    )
       else:
        # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
        # num_processes steps per training step
        a						    = AcceleratorState().num_processes
        for _ in range(__UpperCAmelCase    ):
         # Special case when using OneCycle and `drop_last` was not used
         if hasattr(self.scheduler					,   '''total_steps'''    ):
          if self.scheduler._step_count <= self.scheduler.total_steps:
           self.scheduler.step(*__UpperCAmelCase					,   **__UpperCAmelCase    )
         else:
          self.scheduler.step(*__UpperCAmelCase					,   **__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Optional[int]    )       ->List[str]:




       """simple docstring"""


       return self.scheduler.get_last_lr()







      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->Any:




       """simple docstring"""


       return self.scheduler.state_dict()







      def  __lowerCAmelCase   ( self :		Union[str, Any]					,   __UpperCAmelCase :		Any    )       ->Tuple:




       """simple docstring"""


       self.scheduler.load_state_dict(__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Dict    )       ->Optional[int]:




       """simple docstring"""


       return self.scheduler.get_lr()







      def  __lowerCAmelCase   ( self :		int					,   *__UpperCAmelCase :		List[str]					,   **__UpperCAmelCase :		Optional[Any]    )       ->List[str]:




       """simple docstring"""


       return self.scheduler.print_lr(*__UpperCAmelCase					,   **__UpperCAmelCase    )


