

"""simple docstring"""





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 							SCREAMING_SNAKE_CASE   (						nn.Module			):






					"""simple docstring"""



					def __init__(			self     :							Optional[int]						):
						super().__init__()
						lowerCAmelCase__      :		List[str]    							=      nn.Linear(3  ,4						)
						lowerCAmelCase__      :		List[str]    							=      nn.BatchNormad(4						)
						lowerCAmelCase__      :		Union[str, Any]    							=      nn.Linear(4  ,5						)







					def 			__lowerCAmelCase					(			self     :							Optional[int]  ,lowercase_     :							Any						):
						return self.lineara(self.batchnorm(self.lineara(lowercase_						)						)						)







class 							SCREAMING_SNAKE_CASE   (						a_			):






					"""simple docstring"""



					def 			__lowerCAmelCase					(			self     :							Union[str, Any]  ,lowercase_     :							Tuple  ,*lowercase_     :							Optional[int]  ,**lowercase_     :							Union[str, Any]						):
						return (args[0] + 1,) + args[1:], kwargs







class 							SCREAMING_SNAKE_CASE   (						a_			):






					"""simple docstring"""



					def 			__lowerCAmelCase					(			self     :							str  ,lowercase_     :							int  ,lowercase_     :							Dict						):
						return output + 1







