


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

__snake_case				=				[
    "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",
]

__snake_case				=				[
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]
def 					A_      (      _lowerCAmelCase   :    Any,   _lowerCAmelCase   :    List[Any]				):






     """simple docstring"""

     _a					= {
         '''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
     _a					= int(re.match(R'''.*layer_(\d*).*''',   __lowerCAmelCase				)[1]				)
     layer_number -= 3
     return f'h.{layer_number}.' + key
def 					A_      (      _lowerCAmelCase   :    Dict				):






     """simple docstring"""

     if dtype == torch.bool:
          return 1 / 8
     _a					= re.search(R'''[^\d](\d+)$''',   str(__lowerCAmelCase				)				)
     if bit_search is None:
          raise ValueError(f'`dtype` is not a valid dtype: {dtype}.'				)
     _a					= int(bit_search.groups()[0]				)
     return bit_size // 8
def 					A_      (      _lowerCAmelCase   :    Optional[Any],   _lowerCAmelCase   :    Union[str, Any],   _lowerCAmelCase   :    List[Any],   _lowerCAmelCase   :    List[Any],   _lowerCAmelCase   :    Optional[Any]				):






     """simple docstring"""

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

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

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

          _a					= None

          _a					= BloomConfig()

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

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

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

                    if tensors is None:
                         _a					= temp
                    else:
                         for key in tensors.keys():
                              if any(key.endswith(__lowerCAmelCase				) 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
                                   _a					= 1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN				) else 0
                                   # We concatenate these weights accross TP ranks
                                   _a					= torch.cat([tensors[key], temp[key]],   dim=__lowerCAmelCase				)

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

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

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

          _a					= os.listdir(__lowerCAmelCase				)
          _a					= sorted(filter(lambda _lowerCAmelCase				: s.startswith('''layer'''				) and "model_00" in s,   __lowerCAmelCase				)				)

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

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

                    if tensors is None:
                         _a					= 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(__lowerCAmelCase				) for end in WEIGHTS_TO_AVERAGE_ENDSWITH				):
                                   tensors[key] += temp[key]
                              else:
                                   # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
                                   _a					= 1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN				) else 0
                                   # We concatenate these weights accross TP ranks
                                   _a					= torch.cat([tensors[key], temp[key]],   dim=__lowerCAmelCase				)

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

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

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

          # Save pytorch-model
          os.makedirs(__lowerCAmelCase,   exist_ok=__lowerCAmelCase				)
          _a					= pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
          _a					= 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:
               _a					= model.to(config.torch_dtype				)
          torch.save(model.state_dict(),   __lowerCAmelCase				)
          print(f'Save configuration file to {pytorch_config_dump_path}'				)
          with open(__lowerCAmelCase,   '''w''',   encoding='''utf-8'''				) as f:
               f.write(config.to_json_string()				)


if __name__ == "__main__":
   __snake_case				=				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''',
   )
   __snake_case				=				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,
   )

import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
					from sagemaker import Session, TrainingJobAnalytics
					from sagemaker.huggingface import HuggingFace

@pytest.mark.skipif(
    literal_eval(os.getenv("""TEST_SAGEMAKER"""  ,	"""False"""			)			) is not True  ,	reason="""Skipping test because should only be run when releasing minor transformers version"""  ,	)
@pytest.mark.usefixtures("""sm_env"""			)
@parameterized_class(
    [
        {
            """framework""": """pytorch""",
            """script""": """run_glue_model_parallelism.py""",
            """model_name_or_path""": """roberta-large""",
            """instance_type""": """ml.p3dn.24xlarge""",
            """results""": {"""train_runtime""": 1_600, """eval_accuracy""": 0.3, """eval_loss""": 1.2},
        },
        {
            """framework""": """pytorch""",
            """script""": """run_glue.py""",
            """model_name_or_path""": """roberta-large""",
            """instance_type""": """ml.p3dn.24xlarge""",
            """results""": {"""train_runtime""": 1_600, """eval_accuracy""": 0.3, """eval_loss""": 1.2},
        },
    ]			)
