





'''simple docstring'''






import gc
import importlib.metadata
import tempfile
import unittest

from packaging import version

from transformers import (
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    BitsAndBytesConfig,
    pipeline,
)
from transformers.testing_utils import (
    is_torch_available,
    require_accelerate,
    require_bitsandbytes,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)





def       lowercase_      (		lowercase__						)   ->Optional[Any]:
						if model.config.model_type == "gpt2":
												return model.transformer.h[0].mlp.c_fc
						return model.transformer.h[0].mlp.dense_ah_to_h


if is_torch_available():
		import torch
		import torch.nn as nn







		class 			lowerCamelCase					(						nn.Module  ):


							def __init__(					self	:  List[str]       ,     __snake_case	:  Tuple       ,     __snake_case	:  Any       ):





													'''simple docstring'''





													super().__init__()
													_snake_case:   str  					=		module
													_snake_case:   str  					=		nn.Sequential(
													    nn.Linear(module.in_features       ,     _lowerCamelCase       ,     bias=_lowerCamelCase       )       ,     nn.Linear(_lowerCamelCase       ,     module.out_features       ,     bias=_lowerCamelCase       )       ,     )
													_snake_case:   Optional[int]  					=		(2.0 / (5 * min(module.in_features       ,     module.out_features       ))) ** 0.5
													nn.init.normal_(self.adapter[0].weight       ,     std=_lowerCamelCase       )
													nn.init.zeros_(self.adapter[1].weight       )
													self.adapter.to(module.weight.device       )




							def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ,     __snake_case	:  List[str]       ,     *__snake_case	:  Optional[int]       ,     **__snake_case	:  Any       ):





													'''simple docstring'''





													return self.module(_lowerCamelCase       ,     *_lowerCamelCase       ,     **_lowerCamelCase       ) + self.adapter(_lowerCamelCase       )







@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 			lowerCamelCase					(						unittest.TestCase  ):
					_SCREAMING_SNAKE_CASE							   =    "bigscience/bloom-1b7"

					# Constant values
					_SCREAMING_SNAKE_CASE							   =    2.109659552692574

					_SCREAMING_SNAKE_CASE							   =    "Hello my name is"
					_SCREAMING_SNAKE_CASE							   =    set()
					EXPECTED_OUTPUTS.add("Hello my name is John and I am a professional photographer. I"  )
					EXPECTED_OUTPUTS.add("Hello my name is John.\nI am a friend of your father.\n"  )
					EXPECTED_OUTPUTS.add("Hello my name is John Doe, I am a student at the University"  )
					_SCREAMING_SNAKE_CASE							   =    10


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											_snake_case:   List[Any]  					=		AutoTokenizer.from_pretrained(self.model_name       )







