



"""simple docstring"""


import importlib.metadata
import warnings
from copy import deepcopy

from packaging import version

from ..utils import logging
from .import_utils import is_accelerate_available, is_bitsandbytes_available


if is_bitsandbytes_available():
				import bitsandbytes as bnb
				import torch
				import torch.nn as nn

				from ..pytorch_utils import ConvaD

if is_accelerate_available():
				from accelerate import init_empty_weights
				from accelerate.utils import find_tied_parameters

snake_case_    					=	logging.get_logger(__name__)

def 			_lowerCAmelCase				(       lowercase_		,     lowercase_		,     lowercase_		,     lowercase_=None		,     lowercase_=None						):
						# Recurse if needed
						if "." in tensor_name:
												UpperCAmelCase					=			tensor_name.split('.'						)
												for split in splits[:-1]:
																		UpperCAmelCase					=			getattr(lowercase_		,     lowercase_						)
																		if new_module is None:
																								raise ValueError(F"""{module} has no attribute {split}."""						)
																		UpperCAmelCase					=			new_module
												UpperCAmelCase					=			splits[-1]

						if tensor_name not in module._parameters and tensor_name not in module._buffers:
												raise ValueError(F"""{module} does not have a parameter or a buffer named {tensor_name}."""						)
						UpperCAmelCase					=			tensor_name in module._buffers
						UpperCAmelCase					=			getattr(lowercase_		,     lowercase_						)

						if old_value.device == torch.device('meta'						) and device not in ["meta", torch.device('meta'						)] and value is None:
												raise ValueError(F"""{tensor_name} is on the meta device, we need a `value` to put in on {device}."""						)

						UpperCAmelCase					=			False
						UpperCAmelCase					=			False
						if is_buffer or not is_bitsandbytes_available():
												UpperCAmelCase					=			False
												UpperCAmelCase					=			False
						else:
												UpperCAmelCase					=			hasattr(bnb.nn		,     'Params4bit'						) and isinstance(module._parameters[tensor_name]		,     bnb.nn.Paramsabit						)
												UpperCAmelCase					=			isinstance(module._parameters[tensor_name]		,     bnb.nn.IntaParams						)

						if is_abit or is_abit:
												UpperCAmelCase					=			module._parameters[tensor_name]
												if param.device.type != "cuda":
																		if value is None:
																								UpperCAmelCase					=			old_value.to(lowercase_						)
																		elif isinstance(lowercase_		,     torch.Tensor						):
																								UpperCAmelCase					=			value.to('cpu'						)
																								if value.dtype == torch.inta:
																														UpperCAmelCase					=			version.parse(importlib.metadata.version('bitsandbytes'						)						) > version.parse(
																														    '0.37.2'						)
																														if not is_abit_serializable:
																																				raise ValueError(
																																				    'Detected int8 weights but the version of bitsandbytes is not compatible with int8 serialization. '
																																				    'Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`.'						)
																		else:
																								UpperCAmelCase					=			torch.tensor(lowercase_		,     device='cpu'						)

																		# Support models using `Conv1D` in place of `nn.Linear` (e.g. gpt2) by transposing the weight matrix prior to quantization.
																		# Since weights are saved in the correct "orientation", we skip transposing when loading.
																		if issubclass(module.source_cls		,     lowercase_						) and fpaa_statistics is None:
																								UpperCAmelCase					=			new_value.T

																		UpperCAmelCase					=			old_value.__dict__
																		if is_abit:
																								UpperCAmelCase					=			bnb.nn.IntaParams(lowercase_		,     requires_grad=lowercase_		,     **lowercase_						).to(lowercase_						)
																		elif is_abit:
																								UpperCAmelCase					=			bnb.nn.Paramsabit(lowercase_		,     requires_grad=lowercase_		,     **lowercase_						).to(lowercase_						)

																		UpperCAmelCase					=			new_value
																		if fpaa_statistics is not None:
																								setattr(module.weight		,     'SCB'		,     fpaa_statistics.to(lowercase_						)						)

						else:
												if value is None:
																		UpperCAmelCase					=			old_value.to(lowercase_						)
												elif isinstance(lowercase_		,     torch.Tensor						):
																		UpperCAmelCase					=			value.to(lowercase_						)
												else:
																		UpperCAmelCase					=			torch.tensor(lowercase_		,     device=lowercase_						)

												if is_buffer:
																		UpperCAmelCase					=			new_value
												else:
																		UpperCAmelCase					=			nn.Parameter(lowercase_		,     requires_grad=old_value.requires_grad						)
																		UpperCAmelCase					=			new_value

