





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  __magic_name__  (				lowercase				)				-> List[str]:

							"""simple docstring"""



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






									'''simple docstring'''






									def __init__(							self,   snake_case__,   snake_case__  )       -> Tuple:






																"""simple docstring"""
																super().__init__()
																lowercase_	:   Union[str, Any]  				=     module
																lowercase_	:   str  				=     nn.Sequential(
																    nn.Linear(module.in_features,   snake_case__,   bias=snake_case__  ),   nn.Linear(snake_case__,   module.out_features,   bias=snake_case__  ),   )
																lowercase_	:   Tuple  				=     (2.0 / (5 * min(module.in_features,   module.out_features  ))) ** 0.5
																nn.init.normal_(self.adapter[0].weight,   std=snake_case__  )
																nn.init.zeros_(self.adapter[1].weight  )
																self.adapter.to(module.weight.device  )






									def 						snake_case__				(							self,   snake_case__,   *snake_case__,   **snake_case__  )       -> Tuple:






																"""simple docstring"""
																return self.module(snake_case__,   *snake_case__,   **snake_case__  ) + self.adapter(snake_case__  )
@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class     UpperCamelCase__					(     unittest.TestCase				):






		'''simple docstring'''






		__a       :       List[Any]       =	"""bigscience/bloom-1b7"""

		# Constant values
		__a       :       List[Any]       =	2.1_0_9_6_5_9_5_5_2_6_9_2_5_7_4

		__a       :       Optional[Any]       =	"""Hello my name is"""
		__a       :       List[Any]       =	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"""				)
		__a       :       Dict       =	10






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






									"""simple docstring"""
									# Models and tokenizer
									lowercase_	:   Tuple  				=     AutoTokenizer.from_pretrained(self.model_name  )




class     UpperCamelCase__					(     lowerCamelCase__				):






		'''simple docstring'''






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






									"""simple docstring"""
									super().setUp()

									# Models and tokenizer
									lowercase_	:   Tuple  				=     AutoModelForCausalLM.from_pretrained(
									    self.model_name,   torch_dtype=torch.floataa,   device_map="""auto"""  )
									lowercase_	:   Any  				=     AutoModelForCausalLM.from_pretrained(self.model_name,   load_in_abit=snake_case__,   device_map="""auto"""  )






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






									"""simple docstring"""
									del self.model_fpaa
									del self.model_abit

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






		def 						snake_case__				(							self  )       -> Dict:






									"""simple docstring"""
									lowercase_	:   List[Any]  				=     self.model_abit.config

									self.assertTrue(hasattr(snake_case__,   """quantization_config"""  )  )

									lowercase_	:   List[Any]  				=     config.to_dict()
									lowercase_	:   Dict  				=     config.to_diff_dict()

									lowercase_	:   List[str]  				=     config.to_json_string()






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






									"""simple docstring"""
									from bitsandbytes.nn import Paramsabit

									lowercase_	:   List[Any]  				=     self.model_fpaa.get_memory_footprint()
									lowercase_	:   Dict  				=     self.model_abit.get_memory_footprint()

									self.assertAlmostEqual(mem_fpaa / mem_abit,   self.EXPECTED_RELATIVE_DIFFERENCE  )
									lowercase_	:   List[Any]  				=     get_some_linear_layer(self.model_abit  )
									self.assertTrue(linear.weight.__class__ == Paramsabit  )






		def 						snake_case__				(							self  )       -> Tuple:






									"""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(snake_case__,   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 						snake_case__				(							self  )       -> Optional[Any]:






									"""simple docstring"""
									lowercase_	:   List[str]  				=     self.tokenizer(self.input_text,   return_tensors="""pt"""  )
									lowercase_	:   Optional[Any]  				=     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=snake_case__  ),   self.EXPECTED_OUTPUTS  )






		def 						snake_case__				(							self  )       -> Dict:






									"""simple docstring"""
									lowercase_	:   List[Any]  				=     BitsAndBytesConfig()
									lowercase_	:   List[str]  				=     True

									lowercase_	:   List[str]  				=     AutoModelForCausalLM.from_pretrained(
									    self.model_name,   quantization_config=snake_case__,   device_map="""auto"""  )

									lowercase_	:   Tuple  				=     self.tokenizer(self.input_text,   return_tensors="""pt"""  )
									lowercase_	:   Any  				=     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=snake_case__  ),   self.EXPECTED_OUTPUTS  )






		def 						snake_case__				(							self  )       -> List[Any]:






									"""simple docstring"""
									with self.assertRaises(snake_case__  ), tempfile.TemporaryDirectory() as tmpdirname:
																self.model_abit.save_pretrained(snake_case__  )






		def 						snake_case__				(							self  )       -> List[Any]:






									"""simple docstring"""
									lowercase_	:   Any  				=     BitsAndBytesConfig()

									with self.assertRaises(snake_case__  ):
																lowercase_	:   List[str]  				=     AutoModelForCausalLM.from_pretrained(
																    self.model_name,   quantization_config=snake_case__,   load_in_abit=snake_case__,   device_map="""auto""",   bnb_abit_quant_type="""nf4""",   )






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






									"""simple docstring"""
									with self.assertRaises(snake_case__  ):
																# Tries with `str`
																self.model_abit.to("""cpu"""  )

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

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

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

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

									# Test if we did not break anything
									lowercase_	:   Optional[Any]  				=     self.tokenizer(self.input_text,   return_tensors="""pt"""  )

									lowercase_	:   int  				=     self.model_fpaa.to(torch.floataa  )
									lowercase_	:   Any  				=     self.model_fpaa.generate(input_ids=encoded_input["""input_ids"""].to(0  ),   max_new_tokens=10  )

									# Check this does not throw an error
									lowercase_	:   Dict  				=     self.model_fpaa.to("""cpu"""  )

									# Check this does not throw an error
									lowercase_	:   List[Any]  				=     self.model_fpaa.half()

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






		def 						snake_case__				(							self  )       -> List[Any]:






									"""simple docstring"""
									lowercase_	:   Optional[int]  				=     AutoModelForSeqaSeqLM.from_pretrained("""t5-small""",   load_in_abit=snake_case__,   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     UpperCamelCase__					(     unittest.TestCase				):






		'''simple docstring'''






		@classmethod
		def 						snake_case__				(							cls  )       -> Tuple:






									"""simple docstring"""
									lowercase_	:   int  				=     """t5-small"""
									lowercase_	:   Any  				=     """google/flan-t5-small"""  # flan-t5 uses dense-act instead of dense-relu-dense
									lowercase_	:   Optional[Any]  				=     AutoTokenizer.from_pretrained(cls.model_name  )
									lowercase_	:   Optional[Any]  				=     """Translate in German: Hello, my dog is cute"""






		def 						snake_case__				(							self  )       -> List[Any]:






									"""simple docstring"""
									gc.collect()
									torch.cuda.empty_cache()






		def 						snake_case__				(							self  )       -> Any:






									"""simple docstring"""
									from transformers import TaForConditionalGeneration

									lowercase_	:   List[str]  				=     TaForConditionalGeneration._keep_in_fpaa_modules
									lowercase_	:   str  				=     None

									# test with `t5-small`
									lowercase_	:   int  				=     TaForConditionalGeneration.from_pretrained(self.model_name,   load_in_abit=snake_case__,   device_map="""auto"""  )
									lowercase_	:   str  				=     self.tokenizer(self.input_text,   return_tensors="""pt"""  ).to(0  )
									lowercase_	:   Optional[int]  				=     model.generate(**snake_case__  )

									# test with `flan-t5-small`
									lowercase_	:   str  				=     TaForConditionalGeneration.from_pretrained(
									    self.dense_act_model_name,   load_in_abit=snake_case__,   device_map="""auto"""  )
									lowercase_	:   int  				=     self.tokenizer(self.input_text,   return_tensors="""pt"""  ).to(0  )
									lowercase_	:   Dict  				=     model.generate(**snake_case__  )
									lowercase_	:   int  				=     modules






		def 						snake_case__				(							self  )       -> str:






									"""simple docstring"""
									import bitsandbytes as bnb

									from transformers import TaForConditionalGeneration

									# test with `t5-small`
									lowercase_	:   Tuple  				=     TaForConditionalGeneration.from_pretrained(self.model_name,   load_in_abit=snake_case__,   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  )  )

									lowercase_	:   Tuple  				=     self.tokenizer(self.input_text,   return_tensors="""pt"""  ).to(0  )
									lowercase_	:   Optional[int]  				=     model.generate(**snake_case__  )

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