class 			lowerCamelCase					(						lowerCAmelCase__  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  str       ):





											'''simple docstring'''





											super().setUp()

											# Models and tokenizer
											_snake_case:   Optional[int]  					=		AutoModelForCausalLM.from_pretrained(
											    self.model_name       ,     torch_dtype=torch.floataa       ,     device_map='auto'       )
											_snake_case:   List[Any]  					=		AutoModelForCausalLM.from_pretrained(self.model_name       ,     load_in_abit=_lowerCamelCase       ,     device_map='auto'       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											del self.model_fpaa
											del self.model_abit

											gc.collect()
											torch.cuda.empty_cache()


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case:   Optional[int]  					=		self.model_abit.config

											self.assertTrue(hasattr(_lowerCamelCase       ,     'quantization_config'       )       )

											_snake_case:   Tuple  					=		config.to_dict()
											_snake_case:   Dict  					=		config.to_diff_dict()

											_snake_case:   Any  					=		config.to_json_string()


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											from bitsandbytes.nn import Paramsabit

											_snake_case:   Optional[int]  					=		self.model_fpaa.get_memory_footprint()
											_snake_case:   int  					=		self.model_abit.get_memory_footprint()

											self.assertAlmostEqual(mem_fpaa / mem_abit       ,     self.EXPECTED_RELATIVE_DIFFERENCE       )
											_snake_case:   Dict  					=		get_some_linear_layer(self.model_abit       )
											self.assertTrue(linear.weight.__class__ == Paramsabit       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ):





											'''simple docstring'''





											from transformers import TaPreTrainedModel

											self.model_fpaa.get_memory_footprint()
											self.model_abit.get_memory_footprint()

											for name, module in self.model_abit.named_modules():
																	if isinstance(_lowerCamelCase       ,     torch.nn.Linear       ):
																							if name not in ["lm_head"] + TaPreTrainedModel._keep_in_fpaa_modules:
																													# 4-bit parameters are packed in uint8 variables
																													self.assertTrue(module.weight.dtype == torch.uinta       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[int]       ):





											'''simple docstring'''





											_snake_case:   Any  					=		self.tokenizer(self.input_text       ,     return_tensors='pt'       )
											_snake_case:   str  					=		self.model_abit.generate(input_ids=encoded_input['input_ids'].to(0       )       ,     max_new_tokens=10       )

											self.assertIn(self.tokenizer.decode(output_sequences[0]       ,     skip_special_tokens=_lowerCamelCase       )       ,     self.EXPECTED_OUTPUTS       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ):





											'''simple docstring'''





											_snake_case:   Optional[int]  					=		BitsAndBytesConfig()
											_snake_case:   int  					=		True

											_snake_case:   str  					=		AutoModelForCausalLM.from_pretrained(
											    self.model_name       ,     quantization_config=_lowerCamelCase       ,     device_map='auto'       )

											_snake_case:   Any  					=		self.tokenizer(self.input_text       ,     return_tensors='pt'       )
											_snake_case:   Dict  					=		model_abit_from_config.generate(
											    input_ids=encoded_input['input_ids'].to(0       )       ,     max_new_tokens=10       )

											self.assertIn(self.tokenizer.decode(output_sequences[0]       ,     skip_special_tokens=_lowerCamelCase       )       ,     self.EXPECTED_OUTPUTS       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ):





											'''simple docstring'''





											with self.assertRaises(_lowerCamelCase       ), tempfile.TemporaryDirectory() as tmpdirname:
																	self.model_abit.save_pretrained(_lowerCamelCase       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case:   int  					=		BitsAndBytesConfig()

											with self.assertRaises(_lowerCamelCase       ):
																	_snake_case:   Optional[Any]  					=		AutoModelForCausalLM.from_pretrained(
																	    self.model_name       ,     quantization_config=_lowerCamelCase       ,     load_in_abit=_lowerCamelCase       ,     device_map='auto'       ,     bnb_abit_quant_type='nf4'       ,     )


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											with self.assertRaises(_lowerCamelCase       ):
																	# Tries with `str`
																	self.model_abit.to('cpu'       )

											with self.assertRaises(_lowerCamelCase       ):
																	# Tries with a `dtype``
																	self.model_abit.to(torch.floataa       )

											with self.assertRaises(_lowerCamelCase       ):
																	# Tries with a `device`
																	self.model_abit.to(torch.device('cuda:0'       )       )

											with self.assertRaises(_lowerCamelCase       ):
																	# Tries with a `device`
																	self.model_abit.float()

											with self.assertRaises(_lowerCamelCase       ):
																	# Tries with a `device`
																	self.model_abit.half()

											# Test if we did not break anything
											_snake_case:   List[Any]  					=		self.tokenizer(self.input_text       ,     return_tensors='pt'       )

											_snake_case:   Dict  					=		self.model_fpaa.to(torch.floataa       )
											_snake_case:   Optional[int]  					=		self.model_fpaa.generate(input_ids=encoded_input['input_ids'].to(0       )       ,     max_new_tokens=10       )

											# Check this does not throw an error
											_snake_case:   int  					=		self.model_fpaa.to('cpu'       )

											# Check this does not throw an error
											_snake_case:   Optional[int]  					=		self.model_fpaa.half()

											# Check this does not throw an error
											_snake_case:   Dict  					=		self.model_fpaa.float()




					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ):





											'''simple docstring'''





											_snake_case:   Dict  					=		AutoModelForSeqaSeqLM.from_pretrained('t5-small'       ,     load_in_abit=_lowerCamelCase       ,     device_map='auto'       )
											self.assertTrue(model.decoder.block[0].layer[2].DenseReluDense.wo.weight.dtype == torch.floataa       )







@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 			lowerCamelCase					(						unittest.TestCase  ):


					@classmethod
					def    SCREAMING_SNAKE_CASE_		(					cls	:  int       ):





											'''simple docstring'''





											_snake_case:   List[str]  					=		't5-small'
											_snake_case:   List[str]  					=		'google/flan-t5-small'  # flan-t5 uses dense-act instead of dense-relu-dense
											_snake_case:   Any  					=		AutoTokenizer.from_pretrained(cls.model_name       )
											_snake_case:   Tuple  					=		'Translate in German: Hello, my dog is cute'


					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ):





											'''simple docstring'''





											gc.collect()
											torch.cuda.empty_cache()


					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ):





											'''simple docstring'''





											from transformers import TaForConditionalGeneration

											_snake_case:   Optional[int]  					=		TaForConditionalGeneration._keep_in_fpaa_modules
											_snake_case:   List[str]  					=		None

											# test with `t5-small`
											_snake_case:   Union[str, Any]  					=		TaForConditionalGeneration.from_pretrained(self.model_name       ,     load_in_abit=_lowerCamelCase       ,     device_map='auto'       )
											_snake_case:   Optional[Any]  					=		self.tokenizer(self.input_text       ,     return_tensors='pt'       ).to(0       )
											_snake_case:   Optional[Any]  					=		model.generate(**_lowerCamelCase       )

											# test with `flan-t5-small`
											_snake_case:   List[str]  					=		TaForConditionalGeneration.from_pretrained(
											    self.dense_act_model_name       ,     load_in_abit=_lowerCamelCase       ,     device_map='auto'       )
											_snake_case:   Optional[int]  					=		self.tokenizer(self.input_text       ,     return_tensors='pt'       ).to(0       )
											_snake_case:   Union[str, Any]  					=		model.generate(**_lowerCamelCase       )
											_snake_case:   Any  					=		modules




					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											import bitsandbytes as bnb

											from transformers import TaForConditionalGeneration

											# test with `t5-small`
											_snake_case:   int  					=		TaForConditionalGeneration.from_pretrained(self.model_name       ,     load_in_abit=_lowerCamelCase       ,     device_map='auto'       )

											# there was a bug with decoders - this test checks that it is fixed
											self.assertTrue(isinstance(model.decoder.block[0].layer[0].SelfAttention.q       ,     bnb.nn.Linearabit       )       )

											_snake_case:   Optional[Any]  					=		self.tokenizer(self.input_text       ,     return_tensors='pt'       ).to(0       )
											_snake_case:   int  					=		model.generate(**_lowerCamelCase       )

											# test with `flan-t5-small`
											_snake_case:   List[Any]  					=		TaForConditionalGeneration.from_pretrained(
											    self.dense_act_model_name       ,     load_in_abit=_lowerCamelCase       ,     device_map='auto'       )
											_snake_case:   Optional[int]  					=		self.tokenizer(self.input_text       ,     return_tensors='pt'       ).to(0       )
											_snake_case:   List[Any]  					=		model.generate(**_lowerCamelCase       )