def 			_lowerCAmelCase				(       lowercase_		,     lowercase_=None		,     lowercase_=None		,     lowercase_=None		,     lowercase_=False						):
						for name, module in model.named_children():
												if current_key_name is None:
																		UpperCAmelCase					=			[]
												current_key_name.append(lowercase_						)

												if (isinstance(lowercase_		,     nn.Linear						) or isinstance(lowercase_		,     lowercase_						)) and name not in modules_to_not_convert:
																		# Check if the current key is not in the `modules_to_not_convert`
																		if not any(key in '.'.join(lowercase_						) for key in modules_to_not_convert						):
																								with init_empty_weights():
																														if isinstance(lowercase_		,     lowercase_						):
																																				UpperCAmelCase				,   UpperCAmelCase					=			module.weight.shape
																														else:
																																				UpperCAmelCase					=			module.in_features
																																				UpperCAmelCase					=			module.out_features

																														if quantization_config.quantization_method() == "llm_int8":
																																				UpperCAmelCase					=			bnb.nn.LinearabitLt(
																																				    lowercase_		,     lowercase_		,     module.bias is not None		,     has_fpaa_weights=quantization_config.llm_inta_has_fpaa_weight		,     threshold=quantization_config.llm_inta_threshold		,     )
																																				UpperCAmelCase					=			True
																														else:
																																				if (
																																				    quantization_config.llm_inta_skip_modules is not None
																																				    and name in quantization_config.llm_inta_skip_modules
																																				):
																																										pass
																																				else:
																																										UpperCAmelCase					=			bnb.nn.Linearabit(
																																										    lowercase_		,     lowercase_		,     module.bias is not None		,     quantization_config.bnb_abit_compute_dtype		,     compress_statistics=quantization_config.bnb_abit_use_double_quant		,     quant_type=quantization_config.bnb_abit_quant_type		,     )
																																										UpperCAmelCase					=			True
                    # Store the module class in case we need to transpose the weight later
																														UpperCAmelCase					=			type(lowercase_						)
																														# Force requires grad to False to avoid unexpected errors
																														model._modules[name].requires_grad_(lowercase_						)
												if len(list(module.children()						)						) > 0:
																		UpperCAmelCase				,   UpperCAmelCase					=			_replace_with_bnb_linear(
																		    lowercase_		,     lowercase_		,     lowercase_		,     lowercase_		,     has_been_replaced=lowercase_		,     )
												# Remove the last key for recursion
												current_key_name.pop(-1						)
						return model, has_been_replaced

def 			_lowerCAmelCase				(       lowercase_		,     lowercase_=None		,     lowercase_=None		,     lowercase_=None						):
						UpperCAmelCase					=			['lm_head'] if modules_to_not_convert is None else modules_to_not_convert
						UpperCAmelCase				,   UpperCAmelCase					=			_replace_with_bnb_linear(
						    lowercase_		,     lowercase_		,     lowercase_		,     lowercase_						)

						if not has_been_replaced:
												logger.warning(
												    'You are loading your model in 8bit or 4bit but no linear modules were found in your model.'
												    ' Please double check your model architecture, or submit an issue on github if you think this is'
												    ' a bug.'						)

						return model

def 			_lowerCAmelCase				(       *lowercase_		,     **lowercase_						):
						warnings.warn(
						    '`replace_8bit_linear` will be deprecated in a future version, please use `replace_with_bnb_linear` instead'		,     lowercase_		,     )
						return replace_with_bnb_linear(*lowercase_		,     **lowercase_						)

def 			_lowerCAmelCase				(       *lowercase_		,     **lowercase_						):
						warnings.warn(
						    '`set_module_8bit_tensor_to_device` will be deprecated in a future version, please use `set_module_quantized_tensor_to_device` instead'		,     lowercase_		,     )
						return set_module_quantized_tensor_to_device(*lowercase_		,     **lowercase_						)






def 			_lowerCAmelCase				(       lowercase_						):
						UpperCAmelCase					=			deepcopy(lowercase_						)  # this has 0 cost since it is done inside `init_empty_weights` context manager`
						tied_model.tie_weights()

						UpperCAmelCase					=			find_tied_parameters(lowercase_						)
						# For compatibility with Accelerate < 0.18
						if isinstance(lowercase_		,     lowercase_						):
												UpperCAmelCase					=			sum(list(tied_params.values()						)		,     []						) + list(tied_params.keys()						)
						else:
												UpperCAmelCase					=			sum(lowercase_		,     []						)
						UpperCAmelCase					=			len(lowercase_						) > 0

						# Check if it is a base model
						UpperCAmelCase					=			not hasattr(lowercase_		,     model.base_model_prefix						)

						# Ignore this for base models (BertModel, GPT2Model, etc.)
						if (not has_tied_params) and is_base_model:
												return []

						# otherwise they have an attached head
						UpperCAmelCase					=			list(model.named_children()						)
						UpperCAmelCase					=			[list_modules[-1][0]]

						# add last module together with tied weights
						UpperCAmelCase					=			set(lowercase_						) - set(lowercase_						)
						UpperCAmelCase					=			list(set(lowercase_						)						) + list(lowercase_						)

						# remove ".weight" from the keys
						UpperCAmelCase					=			['.weight', '.bias']
						UpperCAmelCase					=			[]
						for name in list_untouched:
												for name_to_remove in names_to_remove:
																		if name_to_remove in name:
																								UpperCAmelCase					=			name.replace(lowercase_		,     ''						)
												filtered_module_names.append(lowercase_						)

						return filtered_module_names










"""simple docstring"""


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

import torch_xla.distributed.xla_multiprocessing as xmp

def 			_lowerCAmelCase				(       ):
						UpperCAmelCase					=			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 			_lowerCAmelCase				(       ):
						UpperCAmelCase					=			parse_args()

						# Import training_script as a module.
						UpperCAmelCase					=			Path(args.training_script						)
						sys.path.append(str(script_fpath.parent.resolve()						)						)
						UpperCAmelCase					=			script_fpath.stem
						UpperCAmelCase					=			importlib.import_module(lowercase_						)

						# Patch sys.argv
						UpperCAmelCase					=			[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()





