





'''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        UpperCamelCase__				(		lowerCAmelCase  ):





      """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       ):




            def __init__(      self      :			Any      ,      __snake_case      :			nn.Module      ,      __snake_case      :			int  )    ->				str:
                  super().__init__()
                  _lowerCAmelCase												=							module
                  _lowerCAmelCase												=							nn.Sequential(
                      nn.Linear(module.in_features      ,      __snake_case      ,      bias=__snake_case  )      ,      nn.Linear(__snake_case      ,      module.out_features      ,      bias=__snake_case  )      ,      )
                  _lowerCAmelCase												=							(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     lowercase__   (      self      :			Optional[Any]      ,      __snake_case      :			Union[str, Any]      ,      *__snake_case      :			Optional[int]      ,      **__snake_case      :			Optional[Any]  )    ->				Dict:
                  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       ):
      # We keep the constants inside the init function and model loading inside setUp function

      # We need to test on relatively large models (aka >1b parameters otherwise the quantiztion may not work as expected)
      # Therefore here we use only bloom-1b3 to test our module
      _lowercase:      Optional[int]											=       '''bigscience/bloom-1b7'''

      # Constant values
      _lowercase:      Union[str, Any]											=       2.109659552692574

      _lowercase:      Union[str, Any]											=       '''Hello my name is'''
      _lowercase:      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'''       )
      _lowercase:      int											=       10




      def     lowercase__   (      self      :			str  )    ->				Optional[Any]:
            # Models and tokenizer
            _lowerCAmelCase												=							AutoTokenizer.from_pretrained(self.model_name  )




