



"""simple docstring"""

# 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_gpt2/convert_megatron_gpt2_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 AutoTokenizer, GPTaConfig




def __lowercase			(							_a		,			_a		,			_a=0    ):
    # Format the message.
    if name is None:
        snake_case_ :     str				=	None
    else:
        snake_case_ :     Optional[int]				=	'''.''' * max(0		,			spaces - 2    ) + '''# {:''' + str(50 - spaces    ) + '''s}'''
        snake_case_ :     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 __lowercase			(							_a		,			_a		,			_a		,			_a		,			_a    ):
    # Permutes layout of param tensor to [num_splits * num_heads * hidden_size, :]
    # for compatibility with later versions of NVIDIA Megatron-LM.
    # The inverse operation is performed inside Megatron-LM to read checkpoints:
    # https://github.com/NVIDIA/Megatron-LM/blob/v2.4/megatron/checkpointing.py#L209
    # If param is the weight tensor of the self-attention block, the returned tensor
    # will have to be transposed one more time to be read by HuggingFace GPT2.
    snake_case_ :     Tuple				=	param.size()
    if checkpoint_version == 1.0:
        # version 1.0 stores [num_heads * hidden_size * num_splits, :]
        snake_case_ :     str				=	(num_heads, hidden_size, num_splits) + input_shape[1:]
        snake_case_ :     Any				=	param.view(*_a    )
        snake_case_ :     Tuple				=	param.transpose(0		,			2    )
        snake_case_ :     int				=	param.transpose(1		,			2    ).contiguous()
    elif checkpoint_version >= 2.0:
        # other versions store [num_heads * num_splits * hidden_size, :]
        snake_case_ :     Tuple				=	(num_heads, num_splits, hidden_size) + input_shape[1:]
        snake_case_ :     Any				=	param.view(*_a    )
        snake_case_ :     Optional[Any]				=	param.transpose(0		,			1    ).contiguous()
    snake_case_ :     Union[str, Any]				=	param.view(*_a    )
    return param




