
'''simple docstring'''


import argparse
import gc
import json
import os

import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils.deepspeed import DummyOptim, DummyScheduler


snake_case_ :  Optional[int]     =					16
snake_case_ :  str     =					32



def  __snake_case     (    _UpperCAmelCase					:       List[str]):
   return int(x / 2**20)


class 						lowercase__	:






   '''simple docstring'''







   def __enter__(			self   ):

      '''simple docstring'''







      gc.collect()
      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
      UpperCamelCase       =  torch.cuda.memory_allocated()
      return self







   def __exit__(			self ,   *lowerCamelCase__   ):

      '''simple docstring'''







      gc.collect()
      torch.cuda.empty_cache()
      UpperCamelCase       =  torch.cuda.memory_allocated()
      UpperCamelCase       =  torch.cuda.max_memory_allocated()
      UpperCamelCase       =  bamb(self.end - self.begin   )
      UpperCamelCase       =  bamb(self.peak - self.begin   )
      # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")




def  __snake_case     (    _UpperCAmelCase					:       Accelerator,      _UpperCAmelCase					:       int = 16,      _UpperCAmelCase					:       str = "bert-base-cased",      _UpperCAmelCase					:       int = 320,      _UpperCAmelCase					:       int = 160,      ):
   UpperCamelCase       =  AutoTokenizer.from_pretrained(_UpperCAmelCase)
   UpperCamelCase       =  load_dataset(
       '''glue''',      '''mrpc''',      split={'''train''': f'train[:{n_train}]', '''validation''': f'validation[:{n_val}]'})

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

   # Apply the method we just defined to all the examples in all the splits of the dataset
   UpperCamelCase       =  datasets.map(
       _UpperCAmelCase,      batched=_UpperCAmelCase,      remove_columns=['''idx''', '''sentence1''', '''sentence2'''],      load_from_cache_file=_UpperCAmelCase)

   # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
   # transformers library
   UpperCamelCase       =  tokenized_datasets.rename_column('''label''',      '''labels''')

   def collate_fn(_UpperCAmelCase					:       Dict):
      # On TPU it's best to pad everything to the same length or training will be very slow.
      if accelerator.distributed_type == DistributedType.TPU:
         return tokenizer.pad(_UpperCAmelCase,      padding='''max_length''',      max_length=128,      return_tensors='''pt''')
      return tokenizer.pad(_UpperCAmelCase,      padding='''longest''',      return_tensors='''pt''')

   # Instantiate dataloaders.
   UpperCamelCase       =  DataLoader(
       tokenized_datasets['''train'''],      shuffle=_UpperCAmelCase,      collate_fn=_UpperCAmelCase,      batch_size=_UpperCAmelCase)
   UpperCamelCase       =  DataLoader(
       tokenized_datasets['''validation'''],      shuffle=_UpperCAmelCase,      collate_fn=_UpperCAmelCase,      batch_size=_UpperCAmelCase)

   return train_dataloader, eval_dataloader



