

import argparse
import gc
import os
import threading

import evaluate
import psutil
import torch
from datasets import load_dataset
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig
from torch.utils.data import DataLoader
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    get_linear_schedule_with_warmup,
    set_seed,
)

from accelerate import Accelerator, DistributedType, FullyShardedDataParallelPlugin


########################################################################
# This is a fully working simple example to use Accelerate
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#   - FSDP
#
# This example also demonstrates the checkpointing and sharding capabilities
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


__snake_case							:     int 						=	16
__snake_case							:     int 						=	32


def 			_UpperCAmelCase							(		a__):

			'''simple docstring'''



			return int(x / 2**2_0)



class    A__:
						"""simple docstring"""






						def __enter__(  self  )      ->				Dict:
									gc.collect()
									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
									a_		:		Any				    =				torch.cuda.memory_allocated()
									a_		:		Dict				    =				psutil.Process()

									a_		:		Tuple				    =				self.cpu_mem_used()
									a_		:		List[str]				    =				True
									a_		:		Optional[Any]				    =				threading.Thread(target=self.peak_monitor_func  )
									a_		:		int				    =				True
									peak_monitor_thread.start()
									return self





						def 	UpperCamelCase__			(  self  )      ->				Dict:
									return self.process.memory_info().rss





						def 	UpperCamelCase__			(  self  )      ->				Optional[Any]:
									a_		:		Any				    =				-1

									while True:
												a_		:		List[str]				    =				max(self.cpu_mem_used()       ,     self.cpu_peak  )

												# can't sleep or will not catch the peak right (this comment is here on purpose)
												# time.sleep(0.001) # 1msec

												if not self.peak_monitoring:
															break





						def __exit__(  self       ,     *_lowercase  )      ->				List[Any]:
									a_		:		Dict				    =				False

									gc.collect()
									torch.cuda.empty_cache()
									a_		:		Dict				    =				torch.cuda.memory_allocated()
									a_		:		Any				    =				torch.cuda.max_memory_allocated()
									a_		:		Any				    =				bamb(self.end - self.begin  )
									a_		:		Dict				    =				bamb(self.peak - self.begin  )

									a_		:		Optional[Any]				    =				self.cpu_mem_used()
									a_		:		Union[str, Any]				    =				bamb(self.cpu_end - self.cpu_begin  )
									a_		:		List[Any]				    =				bamb(self.cpu_peak - self.cpu_begin  )
									# print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")


# For testing only
if os.environ.get("""TESTING_MOCKED_DATALOADERS""", None) == "1":
	from accelerate.test_utils.training import mocked_dataloaders

	__snake_case							:     List[str] 						=	mocked_dataloaders  # noqa: F811