def __lowercase			(							_a		,			_a		,			_a    ):
    # The converted output model.
    snake_case_ :     List[str]				=	{}

    # old versions did not store training args
    snake_case_ :     List[str]				=	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))

        snake_case_ :     int				=	ds_args.padded_vocab_size
        snake_case_ :     List[Any]				=	ds_args.max_position_embeddings
        snake_case_ :     List[Any]				=	ds_args.hidden_size
        snake_case_ :     str				=	ds_args.num_layers
        snake_case_ :     int				=	ds_args.num_attention_heads
        snake_case_ :     int				=	ds_args.ffn_hidden_size
        # pprint(config)

    # The number of heads.
    snake_case_ :     Any				=	config.n_head
    # The hidden_size per head.
    snake_case_ :     Any				=	config.n_embd // config.n_head
    # Megatron-LM checkpoint version
    if "checkpoint_version" in input_state_dict.keys():
        snake_case_ :     Any				=	input_state_dict['''checkpoint_version''']
    else:
        snake_case_ :     str				=	0.0

    # The model.
    snake_case_ :     int				=	input_state_dict['''model''']
    # The language model.
    snake_case_ :     Union[str, Any]				=	model['''language_model''']
    # The embeddings.
    snake_case_ :     str				=	lm['''embedding''']

    # The word embeddings.
    snake_case_ :     Dict				=	embeddings['''word_embeddings''']['''weight''']
    # Truncate the embedding table to vocab_size rows.
    snake_case_ :     List[str]				=	word_embeddings[: config.vocab_size, :]
    snake_case_ :     int				=	word_embeddings

    # The position embeddings.
    snake_case_ :     List[str]				=	embeddings['''position_embeddings''']['''weight''']
    # Read the causal mask dimension (seqlen). [max_sequence_length, hidden_size]
    snake_case_ :     str				=	pos_embeddings.size(0    )
    if n_positions != config.n_positions:
        raise ValueError(
            f"pos_embeddings.max_sequence_length={n_positions} and config.n_positions={config.n_positions} don't match"    )
    # Store the position embeddings.
    snake_case_ :     Dict				=	pos_embeddings

    # The transformer.
    snake_case_ :     Union[str, Any]				=	lm['''transformer'''] if '''transformer''' in lm.keys() else lm['''encoder''']

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

    # The simple map of names for "automated" rules.
    snake_case_ :     Dict				=	{
        '''attention.dense''': '''.attn.c_proj.''',
        '''self_attention.dense''': '''.attn.c_proj.''',
        '''mlp.dense_h_to_4h''': '''.mlp.c_fc.''',
        '''mlp.dense_4h_to_h''': '''.mlp.c_proj.''',
    }

    # Extract the layers.
    for key, val in transformer.items():
        # Match the name.
        snake_case_ :     Union[str, Any]				=	layer_re.match(_a    )

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

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

        # The name of the layer.
        snake_case_ :     Any				=	f"transformer.h.{layer_idx}"

        # For layernorm(s), simply store the layer norm.
        if op_name.endswith('''layernorm'''    ):
            snake_case_ :     str				=	'''ln_1''' if op_name.startswith('''input'''    ) else '''ln_2'''
            snake_case_ :     Optional[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":
            # Insert a tensor of 1x1xDxD bias.
            snake_case_ :     Any				=	torch.tril(torch.ones((n_positions, n_positions)		,			dtype=torch.floataa    )    ).view(
                1		,			1		,			_a		,			_a    )
            snake_case_ :     List[Any]				=	causal_mask

            # Insert a "dummy" tensor for masked_bias.
            snake_case_ :     Optional[Any]				=	torch.tensor(-1E4		,			dtype=torch.floataa    )
            snake_case_ :     List[str]				=	masked_bias

            snake_case_ :     Any				=	fix_query_key_value_ordering(_a		,			_a		,			3		,			_a		,			_a    )
            # Megatron stores (3*D) x D but transformers-GPT2 expects D x 3*D.
            snake_case_ :     List[Any]				=	out_val.transpose(0		,			1    ).contiguous()
            # Store.
            snake_case_ :     Tuple				=	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":
            snake_case_ :     Tuple				=	fix_query_key_value_ordering(_a		,			_a		,			3		,			_a		,			_a    )
            # Store. No change of shape.
            snake_case_ :     List[str]				=	out_val

        # Transpose the weights.
        elif weight_or_bias == "weight":
            snake_case_ :     Union[str, Any]				=	megatron_to_transformers[op_name]
            snake_case_ :     List[Any]				=	val.transpose(0		,			1    )

        # Copy the bias.
        elif weight_or_bias == "bias":
            snake_case_ :     Optional[int]				=	megatron_to_transformers[op_name]
            snake_case_ :     Union[str, Any]				=	val

    # DEBUG.
    assert config.n_layer == layer_idx + 1

    # The final layernorm.
    snake_case_ :     Optional[int]				=	transformer['''final_layernorm.weight''']
    snake_case_ :     int				=	transformer['''final_layernorm.bias''']

    # For LM head, transformers' wants the matrix to weight embeddings.
    snake_case_ :     str				=	word_embeddings

    # It should be done!
    return output_state_dict




def __lowercase			(							):
    # Create the argument parser.
    snake_case_ :     Dict				=	argparse.ArgumentParser()
    parser.add_argument('''--print-checkpoint-structure'''		,			action='''store_true'''    )
    parser.add_argument(
        '''path_to_checkpoint'''		,			type=_a		,			help='''Path to the checkpoint file (.zip archive or direct .pt file)'''		,			)
    parser.add_argument(
        '''--config_file'''		,			default=''''''		,			type=_a		,			help='''An optional config json file describing the pre-trained model.'''		,			)
    snake_case_ :     List[str]				=	parser.parse_args()

    # Extract the basename.
    snake_case_ :     int				=	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:
                snake_case_ :     str				=	torch.load(_a		,			map_location='''cpu'''    )
    else:
        snake_case_ :     Any				=	torch.load(args.path_to_checkpoint		,			map_location='''cpu'''    )

    snake_case_ :     Tuple				=	input_state_dict.get('''args'''		,			_a    )

    # Read the config, or default to the model released by NVIDIA.
    if args.config_file == "":
        if ds_args is not None:
            if ds_args.bias_gelu_fusion:
                snake_case_ :     List[Any]				=	'''gelu_fast'''
            elif ds_args.openai_gelu:
                snake_case_ :     int				=	'''gelu_new'''
            else:
                snake_case_ :     List[str]				=	'''gelu'''
        else:
            # in the very early days this used to be "gelu_new"
            snake_case_ :     List[str]				=	'''gelu_new'''

        # Spell out all parameters in case the defaults change.
        snake_case_ :     str				=	GPTaConfig(
            vocab_size=50_257		,			n_positions=1_024		,			n_embd=1_024		,			n_layer=24		,			n_head=16		,			n_inner=4_096		,			activation_function=_a		,			resid_pdrop=0.1		,			embd_pdrop=0.1		,			attn_pdrop=0.1		,			layer_norm_epsilon=1E-5		,			initializer_range=0.02		,			summary_type='''cls_index'''		,			summary_use_proj=_a		,			summary_activation=_a		,			summary_proj_to_labels=_a		,			summary_first_dropout=0.1		,			scale_attn_weights=_a		,			use_cache=_a		,			bos_token_id=50_256		,			eos_token_id=50_256		,			)
    else:
        snake_case_ :     List[Any]				=	GPTaConfig.from_json_file(args.config_file    )

    snake_case_ :     str				=	['''GPT2LMHeadModel''']

    # Convert.
    print('''Converting'''    )
    snake_case_ :     Union[str, Any]				=	convert_megatron_checkpoint(_a		,			_a		,			_a    )

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

    # Add tokenizer class info to config
    # see https://github.com/huggingface/transformers/issues/13906)
    if ds_args is not None:
        snake_case_ :     Dict				=	ds_args.tokenizer_type
        if tokenizer_type == "GPT2BPETokenizer":
            snake_case_ :     List[Any]				=	'''gpt2'''
        elif tokenizer_type == "PretrainedFromHF":
            snake_case_ :     int				=	ds_args.tokenizer_name_or_path
        else:
            raise ValueError(f"Unrecognized tokenizer_type {tokenizer_type}"    )
    else:
        snake_case_ :     str				=	'''gpt2'''

    snake_case_ :     Optional[int]				=	AutoTokenizer.from_pretrained(_a    )
    snake_case_ :     Optional[int]				=	type(_a    ).__name__
    snake_case_ :     List[str]				=	tokenizer_class

    # Store the config to file.
    print('''Saving config'''    )
    config.save_pretrained(_a    )

    # Save tokenizer based on args
    print(f"Adding {tokenizer_class} tokenizer files"    )
    tokenizer.save_pretrained(_a    )

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


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

if __name__ == "__main__":
     main()

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










"""simple docstring"""

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()

lowercase__    :  Dict							    =	[
    '''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''',
]

lowercase__    :  str							    =	[
    '''mlp.dense_4h_to_h.weight''',
    '''self_attention.dense.weight''',
]




def __lowercase			(							_a		,			_a    ):
    snake_case_ :     Optional[int]				=	{
        '''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
    snake_case_ :     List[Any]				=	int(re.match(r'''.*layer_(\d*).*'''		,			_a    )[1]    )
    layer_number -= 3
    return f"h.{layer_number}." + key




def __lowercase			(							_a    ):
    if dtype == torch.bool:
        return 1 / 8
    snake_case_ :     Dict				=	re.search(r'''[^\d](\d+)$'''		,			str(_a    )    )
    if bit_search is None:
        raise ValueError(f"`dtype` is not a valid dtype: {dtype}."    )
    snake_case_ :     Optional[int]				=	int(bit_search.groups()[0]    )
    return bit_size // 8




def __lowercase			(							_a		,			_a		,			_a		,			_a		,			_a    ):
    # Construct model
    if bloom_config_file == "":
        snake_case_ :     int				=	BloomConfig()
    else:
        snake_case_ :     List[str]				=	BloomConfig.from_json_file(_a    )

    if shard_model:
        snake_case_ :     List[str]				=	os.listdir(_a    )
        snake_case_ :     int				=	sorted(filter(lambda _a    : s.startswith('''layer'''    ) and "model_00" in s		,			_a    )    )

        snake_case_ :     List[str]				=	{'''weight_map''': {}, '''metadata''': {}}
        snake_case_ :     Any				=	0

        snake_case_ :     Union[str, Any]				=	None

        snake_case_ :     List[str]				=	BloomConfig()

        for j, file in enumerate(_a    ):
            print('''Processing file: {}'''.format(_a    )    )
            snake_case_ :     Dict				=	None

            for i in range(_a    ):
                # load all TP files
                snake_case_ :     Union[str, Any]				=	file.replace('''model_00'''		,			f"model_0{i}"    )
                snake_case_ :     List[str]				=	torch.load(os.path.join(_a		,			_a    )		,			map_location='''cpu'''    )

                # Rename keys in the transformers names
                snake_case_ :     str				=	list(temp.keys()    )
                for key in keys:
                    snake_case_ :     Any				=	temp.pop(_a    )

                if tensors is None:
                    snake_case_ :     Any				=	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
                            snake_case_ :     Tuple				=	1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN    ) else 0
                            # We concatenate these weights accross TP ranks
                            snake_case_ :     List[str]				=	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    ):
                    snake_case_ :     Any				=	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():
                snake_case_ :     List[str]				=	tensors[key]
                total_size += value.numel() * get_dtype_size(value.dtype    )
                if key not in index_dict["weight_map"]:
                    snake_case_ :     List[str]				=	'''pytorch_model_{}-of-{}.bin'''.format(
                        str(j + 1    ).zfill(5    )		,			str(len(_a    )    ).zfill(5    )    )

        snake_case_ :     int				=	BloomConfig()
        snake_case_ :     Any				=	pytorch_dump_folder_path + '''/''' + CONFIG_NAME
        snake_case_ :     Dict				=	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:
            snake_case_ :     Tuple				=	json.dumps(_a		,			indent=2		,			sort_keys=_a    ) + '''\n'''
            f.write(_a    )
    else:
        snake_case_ :     Union[str, Any]				=	BloomModel(_a    )

        snake_case_ :     List[str]				=	os.listdir(_a    )
        snake_case_ :     Dict				=	sorted(filter(lambda _a    : s.startswith('''layer'''    ) and "model_00" in s		,			_a    )    )

        snake_case_ :     List[Any]				=	None
        for i, file in enumerate(_a    ):
            snake_case_ :     Optional[Any]				=	None
            for i in range(_a    ):
                # load all TP files
                snake_case_ :     List[str]				=	file.replace('''model_00'''		,			f"model_0{i}"    )
                snake_case_ :     Optional[Any]				=	torch.load(os.path.join(_a		,			_a    )		,			map_location='''cpu'''    )

                # Rename keys in the transformers names
                snake_case_ :     str				=	list(temp.keys()    )
                for key in keys:
                    snake_case_ :     str				=	temp.pop(_a    )

                if tensors is None:
                    snake_case_ :     int				=	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
                            snake_case_ :     Tuple				=	1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN    ) else 0
                            # We concatenate these weights accross TP ranks
                            snake_case_ :     Optional[Any]				=	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    ):
                    snake_case_ :     Union[str, Any]				=	tensors[key] / pretraining_tp

            snake_case_ :     Any				=	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:
                snake_case_ :     Optional[int]				=	set(other_keys.missing_keys    )
            else:
                snake_case_ :     Tuple				=	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    )
        snake_case_ :     List[str]				=	pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
        snake_case_ :     Optional[Any]				=	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:
            snake_case_ :     Optional[Any]				=	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__":
     lowercase__    :  str							    =	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''',
     )
     lowercase__    :  List[Any]							    =	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,
     )