class 			lowerCamelCase					(						lowerCAmelCase__  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											super().setUp()
											# model_name
											_snake_case:   Dict  					=		'bigscience/bloom-560m'
											_snake_case:   Dict  					=		't5-small'

											# Different types of model

											_snake_case:   List[Any]  					=		AutoModel.from_pretrained(self.model_name       ,     load_in_abit=_lowerCamelCase       ,     device_map='auto'       )
											# Sequence classification model
											_snake_case:   Union[str, Any]  					=		AutoModelForSequenceClassification.from_pretrained(
											    self.model_name       ,     load_in_abit=_lowerCamelCase       ,     device_map='auto'       )
											# CausalLM model
											_snake_case:   str  					=		AutoModelForCausalLM.from_pretrained(self.model_name       ,     load_in_abit=_lowerCamelCase       ,     device_map='auto'       )
											# Seq2seq model
											_snake_case:   Any  					=		AutoModelForSeqaSeqLM.from_pretrained(
											    self.seq_to_seq_name       ,     load_in_abit=_lowerCamelCase       ,     device_map='auto'       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[int]       ):





											'''simple docstring'''





											del self.base_model
											del self.sequence_model
											del self.model_abit
											del self.seq_to_seq_model

											gc.collect()
											torch.cuda.empty_cache()




					def    SCREAMING_SNAKE_CASE_		(					self	:  str       ):





											'''simple docstring'''





											from bitsandbytes.nn import Paramsabit

											self.assertTrue(self.base_model.h[-1].mlp.dense_ah_to_h.weight.__class__ == Paramsabit       )

											# Other heads should be nn.Parameter
											self.assertTrue(self.model_abit.lm_head.weight.__class__ == torch.nn.Parameter       )
											self.assertTrue(self.sequence_model.score.weight.__class__ == torch.nn.Parameter       )
											self.assertTrue(self.seq_to_seq_model.lm_head.weight.__class__ == torch.nn.Parameter       )







class 			lowerCamelCase					(						lowerCAmelCase__  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											super().setUp()


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





											del self.pipe

											gc.collect()
											torch.cuda.empty_cache()




					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		pipeline(
											    'text-generation'       ,     model=self.model_name       ,     model_kwargs={'device_map': 'auto', 'load_in_4bit': True, 'torch_dtype': torch.floataa}       ,     max_new_tokens=self.MAX_NEW_TOKENS       ,     )

											# Real second forward pass
											_snake_case:   Dict  					=		self.pipe(self.input_text       )
											self.assertIn(pipeline_output[0]['generated_text']       ,     self.EXPECTED_OUTPUTS       )







@require_torch_multi_gpu
class 			lowerCamelCase					(						lowerCAmelCase__  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ):





											'''simple docstring'''





											super().setUp()




					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[int]       ):





											'''simple docstring'''





											_snake_case:   str  					=		AutoModelForCausalLM.from_pretrained(
											    self.model_name       ,     load_in_abit=_lowerCamelCase       ,     device_map='balanced'       )

											# Check correct device map
											self.assertEqual(set(model_parallel.hf_device_map.values()       )       ,     {0, 1}       )

											# Check that inference pass works on the model
											_snake_case:   Dict  					=		self.tokenizer(self.input_text       ,     return_tensors='pt'       )

											# Second real batch
											_snake_case:   int  					=		model_parallel.generate(input_ids=encoded_input['input_ids'].to(0       )       ,     max_new_tokens=10       )
											self.assertIn(self.tokenizer.decode(output_parallel[0]       ,     skip_special_tokens=_lowerCamelCase       )       ,     self.EXPECTED_OUTPUTS       )







class 			lowerCamelCase					(						lowerCAmelCase__  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		'facebook/opt-350m'
											super().setUp()




					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ):





											'''simple docstring'''





											if version.parse(importlib.metadata.version('bitsandbytes'       )       ) < version.parse('0.37.0'       ):
																	return

											# Step 1: freeze all parameters
											_snake_case:   Union[str, Any]  					=		AutoModelForCausalLM.from_pretrained(self.model_name       ,     load_in_abit=_lowerCamelCase       )

											self.assertEqual(set(model.hf_device_map.values()       )       ,     {torch.cuda.current_device()}       )

											for param in model.parameters():
																	_snake_case:   Union[str, Any]  					=		False  # freeze the model - train adapters later
																	if param.ndim == 1:
																							# cast the small parameters (e.g. layernorm) to fp32 for stability
																							_snake_case:   Optional[Any]  					=		param.data.to(torch.floataa       )

        # Step 2: add adapters
											for _, module in model.named_modules():
																	if "OPTAttention" in repr(type(_lowerCamelCase       )       ):
																							_snake_case:   Union[str, Any]  					=		LoRALayer(module.q_proj       ,     rank=16       )
																							_snake_case:   Dict  					=		LoRALayer(module.k_proj       ,     rank=16       )
																							_snake_case:   Any  					=		LoRALayer(module.v_proj       ,     rank=16       )

        # Step 3: dummy batch
											_snake_case:   Union[str, Any]  					=		self.tokenizer('Test batch '       ,     return_tensors='pt'       ).to(0       )

											# Step 4: Check if the gradient is not None
											with torch.cuda.amp.autocast():
																	_snake_case:   List[str]  					=		model.forward(**_lowerCamelCase       )
																	out.logits.norm().backward()

											for module in model.modules():
																	if isinstance(_lowerCamelCase       ,     _lowerCamelCase       ):
																							self.assertTrue(module.adapter[1].weight.grad is not None       )
																							self.assertTrue(module.adapter[1].weight.grad.norm().item() > 0       )
																	elif isinstance(_lowerCamelCase       ,     nn.Embedding       ):
																							self.assertTrue(module.weight.grad is None       )







class 			lowerCamelCase					(						lowerCAmelCase__  ):
					_SCREAMING_SNAKE_CASE							   =    "gpt2-xl"
					_SCREAMING_SNAKE_CASE							   =    3.3191854854152187












'''simple docstring'''






import torch





def       lowercase_      (		)   ->Dict:
						if torch.cuda.is_available():
												_snake_case:   Any  					=		torch.cuda.device_count()
						else:
												_snake_case:   Optional[int]  					=		0
						print(F'''Successfully ran on {num_gpus} GPUs'''						)


if __name__ == "__main__":
		main()





