



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 __lowerCAmelCase (     SCREAMING_SNAKE_CASE_  ):
			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 				_snake_case			(					nn.Module):




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




											def 							A__      (						self					:						List[str],					__lowercase					:						Any,					*__lowercase					:						Any,					**__lowercase					:						Dict						):
														return self.module(__lowercase,					*__lowercase,					**__lowercase						) + self.adapter(__lowercase						)
@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 				_snake_case			(					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
				UpperCamelCase__				:					int     		="""bigscience/bloom-1b7"""

				# Constant values
				UpperCamelCase__				:					List[str]     		=2.109_659_552_692_574

				UpperCamelCase__				:					str     		="""Hello my name is"""
				UpperCamelCase__				:					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""")
				UpperCamelCase__				:					Optional[Any]     		=1_0




				def 							A__      (						self					:						Optional[int]						):
							# Models and tokenizer
							lowercase__					  =				AutoTokenizer.from_pretrained(self.model_name						)
class 				_snake_case			(					lowercase__):




				def 							A__      (						self					:						Union[str, Any]						):
							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=__lowercase,					device_map="auto"						)




				def 							A__      (						self					:						Union[str, Any]						):
							del self.model_fpaa
							del self.model_abit

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




				def 							A__      (						self					:						int						):
							lowercase__					  =				self.model_abit.config

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

							lowercase__					  =				config.to_dict()
							lowercase__					  =				config.to_diff_dict()

							lowercase__					  =				config.to_json_string()




				def 							A__      (						self					:						List[str]						):
							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 							A__      (						self					:						Optional[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(__lowercase,					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 							A__      (						self					:						Optional[int]						):
							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=__lowercase						),					self.EXPECTED_OUTPUTS						)




				def 							A__      (						self					:						str						):
							lowercase__					  =				BitsAndBytesConfig()
							lowercase__					  =				True

							lowercase__					  =				AutoModelForCausalLM.from_pretrained(
							    self.model_name,					quantization_config=__lowercase,					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=__lowercase						),					self.EXPECTED_OUTPUTS						)




				def 							A__      (						self					:						int						):
							with self.assertRaises(__lowercase						), tempfile.TemporaryDirectory() as tmpdirname:
										self.model_abit.save_pretrained(__lowercase						)




				def 							A__      (						self					:						str						):
							lowercase__					  =				BitsAndBytesConfig()

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




				def 							A__      (						self					:						Tuple						):
							with self.assertRaises(__lowercase						):
										# Tries with `str`
										self.model_abit.to("cpu"						)

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

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

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

							with self.assertRaises(__lowercase						):
										# 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 							A__      (						self					:						Any						):
							lowercase__					  =				AutoModelForSeqaSeqLM.from_pretrained("t5-small",					load_in_abit=__lowercase,					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 				_snake_case			(					unittest.TestCase):




				@classmethod
				def 							A__      (						cls					:						Union[str, Any]						):
							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 							A__      (						self					:						List[str]						):
							gc.collect()
							torch.cuda.empty_cache()




				def 							A__      (						self					:						Dict						):
							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=__lowercase,					device_map="auto"						)
							lowercase__					  =				self.tokenizer(self.input_text,					return_tensors="pt"						).to(0						)
							lowercase__					  =				model.generate(**__lowercase						)

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




				def 							A__      (						self					:						int						):
							import bitsandbytes as bnb

							from transformers import TaForConditionalGeneration

							# test with `t5-small`
							lowercase__					  =				TaForConditionalGeneration.from_pretrained(self.model_name,					load_in_abit=__lowercase,					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(**__lowercase						)

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




				def 							A__      (						self					:						Tuple						):
							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=__lowercase,					device_map="auto"						)
							# Sequence classification model
							lowercase__					  =				AutoModelForSequenceClassification.from_pretrained(
							    self.model_name,					load_in_abit=__lowercase,					device_map="auto"						)
							# CausalLM model
							lowercase__					  =				AutoModelForCausalLM.from_pretrained(self.model_name,					load_in_abit=__lowercase,					device_map="auto"						)
							# Seq2seq model
							lowercase__					  =				AutoModelForSeqaSeqLM.from_pretrained(
							    self.seq_to_seq_name,					load_in_abit=__lowercase,					device_map="auto"						)




				def 							A__      (						self					:						List[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 							A__      (						self					:						Optional[int]						):
							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 				_snake_case			(					lowercase__):




				def 							A__      (						self					:						Tuple						):
							super().setUp()




				def 							A__      (						self					:						List[Any]						):
							del self.pipe

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




				def 							A__      (						self					:						Union[str, Any]						):
							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 				_snake_case			(					lowercase__):




				def 							A__      (						self					:						Any						):
							super().setUp()




				def 							A__      (						self					:						str						):
							lowercase__					  =				AutoModelForCausalLM.from_pretrained(
							    self.model_name,					load_in_abit=__lowercase,					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=__lowercase						),					self.EXPECTED_OUTPUTS						)
class 				_snake_case			(					lowercase__):




				def 							A__      (						self					:						Optional[Any]						):
							lowercase__					  =				"facebook/opt-350m"
							super().setUp()




				def 							A__      (						self					:						Dict						):
							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=__lowercase						)

							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(__lowercase						)						):
													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(**__lowercase						)
										out.logits.norm().backward()

							for module in model.modules():
										if isinstance(__lowercase,					__lowercase						):
													self.assertTrue(module.adapter[1].weight.grad is not None						)
													self.assertTrue(module.adapter[1].weight.grad.norm().item() > 0						)
										elif isinstance(__lowercase,					nn.Embedding						):
													self.assertTrue(module.weight.grad is None						)
class 				_snake_case			(					lowercase__):
				UpperCamelCase__				:					Union[str, Any]     		="""gpt2-xl"""
				UpperCamelCase__				:					Any     		=3.3_191_854_854_152_187









# Copyright 2021 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.

from argparse import ArgumentParser

from accelerate.commands.config import get_config_parser
from accelerate.commands.env import env_command_parser
from accelerate.commands.launch import launch_command_parser
from accelerate.commands.test import test_command_parser
from accelerate.commands.tpu import tpu_command_parser



def __lowerCAmelCase (     ):
			lowercase__					  =				ArgumentParser("Accelerate CLI tool"						,  usage="accelerate <command> [<args>]"						,  allow_abbrev=SCREAMING_SNAKE_CASE_  )
			lowercase__					  =				parser.add_subparsers(help="accelerate command helpers"  )

			# Register commands
			get_config_parser(subparsers=SCREAMING_SNAKE_CASE_  )
			env_command_parser(subparsers=SCREAMING_SNAKE_CASE_  )
			launch_command_parser(subparsers=SCREAMING_SNAKE_CASE_  )
			tpu_command_parser(subparsers=SCREAMING_SNAKE_CASE_  )
			test_command_parser(subparsers=SCREAMING_SNAKE_CASE_  )

			# Let's go
			lowercase__					  =				parser.parse_args()

			if not hasattr(SCREAMING_SNAKE_CASE_						,  "func"  ):
						parser.print_help()
						exit(1  )

			# Run
			args.func(SCREAMING_SNAKE_CASE_  )


if __name__ == "__main__":
							main()