class     UpperCamelCase__					(     lowerCamelCase__				):






		'''simple docstring'''






		def 						snake_case__				(							self  )       -> List[Any]:






									"""simple docstring"""
									super().setUp()
									# model_name
									lowercase_	:   List[str]  				=     """bigscience/bloom-560m"""
									lowercase_	:   List[str]  				=     """t5-small"""

									# Different types of model

									lowercase_	:   List[Any]  				=     AutoModel.from_pretrained(self.model_name,   load_in_abit=snake_case__,   device_map="""auto"""  )
									# Sequence classification model
									lowercase_	:   str  				=     AutoModelForSequenceClassification.from_pretrained(
									    self.model_name,   load_in_abit=snake_case__,   device_map="""auto"""  )
									# CausalLM model
									lowercase_	:   Any  				=     AutoModelForCausalLM.from_pretrained(self.model_name,   load_in_abit=snake_case__,   device_map="""auto"""  )
									# Seq2seq model
									lowercase_	:   Optional[int]  				=     AutoModelForSeqaSeqLM.from_pretrained(
									    self.seq_to_seq_name,   load_in_abit=snake_case__,   device_map="""auto"""  )






		def 						snake_case__				(							self  )       -> List[str]:






									"""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 						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     UpperCamelCase__					(     lowerCamelCase__				):






		'''simple docstring'''






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






									"""simple docstring"""
									super().setUp()






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






									"""simple docstring"""
									del self.pipe

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






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






									"""simple docstring"""
									lowercase_	:   List[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
									lowercase_	:   Union[str, Any]  				=     self.pipe(self.input_text  )
									self.assertIn(pipeline_output[0]["""generated_text"""],   self.EXPECTED_OUTPUTS  )




@require_torch_multi_gpu
class     UpperCamelCase__					(     lowerCamelCase__				):






		'''simple docstring'''






		def 						snake_case__				(							self  )       -> Dict:






									"""simple docstring"""
									super().setUp()






		def 						snake_case__				(							self  )       -> Optional[int]:






									"""simple docstring"""
									lowercase_	:   Union[str, Any]  				=     AutoModelForCausalLM.from_pretrained(
									    self.model_name,   load_in_abit=snake_case__,   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
									lowercase_	:   Tuple  				=     self.tokenizer(self.input_text,   return_tensors="""pt"""  )

									# Second real batch
									lowercase_	:   Dict  				=     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=snake_case__  ),   self.EXPECTED_OUTPUTS  )




class     UpperCamelCase__					(     lowerCamelCase__				):






		'''simple docstring'''






		def 						snake_case__				(							self  )       -> int:






									"""simple docstring"""
									lowercase_	:   Optional[Any]  				=     """facebook/opt-350m"""
									super().setUp()






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






									"""simple docstring"""
									if version.parse(importlib.metadata.version("""bitsandbytes"""  )  ) < version.parse("""0.37.0"""  ):
																return

									# Step 1: freeze all parameters
									lowercase_	:   Optional[Any]  				=     AutoModelForCausalLM.from_pretrained(self.model_name,   load_in_abit=snake_case__  )

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

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

        # Step 2: add adapters
									for _, module in model.named_modules():
																if "OPTAttention" in repr(type(snake_case__  )  ):
																							lowercase_	:   List[str]  				=     LoRALayer(module.q_proj,   rank=16  )
																							lowercase_	:   str  				=     LoRALayer(module.k_proj,   rank=16  )
																							lowercase_	:   str  				=     LoRALayer(module.v_proj,   rank=16  )

        # Step 3: dummy batch
									lowercase_	:   Optional[Any]  				=     self.tokenizer("""Test batch """,   return_tensors="""pt"""  ).to(0  )

									# Step 4: Check if the gradient is not None
									with torch.cuda.amp.autocast():
																lowercase_	:   Optional[int]  				=     model.forward(**snake_case__  )
																out.logits.norm().backward()

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




class     UpperCamelCase__					(     lowerCamelCase__				):






		'''simple docstring'''






		__a       :       Optional[Any]       =	"""gpt2-xl"""
		__a       :       str       =	3.3_1_9_1_8_5_4_8_5_4_1_5_2_1_8_7






import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp





def  __magic_name__  (				)				-> str:

							"""simple docstring"""



							lowercase_	:   Optional[int]  				=     ArgumentParser(
							    description=(
							        """PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes"""
							    )				)

							# Optional arguments for the launch helper
							parser.add_argument("""--num_cores"""				,       type=lowercase				,       default=1				,       help="""Number of TPU cores to use (1 or 8)."""				)

							# positional
							parser.add_argument(
							    """training_script"""				,       type=lowercase				,       help=(
							        """The full path to the single TPU training """
							        """program/script to be launched in parallel, """
							        """followed by all the arguments for the """
							        """training script"""
							    )				,       )

							# rest from the training program
							parser.add_argument("""training_script_args"""				,       nargs=lowercase				)

							return parser.parse_args()





def  __magic_name__  (				)				-> List[Any]:

							"""simple docstring"""



							lowercase_	:   Union[str, Any]  				=     parse_args()

							# Import training_script as a module.
							lowercase_	:   List[Any]  				=     Path(args.training_script				)
							sys.path.append(str(script_fpath.parent.resolve()				)				)
							lowercase_	:   Tuple  				=     script_fpath.stem
							lowercase_	:   List[str]  				=     importlib.import_module(lowercase				)

							# Patch sys.argv
							lowercase_	:   Optional[int]  				=     [args.training_script] + args.training_script_args + ["""--tpu_num_cores""", str(args.num_cores				)]

							xmp.spawn(mod._mp_fn				,       args=()				,       nprocs=args.num_cores				)


if __name__ == "__main__":
							main()