


import argparse
import json
import os
import re

import torch

from transformers import BloomConfig, BloomModel
from transformers.file_utils import CONFIG_NAME, WEIGHTS_NAME
from transformers.utils import logging


logging.set_verbosity_info()

A  :     Tuple    					=    [
    "word_embeddings_layernorm.weight",
    "word_embeddings_layernorm.bias",
    "input_layernorm.weight",
    "input_layernorm.bias",
    "post_attention_layernorm.weight",
    "post_attention_layernorm.bias",
    "self_attention.dense.bias",
    "mlp.dense_4h_to_h.bias",
    "ln_f.weight",
    "ln_f.bias",
]

A  :     Dict    					=    [
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]


def _lowerCamelCase			(	_UpperCamelCase					,				_UpperCamelCase   ):

			'''simple docstring'''




			__lowerCAmelCase       					=    {
			    '''word_embeddings.weight''': '''word_embeddings.weight''',
			    '''word_embeddings.norm.weight''': '''word_embeddings_layernorm.weight''',
			    '''word_embeddings.norm.bias''': '''word_embeddings_layernorm.bias''',
			    '''weight''': '''ln_f.weight''',
			    '''bias''': '''ln_f.bias''',
			}

			if key in layer_rename_map:
						return layer_rename_map[key]

			# Handle transformer blocks
			__lowerCAmelCase       					=    int(re.match(R".*layer_(\d*).*"					,				a_   )[1]   )
			layer_number -= 3
			return f"h.{layer_number}." + key


def _lowerCamelCase			(	_UpperCamelCase   ):

			'''simple docstring'''




			if dtype == torch.bool:
						return 1 / 8
			__lowerCAmelCase       					=    re.search(R"[^\d](\d+)$"					,				str(a_   )   )
			if bit_search is None:
						raise ValueError(f"`dtype` is not a valid dtype: {dtype}."   )
			__lowerCAmelCase       					=    int(bit_search.groups()[0]   )
			return bit_size // 8