class    lowerCAmelCase__     (	unittest.TestCase			):






		def  __A       (		self		:							Optional[int]  )     ->							List[Any]:
						if self.framework == "pytorch":
										subprocess.run(
										    f'''cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py'''.split()							,			encoding='''utf-8'''							,			check=SCREAMING_SNAKE_CASE__							,			)
						assert hasattr(self							,			'''env'''  )






		def  __A       (		self		:							List[str]							,			SCREAMING_SNAKE_CASE__		:							List[Any]  )     ->							Optional[Any]:
						# configuration for running training on smdistributed Model Parallel
						__lowerCamelCase    					=	{
						    '''enabled''': True,
						    '''processes_per_host''': 8,
						}
						__lowerCamelCase    					=	{
						    '''enabled''': True,
						    '''parameters''': {
						        '''microbatches''': 4,
						        '''placement_strategy''': '''spread''',
						        '''pipeline''': '''interleaved''',
						        '''optimize''': '''speed''',
						        '''partitions''': 4,
						        '''ddp''': True,
						    },
						}

						__lowerCamelCase    					=	{'''smdistributed''': {'''modelparallel''': smp_options}, '''mpi''': mpi_options}

						__lowerCamelCase    					=	'''trainer''' if self.script == '''run_glue.py''' else '''smtrainer'''
						# creates estimator
						return HuggingFace(
						    entry_point=self.script							,			source_dir=self.env.test_path							,			role=self.env.role							,			image_uri=self.env.image_uri							,			base_job_name=f'''{self.env.base_job_name}-{instance_count}-smp-{name_extension}'''							,			instance_count=SCREAMING_SNAKE_CASE__							,			instance_type=self.instance_type							,			debugger_hook_config=SCREAMING_SNAKE_CASE__							,			hyperparameters={
						        **self.env.hyperparameters,
						        '''model_name_or_path''': self.model_name_or_path,
						        '''max_steps''': 5_00,
						    }							,			metric_definitions=self.env.metric_definitions							,			distribution=SCREAMING_SNAKE_CASE__							,			py_version='''py36'''							,			)






		def  __A       (		self		:							List[Any]							,			SCREAMING_SNAKE_CASE__		:							List[Any]  )     ->							List[str]:
						TrainingJobAnalytics(SCREAMING_SNAKE_CASE__  ).export_csv(f'''{self.env.test_path}/{job_name}_metrics.csv'''  )







		@parameterized.expand([(1,)]  )
		def  __A       (		self		:							Optional[Any]							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]  )     ->							Optional[int]:
						# create estimator
						__lowerCamelCase    					=	self.create_estimator(SCREAMING_SNAKE_CASE__  )

						# run training
						estimator.fit()

						# result dataframe
						__lowerCamelCase    					=	TrainingJobAnalytics(estimator.latest_training_job.name  ).dataframe()

						# extract kpis
						__lowerCamelCase    					=	list(result_metrics_df[result_metrics_df.metric_name == '''eval_accuracy''']['''value''']  )
						__lowerCamelCase    					=	list(result_metrics_df[result_metrics_df.metric_name == '''eval_loss''']['''value''']  )
						# get train time from SageMaker job, this includes starting, preprocessing, stopping
						__lowerCamelCase    					=	(
						    Session().describe_training_job(estimator.latest_training_job.name  ).get('''TrainingTimeInSeconds'''							,			99_99_99  )
						)

						# assert kpis
						assert train_runtime <= self.results["train_runtime"]
						assert all(t >= self.results['''eval_accuracy'''] for t in eval_accuracy  )
						assert all(t <= self.results['''eval_loss'''] for t in eval_loss  )

						# dump tests result into json file to share in PR
						with open(f'''{estimator.latest_training_job.name}.json'''							,			'''w'''  ) as outfile:
										json.dump({'''train_time''': train_runtime, '''eval_accuracy''': eval_accuracy, '''eval_loss''': eval_loss}							,			SCREAMING_SNAKE_CASE__  )