class 			UpperCAmelCase					(      snake_case_       ):




      def     lowercase__   (      self      :			Optional[int]  )    ->				List[str]:
            super().setUp()

            # Models and tokenizer
            _lowerCAmelCase												=							AutoModelForCausalLM.from_pretrained(
                self.model_name      ,      torch_dtype=torch.floataa      ,      device_map="""auto"""  )
            _lowerCAmelCase												=							AutoModelForCausalLM.from_pretrained(self.model_name      ,      load_in_abit=__snake_case      ,      device_map="""auto"""  )




      def     lowercase__   (      self      :			List[str]  )    ->				Tuple:
            del self.model_fpaa
            del self.model_abit

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




      def     lowercase__   (      self      :			Dict  )    ->				int:
            _lowerCAmelCase												=							self.model_abit.config

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

            _lowerCAmelCase												=							config.to_dict()
            _lowerCAmelCase												=							config.to_diff_dict()

            _lowerCAmelCase												=							config.to_json_string()




      def     lowercase__   (      self      :			Optional[Any]  )    ->				List[Any]:
            from bitsandbytes.nn import Paramsabit

            _lowerCAmelCase												=							self.model_fpaa.get_memory_footprint()
            _lowerCAmelCase												=							self.model_abit.get_memory_footprint()

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




      def     lowercase__   (      self      :			Dict  )    ->				Union[str, Any]:
            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     lowercase__   (      self      :			Union[str, Any]  )    ->				str:
            _lowerCAmelCase												=							self.tokenizer(self.input_text      ,      return_tensors="""pt"""  )
            _lowerCAmelCase												=							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     lowercase__   (      self      :			Union[str, Any]  )    ->				Optional[int]:
            _lowerCAmelCase												=							BitsAndBytesConfig()
            _lowerCAmelCase												=							True

            _lowerCAmelCase												=							AutoModelForCausalLM.from_pretrained(
                self.model_name      ,      quantization_config=__snake_case      ,      device_map="""auto"""  )

            _lowerCAmelCase												=							self.tokenizer(self.input_text      ,      return_tensors="""pt"""  )
            _lowerCAmelCase												=							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     lowercase__   (      self      :			Tuple  )    ->				List[str]:
            with self.assertRaises(__snake_case  ), tempfile.TemporaryDirectory() as tmpdirname:
                  self.model_abit.save_pretrained(__snake_case  )




      def     lowercase__   (      self      :			Union[str, Any]  )    ->				Union[str, Any]:
            _lowerCAmelCase												=							BitsAndBytesConfig()

            with self.assertRaises(__snake_case  ):
                  _lowerCAmelCase												=							AutoModelForCausalLM.from_pretrained(
                      self.model_name      ,      quantization_config=__snake_case      ,      load_in_abit=__snake_case      ,      device_map="""auto"""      ,      bnb_abit_quant_type="""nf4"""      ,      )




      def     lowercase__   (      self      :			Any  )    ->				Dict:
            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
            _lowerCAmelCase												=							self.tokenizer(self.input_text      ,      return_tensors="""pt"""  )

            _lowerCAmelCase												=							self.model_fpaa.to(torch.floataa  )
            _lowerCAmelCase												=							self.model_fpaa.generate(input_ids=encoded_input["""input_ids"""].to(0  )      ,      max_new_tokens=10  )

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

            # Check this does not throw an error
            _lowerCAmelCase												=							self.model_fpaa.half()

            # Check this does not throw an error
            _lowerCAmelCase												=							self.model_fpaa.float()




      def     lowercase__   (      self      :			List[str]  )    ->				Dict:
            _lowerCAmelCase												=							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       ):




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




      def     lowercase__   (      self      :			Any  )    ->				List[str]:
            gc.collect()
            torch.cuda.empty_cache()




      def     lowercase__   (      self      :			str  )    ->				Dict:
            from transformers import TaForConditionalGeneration

            _lowerCAmelCase												=							TaForConditionalGeneration._keep_in_fpaa_modules
            _lowerCAmelCase												=							None

            # test with `t5-small`
            _lowerCAmelCase												=							TaForConditionalGeneration.from_pretrained(self.model_name      ,      load_in_abit=__snake_case      ,      device_map="""auto"""  )
            _lowerCAmelCase												=							self.tokenizer(self.input_text      ,      return_tensors="""pt"""  ).to(0  )
            _lowerCAmelCase												=							model.generate(**__snake_case  )

            # test with `flan-t5-small`
            _lowerCAmelCase												=							TaForConditionalGeneration.from_pretrained(
                self.dense_act_model_name      ,      load_in_abit=__snake_case      ,      device_map="""auto"""  )
            _lowerCAmelCase												=							self.tokenizer(self.input_text      ,      return_tensors="""pt"""  ).to(0  )
            _lowerCAmelCase												=							model.generate(**__snake_case  )
            _lowerCAmelCase												=							modules




      def     lowercase__   (      self      :			Optional[int]  )    ->				List[Any]:
            import bitsandbytes as bnb

            from transformers import TaForConditionalGeneration

            # test with `t5-small`
            _lowerCAmelCase												=							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  )  )

            _lowerCAmelCase												=							self.tokenizer(self.input_text      ,      return_tensors="""pt"""  ).to(0  )
            _lowerCAmelCase												=							model.generate(**__snake_case  )

            # test with `flan-t5-small`
            _lowerCAmelCase												=							TaForConditionalGeneration.from_pretrained(
                self.dense_act_model_name      ,      load_in_abit=__snake_case      ,      device_map="""auto"""  )
            _lowerCAmelCase												=							self.tokenizer(self.input_text      ,      return_tensors="""pt"""  ).to(0  )
            _lowerCAmelCase												=							model.generate(**__snake_case  )




class 			UpperCAmelCase					(      snake_case_       ):




      def     lowercase__   (      self      :			Optional[Any]  )    ->				List[Any]:
            super().setUp()
            # model_name
            _lowerCAmelCase												=							"""bigscience/bloom-560m"""
            _lowerCAmelCase												=							"""t5-small"""

            # Different types of model

            _lowerCAmelCase												=							AutoModel.from_pretrained(self.model_name      ,      load_in_abit=__snake_case      ,      device_map="""auto"""  )
            # Sequence classification model
            _lowerCAmelCase												=							AutoModelForSequenceClassification.from_pretrained(
                self.model_name      ,      load_in_abit=__snake_case      ,      device_map="""auto"""  )
            # CausalLM model
            _lowerCAmelCase												=							AutoModelForCausalLM.from_pretrained(self.model_name      ,      load_in_abit=__snake_case      ,      device_map="""auto"""  )
            # Seq2seq model
            _lowerCAmelCase												=							AutoModelForSeqaSeqLM.from_pretrained(
                self.seq_to_seq_name      ,      load_in_abit=__snake_case      ,      device_map="""auto"""  )




      def     lowercase__   (      self      :			Optional[Any]  )    ->				Optional[Any]:
            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     lowercase__   (      self      :			str  )    ->				Optional[Any]:
            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					(      snake_case_       ):




      def     lowercase__   (      self      :			List[Any]  )    ->				Dict:
            super().setUp()




      def     lowercase__   (      self      :			Union[str, Any]  )    ->				Union[str, Any]:
            del self.pipe

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




      def     lowercase__   (      self      :			Union[str, Any]  )    ->				int:
            _lowerCAmelCase												=							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
            _lowerCAmelCase												=							self.pipe(self.input_text  )
            self.assertIn(pipeline_output[0]["""generated_text"""]      ,      self.EXPECTED_OUTPUTS  )




