


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_				(      UpperCAmelCase__  ):
					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       :   int							,					UpperCamelCase__       :   nn.Module							,					UpperCamelCase__       :   int       ):






																	'''simple docstring'''






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







												def   UpperCAmelCase__       (     self       :   int							,					UpperCamelCase__       :   Any							,					*UpperCamelCase__       :   int							,					**UpperCamelCase__       :   List[Any]       ):






																	'''simple docstring'''






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







@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
					__SCREAMING_SNAKE_CASE       :  Union[str, Any]						=   'bigscience/bloom-1b7'

					# Constant values
					__SCREAMING_SNAKE_CASE       :  Optional[int]						=   2.1_0_9_6_5_9_5_5_2_6_9_2_5_7_4

					__SCREAMING_SNAKE_CASE       :  Dict						=   'Hello my name is'
					__SCREAMING_SNAKE_CASE       :  Optional[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'   )
					__SCREAMING_SNAKE_CASE       :  str						=   10

					def   UpperCAmelCase__       (     self       :   List[str]       ):






										'''simple docstring'''






										lowercase_  							=			AutoTokenizer.from_pretrained(self.model_name       )







class        UpperCamelCase__   (	__magic_name__   ):

					def   UpperCAmelCase__       (     self       :   int       ):






										'''simple docstring'''






										super().setUp()

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

					def   UpperCAmelCase__       (     self       :   Any       ):






										'''simple docstring'''






										del self.model_fpaa
										del self.model_abit

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

					def   UpperCAmelCase__       (     self       :   Optional[int]       ):






										'''simple docstring'''






										lowercase_  							=			self.model_abit.config

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

										lowercase_  							=			config.to_dict()
										lowercase_  							=			config.to_diff_dict()

										lowercase_  							=			config.to_json_string()

					def   UpperCAmelCase__       (     self       :   Any       ):






										'''simple docstring'''






										from bitsandbytes.nn import Paramsabit

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

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

					def   UpperCAmelCase__       (     self       :   Any       ):






										'''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(UpperCamelCase__							,					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   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   List[str]       ):






										'''simple docstring'''






										lowercase_  							=			BitsAndBytesConfig()
										lowercase_  							=			True

										lowercase_  							=			AutoModelForCausalLM.from_pretrained(
										    self.model_name							,					quantization_config=UpperCamelCase__							,					device_map="""auto"""       )

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

					def   UpperCAmelCase__       (     self       :   List[str]       ):






										'''simple docstring'''






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

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






										'''simple docstring'''






										lowercase_  							=			BitsAndBytesConfig()

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

					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






										with self.assertRaises(UpperCamelCase__       ):
															# Tries with `str`
															self.model_abit.to("""cpu"""       )

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

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

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

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

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

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

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

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

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







					def   UpperCAmelCase__       (     self       :   Optional[int]       ):






										'''simple docstring'''






										lowercase_  							=			AutoModelForSeqaSeqLM.from_pretrained("""t5-small"""							,					load_in_abit=UpperCamelCase__							,					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   UpperCAmelCase__       (     cls       :   Any       ):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






										from transformers import TaForConditionalGeneration

										lowercase_  							=			TaForConditionalGeneration._keep_in_fpaa_modules
										lowercase_  							=			None

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

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







					def   UpperCAmelCase__       (     self       :   Any       ):






										'''simple docstring'''






										import bitsandbytes as bnb

										from transformers import TaForConditionalGeneration

										# test with `t5-small`
										lowercase_  							=			TaForConditionalGeneration.from_pretrained(self.model_name							,					load_in_abit=UpperCamelCase__							,					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_  							=			self.tokenizer(self.input_text							,					return_tensors="""pt"""       ).to(0       )
										lowercase_  							=			model.generate(**UpperCamelCase__       )

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







class        UpperCamelCase__   (	__magic_name__   ):

					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''simple docstring'''






										super().setUp()
										# model_name
										lowercase_  							=			"""bigscience/bloom-560m"""
										lowercase_  							=			"""t5-small"""

										# Different types of model

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

					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''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   UpperCAmelCase__       (     self       :   Optional[int]       ):






										'''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__   (	__magic_name__   ):

					def   UpperCAmelCase__       (     self       :   Any       ):






										'''simple docstring'''






										super().setUp()

					def   UpperCAmelCase__       (     self       :   List[str]       ):






										'''simple docstring'''






										del self.pipe

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







					def   UpperCAmelCase__       (     self       :   str       ):






										'''simple docstring'''






										lowercase_  							=			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_  							=			self.pipe(self.input_text       )
										self.assertIn(pipeline_output[0]["""generated_text"""]							,					self.EXPECTED_OUTPUTS       )







@require_torch_multi_gpu
class        UpperCamelCase__   (	__magic_name__   ):

					def   UpperCAmelCase__       (     self       :   int       ):






										'''simple docstring'''






										super().setUp()







					def   UpperCAmelCase__       (     self       :   Dict       ):






										'''simple docstring'''






										lowercase_  							=			AutoModelForCausalLM.from_pretrained(
										    self.model_name							,					load_in_abit=UpperCamelCase__							,					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_  							=			self.tokenizer(self.input_text							,					return_tensors="""pt"""       )

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







class        UpperCamelCase__   (	__magic_name__   ):

					def   UpperCAmelCase__       (     self       :   Dict       ):






										'''simple docstring'''






										lowercase_  							=			"""facebook/opt-350m"""
										super().setUp()







					def   UpperCAmelCase__       (     self       :   List[Any]       ):






										'''simple docstring'''






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

										# Step 1: freeze all parameters
										lowercase_  							=			AutoModelForCausalLM.from_pretrained(self.model_name							,					load_in_abit=UpperCamelCase__       )

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

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

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

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

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

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







class        UpperCamelCase__   (	__magic_name__   ):
					__SCREAMING_SNAKE_CASE       :  Optional[int]						=   'gpt2-xl'
					__SCREAMING_SNAKE_CASE       :  Tuple						=   3.3_1_9_1_8_5_4_8_5_4_1_5_2_1_8_7






# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import platform

import numpy as np
import psutil
import torch

from accelerate import __version__ as version
from accelerate.commands.config import default_config_file, load_config_from_file

from ..utils import is_npu_available, is_xpu_available


def 			UpperCAmelCase_				(      UpperCAmelCase__=None  ):
					if subparsers is not None:
										lowercase_  							=			subparsers.add_parser("""env"""  )
					else:
										lowercase_  							=			argparse.ArgumentParser("""Accelerate env command"""  )

					parser.add_argument(
					    """--config_file"""						,							default=UpperCAmelCase__						,							help="""The config file to use for the default values in the launching script."""  )

					if subparsers is not None:
										parser.set_defaults(func=UpperCAmelCase__  )
					return parser


def 			UpperCAmelCase_				(      UpperCAmelCase__  ):
					lowercase_  							=			torch.__version__
					lowercase_  							=			torch.cuda.is_available()
					lowercase_  							=			is_xpu_available()
					lowercase_  							=			is_npu_available()

					lowercase_  							=			"""Not found"""
					# Get the default from the config file.
					if args.config_file is not None or os.path.isfile(UpperCAmelCase__  ):
										lowercase_  							=			load_config_from_file(args.config_file  ).to_dict()

					lowercase_  							=			{
					    """`Accelerate` version""": version,
					    """Platform""": platform.platform(),
					    """Python version""": platform.python_version(),
					    """Numpy version""": np.__version__,
					    """PyTorch version (GPU?)""": F'''{pt_version} ({pt_cuda_available})''',
					    """PyTorch XPU available""": str(UpperCAmelCase__  ),
					    """PyTorch NPU available""": str(UpperCAmelCase__  ),
					    """System RAM""": F'''{psutil.virtual_memory().total / 1_0_2_4 ** 3:.2f} GB''',
					}
					if pt_cuda_available:
										lowercase_  							=			torch.cuda.get_device_name()

					print("""\nCopy-and-paste the text below in your GitHub issue\n"""  )
					print("""\n""".join([F'''- {prop}: {val}''' for prop, val in info.items()]  )  )

					print("""- `Accelerate` default config:""" if args.config_file is None else """- `Accelerate` config passed:"""  )
					lowercase_  							=			(
					    """\n""".join([F'''\t- {prop}: {val}''' for prop, val in accelerate_config.items()]  )
					    if isinstance(UpperCAmelCase__						,							UpperCAmelCase__  )
					    else F'''\t{accelerate_config}'''
					)
					print(UpperCAmelCase__  )

					lowercase_  							=			accelerate_config

					return info


def 			UpperCAmelCase_				(      ):
					lowercase_  							=			env_command_parser()
					lowercase_  							=			parser.parse_args()
					env_command(UpperCAmelCase__  )
					return 0


if __name__ == "__main__":
							raise SystemExit(main())