def 			_UpperCAmelCase							(		a__       ,					a__):

			'''simple docstring'''



			if os.environ.get("""TESTING_MOCKED_DATALOADERS"""       ,					a__) == "1":
						a_		:		int				    =				2

			# New Code #
			# Pass the advanced FSDP settings not part of the accelerate config by creating fsdp_plugin
			a_		:		Any				    =				FullyShardedDataParallelPlugin(
			    state_dict_config=FullStateDictConfig(offload_to_cpu=a__       ,					ranka_only=a__)       ,					optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=a__       ,					ranka_only=a__)       ,					)

			# Initialize accelerator
			if args.with_tracking:
						a_		:		Dict				    =				Accelerator(
						    cpu=args.cpu       ,					mixed_precision=args.mixed_precision       ,					log_with="""wandb"""       ,					project_dir=args.logging_dir       ,					fsdp_plugin=a__       ,					)
			else:
						a_		:		str				    =				Accelerator(fsdp_plugin=a__)
			accelerator.print(accelerator.distributed_type)

			if hasattr(args.checkpointing_steps       ,					"""isdigit"""):
						if args.checkpointing_steps == "epoch":
									a_		:		Any				    =				args.checkpointing_steps
						elif args.checkpointing_steps.isdigit():
									a_		:		str				    =				int(args.checkpointing_steps)
						else:
									raise ValueError(
									    f'''Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.''')
			else:
						a_		:		List[str]				    =				None
			# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
			a_		:		List[Any]				    =				config["""lr"""]
			a_		:		str				    =				int(config["""num_epochs"""])
			a_		:		Any				    =				int(config["""seed"""])
			a_		:		Tuple				    =				int(config["""batch_size"""])

			# We need to initialize the trackers we use, and also store our configuration
			if args.with_tracking:
						a_		:		List[str]				    =				vars(a__)
						accelerator.init_trackers("""fsdp_glue_no_trainer"""       ,					a__)

			a_		:		Tuple				    =				AutoTokenizer.from_pretrained(args.model_name_or_path)
			a_		:		Optional[Any]				    =				load_dataset("""glue"""       ,					"""mrpc""")
			a_		:		int				    =				evaluate.load("""glue"""       ,					"""mrpc""")

			def tokenize_function(a__):
						# max_length=None => use the model max length (it's actually the default)
						a_		:		Any				    =				tokenizer(examples["""sentence1"""]       ,					examples["""sentence2"""]       ,					truncation=a__       ,					max_length=a__)
						return outputs

			# Apply the method we just defined to all the examples in all the splits of the dataset
			# starting with the main process first:
			with accelerator.main_process_first():
						a_		:		Tuple				    =				datasets.map(
						    a__       ,					batched=a__       ,					remove_columns=["""idx""", """sentence1""", """sentence2"""]       ,					)

			# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
			# transformers library
			a_		:		List[str]				    =				tokenized_datasets.rename_column("""label"""       ,					"""labels""")

			# If the batch size is too big we use gradient accumulation
			a_		:		Tuple				    =				1
			if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
						a_		:		Dict				    =				batch_size // MAX_GPU_BATCH_SIZE
						a_		:		str				    =				MAX_GPU_BATCH_SIZE

			def collate_fn(a__):
						# On TPU it's best to pad everything to the same length or training will be very slow.
						a_		:		Tuple				    =				1_2_8 if accelerator.distributed_type == DistributedType.TPU else None
						# When using mixed precision we want round multiples of 8/16
						if accelerator.mixed_precision == "fp8":
									a_		:		str				    =				1_6
						elif accelerator.mixed_precision != "no":
									a_		:		Any				    =				8
						else:
									a_		:		List[Any]				    =				None

						return tokenizer.pad(
						    a__       ,					padding="""longest"""       ,					max_length=a__       ,					pad_to_multiple_of=a__       ,					return_tensors="""pt"""       ,					)

			# Instantiate dataloaders.
			a_		:		str				    =				DataLoader(
			    tokenized_datasets["""train"""]       ,					shuffle=a__       ,					collate_fn=a__       ,					batch_size=a__)
			a_		:		int				    =				DataLoader(
			    tokenized_datasets["""validation"""]       ,					shuffle=a__       ,					collate_fn=a__       ,					batch_size=a__)

			set_seed(a__)

			# Instantiate the model (we build the model here so that the seed also control new weights initialization)
			a_		:		Dict				    =				AutoModelForSequenceClassification.from_pretrained(
			    args.model_name_or_path       ,					return_dict=a__       ,					low_cpu_mem_usage=a__)

			# New Code #
			# For FSDP feature, it is highly recommended and efficient to prepare the model before creating optimizer
			a_		:		Optional[int]				    =				accelerator.prepare(a__)
			accelerator.print(a__)

			# Instantiate optimizer
			# New Code #
			# For FSDP feature, at present it doesn't support multiple parameter groups,
			# so we need to create a single parameter group for the whole model
			a_		:		Tuple				    =				torch.optim.AdamW(params=model.parameters()       ,					lr=a__       ,					weight_decay=2e-4)

			# Instantiate scheduler
			a_		:		str				    =				get_linear_schedule_with_warmup(
			    optimizer=a__       ,					num_warmup_steps=1_0       ,					num_training_steps=(len(a__) * num_epochs) // gradient_accumulation_steps       ,					)

			# New Code #
			# For FSDP feature, prepare everything except the model as we have already prepared the model
			# before creating the optimizer
			# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
			# prepare method.
			a_    ,			a_    ,			a_    ,			a_		:		Dict				    =				accelerator.prepare(
			    a__       ,					a__       ,					a__       ,					a__)

			a_		:		Dict				    =				0

			# Potentially load in the weights and states from a previous save
			if args.resume_from_checkpoint:
						if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
									accelerator.print(f'''Resumed from checkpoint: {args.resume_from_checkpoint}''')
									accelerator.load_state(args.resume_from_checkpoint)
									a_		:		int				    =				os.path.basename(args.resume_from_checkpoint)
						else:
									# Get the most recent checkpoint
									a_		:		List[str]				    =				[f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
									dirs.sort(key=os.path.getctime)
									a_		:		Dict				    =				dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
						# Extract `epoch_{i}` or `step_{i}`
						a_		:		Union[str, Any]				    =				os.path.splitext(a__)[0]

						if "epoch" in training_difference:
									num_epochs -= int(training_difference.replace("""epoch_"""       ,					""""""))
									a_		:		Any				    =				None
						else:
									a_		:		Optional[int]				    =				int(training_difference.replace("""step_"""       ,					""""""))
									num_epochs -= resume_step // len(a__)
									# If resuming by step, we also need to know exactly how far into the DataLoader we went
									a_		:		Union[str, Any]				    =				(num_epochs * len(a__)) - resume_step

    # Now we train the model
			for epoch in range(a__):
						# New Code #
						# context manager to track the peak memory usage during the training epoch
						with TorchTracemalloc() as tracemalloc:
									model.train()
									if args.with_tracking:
												a_		:		Tuple				    =				0
									for step, batch in enumerate(a__):
												# We need to skip steps until we reach the resumed step
												if args.resume_from_checkpoint and epoch == 0:
															if resume_step is not None and step < resume_step:
																		pass
                # We could avoid this line since we set the accelerator with `device_placement=True`.
												batch.to(accelerator.device)
												a_		:		Any				    =				model(**a__)
												a_		:		Optional[Any]				    =				outputs.loss
												a_		:		Any				    =				loss / gradient_accumulation_steps
												# We keep track of the loss at each epoch
												if args.with_tracking:
															total_loss += loss.detach().float()
												accelerator.backward(a__)
												if step % gradient_accumulation_steps == 0:
															optimizer.step()
															lr_scheduler.step()
															optimizer.zero_grad()
															# accelerator.print(lr_scheduler.get_lr())

												overall_step += 1

												if isinstance(a__       ,					a__):
															a_		:		Dict				    =				f'''step_{overall_step}'''
															if overall_step % checkpointing_steps == 0:
																		if args.output_dir is not None:
																					a_		:		Union[str, Any]				    =				os.path.join(args.output_dir       ,					a__)
																		accelerator.save_state(a__)
        # New Code #
        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
						accelerator.print("""Memory before entering the train : {}""".format(bamb(tracemalloc.begin)))
						accelerator.print("""Memory consumed at the end of the train (end-begin): {}""".format(tracemalloc.used))
						accelerator.print("""Peak Memory consumed during the train (max-begin): {}""".format(tracemalloc.peaked))
						accelerator.print(
						    """Total Peak Memory consumed during the train (max): {}""".format(
						        tracemalloc.peaked + bamb(tracemalloc.begin)))
						# Logging the peak memory usage of the GPU to the tracker
						if args.with_tracking:
									accelerator.log(
									    {
									        """train_total_peak_memory""": tracemalloc.peaked + bamb(tracemalloc.begin),
									    }       ,					step=a__       ,					)

						# New Code #
						# context manager to track the peak memory usage during the evaluation
						with TorchTracemalloc() as tracemalloc:
									model.eval()
									for step, batch in enumerate(a__):
												# We could avoid this line since we set the accelerator with `device_placement=True`.
												batch.to(accelerator.device)
												with torch.no_grad():
															a_		:		str				    =				model(**a__)
												a_		:		Tuple				    =				outputs.logits.argmax(dim=-1)
												a_    ,			a_		:		Dict				    =				accelerator.gather_for_metrics((predictions, batch["""labels"""]))
												metric.add_batch(
												    predictions=a__       ,					references=a__       ,					)

									a_		:		List[str]				    =				metric.compute()
									# Use accelerator.print to print only on the main process.
									accelerator.print(f'''epoch {epoch}:'''       ,					a__)
									if args.with_tracking:
												accelerator.log(
												    {
												        """accuracy""": eval_metric["""accuracy"""],
												        """f1""": eval_metric["""f1"""],
												        """train_loss""": total_loss.item() / len(a__),
												    }       ,					step=a__       ,					)

									if checkpointing_steps == "epoch":
												a_		:		Dict				    =				f'''epoch_{epoch}'''
												if args.output_dir is not None:
															a_		:		Optional[Any]				    =				os.path.join(args.output_dir       ,					a__)
												accelerator.save_state(a__)
        # New Code #
        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
						accelerator.print("""Memory before entering the eval : {}""".format(bamb(tracemalloc.begin)))
						accelerator.print("""Memory consumed at the end of the eval (end-begin): {}""".format(tracemalloc.used))
						accelerator.print("""Peak Memory consumed during the eval (max-begin): {}""".format(tracemalloc.peaked))
						accelerator.print(
						    """Total Peak Memory consumed during the eval (max): {}""".format(tracemalloc.peaked + bamb(tracemalloc.begin)))
						# Logging the peak memory usage of the GPU to the tracker
						if args.with_tracking:
									accelerator.log(
									    {
									        """eval_total_peak_memory""": tracemalloc.peaked + bamb(tracemalloc.begin),
									    }       ,					step=a__       ,					)

			if args.with_tracking:
						accelerator.end_training()


def 			_UpperCAmelCase							(		):

			'''simple docstring'''



			a_		:		Optional[Any]				    =				argparse.ArgumentParser(description="""Simple example of training script.""")
			parser.add_argument(
			    """--mixed_precision"""       ,					type=a__       ,					default=a__       ,					choices=["""no""", """fp16""", """bf16""", """fp8"""]       ,					help="""Whether to use mixed precision. Choose"""
			    """between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."""
			    """and an Nvidia Ampere GPU."""       ,					)
			parser.add_argument("""--cpu"""       ,					action="""store_true"""       ,					help="""If passed, will train on the CPU.""")
			parser.add_argument(
			    """--checkpointing_steps"""       ,					type=a__       ,					default=a__       ,					help="""Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch."""       ,					)
			parser.add_argument(
			    """--resume_from_checkpoint"""       ,					type=a__       ,					default=a__       ,					help="""If the training should continue from a checkpoint folder."""       ,					)
			parser.add_argument(
			    """--with_tracking"""       ,					action="""store_true"""       ,					help="""Whether to load in all available experiment trackers from the environment and use them for logging."""       ,					)
			parser.add_argument(
			    """--output_dir"""       ,					type=a__       ,					default="""."""       ,					help="""Optional save directory where all checkpoint folders will be stored. Default is the current working directory."""       ,					)
			parser.add_argument(
			    """--logging_dir"""       ,					type=a__       ,					default="""logs"""       ,					help="""Location on where to store experiment tracking logs`"""       ,					)
			parser.add_argument(
			    """--model_name_or_path"""       ,					type=a__       ,					help="""Path to pretrained model or model identifier from huggingface.co/models."""       ,					required=a__       ,					)
			a_		:		str				    =				parser.parse_args()
			a_		:		Tuple				    =				{"""lr""": 2e-5, """num_epochs""": 3, """seed""": 4_2, """batch_size""": 1_6}
			training_function(a__       ,					a__)


if __name__ == "__main__":
	main()





# Copyright (c) 2021-, NVIDIA CORPORATION.  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.

####################################################################################################

#
# Note: If when running this conversion script you're getting an exception:
#     ModuleNotFoundError: No module named 'megatron.model.enums'
# you need to tell python where to find the clone of Megatron-LM, e.g.:
#
# cd /tmp
# git clone https://github.com/NVIDIA/Megatron-LM
# PYTHONPATH=/tmp/Megatron-LM python src/transformers/models/megatron_bert/convert_megatron_bert_checkpoint.py ...
#
# if you already have it cloned elsewhere, simply adjust the path to the existing path
#
# If the training was done using a Megatron-LM fork, e.g.,
# https://github.com/microsoft/Megatron-DeepSpeed/ then chances are that you need to have that one
# in your path, i.e., /path/to/Megatron-DeepSpeed/
#

import argparse
import os
import re
import zipfile

import torch

from transformers import MegatronBertConfig


def 			_UpperCAmelCase							(		a__       ,					a__       ,					a__=0):

			'''simple docstring'''



			if name is None:
						a_		:		str				    =				None
			else:
						a_		:		List[str]				    =				""".""" * max(0       ,					spaces - 2) + """# {:""" + str(5_0 - spaces) + """s}"""
						a_		:		Optional[Any]				    =				fmt.format(a__)

			# Print and recurse (if needed).
			if isinstance(a__       ,					a__):
						if msg is not None:
									print(a__)
						for k in val.keys():
									recursive_print(a__       ,					val[k]       ,					spaces + 2)
			elif isinstance(a__       ,					torch.Tensor):
						print(a__       ,					""":"""       ,					val.size())
			else:
						print(a__       ,					""":"""       ,					a__)


def 			_UpperCAmelCase							(		a__       ,					a__       ,					a__       ,					a__       ,					a__):

			'''simple docstring'''



			a_		:		int				    =				param.size()
			if checkpoint_version == 1.0:
						# version 1.0 stores [num_heads * hidden_size * num_splits, :]
						a_		:		List[Any]				    =				(num_heads, hidden_size, num_splits) + input_shape[1:]
						a_		:		Any				    =				param.view(*a__)
						a_		:		Any				    =				param.transpose(0       ,					2)
						a_		:		Any				    =				param.transpose(1       ,					2).contiguous()
			elif checkpoint_version >= 2.0:
						# other versions store [num_heads * num_splits * hidden_size, :]
						a_		:		Optional[Any]				    =				(num_heads, num_splits, hidden_size) + input_shape[1:]
						a_		:		Any				    =				param.view(*a__)
						a_		:		Optional[Any]				    =				param.transpose(0       ,					1).contiguous()
			a_		:		Union[str, Any]				    =				param.view(*a__)
			return param


def 			_UpperCAmelCase							(		a__       ,					a__       ,					a__):

			'''simple docstring'''



			a_		:		Dict				    =				{}

			# old versions did not store training args
			a_		:		Union[str, Any]				    =				input_state_dict.get("""args"""       ,					a__)
			if ds_args is not None:
						# do not make the user write a config file when the exact dimensions/sizes are already in the checkpoint
						# from pprint import pprint
						# pprint(vars(ds_args))

						a_		:		Optional[Any]				    =				ds_args.tokenizer_type
						a_		:		str				    =				ds_args.padded_vocab_size
						a_		:		Optional[Any]				    =				ds_args.max_position_embeddings
						a_		:		Optional[Any]				    =				ds_args.hidden_size
						a_		:		Optional[Any]				    =				ds_args.num_layers
						a_		:		Optional[int]				    =				ds_args.num_attention_heads
						a_		:		List[Any]				    =				ds_args.ffn_hidden_size if """ffn_hidden_size""" in ds_args else 4 * ds_args.hidden_size
						# pprint(config)

			# The number of heads.
			a_		:		Tuple				    =				config.num_attention_heads
			# The hidden_size per head.
			a_		:		List[str]				    =				config.hidden_size // heads
			# Megatron-LM checkpoint version
			if "checkpoint_version" in input_state_dict.keys():
						a_		:		Optional[Any]				    =				input_state_dict["""checkpoint_version"""]
			else:
						a_		:		Tuple				    =				0.0

			# The model.
			a_		:		Optional[Any]				    =				input_state_dict["""model"""]
			# The language model.
			a_		:		str				    =				model["""language_model"""]
			# The embeddings.
			a_		:		Dict				    =				lm["""embedding"""]

			# The word embeddings.
			a_		:		str				    =				embeddings["""word_embeddings"""]["""weight"""]
			# Truncate the embedding table to vocab_size rows.
			a_		:		Dict				    =				word_embeddings[: config.vocab_size, :]
			# Store the word embeddings.
			a_		:		str				    =				word_embeddings

			# The position embeddings.
			a_		:		Tuple				    =				embeddings["""position_embeddings"""]["""weight"""]
			assert pos_embeddings.size(0) == config.max_position_embeddings and pos_embeddings.size(1) == config.hidden_size
			# Store the position embeddings.
			a_		:		Optional[Any]				    =				pos_embeddings

			# The token-type embeddings.
			a_		:		Tuple				    =				embeddings["""tokentype_embeddings"""]["""weight"""]
			# Store the position embeddings.
			a_		:		List[Any]				    =				tokentype_embeddings

			# The transformer.
			a_		:		List[str]				    =				lm["""transformer"""] if """transformer""" in lm.keys() else lm["""encoder"""]

			# The regex to extract layer names.
			a_		:		str				    =				re.compile(r"""layers\.(\d+)\.([a-z0-9_.]+)\.([a-z]+)""")

			# The simple map of names for "automated" rules.
			a_		:		List[str]				    =				{
			    """attention.dense""": """.attention.output.dense.""",
			    """self_attention.dense""": """.attention.output.dense.""",
			    """mlp.dense_h_to_4h""": """.intermediate.dense.""",
			    """mlp.dense_4h_to_h""": """.output.dense.""",
			}

			# Keep track of the attention/query/value tensor.
			a_		:		Optional[Any]				    =				None

			# Extract the layers.
			for key, val in transformer.items():
						# Match the name.
						a_		:		Optional[Any]				    =				layer_re.match(a__)

						# Stop if that's not a layer
						if m is None:
									break

						# The index of the layer.
						a_		:		Any				    =				int(m.group(1))
						# The name of the operation.
						a_		:		List[Any]				    =				m.group(2)
						# Is it a weight or a bias?
						a_		:		str				    =				m.group(3)

						# The name of the layer.
						a_		:		Optional[Any]				    =				f'''bert.encoder.layer.{layer_idx}'''

						# For layernorm(s), simply store the layer norm.
						if op_name.endswith("""layernorm"""):
									a_		:		Any				    =				"""attention.ln""" if op_name.startswith("""input""") else """ln"""
									a_		:		Union[str, Any]				    =				val

						# Transpose the QKV matrix.
						elif (
						    op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
						) and weight_or_bias == "weight":
									# Make sure the QKV pointer is nil.
									assert attention_qkv_weight is None, ""

									a_		:		int				    =				fix_query_key_value_ordering(a__       ,					a__       ,					3       ,					a__       ,					a__)
									# Store the tensor as we need the bias as well to interleave QKV and biases.
									a_		:		Any				    =				out_val

						# Transpose the bias.
						elif (
						    op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
						) and weight_or_bias == "bias":
									# Make sure we read the weight tensor.
									assert attention_qkv_weight is not None, ""

									# Split the QKV matrix into Q, K and V. Megatron stores Q,K,V interleaved.
									a_		:		Tuple				    =				attention_qkv_weight[0 * config.hidden_size : 1 * config.hidden_size, :]
									a_		:		Tuple				    =				attention_qkv_weight[1 * config.hidden_size : 2 * config.hidden_size, :]
									a_		:		Optional[Any]				    =				attention_qkv_weight[2 * config.hidden_size : 3 * config.hidden_size, :]

									a_		:		Optional[Any]				    =				fix_query_key_value_ordering(a__       ,					a__       ,					3       ,					a__       ,					a__)
									# Split the bias.
									a_		:		Optional[int]				    =				out_val[0 * config.hidden_size : 1 * config.hidden_size]
									a_		:		int				    =				out_val[1 * config.hidden_size : 2 * config.hidden_size]
									a_		:		Tuple				    =				out_val[2 * config.hidden_size : 3 * config.hidden_size]

									# Store.
									a_		:		List[str]				    =				q
									a_		:		Any				    =				q_bias
									a_		:		int				    =				k
									a_		:		Optional[int]				    =				k_bias
									a_		:		Optional[Any]				    =				v
									a_		:		Dict				    =				v_bias

									# Clear the stored tensor.
									a_		:		Union[str, Any]				    =				None

						# Copy weights and biases as is.
						elif weight_or_bias in ["weight", "bias"]:
									a_		:		Any				    =				megatron_to_transformers[op_name]
									a_		:		List[str]				    =				val

    # The final layernorm.
			a_		:		int				    =				transformer["""final_layernorm.weight"""]
			a_		:		List[Any]				    =				transformer["""final_layernorm.bias"""]

			# The pooler.
			a_		:		List[str]				    =				lm["""pooler"""]

			# Store the matrix and the bias.
			a_		:		List[str]				    =				pooler["""dense.weight"""]
			a_		:		List[Any]				    =				pooler["""dense.bias"""]

			# The LM head from Megatron (for RACE).
			a_		:		Union[str, Any]				    =				model["""lm_head"""]

			# The transform matrix.
			a_		:		int				    =				lm_head["""dense.weight"""]
			a_		:		Union[str, Any]				    =				lm_head["""dense.bias"""]

			# The transform LN.
			a_		:		Any				    =				lm_head["""layernorm.weight"""]
			a_		:		List[str]				    =				lm_head["""layernorm.bias"""]

			# For the decoder, we replicate the weights.
			a_		:		Optional[Any]				    =				word_embeddings
			a_		:		int				    =				lm_head["""bias"""]

			# The classifier from Megatron (for MLNI).
			a_		:		str				    =				model["""binary_head"""]

			# Store the classifier.
			a_		:		int				    =				binary_head["""weight"""]
			a_		:		str				    =				binary_head["""bias"""]

			# It should be done!
			return output_state_dict


def 			_UpperCAmelCase							(		):

			'''simple docstring'''



			a_		:		List[Any]				    =				argparse.ArgumentParser()
			parser.add_argument("""--print-checkpoint-structure"""       ,					action="""store_true""")
			parser.add_argument("""path_to_checkpoint"""       ,					type=a__       ,					help="""Path to the ZIP file containing the checkpoint""")
			parser.add_argument(
			    """--config_file"""       ,					default=""""""       ,					type=a__       ,					help="""An optional config json file describing the pre-trained model."""       ,					)
			a_		:		Union[str, Any]				    =				parser.parse_args()

			# Extract the basename.
			a_		:		Any				    =				os.path.dirname(args.path_to_checkpoint)

			# Load the model.
			# the .zip is very optional, let's keep it for backward compatibility
			print(f'''Extracting PyTorch state dictionary from "{args.path_to_checkpoint}"''')
			if args.path_to_checkpoint.endswith(""".zip"""):
						with zipfile.ZipFile(args.path_to_checkpoint       ,					"""r""") as checkpoint:
									with checkpoint.open("""release/mp_rank_00/model_optim_rng.pt""") as pytorch_dict:
												a_		:		int				    =				torch.load(a__       ,					map_location="""cpu""")
			else:
						a_		:		int				    =				torch.load(args.path_to_checkpoint       ,					map_location="""cpu""")

			if args.config_file == "":
						# Default config of megatron-bert 345m
						a_		:		str				    =				MegatronBertConfig()

						# different megatron-bert-*-345m models have different vocab sizes, so override the default
						# config (which is for megatron-bert-cased-345m) with the actual vocab dimension
						a_		:		int				    =				input_state_dict["""model"""]["""lm_head"""]["""bias"""].numel()
			else:
						a_		:		Optional[int]				    =				MegatronBertConfig.from_json_file(args.config_file)

			# Convert.
			print("""Converting""")
			a_		:		int				    =				convert_megatron_checkpoint(a__       ,					a__       ,					a__)

			# Print the structure of converted state dict.
			if args.print_checkpoint_structure:
						recursive_print(a__       ,					a__)

			# Store the config to file.
			print("""Saving config""")
			config.save_pretrained(a__)

			# Store the state_dict to file.
			a_		:		Dict				    =				os.path.join(a__       ,					"""pytorch_model.bin""")
			print(f'''Saving checkpoint to "{output_checkpoint_file}"''')
			torch.save(a__       ,					a__)


####################################################################################################

if __name__ == "__main__":
	main()

####################################################################################################


