


import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu






class   __lowercase				(  nn.Module  ):





      def __init__(		self						:					Dict       )    ->   Any:






         """simple docstring"""







         super().__init__()
         UpperCAmelCase							=     nn.Linear(3					,       4       )
         UpperCAmelCase							=     nn.BatchNormad(4       )
         UpperCAmelCase							=     nn.Linear(4					,       5       )





      def        _lowercase				(		self						:					Optional[int]					,       __lowerCamelCase						:					Dict       )    ->   Any:






         """simple docstring"""







         return self.lineara(self.batchnorm(self.lineara(UpperCamelCase__       )       )       )






class   __lowercase				(  __A  ):





      def        _lowercase				(		self						:					int					,       __lowerCamelCase						:					List[Any]					,       *__lowerCamelCase						:					Union[str, Any]					,       **__lowerCamelCase						:					Tuple       )    ->   Tuple:






         """simple docstring"""







         return (args[0] + 1,) + args[1:], kwargs






class   __lowercase				(  __A  ):





      def        _lowercase				(		self						:					Optional[Any]					,       __lowerCamelCase						:					Any					,       __lowerCamelCase						:					Dict       )    ->   Dict:






         """simple docstring"""







         return output + 1






class   __lowercase				(  unittest.TestCase  ):





      def        _lowercase				(		self						:					Union[str, Any]       )    ->   Dict:






         """simple docstring"""







         UpperCAmelCase							=     ModelForTest()
         UpperCAmelCase							=     ModelHook()

         add_hook_to_module(UpperCamelCase__					,       UpperCamelCase__       )
         self.assertEqual(test_model._hf_hook					,       UpperCamelCase__       )
         self.assertTrue(hasattr(UpperCamelCase__					,       """_old_forward"""       )       )

         # Check adding the hook did not change the name or the signature
         self.assertEqual(test_model.forward.__name__					,       """forward"""       )
         self.assertListEqual(list(inspect.signature(test_model.forward       ).parameters       )					,       ["""x"""]       )

         remove_hook_from_module(UpperCamelCase__       )
         self.assertFalse(hasattr(UpperCamelCase__					,       """_hf_hook"""       )       )
         self.assertFalse(hasattr(UpperCamelCase__					,       """_old_forward"""       )       )





      def        _lowercase				(		self						:					str       )    ->   Optional[Any]:






         """simple docstring"""







         UpperCAmelCase							=     ModelForTest()
         UpperCAmelCase							=     ModelHook()

         add_hook_to_module(UpperCamelCase__					,       UpperCamelCase__       )
         add_hook_to_module(UpperCamelCase__					,       UpperCamelCase__					,       append=UpperCamelCase__       )

         self.assertEqual(isinstance(test_model._hf_hook					,       UpperCamelCase__       )					,       UpperCamelCase__       )
         self.assertEqual(len(test_model._hf_hook.hooks       )					,       2       )
         self.assertTrue(hasattr(UpperCamelCase__					,       """_old_forward"""       )       )

         # Check adding the hook did not change the name or the signature
         self.assertEqual(test_model.forward.__name__					,       """forward"""       )
         self.assertListEqual(list(inspect.signature(test_model.forward       ).parameters       )					,       ["""x"""]       )

         remove_hook_from_module(UpperCamelCase__       )
         self.assertFalse(hasattr(UpperCamelCase__					,       """_hf_hook"""       )       )
         self.assertFalse(hasattr(UpperCamelCase__					,       """_old_forward"""       )       )





      def        _lowercase				(		self						:					str       )    ->   Any:






         """simple docstring"""







         UpperCAmelCase							=     ModelForTest()
         UpperCAmelCase							=     torch.randn(2					,       3       )
         UpperCAmelCase							=     test_model(x + 1       )
         UpperCAmelCase							=     test_model(x + 2       )

         UpperCAmelCase							=     PreForwardHook()
         add_hook_to_module(UpperCamelCase__					,       UpperCamelCase__       )
         UpperCAmelCase							=     test_model(UpperCamelCase__       )
         self.assertTrue(torch.allclose(UpperCamelCase__					,       UpperCamelCase__					,       atol=1e-5       )       )

         # Attaching a hook to a model when it already has one replaces, does not chain
         UpperCAmelCase							=     PreForwardHook()
         add_hook_to_module(UpperCamelCase__					,       UpperCamelCase__       )
         UpperCAmelCase							=     test_model(UpperCamelCase__       )
         self.assertTrue(torch.allclose(UpperCamelCase__					,       UpperCamelCase__					,       atol=1e-5       )       )

         # You need to use the sequential hook to chain two or more hooks
         UpperCAmelCase							=     SequentialHook(PreForwardHook()					,       PreForwardHook()       )
         add_hook_to_module(UpperCamelCase__					,       UpperCamelCase__       )

         UpperCAmelCase							=     test_model(UpperCamelCase__       )
         assert torch.allclose(UpperCamelCase__					,       UpperCamelCase__					,       atol=1e-5       )





      def        _lowercase				(		self						:					List[str]       )    ->   List[Any]:






         """simple docstring"""







         UpperCAmelCase							=     ModelForTest()
         UpperCAmelCase							=     torch.randn(2					,       3       )
         UpperCAmelCase							=     test_model(UpperCamelCase__       )

         UpperCAmelCase							=     PostForwardHook()
         add_hook_to_module(UpperCamelCase__					,       UpperCamelCase__       )
         UpperCAmelCase							=     test_model(UpperCamelCase__       )
         self.assertTrue(torch.allclose(UpperCamelCase__					,       output + 1					,       atol=1e-5       )       )

         # Attaching a hook to a model when it already has one replaces, does not chain
         UpperCAmelCase							=     PostForwardHook()
         add_hook_to_module(UpperCamelCase__					,       UpperCamelCase__       )
         UpperCAmelCase							=     test_model(UpperCamelCase__       )
         self.assertTrue(torch.allclose(UpperCamelCase__					,       output + 1					,       atol=1e-5       )       )

         # You need to use the sequential hook to chain two or more hooks
         UpperCAmelCase							=     SequentialHook(PostForwardHook()					,       PostForwardHook()       )
         add_hook_to_module(UpperCamelCase__					,       UpperCamelCase__       )

         UpperCAmelCase							=     test_model(UpperCamelCase__       )
         assert torch.allclose(UpperCamelCase__					,       output + 2					,       atol=1e-5       )





      def        _lowercase				(		self						:					Optional[Any]       )    ->   str:






         """simple docstring"""







         UpperCAmelCase							=     ModelForTest()
         UpperCAmelCase							=     torch.randn(2					,       3       )
         UpperCAmelCase							=     test_model(UpperCamelCase__       )

         UpperCAmelCase							=     PostForwardHook()
         add_hook_to_module(UpperCamelCase__					,       UpperCamelCase__       )
         UpperCAmelCase							=     test_model(UpperCamelCase__       )
         self.assertTrue(torch.allclose(UpperCamelCase__					,       output + 1       )       )
         self.assertTrue(outputa.requires_grad       )

         UpperCAmelCase							=     True
         UpperCAmelCase							=     test_model(UpperCamelCase__       )
         self.assertFalse(outputa.requires_grad       )





      @require_multi_gpu
      def        _lowercase				(		self						:					Optional[Any]       )    ->   Tuple:






         """simple docstring"""







         UpperCAmelCase							=     ModelForTest()
         # Everything is on CPU
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )

         # This will move each submodule on different devices
         add_hook_to_module(model.lineara					,       AlignDevicesHook(execution_device=0       )       )
         add_hook_to_module(model.batchnorm					,       AlignDevicesHook(execution_device=0       )       )
         add_hook_to_module(model.lineara					,       AlignDevicesHook(execution_device=1       )       )

         self.assertEqual(model.lineara.weight.device					,       torch.device(0       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device(0       )       )
         self.assertEqual(model.batchnorm.running_mean.device					,       torch.device(0       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device(1       )       )

         # We can still make a forward pass. The input does not need to be on any particular device
         UpperCAmelCase							=     torch.randn(2					,       3       )
         UpperCAmelCase							=     model(UpperCamelCase__       )
         self.assertEqual(output.device					,       torch.device(1       )       )

         # We can add a general hook to put back output on same device as input.
         add_hook_to_module(UpperCamelCase__					,       AlignDevicesHook(io_same_device=UpperCamelCase__       )       )
         UpperCAmelCase							=     torch.randn(2					,       3       ).to(0       )
         UpperCAmelCase							=     model(UpperCamelCase__       )
         self.assertEqual(output.device					,       torch.device(0       )       )





      def        _lowercase				(		self						:					Union[str, Any]       )    ->   Optional[int]:






         """simple docstring"""







         UpperCAmelCase							=     ModelForTest()

         # Everything is on CPU
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )

         # This will move each submodule on different devices
         UpperCAmelCase							=     {'execution_device': 0 if torch.cuda.is_available() else 'cpu', 'offload': True}

         add_hook_to_module(model.lineara					,       AlignDevicesHook(**UpperCamelCase__       )       )
         add_hook_to_module(model.batchnorm					,       AlignDevicesHook(**UpperCamelCase__       )       )
         add_hook_to_module(model.lineara					,       AlignDevicesHook(**UpperCamelCase__       )       )

         # Parameters have been offloaded, so on the meta device
         self.assertEqual(model.lineara.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""meta"""       )       )
         # Buffers are not included in the offload by default, so are on the execution device
         UpperCAmelCase							=     torch.device(hook_kwargs["""execution_device"""]       )
         self.assertEqual(model.batchnorm.running_mean.device					,       UpperCamelCase__       )

         UpperCAmelCase							=     torch.randn(2					,       3       )
         UpperCAmelCase							=     model(UpperCamelCase__       )
         self.assertEqual(output.device					,       UpperCamelCase__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_module(model.lineara       )
         remove_hook_from_module(model.batchnorm       )
         remove_hook_from_module(model.lineara       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )

         # Now test with buffers included in the offload
         UpperCAmelCase							=     {
             'execution_device': 0 if torch.cuda.is_available() else 'cpu',
             'offload': True,
             'offload_buffers': True,
         }

         add_hook_to_module(model.lineara					,       AlignDevicesHook(**UpperCamelCase__       )       )
         add_hook_to_module(model.batchnorm					,       AlignDevicesHook(**UpperCamelCase__       )       )
         add_hook_to_module(model.lineara					,       AlignDevicesHook(**UpperCamelCase__       )       )

         # Parameters have been offloaded, so on the meta device, buffers included
         self.assertEqual(model.lineara.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.batchnorm.running_mean.device					,       torch.device("""meta"""       )       )

         UpperCAmelCase							=     torch.randn(2					,       3       )
         UpperCAmelCase							=     model(UpperCamelCase__       )
         self.assertEqual(output.device					,       UpperCamelCase__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_module(model.lineara       )
         remove_hook_from_module(model.batchnorm       )
         remove_hook_from_module(model.lineara       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )





      def        _lowercase				(		self						:					List[str]       )    ->   Any:






         """simple docstring"""







         UpperCAmelCase							=     ModelForTest()

         # Everything is on CPU
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )

         # This will move each submodule on different devices
         UpperCAmelCase							=     0 if torch.cuda.is_available() else 'cpu'
         attach_align_device_hook(UpperCamelCase__					,       execution_device=UpperCamelCase__					,       offload=UpperCamelCase__       )

         # Parameters have been offloaded, so on the meta device
         self.assertEqual(model.lineara.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""meta"""       )       )
         # Buffers are not included in the offload by default, so are on the execution device
         UpperCAmelCase							=     torch.device(UpperCamelCase__       )
         self.assertEqual(model.batchnorm.running_mean.device					,       UpperCamelCase__       )

         UpperCAmelCase							=     torch.randn(2					,       3       )
         UpperCAmelCase							=     model(UpperCamelCase__       )
         self.assertEqual(output.device					,       UpperCamelCase__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_submodules(UpperCamelCase__       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )

         # Now test with buffers included in the offload
         attach_align_device_hook(UpperCamelCase__					,       execution_device=UpperCamelCase__					,       offload=UpperCamelCase__					,       offload_buffers=UpperCamelCase__       )

         # Parameters have been offloaded, so on the meta device, buffers included
         self.assertEqual(model.lineara.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.batchnorm.running_mean.device					,       torch.device("""meta"""       )       )

         UpperCAmelCase							=     torch.randn(2					,       3       )
         UpperCAmelCase							=     model(UpperCamelCase__       )
         self.assertEqual(output.device					,       UpperCamelCase__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_submodules(UpperCamelCase__       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )





      def        _lowercase				(		self						:					Dict       )    ->   Union[str, Any]:






         """simple docstring"""







         UpperCAmelCase							=     ModelForTest()

         # Everything is on CPU
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )

         # This will move each submodule on different devices
         UpperCAmelCase							=     0 if torch.cuda.is_available() else 'cpu'
         attach_align_device_hook(
             UpperCamelCase__					,       execution_device=UpperCamelCase__					,       offload=UpperCamelCase__					,       weights_map=model.state_dict()       )

         # Parameters have been offloaded, so on the meta device
         self.assertEqual(model.lineara.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""meta"""       )       )
         # Buffers are not included in the offload by default, so are on the execution device
         UpperCAmelCase							=     torch.device(UpperCamelCase__       )
         self.assertEqual(model.batchnorm.running_mean.device					,       UpperCamelCase__       )

         UpperCAmelCase							=     torch.randn(2					,       3       )
         UpperCAmelCase							=     model(UpperCamelCase__       )
         self.assertEqual(output.device					,       UpperCamelCase__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_submodules(UpperCamelCase__       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )

         # Now test with buffers included in the offload
         attach_align_device_hook(
             UpperCamelCase__					,       execution_device=UpperCamelCase__					,       offload=UpperCamelCase__					,       weights_map=model.state_dict()					,       offload_buffers=UpperCamelCase__					,       )

         # Parameters have been offloaded, so on the meta device, buffers included
         self.assertEqual(model.lineara.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""meta"""       )       )
         self.assertEqual(model.batchnorm.running_mean.device					,       torch.device("""meta"""       )       )

         UpperCAmelCase							=     torch.randn(2					,       3       )
         UpperCAmelCase							=     model(UpperCamelCase__       )
         self.assertEqual(output.device					,       UpperCamelCase__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_submodules(UpperCamelCase__       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.batchnorm.weight.device					,       torch.device("""cpu"""       )       )
         self.assertEqual(model.lineara.weight.device					,       torch.device("""cpu"""       )       )







from ...configuration_utils import PretrainedConfig
from ...utils import logging


__a     =			logging.get_logger(__name__)

__a     =			{"""ctrl""": """https://huggingface.co/ctrl/resolve/main/config.json"""}






class   __lowercase				(  __snake_case  ):
      UpperCamelCase								=       '''ctrl'''
      UpperCamelCase								=       ['''past_key_values''']
      UpperCamelCase								=       {
          '''max_position_embeddings''': '''n_positions''',
          '''hidden_size''': '''n_embd''',
          '''num_attention_heads''': '''n_head''',
          '''num_hidden_layers''': '''n_layer''',
      }





      def __init__(		self						:					str					,       __lowerCamelCase						:					Optional[int]=2_4_6_5_3_4					,       __lowerCamelCase						:					Union[str, Any]=2_5_6					,       __lowerCamelCase						:					int=1_2_8_0					,       __lowerCamelCase						:					Optional[Any]=8_1_9_2					,       __lowerCamelCase						:					List[str]=4_8					,       __lowerCamelCase						:					Dict=1_6					,       __lowerCamelCase						:					Union[str, Any]=0.1					,       __lowerCamelCase						:					Union[str, Any]=0.1					,       __lowerCamelCase						:					Union[str, Any]=1e-6					,       __lowerCamelCase						:					str=0.02					,       __lowerCamelCase						:					Union[str, Any]=True					,       **__lowerCamelCase						:					List[str]					,       )    ->   Optional[Any]:






         """simple docstring"""







         UpperCAmelCase							=     vocab_size
         UpperCAmelCase							=     n_positions
         UpperCAmelCase							=     n_embd
         UpperCAmelCase							=     n_layer
         UpperCAmelCase							=     n_head
         UpperCAmelCase							=     dff
         UpperCAmelCase							=     resid_pdrop
         UpperCAmelCase							=     embd_pdrop
         UpperCAmelCase							=     layer_norm_epsilon
         UpperCAmelCase							=     initializer_range

         UpperCAmelCase							=     use_cache

         super().__init__(**__lowerCamelCase       )