def  __snake_case     (    _UpperCAmelCase					:       Union[str, Any],      _UpperCAmelCase					:       List[Any]):
   # Initialize accelerator
   UpperCamelCase       =  Accelerator()

   # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
   UpperCamelCase       =  config['''lr''']
   UpperCamelCase       =  int(config['''num_epochs'''])
   UpperCamelCase       =  int(config['''seed'''])
   UpperCamelCase       =  int(config['''batch_size'''])
   UpperCamelCase       =  args.model_name_or_path

   set_seed(_UpperCAmelCase)
   UpperCamelCase     , UpperCamelCase       =  get_dataloaders(_UpperCAmelCase,      _UpperCAmelCase,      _UpperCAmelCase,      args.n_train,      args.n_val)

   # Instantiate the model (we build the model here so that the seed also control new weights initialization)
   UpperCamelCase       =  AutoModelForSequenceClassification.from_pretrained(_UpperCAmelCase,      return_dict=_UpperCAmelCase)

   # Instantiate optimizer
   UpperCamelCase       =  (
       AdamW
       if accelerator.state.deepspeed_plugin is None
       or '''optimizer''' not in accelerator.state.deepspeed_plugin.deepspeed_config
       else DummyOptim
   )
   UpperCamelCase       =  optimizer_cls(params=model.parameters(),      lr=_UpperCAmelCase)

   if accelerator.state.deepspeed_plugin is not None:
      UpperCamelCase       =  accelerator.state.deepspeed_plugin.deepspeed_config[
          '''gradient_accumulation_steps'''
      ]
   else:
      UpperCamelCase       =  1
   UpperCamelCase       =  (len(_UpperCAmelCase) * num_epochs) // gradient_accumulation_steps

   # Instantiate scheduler
   if (
       accelerator.state.deepspeed_plugin is None
       or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
   ):
      UpperCamelCase       =  get_linear_schedule_with_warmup(
          optimizer=_UpperCAmelCase,      num_warmup_steps=0,      num_training_steps=_UpperCAmelCase,      )
   else:
      UpperCamelCase       =  DummyScheduler(_UpperCAmelCase,      total_num_steps=_UpperCAmelCase,      warmup_num_steps=0)

   # Prepare everything
   # 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.
   UpperCamelCase     , UpperCamelCase     , UpperCamelCase     , UpperCamelCase     , UpperCamelCase       =  accelerator.prepare(
       _UpperCAmelCase,      _UpperCAmelCase,      _UpperCAmelCase,      _UpperCAmelCase,      _UpperCAmelCase)

   # We need to keep track of how many total steps we have iterated over
   UpperCamelCase       =  0
   # We also need to keep track of the stating epoch so files are named properly
   UpperCamelCase       =  0

   # Now we train the model
   UpperCamelCase       =  {}
   for epoch in range(_UpperCAmelCase,      _UpperCAmelCase):
      with TorchTracemalloc() as tracemalloc:
         model.train()
         for step, batch in enumerate(_UpperCAmelCase):
            UpperCamelCase       =  model(**_UpperCAmelCase)
            UpperCamelCase       =  outputs.loss
            UpperCamelCase       =  loss / gradient_accumulation_steps
            accelerator.backward(_UpperCAmelCase)
            if step % gradient_accumulation_steps == 0:
               optimizer.step()
               lr_scheduler.step()
               optimizer.zero_grad()

            overall_step += 1

        # 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)))
      UpperCamelCase       =  tracemalloc.peaked + bamb(tracemalloc.begin)
      if args.peak_memory_upper_bound is not None:
         assert (
             train_total_peak_memory[f'epoch-{epoch}'] <= args.peak_memory_upper_bound
         ), "Peak memory usage exceeded the upper bound"

   accelerator.wait_for_everyone()
   if accelerator.is_main_process:
      with open(os.path.join(args.output_dir,      '''peak_memory_utilization.json'''),      '''w''') as f:
         json.dump(_UpperCAmelCase,      _UpperCAmelCase)



def  __snake_case     (    ):
   UpperCamelCase       =  argparse.ArgumentParser(description='''Simple example of training script tracking peak GPU memory usage.''')
   parser.add_argument(
       '''--model_name_or_path''',      type=_UpperCAmelCase,      default='''bert-base-cased''',      help='''Path to pretrained model or model identifier from huggingface.co/models.''',      required=_UpperCAmelCase,      )
   parser.add_argument(
       '''--output_dir''',      type=_UpperCAmelCase,      default='''.''',      help='''Optional save directory where all checkpoint folders will be stored. Default is the current working directory.''',      )
   parser.add_argument(
       '''--peak_memory_upper_bound''',      type=_UpperCAmelCase,      default=_UpperCAmelCase,      help='''The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value.''',      )
   parser.add_argument(
       '''--n_train''',      type=_UpperCAmelCase,      default=320,      help='''Number of training examples to use.''',      )
   parser.add_argument(
       '''--n_val''',      type=_UpperCAmelCase,      default=160,      help='''Number of validation examples to use.''',      )
   parser.add_argument(
       '''--num_epochs''',      type=_UpperCAmelCase,      default=1,      help='''Number of train epochs.''',      )
   UpperCamelCase       =  parser.parse_args()
   UpperCamelCase       =  {'''lr''': 2E-5, '''num_epochs''': args.num_epochs, '''seed''': 42, '''batch_size''': 16}
   training_function(_UpperCAmelCase,      _UpperCAmelCase)


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

snake_case_ :  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',
]

snake_case_ :  Union[str, Any]     =					[
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]



def  __snake_case     (    _UpperCAmelCase					:       List[str],      _UpperCAmelCase					:       List[str]):
   UpperCamelCase       =  {
       '''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
   UpperCamelCase       =  int(re.match(R'''.*layer_(\d*).*''',      _UpperCAmelCase)[1])
   layer_number -= 3
   return f'h.{layer_number}.' + key



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



def  __snake_case     (    _UpperCAmelCase					:       int,      _UpperCAmelCase					:       Optional[Any],      _UpperCAmelCase					:       str,      _UpperCAmelCase					:       List[Any],      _UpperCAmelCase					:       Optional[int]):
   # Construct model
   if bloom_config_file == "":
      UpperCamelCase       =  BloomConfig()
   else:
      UpperCamelCase       =  BloomConfig.from_json_file(_UpperCAmelCase)

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

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

      UpperCamelCase       =  None

      UpperCamelCase       =  BloomConfig()

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

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

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

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

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

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

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

      UpperCamelCase       =  os.listdir(_UpperCAmelCase)
      UpperCamelCase       =  sorted(filter(lambda _UpperCAmelCase: s.startswith('''layer''') and "model_00" in s,      _UpperCAmelCase))

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

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

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

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

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

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

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


if __name__ == "__main__":
   snake_case_ :  Tuple     =					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_ :  List[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,
   )