@require_torch_multi_gpu
class 			UpperCAmelCase					(      snake_case_       ):




      def     lowercase__   (      self      :			int  )    ->				Union[str, Any]:
            super().setUp()




      def     lowercase__   (      self      :			Union[str, Any]  )    ->				List[str]:
            _lowerCAmelCase												=							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
            _lowerCAmelCase												=							self.tokenizer(self.input_text      ,      return_tensors="""pt"""  )

            # Second real batch
            _lowerCAmelCase												=							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					(      snake_case_       ):




      def     lowercase__   (      self      :			Tuple  )    ->				Union[str, Any]:
            _lowerCAmelCase												=							"""facebook/opt-350m"""
            super().setUp()




      def     lowercase__   (      self      :			Optional[Any]  )    ->				Dict:
            if version.parse(importlib.metadata.version("""bitsandbytes"""  )  ) < version.parse("""0.37.0"""  ):
                  return

            # Step 1: freeze all parameters
            _lowerCAmelCase												=							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():
                  _lowerCAmelCase												=							False  # freeze the model - train adapters later
                  if param.ndim == 1:
                        # cast the small parameters (e.g. layernorm) to fp32 for stability
                        _lowerCAmelCase												=							param.data.to(torch.floataa  )

        # Step 2: add adapters
            for _, module in model.named_modules():
                  if "OPTAttention" in repr(type(__snake_case  )  ):
                        _lowerCAmelCase												=							LoRALayer(module.q_proj      ,      rank=16  )
                        _lowerCAmelCase												=							LoRALayer(module.k_proj      ,      rank=16  )
                        _lowerCAmelCase												=							LoRALayer(module.v_proj      ,      rank=16  )

        # Step 3: dummy batch
            _lowerCAmelCase												=							self.tokenizer("""Test batch """      ,      return_tensors="""pt"""  ).to(0  )

            # Step 4: Check if the gradient is not None
            with torch.cuda.amp.autocast():
                  _lowerCAmelCase												=							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					(      snake_case_       ):
      _lowercase:      Optional[Any]											=       '''gpt2-xl'''
      _lowercase:      Tuple											=       3.3191854854152187













'''simple docstring'''






from manim import *