class 							SCREAMING_SNAKE_CASE   (						unittest.TestCase			):






					"""simple docstring"""



					def 			__lowerCAmelCase					(			self     :							int						):
						lowerCAmelCase__      :		List[str]    							=      ModelForTest()
						lowerCAmelCase__      :		List[str]    							=      ModelHook()

						add_hook_to_module(lowercase_  ,lowercase_						)
						self.assertEqual(test_model._hf_hook  ,lowercase_						)
						self.assertTrue(hasattr(lowercase_  ,'''_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(lowercase_						)
						self.assertFalse(hasattr(lowercase_  ,'''_hf_hook'''						)						)
						self.assertFalse(hasattr(lowercase_  ,'''_old_forward'''						)						)







					def 			__lowerCAmelCase					(			self     :							Optional[Any]						):
						lowerCAmelCase__      :		str    							=      ModelForTest()
						lowerCAmelCase__      :		List[str]    							=      ModelHook()

						add_hook_to_module(lowercase_  ,lowercase_						)
						add_hook_to_module(lowercase_  ,lowercase_  ,append=lowercase_						)

						self.assertEqual(isinstance(test_model._hf_hook  ,lowercase_						)  ,lowercase_						)
						self.assertEqual(len(test_model._hf_hook.hooks						)  ,2						)
						self.assertTrue(hasattr(lowercase_  ,'''_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(lowercase_						)
						self.assertFalse(hasattr(lowercase_  ,'''_hf_hook'''						)						)
						self.assertFalse(hasattr(lowercase_  ,'''_old_forward'''						)						)







					def 			__lowerCAmelCase					(			self     :							Dict						):
						lowerCAmelCase__      :		Union[str, Any]    							=      ModelForTest()
						lowerCAmelCase__      :		List[str]    							=      torch.randn(2  ,3						)
						lowerCAmelCase__      :		List[Any]    							=      test_model(x + 1						)
						lowerCAmelCase__      :		List[Any]    							=      test_model(x + 2						)

						lowerCAmelCase__      :		List[Any]    							=      PreForwardHook()
						add_hook_to_module(lowercase_  ,lowercase_						)
						lowerCAmelCase__      :		int    							=      test_model(lowercase_						)
						self.assertTrue(torch.allclose(lowercase_  ,lowercase_  ,atol=1E-5						)						)

						# Attaching a hook to a model when it already has one replaces, does not chain
						lowerCAmelCase__      :		Dict    							=      PreForwardHook()
						add_hook_to_module(lowercase_  ,lowercase_						)
						lowerCAmelCase__      :		Optional[Any]    							=      test_model(lowercase_						)
						self.assertTrue(torch.allclose(lowercase_  ,lowercase_  ,atol=1E-5						)						)

						# You need to use the sequential hook to chain two or more hooks
						lowerCAmelCase__      :		Union[str, Any]    							=      SequentialHook(PreForwardHook()  ,PreForwardHook()						)
						add_hook_to_module(lowercase_  ,lowercase_						)

						lowerCAmelCase__      :		Optional[Any]    							=      test_model(lowercase_						)
						assert torch.allclose(lowercase_  ,lowercase_  ,atol=1E-5						)







					def 			__lowerCAmelCase					(			self     :							Tuple						):
						lowerCAmelCase__      :		Union[str, Any]    							=      ModelForTest()
						lowerCAmelCase__      :		Dict    							=      torch.randn(2  ,3						)
						lowerCAmelCase__      :		Optional[Any]    							=      test_model(lowercase_						)

						lowerCAmelCase__      :		Optional[Any]    							=      PostForwardHook()
						add_hook_to_module(lowercase_  ,lowercase_						)
						lowerCAmelCase__      :		Optional[Any]    							=      test_model(lowercase_						)
						self.assertTrue(torch.allclose(lowercase_  ,output + 1  ,atol=1E-5						)						)

						# Attaching a hook to a model when it already has one replaces, does not chain
						lowerCAmelCase__      :		Any    							=      PostForwardHook()
						add_hook_to_module(lowercase_  ,lowercase_						)
						lowerCAmelCase__      :		Dict    							=      test_model(lowercase_						)
						self.assertTrue(torch.allclose(lowercase_  ,output + 1  ,atol=1E-5						)						)

						# You need to use the sequential hook to chain two or more hooks
						lowerCAmelCase__      :		Union[str, Any]    							=      SequentialHook(PostForwardHook()  ,PostForwardHook()						)
						add_hook_to_module(lowercase_  ,lowercase_						)

						lowerCAmelCase__      :		str    							=      test_model(lowercase_						)
						assert torch.allclose(lowercase_  ,output + 2  ,atol=1E-5						)







					def 			__lowerCAmelCase					(			self     :							Dict						):
						lowerCAmelCase__      :		Any    							=      ModelForTest()
						lowerCAmelCase__      :		List[Any]    							=      torch.randn(2  ,3						)
						lowerCAmelCase__      :		Optional[int]    							=      test_model(lowercase_						)

						lowerCAmelCase__      :		Optional[int]    							=      PostForwardHook()
						add_hook_to_module(lowercase_  ,lowercase_						)
						lowerCAmelCase__      :		str    							=      test_model(lowercase_						)
						self.assertTrue(torch.allclose(lowercase_  ,output + 1						)						)
						self.assertTrue(outputa.requires_grad						)

						lowerCAmelCase__      :		List[Any]    							=      True
						lowerCAmelCase__      :		Optional[Any]    							=      test_model(lowercase_						)
						self.assertFalse(outputa.requires_grad						)







					@require_multi_gpu
					def 			__lowerCAmelCase					(			self     :							Optional[int]						):
						lowerCAmelCase__      :		List[str]    							=      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
						lowerCAmelCase__      :		Any    							=      torch.randn(2  ,3						)
						lowerCAmelCase__      :		Any    							=      model(lowercase_						)
						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(lowercase_  ,AlignDevicesHook(io_same_device=lowercase_						)						)
						lowerCAmelCase__      :		Any    							=      torch.randn(2  ,3						).to(0						)
						lowerCAmelCase__      :		Optional[int]    							=      model(lowercase_						)
						self.assertEqual(output.device  ,torch.device(0						)						)







					def 			__lowerCAmelCase					(			self     :							Tuple						):
						lowerCAmelCase__      :		Any    							=      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
						lowerCAmelCase__      :		List[str]    							=      {'''execution_device''': 0 if torch.cuda.is_available() else '''cpu''', '''offload''': True}

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

						# 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
						lowerCAmelCase__      :		Optional[int]    							=      torch.device(hook_kwargs['''execution_device''']						)
						self.assertEqual(model.batchnorm.running_mean.device  ,lowercase_						)

						lowerCAmelCase__      :		Optional[int]    							=      torch.randn(2  ,3						)
						lowerCAmelCase__      :		Optional[int]    							=      model(lowercase_						)
						self.assertEqual(output.device  ,lowercase_						)

						# 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
						lowerCAmelCase__      :		Optional[int]    							=      {
						    '''execution_device''': 0 if torch.cuda.is_available() else '''cpu''',
						    '''offload''': True,
						    '''offload_buffers''': True,
						}

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

						# 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'''						)						)

						lowerCAmelCase__      :		List[str]    							=      torch.randn(2  ,3						)
						lowerCAmelCase__      :		List[str]    							=      model(lowercase_						)
						self.assertEqual(output.device  ,lowercase_						)

						# 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 			__lowerCAmelCase					(			self     :							Optional[int]						):
						lowerCAmelCase__      :		Optional[Any]    							=      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
						lowerCAmelCase__      :		Optional[int]    							=      0 if torch.cuda.is_available() else '''cpu'''
						attach_align_device_hook(lowercase_  ,execution_device=lowercase_  ,offload=lowercase_						)

						# 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
						lowerCAmelCase__      :		Any    							=      torch.device(lowercase_						)
						self.assertEqual(model.batchnorm.running_mean.device  ,lowercase_						)

						lowerCAmelCase__      :		str    							=      torch.randn(2  ,3						)
						lowerCAmelCase__      :		List[Any]    							=      model(lowercase_						)
						self.assertEqual(output.device  ,lowercase_						)

						# Removing hooks loads back the weights in the model.
						remove_hook_from_submodules(lowercase_						)
						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(lowercase_  ,execution_device=lowercase_  ,offload=lowercase_  ,offload_buffers=lowercase_						)

						# 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'''						)						)

						lowerCAmelCase__      :		List[str]    							=      torch.randn(2  ,3						)
						lowerCAmelCase__      :		Dict    							=      model(lowercase_						)
						self.assertEqual(output.device  ,lowercase_						)

						# Removing hooks loads back the weights in the model.
						remove_hook_from_submodules(lowercase_						)
						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 			__lowerCAmelCase					(			self     :							int						):
						lowerCAmelCase__      :		Union[str, Any]    							=      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
						lowerCAmelCase__      :		Tuple    							=      0 if torch.cuda.is_available() else '''cpu'''
						attach_align_device_hook(
						    lowercase_  ,execution_device=lowercase_  ,offload=lowercase_  ,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
						lowerCAmelCase__      :		Optional[int]    							=      torch.device(lowercase_						)
						self.assertEqual(model.batchnorm.running_mean.device  ,lowercase_						)

						lowerCAmelCase__      :		Optional[Any]    							=      torch.randn(2  ,3						)
						lowerCAmelCase__      :		int    							=      model(lowercase_						)
						self.assertEqual(output.device  ,lowercase_						)

						# Removing hooks loads back the weights in the model.
						remove_hook_from_submodules(lowercase_						)
						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(
						    lowercase_  ,execution_device=lowercase_  ,offload=lowercase_  ,weights_map=model.state_dict()  ,offload_buffers=lowercase_  ,)

						# 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'''						)						)

						lowerCAmelCase__      :		List[str]    							=      torch.randn(2  ,3						)
						lowerCAmelCase__      :		Any    							=      model(lowercase_						)
						self.assertEqual(output.device  ,lowercase_						)

						# Removing hooks loads back the weights in the model.
						remove_hook_from_submodules(lowercase_						)
						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'''						)						)







"""simple docstring"""





from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


__UpperCamelCase    : List[Any]				 =						logging.get_logger(__name__)

__UpperCamelCase    : Any				 =						{
    '''t5-small''': '''https://huggingface.co/t5-small/resolve/main/config.json''',
    '''t5-base''': '''https://huggingface.co/t5-base/resolve/main/config.json''',
    '''t5-large''': '''https://huggingface.co/t5-large/resolve/main/config.json''',
    '''t5-3b''': '''https://huggingface.co/t5-3b/resolve/main/config.json''',
    '''t5-11b''': '''https://huggingface.co/t5-11b/resolve/main/config.json''',
}







class 							SCREAMING_SNAKE_CASE   (						a_			):






					"""simple docstring"""



					lowercase__												=			"t5"
					lowercase__												=			["past_key_values"]
					lowercase__												=			{"hidden_size": "d_model", "num_attention_heads": "num_heads", "num_hidden_layers": "num_layers"}







					def __init__(			self     :							Tuple  ,lowercase_     :							Tuple=3_2_1_2_8  ,lowercase_     :							Union[str, Any]=5_1_2  ,lowercase_     :							Dict=6_4  ,lowercase_     :							Any=2_0_4_8  ,lowercase_     :							Dict=6  ,lowercase_     :							Dict=None  ,lowercase_     :							Dict=8  ,lowercase_     :							Optional[int]=3_2  ,lowercase_     :							List[Any]=1_2_8  ,lowercase_     :							Any=0.1  ,lowercase_     :							Optional[Any]=1E-6  ,lowercase_     :							Any=1.0  ,lowercase_     :							Dict="relu"  ,lowercase_     :							Tuple=True  ,lowercase_     :							List[Any]=True  ,lowercase_     :							Tuple=0  ,lowercase_     :							Optional[Any]=1  ,**lowercase_     :							Any  ,):
						lowerCAmelCase__      :		Tuple    							=      vocab_size
						lowerCAmelCase__      :		List[str]    							=      d_model
						lowerCAmelCase__      :		Union[str, Any]    							=      d_kv
						lowerCAmelCase__      :		List[Any]    							=      d_ff
						lowerCAmelCase__      :		int    							=      num_layers
						lowerCAmelCase__      :		Union[str, Any]    							=      (
						    num_decoder_layers if num_decoder_layers is not None else self.num_layers
						)  # default = symmetry
						lowerCAmelCase__      :		Any    							=      num_heads
						lowerCAmelCase__      :		Optional[int]    							=      relative_attention_num_buckets
						lowerCAmelCase__      :		Union[str, Any]    							=      relative_attention_max_distance
						lowerCAmelCase__      :		List[Any]    							=      dropout_rate
						lowerCAmelCase__      :		str    							=      layer_norm_epsilon
						lowerCAmelCase__      :		int    							=      initializer_factor
						lowerCAmelCase__      :		Dict    							=      feed_forward_proj
						lowerCAmelCase__      :		int    							=      use_cache

						lowerCAmelCase__      :		int    							=      self.feed_forward_proj.split('''-'''						)
						lowerCAmelCase__      :		str    							=      act_info[-1]
						lowerCAmelCase__      :		Tuple    							=      act_info[0] == '''gated'''

						if len(lowercase_						) > 1 and act_info[0] != "gated" or len(lowercase_						) > 2:
							raise ValueError(
							    F'`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer.'
							    '''Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. '''
							    '''\'gated-gelu\' or \'relu\''''						)

						# for backwards compatibility
						if feed_forward_proj == "gated-gelu":
							lowerCAmelCase__      :		Tuple    							=      '''gelu_new'''

						super().__init__(
						    pad_token_id=lowercase_  ,eos_token_id=lowercase_  ,is_encoder_decoder=lowercase_  ,**lowercase_  ,)







class 							SCREAMING_SNAKE_CASE   (						a_			):






					"""simple docstring"""



					@property
					def 			__lowerCAmelCase					(			self     :							List[Any]						):
						lowerCAmelCase__      :		int    							=      {
						    '''input_ids''': {0: '''batch''', 1: '''encoder_sequence'''},
						    '''attention_mask''': {0: '''batch''', 1: '''encoder_sequence'''},
						}
						if self.use_past:
							lowerCAmelCase__      :		Union[str, Any]    							=      '''past_encoder_sequence + sequence'''
							lowerCAmelCase__      :		List[Any]    							=      {0: '''batch'''}
							lowerCAmelCase__      :		Union[str, Any]    							=      {0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
						else:
							lowerCAmelCase__      :		Union[str, Any]    							=      {0: '''batch''', 1: '''decoder_sequence'''}
							lowerCAmelCase__      :		Any    							=      {0: '''batch''', 1: '''decoder_sequence'''}

						if self.use_past:
							self.fill_with_past_key_values_(lowercase_  ,direction='''inputs'''						)

						return common_inputs







					@property
					def 			__lowerCAmelCase					(			self     :							Optional[int]						):
						return 1_3