def _lowerCamelCase			(	_UpperCamelCase					,				_UpperCamelCase					,				_UpperCamelCase					,				_UpperCamelCase					,				_UpperCamelCase   ):

			'''simple docstring'''




			if bloom_config_file == "":
						__lowerCAmelCase       					=    BloomConfig()
			else:
						__lowerCAmelCase       					=    BloomConfig.from_json_file(a_   )

			if shard_model:
						__lowerCAmelCase       					=    os.listdir(a_   )
						__lowerCAmelCase       					=    sorted(filter(lambda _UpperCamelCase   : s.startswith("layer"   ) and "model_00" in s					,				a_   )   )

						__lowerCAmelCase       					=    {'''weight_map''': {}, '''metadata''': {}}
						__lowerCAmelCase       					=    0

						__lowerCAmelCase       					=    None

						__lowerCAmelCase       					=    BloomConfig()

						for j, file in enumerate(a_   ):
									print("Processing file: {}".format(a_   )   )
									__lowerCAmelCase       					=    None

									for i in range(a_   ):
												# load all TP files
												__lowerCAmelCase       					=    file.replace("model_00"					,				f"model_0{i}"   )
												__lowerCAmelCase       					=    torch.load(os.path.join(a_					,				a_   )					,				map_location="cpu"   )

												# Rename keys in the transformers names
												__lowerCAmelCase       					=    list(temp.keys()   )
												for key in keys:
															__lowerCAmelCase       					=    temp.pop(a_   )

												if tensors is None:
															__lowerCAmelCase       					=    temp
												else:
															for key in tensors.keys():
																		if any(key.endswith(a_   ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH   ):
																					# We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
																					tensors[key] += temp[key]
																		else:
																					# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																					__lowerCAmelCase       					=    1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN   ) else 0
																					# We concatenate these weights accross TP ranks
																					__lowerCAmelCase       					=    torch.cat([tensors[key], temp[key]]					,				dim=a_   )

            # Divide by the number of TP the weights we want to average
									for key in tensors.keys():
												if any(key.endswith(a_   ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH   ):
															__lowerCAmelCase       					=    tensors[key] / pretraining_tp
									torch.save(
									    a_					,				os.path.join(
									        a_					,				"pytorch_model_{}-of-{}.bin".format(str(j + 1   ).zfill(5   )					,				str(len(a_   )   ).zfill(5   )   )					,				)					,				)

									for key in tensors.keys():
												__lowerCAmelCase       					=    tensors[key]
												total_size += value.numel() * get_dtype_size(value.dtype   )
												if key not in index_dict["weight_map"]:
															__lowerCAmelCase       					=    '''pytorch_model_{}-of-{}.bin'''.format(
															    str(j + 1   ).zfill(5   )					,				str(len(a_   )   ).zfill(5   )   )

						__lowerCAmelCase       					=    BloomConfig()
						__lowerCAmelCase       					=    pytorch_dump_folder_path + '''/''' + CONFIG_NAME
						__lowerCAmelCase       					=    total_size
						with open(a_					,				"w"					,				encoding="utf-8"   ) as f:
									f.write(config.to_json_string()   )
						with open(os.path.join(a_					,				WEIGHTS_NAME + ".index.json"   )					,				"w"					,				encoding="utf-8"   ) as f:
									__lowerCAmelCase       					=    json.dumps(a_					,				indent=2					,				sort_keys=a_   ) + '''\n'''
									f.write(a_   )
			else:
						__lowerCAmelCase       					=    BloomModel(a_   )

						__lowerCAmelCase       					=    os.listdir(a_   )
						__lowerCAmelCase       					=    sorted(filter(lambda _UpperCamelCase   : s.startswith("layer"   ) and "model_00" in s					,				a_   )   )

						__lowerCAmelCase       					=    None
						for i, file in enumerate(a_   ):
									__lowerCAmelCase       					=    None
									for i in range(a_   ):
												# load all TP files
												__lowerCAmelCase       					=    file.replace("model_00"					,				f"model_0{i}"   )
												__lowerCAmelCase       					=    torch.load(os.path.join(a_					,				a_   )					,				map_location="cpu"   )

												# Rename keys in the transformers names
												__lowerCAmelCase       					=    list(temp.keys()   )
												for key in keys:
															__lowerCAmelCase       					=    temp.pop(a_   )

												if tensors is None:
															__lowerCAmelCase       					=    temp
												else:
															for key in tensors.keys():
																		# We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
																		if any(key.endswith(a_   ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH   ):
																					tensors[key] += temp[key]
																		else:
																					# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																					__lowerCAmelCase       					=    1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN   ) else 0
																					# We concatenate these weights accross TP ranks
																					__lowerCAmelCase       					=    torch.cat([tensors[key], temp[key]]					,				dim=a_   )

            # Divide by the number of TP the weights we want to average
									for key in tensors.keys():
												if any(key.endswith(a_   ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH   ):
															__lowerCAmelCase       					=    tensors[key] / pretraining_tp

									__lowerCAmelCase       					=    model.load_state_dict(a_					,				strict=a_   )
									assert not other_keys.unexpected_keys, f"The keys {other_keys.unexpected_keys} are unexpected"
									if missing_keys is None:
												__lowerCAmelCase       					=    set(other_keys.missing_keys   )
									else:
												__lowerCAmelCase       					=    missing_keys.intersection(set(other_keys.missing_keys   )   )

						assert not missing_keys, f"The keys {missing_keys} are missing"

						# Save pytorch-model
						os.makedirs(a_					,				exist_ok=a_   )
						__lowerCAmelCase       					=    pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
						__lowerCAmelCase       					=    pytorch_dump_folder_path + '''/''' + CONFIG_NAME
						print(f"Save PyTorch model to {pytorch_weights_dump_path} with dtype {config.torch_dtype}"   )
						if config.torch_dtype is not None:
									__lowerCAmelCase       					=    model.to(config.torch_dtype   )
						torch.save(model.state_dict()					,				a_   )
						print(f"Save configuration file to {pytorch_config_dump_path}"   )
						with open(a_					,				"w"					,				encoding="utf-8"   ) as f:
									f.write(config.to_json_string()   )


if __name__ == "__main__":
	A  :     Any    					=    argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    "--bloom_checkpoint_path",
	    default=None,
	    type=str,
	    required=True,
	    help="Path to the Megatron-LM checkpoint path.",
	)
	parser.add_argument(
	    "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
	)
	parser.add_argument(
	    "--bloom_config_file",
	    default="",
	    type=str,
	    help=(
	        "An optional config json file corresponding to the pre-trained model. \n"
	        "This specifies the model architecture."
	    ),
	)
	parser.add_argument(
	    "--shard_model",
	    action="store_true",
	    help="An optional setting to shard the output model \nThis enables sharding the converted checkpoint",
	)
	parser.add_argument(
	    "--pretraining_tp",
	    default=4,
	    type=int,
	    help="Pretraining TP rank that has been used when training the model in Megatron-LM \n",
	)
	A  :     str    					=    parser.parse_args()
	convert_bloom_checkpoint_to_pytorch(
	    args.bloom_checkpoint_path,
	    args.bloom_config_file,
	    args.pytorch_dump_folder_path,
	    args.shard_model,
	    args.pretraining_tp,
	)








"""simple docstring"""


import gc
import threading
import time

import psutil
import torch
class 						_UpperCamelCase						:
							'''simple docstring'''







							def __init__(  self			):
										__lowerCAmelCase       					=    psutil.Process()
										__lowerCAmelCase       					=    False

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    -1

										while True:
													__lowerCAmelCase       					=    max(self.process.memory_info().rss					,     self.cpu_memory_peak			)

													# can't sleep or will not catch the peak right (this comment is here on purpose)
													if not self.peak_monitoring:
																break

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    True
										__lowerCAmelCase       					=    threading.Thread(target=self.peak_monitor			)
										__lowerCAmelCase       					=    True
										self.thread.start()



							def     snake_case       (  self			):
										__lowerCAmelCase       					=    False
										self.thread.join()
										return self.cpu_memory_peak
A  :     Any    					=    PeakCPUMemory()


def _lowerCamelCase			(	):

			'''simple docstring'''




			__lowerCAmelCase       					=    {"time": time.time()}

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

			# CPU mem
			__lowerCAmelCase       					=    psutil.Process().memory_info().rss
			cpu_peak_tracker.start()

			# GPU mem
			for i in range(torch.cuda.device_count()   ):
						__lowerCAmelCase       					=    torch.cuda.memory_allocated(_UpperCamelCase   )
			torch.cuda.reset_peak_memory_stats()

			return measures


def _lowerCamelCase			(	_UpperCamelCase   ):

			'''simple docstring'''




			__lowerCAmelCase       					=    {"time": time.time() - start_measures["time"]}

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

			# CPU mem
			__lowerCAmelCase       					=    (psutil.Process().memory_info().rss - start_measures["cpu"]) / 2**20
			__lowerCAmelCase       					=    (cpu_peak_tracker.stop() - start_measures["cpu"]) / 2**20

			# GPU mem
			for i in range(torch.cuda.device_count()   ):
						__lowerCAmelCase       					=    (torch.cuda.memory_allocated(_UpperCamelCase   ) - start_measures[str(_UpperCamelCase   )]) / 2**20
						__lowerCAmelCase       					=    (torch.cuda.max_memory_allocated(_UpperCamelCase   ) - start_measures[str(_UpperCamelCase   )]) / 2**20

			return measures


def _lowerCamelCase			(	_UpperCamelCase					,				_UpperCamelCase   ):

			'''simple docstring'''




			print(f"{description}:"   )
			print(f"- Time: {measures['time']:.2f}s"   )
			for i in range(torch.cuda.device_count()   ):
						print(f"- GPU {i} allocated: {measures[str(_UpperCamelCase   )]:.2f}MiB"   )
						__lowerCAmelCase       					=    measures[f"{i}-peak"]
						print(f"- GPU {i} peak: {peak:.2f}MiB"   )
			print(f"- CPU RAM allocated: {measures['cpu']:.2f}MiB"   )
			print(f"- CPU RAM peak: {measures['cpu-peak']:.2f}MiB"   )