class 			UpperCAmelCase					(      snake_case_       ):




      def     lowercase__   (      self      :			Union[str, Any]  )    ->				str:
            _lowerCAmelCase												=							Rectangle(height=0.5      ,      width=0.5  )
            _lowerCAmelCase												=							Rectangle(height=0.46      ,      width=0.46  ).set_stroke(width=0  )

            _lowerCAmelCase												=							[mem.copy() for i in range(6  )]
            _lowerCAmelCase												=							[mem.copy() for i in range(6  )]
            _lowerCAmelCase												=							VGroup(*__snake_case  ).arrange(__snake_case      ,      buff=0  )
            _lowerCAmelCase												=							VGroup(*__snake_case  ).arrange(__snake_case      ,      buff=0  )
            _lowerCAmelCase												=							VGroup(__snake_case      ,      __snake_case  ).arrange(__snake_case      ,      buff=0  )
            _lowerCAmelCase												=							Text("""CPU"""      ,      font_size=24  )
            _lowerCAmelCase												=							Group(__snake_case      ,      __snake_case  ).arrange(__snake_case      ,      buff=0.5      ,      aligned_edge=__snake_case  )
            cpu.move_to([-2.5, -0.5, 0]  )
            self.add(__snake_case  )

            _lowerCAmelCase												=							[mem.copy() for i in range(4  )]
            _lowerCAmelCase												=							VGroup(*__snake_case  ).arrange(__snake_case      ,      buff=0  )
            _lowerCAmelCase												=							Text("""GPU"""      ,      font_size=24  )
            _lowerCAmelCase												=							Group(__snake_case      ,      __snake_case  ).arrange(__snake_case      ,      buff=0.5      ,      aligned_edge=__snake_case  )
            gpu.move_to([-1, -1, 0]  )
            self.add(__snake_case  )

            _lowerCAmelCase												=							[mem.copy() for i in range(6  )]
            _lowerCAmelCase												=							VGroup(*__snake_case  ).arrange(__snake_case      ,      buff=0  )

            _lowerCAmelCase												=							Text("""Model"""      ,      font_size=24  )
            _lowerCAmelCase												=							Group(__snake_case      ,      __snake_case  ).arrange(__snake_case      ,      buff=0.5      ,      aligned_edge=__snake_case  )
            model.move_to([3, -1.0, 0]  )
            self.add(__snake_case  )

            _lowerCAmelCase												=							[]
            for i, rect in enumerate(__snake_case  ):
                  rect.set_stroke(__snake_case  )
                  # target = fill.copy().set_fill(YELLOW, opacity=0.7)
                  # target.move_to(rect)
                  # self.add(target)

                  _lowerCAmelCase												=							Rectangle(height=0.46 / 4      ,      width=0.46 / 3  ).set_stroke(width=0.0  ).set_fill(__snake_case      ,      opacity=0.7  )

                  if i == 0:
                        cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT  )      ,      buff=0.02      ,      direction=__snake_case  )
                        cpu_target.set_x(cpu_target.get_x() + 0.1  )
                  elif i == 3:
                        cpu_target.next_to(cpu_targs[0]      ,      direction=__snake_case      ,      buff=0.0  )
                  else:
                        cpu_target.next_to(cpu_targs[i - 1]      ,      direction=__snake_case      ,      buff=0.0  )
                  self.add(__snake_case  )
                  cpu_targs.append(__snake_case  )

            _lowerCAmelCase												=							[mem.copy() for i in range(6  )]
            _lowerCAmelCase												=							VGroup(*__snake_case  ).arrange(__snake_case      ,      buff=0  )

            _lowerCAmelCase												=							Text("""Loaded Checkpoint"""      ,      font_size=24  )
            _lowerCAmelCase												=							Group(__snake_case      ,      __snake_case  ).arrange(__snake_case      ,      aligned_edge=__snake_case      ,      buff=0.4  )
            checkpoint.move_to([3, 0.5, 0]  )

            _lowerCAmelCase												=							Square(side_length=2.2  )
            key.move_to([-5, 2, 0]  )

            _lowerCAmelCase												=							MarkupText(
                f"<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model"      ,      font_size=18      ,      )

            key_text.move_to([-5, 2.4, 0]  )

            self.add(__snake_case      ,      __snake_case  )

            _lowerCAmelCase												=							MarkupText(
                f"<span fgcolor='{BLUE}'>●</span> Checkpoint"      ,      font_size=18      ,      )

            blue_text.next_to(__snake_case      ,      DOWN * 2.4      ,      aligned_edge=key_text.get_left()  )

            _lowerCAmelCase												=							MarkupText(
                f"Next, a <i><span fgcolor=\"{BLUE}\">second</span></i> model is loaded into memory,\nwith the weights of a <span fgcolor=\"{BLUE}\">single shard</span>."      ,      font_size=24      ,      )
            step_a.move_to([2, 2, 0]  )
            self.play(Write(__snake_case  )      ,      Write(__snake_case  )  )

            self.play(Write(__snake_case      ,      run_time=1  )      ,      Create(__snake_case      ,      run_time=1  )  )

            _lowerCAmelCase												=							[]
            _lowerCAmelCase												=							[]
            for i, rect in enumerate(__snake_case  ):
                  _lowerCAmelCase												=							fill.copy().set_fill(__snake_case      ,      opacity=0.7  )
                  target.move_to(__snake_case  )
                  first_animations.append(GrowFromCenter(__snake_case      ,      run_time=1  )  )

                  _lowerCAmelCase												=							target.copy()
                  cpu_target.generate_target()
                  if i < 5:
                        cpu_target.target.move_to(cpu_left_col_base[i + 1]  )
                  else:
                        cpu_target.target.move_to(cpu_right_col_base[i - 5]  )
                  second_animations.append(MoveToTarget(__snake_case      ,      run_time=1.5  )  )

            self.play(*__snake_case  )
            self.play(*__snake_case  )
            self.wait()






