"""simple docstring"""




import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import (
          OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
          OpenAIGPTConfig,
          OpenAIGPTDoubleHeadsModel,
          OpenAIGPTForSequenceClassification,
          OpenAIGPTLMHeadModel,
          OpenAIGPTModel,
      )
class 						__A							:






     def __init__(		self    ,					a__    ,					a__=13    ,					a__=7    ,					a__=True    ,					a__=True    ,					a__=True    ,					a__=99    ,					a__=32    ,					a__=5    ,					a__=4    ,					a__=37    ,					a__="gelu"    ,					a__=0.1    ,					a__=0.1    ,					a__=512    ,					a__=16    ,					a__=2    ,					a__=0.0_2    ,					a__=3    ,					a__=4    ,					a__=None    ,					):
       _lowerCAmelCase						:	Optional[int]        =					parent
       _lowerCAmelCase						:	Union[str, Any]        =					batch_size
       _lowerCAmelCase						:	List[Any]        =					seq_length
       _lowerCAmelCase						:	Optional[int]        =					is_training
       _lowerCAmelCase						:	Optional[Any]        =					use_token_type_ids
       _lowerCAmelCase						:	Optional[int]        =					use_labels
       _lowerCAmelCase						:	Tuple        =					vocab_size
       _lowerCAmelCase						:	Any        =					hidden_size
       _lowerCAmelCase						:	Dict        =					num_hidden_layers
       _lowerCAmelCase						:	Dict        =					num_attention_heads
       _lowerCAmelCase						:	Optional[Any]        =					intermediate_size
       _lowerCAmelCase						:	Optional[Any]        =					hidden_act
       _lowerCAmelCase						:	List[str]        =					hidden_dropout_prob
       _lowerCAmelCase						:	List[str]        =					attention_probs_dropout_prob
       _lowerCAmelCase						:	Union[str, Any]        =					max_position_embeddings
       _lowerCAmelCase						:	Optional[Any]        =					type_vocab_size
       _lowerCAmelCase						:	List[str]        =					type_sequence_label_size
       _lowerCAmelCase						:	Dict        =					initializer_range
       _lowerCAmelCase						:	List[Any]        =					num_labels
       _lowerCAmelCase						:	Any        =					num_choices
       _lowerCAmelCase						:	Optional[int]        =					scope
       _lowerCAmelCase						:	str        =					self.vocab_size - 1






     def    __A  (		self  ):
       _lowerCAmelCase						:	List[str]        =					ids_tensor([self.batch_size, self.seq_length]    ,					self.vocab_size  )

       _lowerCAmelCase						:	Optional[Any]        =					None
       if self.use_token_type_ids:
         _lowerCAmelCase						:	Tuple        =					ids_tensor([self.batch_size, self.seq_length]    ,					self.type_vocab_size  )

       _lowerCAmelCase						:	str        =					None
       _lowerCAmelCase						:	Any        =					None
       _lowerCAmelCase						:	List[str]        =					None
       if self.use_labels:
         _lowerCAmelCase						:	List[Any]        =					ids_tensor([self.batch_size]    ,					self.type_sequence_label_size  )
         _lowerCAmelCase						:	Tuple        =					ids_tensor([self.batch_size, self.seq_length]    ,					self.num_labels  )
         _lowerCAmelCase						:	List[Any]        =					ids_tensor([self.batch_size]    ,					self.num_choices  )

       _lowerCAmelCase						:	Tuple        =					OpenAIGPTConfig(
           vocab_size=self.vocab_size    ,					n_embd=self.hidden_size    ,					n_layer=self.num_hidden_layers    ,					n_head=self.num_attention_heads    ,					n_positions=self.max_position_embeddings    ,					pad_token_id=self.pad_token_id    ,					)

       _lowerCAmelCase						:	str        =					ids_tensor([self.num_hidden_layers, self.num_attention_heads]    ,					2  )

       return (
           config,
           input_ids,
           head_mask,
           token_type_ids,
           sequence_labels,
           token_labels,
           choice_labels,
       )






     def    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					*a__  ):
       _lowerCAmelCase						:	int        =					OpenAIGPTModel(config=a__  )
       model.to(a__  )
       model.eval()

       _lowerCAmelCase						:	Optional[Any]        =					model(a__    ,					token_type_ids=a__    ,					head_mask=a__  )
       _lowerCAmelCase						:	Optional[Any]        =					model(a__    ,					token_type_ids=a__  )
       _lowerCAmelCase						:	int        =					model(a__  )

       self.parent.assertEqual(result.last_hidden_state.shape    ,					(self.batch_size, self.seq_length, self.hidden_size)  )






     def    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					*a__  ):
       _lowerCAmelCase						:	Tuple        =					OpenAIGPTLMHeadModel(a__  )
       model.to(a__  )
       model.eval()

       _lowerCAmelCase						:	Union[str, Any]        =					model(a__    ,					token_type_ids=a__    ,					labels=a__  )
       self.parent.assertEqual(result.loss.shape    ,					()  )
       self.parent.assertEqual(result.logits.shape    ,					(self.batch_size, self.seq_length, self.vocab_size)  )






     def    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					*a__  ):
       _lowerCAmelCase						:	Optional[Any]        =					OpenAIGPTDoubleHeadsModel(a__  )
       model.to(a__  )
       model.eval()

       _lowerCAmelCase						:	Any        =					model(a__    ,					token_type_ids=a__    ,					labels=a__  )
       self.parent.assertEqual(result.loss.shape    ,					()  )
       self.parent.assertEqual(result.logits.shape    ,					(self.batch_size, self.seq_length, self.vocab_size)  )






     def    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					*a__  ):
       _lowerCAmelCase						:	Optional[int]        =					self.num_labels
       _lowerCAmelCase						:	Any        =					OpenAIGPTForSequenceClassification(a__  )
       model.to(a__  )
       model.eval()

       _lowerCAmelCase						:	int        =					ids_tensor([self.batch_size]    ,					self.type_sequence_label_size  )
       _lowerCAmelCase						:	Any        =					model(a__    ,					token_type_ids=a__    ,					labels=a__  )
       self.parent.assertEqual(result.logits.shape    ,					(self.batch_size, self.num_labels)  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	Tuple        =					self.prepare_config_and_inputs()
       (
           _lowerCAmelCase     
       )						:	List[str]        =					config_and_inputs
       _lowerCAmelCase						:	List[Any]        =					{
           """input_ids""": input_ids,
           """token_type_ids""": token_type_ids,
           """head_mask""": head_mask,
       }

       return config, inputs_dict







@require_torch
class 						__A							(			SCREAMING_SNAKE_CASE_   ,     SCREAMING_SNAKE_CASE_   ,     SCREAMING_SNAKE_CASE_   ,     unittest.TestCase		):
     _UpperCamelCase							:      Any		 =		(
         (OpenAIGPTModel, OpenAIGPTLMHeadModel, OpenAIGPTDoubleHeadsModel, OpenAIGPTForSequenceClassification)
         if is_torch_available()
         else ()
     )
     _UpperCamelCase							:      List[str]		 =		(
         (OpenAIGPTLMHeadModel,) if is_torch_available() else ()
     )  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
     _UpperCamelCase							:      Tuple		 =		(
         {
             "feature-extraction": OpenAIGPTModel,
             "text-classification": OpenAIGPTForSequenceClassification,
             "text-generation": OpenAIGPTLMHeadModel,
             "zero-shot": OpenAIGPTForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )






     def    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__  ):
       if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
         # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
         # `OpenAIGPTConfig` was never used in pipeline tests, either because of a missing checkpoint or because a
         # tiny config could not be created.
         return True

       return False






     def    __A  (		self    ,					a__    ,					a__    ,					a__=False  ):
       _lowerCAmelCase						:	Tuple        =					super()._prepare_for_class(a__    ,					a__    ,					return_labels=a__  )

       if return_labels:
         if model_class.__name__ == "OpenAIGPTDoubleHeadsModel":
           _lowerCAmelCase						:	int        =					torch.zeros(
               (self.model_tester.batch_size, self.model_tester.num_choices, self.model_tester.seq_length)    ,					dtype=torch.long    ,					device=a__    ,					)
           _lowerCAmelCase						:	str        =					inputs_dict["""labels"""]
           _lowerCAmelCase						:	Tuple        =					inputs_dict["""labels"""]
           _lowerCAmelCase						:	Optional[int]        =					torch.zeros(
               (self.model_tester.batch_size, self.model_tester.num_choices)    ,					dtype=torch.long    ,					device=a__    ,					)
           _lowerCAmelCase						:	Any        =					torch.zeros(
               self.model_tester.batch_size    ,					dtype=torch.long    ,					device=a__  )
       return inputs_dict






     def    __A  (		self  ):
       _lowerCAmelCase						:	Any        =					OpenAIGPTModelTester(self  )
       _lowerCAmelCase						:	str        =					ConfigTester(self    ,					config_class=a__    ,					n_embd=37  )






     def    __A  (		self  ):
       self.config_tester.run_common_tests()






     def    __A  (		self  ):
       _lowerCAmelCase						:	str        =					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_openai_gpt_model(*a__  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	Optional[int]        =					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_lm_head_model(*a__  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	Any        =					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_double_lm_head_model(*a__  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	Tuple        =					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*a__  )






     @slow
     def    __A  (		self  ):
       for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         _lowerCAmelCase						:	int        =					OpenAIGPTModel.from_pretrained(a__  )
         self.assertIsNotNone(a__  )







@require_torch
class 						__A							(			unittest.TestCase		):






     @slow
     def    __A  (		self  ):
       _lowerCAmelCase						:	str        =					OpenAIGPTLMHeadModel.from_pretrained("""openai-gpt"""  )
       model.to(a__  )
       _lowerCAmelCase						:	Tuple        =					torch.tensor([[481, 4735, 544]]    ,					dtype=torch.long    ,					device=a__  )  # the president is
       _lowerCAmelCase						:	Any        =					[
           481,
           4735,
           544,
           246,
           963,
           870,
           762,
           239,
           244,
           40477,
           244,
           249,
           719,
           881,
           487,
           544,
           240,
           244,
           603,
           481,
       ]  # the president is a very good man. " \n " i\'m sure he is, " said the

       _lowerCAmelCase						:	Union[str, Any]        =					model.generate(a__    ,					do_sample=a__  )
       self.assertListEqual(output_ids[0].tolist()    ,					a__  )





"""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 							SCREAMING_SNAKE_CASE					(      )	->				Tuple:
  _lowerCAmelCase						:	Tuple        =					torch.nn.Linear(2      ,4     )
  _lowerCAmelCase						:	Union[str, Any]        =					torch.optim.AdamW(model.parameters()      ,lr=1.0     )
  _lowerCAmelCase						:	Tuple        =					torch.optim.lr_scheduler.OneCycleLR(_lowerCamelCase      ,max_lr=0.01      ,steps_per_epoch=2      ,epochs=1     )
  _lowerCAmelCase						:	Tuple        =					DataLoader(TensorDataset(torch.tensor([1, 2, 3]     )     )     )
  _lowerCAmelCase						:	List[Any]        =					DataLoader(TensorDataset(torch.tensor([4, 5, 6]     )     )     )

  return model, optimizer, scheduler, train_dl, valid_dl





def 							SCREAMING_SNAKE_CASE					(      _lowerCamelCase						:	Any     )	->				int:
  return (model.weight.abs().sum() + model.bias.abs().sum()).item()





def 							SCREAMING_SNAKE_CASE					(      _lowerCamelCase						:	List[str]     )	->				Any:
  _lowerCAmelCase						:	List[str]        =					torch.nn.Linear(*tuple(model.weight.T.shape     )     ).state_dict()
  model.load_state_dict(_lowerCamelCase     )






class 						__A							(			SCREAMING_SNAKE_CASE_		):






     @require_cuda
     def    __A  (		self  ):
       _lowerCAmelCase						:	Union[str, Any]        =					Accelerator()
       assert PartialState._shared_state["_cpu"] is False
       assert PartialState._shared_state["device"].type == "cuda"
       with self.assertRaises(a__  ):
         _lowerCAmelCase						:	Tuple        =					Accelerator(cpu=a__  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	Dict        =					Accelerator()
       _lowerCAmelCase						:	Any        =					GradientState()
       assert state.num_steps == 1
       _lowerCAmelCase						:	Optional[int]        =					4
       assert state.num_steps == 4

       assert state.sync_gradients is True
       _lowerCAmelCase						:	Dict        =					False
       assert state.sync_gradients is False
       GradientState._reset_state()






     def    __A  (		self  ):
       _lowerCAmelCase						:	Optional[int]        =					Accelerator()
       _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase						:	Any        =					create_components()

       (
           (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     
       )						:	int        =					accelerator.prepare(a__    ,					a__    ,					a__    ,					a__    ,					a__  )

       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    __A  (		self  ):
       _lowerCAmelCase						:	Optional[Any]        =					Accelerator()
       _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase						:	List[Any]        =					create_components()
       accelerator.prepare(a__    ,					a__    ,					a__    ,					a__    ,					a__  )
       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    __A  (		self  ):
       PartialState._reset_state()

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

       with patch("""torch.cuda.set_device"""    ,					a__  ), patch_environment(ACCELERATE_TORCH_DEVICE="""cuda:64"""  ):
         _lowerCAmelCase						:	Dict        =					Accelerator()
         self.assertEqual(str(accelerator.state.device  )    ,					"""cuda:64"""  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	Any        =					Accelerator()
       _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase						:	List[str]        =					create_components()
       accelerator.prepare(a__    ,					a__    ,					a__    ,					a__    ,					a__  )

       _lowerCAmelCase						:	List[Any]        =					get_signature(a__  )

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

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

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






     def    __A  (		self  ):
       _lowerCAmelCase						:	str        =					Accelerator()
       _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase						:	str        =					create_components()
       accelerator.prepare(a__    ,					a__    ,					a__    ,					a__    ,					a__  )

       _lowerCAmelCase						:	Optional[Any]        =					get_signature(a__  )

       # saving hook
       def save_config(a__    ,					a__    ,					a__  ):
         _lowerCAmelCase						:	Dict        =					{"""class_name""": models[0].__class__.__name__}

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

        # loading hook
       def load_config(a__    ,					a__  ):
         with open(os.path.join(a__    ,					"""data.json"""  )    ,					"""r"""  ) as f:
           _lowerCAmelCase						:	int        =					json.load(a__  )

         _lowerCAmelCase						:	str        =					config["""class_name"""]

       _lowerCAmelCase						:	Union[str, Any]        =					accelerator.register_save_state_pre_hook(a__  )
       _lowerCAmelCase						:	int        =					accelerator.register_load_state_pre_hook(a__  )

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

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

         # random class name to verify correct one is loaded
         _lowerCAmelCase						:	Dict        =					"""random"""

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

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

         # random class name to verify correct one is loaded
         _lowerCAmelCase						:	Optional[Any]        =					"""random"""

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

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






     def    __A  (		self  ):
       _lowerCAmelCase						:	Optional[Any]        =					Accelerator()
       _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase						:	Any        =					create_components()
       _lowerCAmelCase						:	Any        =					None

       # This should work
       _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase						:	int        =					accelerator.prepare(
           a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__  )
       self.assertTrue(dummy_obj is None  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	str        =					Accelerator()
       _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase						:	Any        =					create_components()
       _lowerCAmelCase						:	Optional[int]        =					[1, 2, 3]

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






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

       _lowerCAmelCase						:	List[str]        =					AutoModelForCausalLM.from_pretrained(
           """EleutherAI/gpt-neo-125m"""    ,					load_in_abit=a__    ,					device_map={"""""": 0}    ,					)
       _lowerCAmelCase						:	List[str]        =					Accelerator()

       # This should work
       _lowerCAmelCase						:	List[Any]        =					accelerator.prepare(a__  )






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

       _lowerCAmelCase						:	Any        =					Accelerator()

       with init_empty_weights():
         _lowerCAmelCase						:	Dict        =					AutoModelForCausalLM.from_pretrained(
             """EleutherAI/gpt-neo-125m"""    ,					)
         model.tie_weights()
         _lowerCAmelCase						:	int        =					infer_auto_device_map(a__  )
         _lowerCAmelCase						:	Optional[Any]        =					"""cpu"""

       _lowerCAmelCase						:	Dict        =					AutoModelForCausalLM.from_pretrained(
           """EleutherAI/gpt-neo-125m"""    ,					device_map=a__    ,					load_in_abit=a__    ,					llm_inta_enable_fpaa_cpu_offload=a__  )

       # This should not work and get value error
       with self.assertRaises(a__  ):
         _lowerCAmelCase						:	List[str]        =					accelerator.prepare(a__  )






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

       _lowerCAmelCase						:	Dict        =					{"""distributed_type""": DistributedType.MULTI_GPU}

       with init_empty_weights():
         _lowerCAmelCase						:	Optional[Any]        =					AutoModelForCausalLM.from_pretrained(
             """EleutherAI/gpt-neo-125m"""    ,					)
         model.tie_weights()
         _lowerCAmelCase						:	List[str]        =					infer_auto_device_map(a__  )
         _lowerCAmelCase						:	Union[str, Any]        =					1

       _lowerCAmelCase						:	int        =					AutoModelForCausalLM.from_pretrained(
           """EleutherAI/gpt-neo-125m"""    ,					load_in_abit=a__    ,					device_map=a__    ,					)
       _lowerCAmelCase						:	Tuple        =					Accelerator()

       # This should not work and get value error
       with self.assertRaises(a__  ):
         _lowerCAmelCase						:	Optional[int]        =					accelerator.prepare(a__  )

       PartialState._reset_state()






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

       with init_empty_weights():
         _lowerCAmelCase						:	Dict        =					AutoModelForCausalLM.from_pretrained(
             """EleutherAI/gpt-neo-125m"""    ,					)
         _lowerCAmelCase						:	int        =					infer_auto_device_map(a__  )
         _lowerCAmelCase						:	List[Any]        =					1

       _lowerCAmelCase						:	Optional[Any]        =					AutoModelForCausalLM.from_pretrained(
           """EleutherAI/gpt-neo-125m"""    ,					load_in_abit=a__    ,					device_map=a__    ,					)
       _lowerCAmelCase						:	str        =					Accelerator()

       # This should work
       _lowerCAmelCase						:	str        =					accelerator.prepare(a__  )






     @require_cuda
     def    __A  (		self  ):
       _lowerCAmelCase						:	Union[str, Any]        =					torch.nn.Linear(10    ,					10  )
       _lowerCAmelCase						:	Any        =					torch.optim.SGD(model.parameters()    ,					lr=0.0_1  )
       _lowerCAmelCase						:	List[str]        =					Accelerator(cpu=a__  )
       _lowerCAmelCase						:	Tuple        =					accelerator.prepare(a__  )




