




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
#
########################################################################


__magic_name__     :     Any  							=					16
__magic_name__     :     List[str]  							=					32
def 		a_       (   __lowerCAmelCase							):
  return int(x / 2**20							)





class 						SCREAMING_SNAKE_CASE__		:
     def __enter__(	self  :		str	):



       """simple docstring"""
       gc.collect()
       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
       lowerCAmelCase__								=			torch.cuda.memory_allocated()
       lowerCAmelCase__								=			psutil.Process()

       lowerCAmelCase__								=			self.cpu_mem_used()
       lowerCAmelCase__								=			True
       lowerCAmelCase__								=			threading.Thread(target=self.peak_monitor_func	)
       lowerCAmelCase__								=			True
       peak_monitor_thread.start()
       return self
     def 					A__			(	self  :		Any	):



       """simple docstring"""
       return self.process.memory_info().rss
     def 					A__			(	self  :		List[str]	):



       """simple docstring"""
       lowerCAmelCase__								=			-1

       while True:
         lowerCAmelCase__								=			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  :		int					,	*__lowerCamelCase  :		Optional[Any]	):



       """simple docstring"""
       lowerCAmelCase__								=			False

       gc.collect()
       torch.cuda.empty_cache()
       lowerCAmelCase__								=			torch.cuda.memory_allocated()
       lowerCAmelCase__								=			torch.cuda.max_memory_allocated()
       lowerCAmelCase__								=			bamb(self.end - self.begin	)
       lowerCAmelCase__								=			bamb(self.peak - self.begin	)

       lowerCAmelCase__								=			self.cpu_mem_used()
       lowerCAmelCase__								=			bamb(self.cpu_end - self.cpu_begin	)
       lowerCAmelCase__								=			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

      __magic_name__     :     str  							=					mocked_dataloaders  # noqa: F811
def 		a_       (   __lowerCAmelCase      ,   __lowerCAmelCase							):
  # For testing only
  if os.environ.get('''TESTING_MOCKED_DATALOADERS'''      ,   __lowerCAmelCase							) == "1":
    lowerCAmelCase__								=			2

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

  # Initialize accelerator
  if args.with_tracking:
    lowerCAmelCase__								=			Accelerator(
        cpu=args.cpu      ,   mixed_precision=args.mixed_precision      ,   log_with='''wandb'''      ,   project_dir=args.logging_dir      ,   fsdp_plugin=__lowerCAmelCase      ,   )
  else:
    lowerCAmelCase__								=			Accelerator(fsdp_plugin=__lowerCAmelCase							)
  accelerator.print(accelerator.distributed_type							)

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

  # We need to initialize the trackers we use, and also store our configuration
  if args.with_tracking:
    lowerCAmelCase__								=			vars(__lowerCAmelCase							)
    accelerator.init_trackers('''fsdp_glue_no_trainer'''      ,   __lowerCAmelCase							)

  lowerCAmelCase__								=			AutoTokenizer.from_pretrained(args.model_name_or_path							)
  lowerCAmelCase__								=			load_dataset('''glue'''      ,   '''mrpc'''							)
  lowerCAmelCase__								=			evaluate.load('''glue'''      ,   '''mrpc'''							)

  def tokenize_function(__lowerCAmelCase							):
    # max_length=None => use the model max length (it's actually the default)
    lowerCAmelCase__								=			tokenizer(examples['''sentence1''']      ,   examples['''sentence2''']      ,   truncation=__lowerCAmelCase      ,   max_length=__lowerCAmelCase							)
    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():
    lowerCAmelCase__								=			datasets.map(
        __lowerCAmelCase      ,   batched=__lowerCAmelCase      ,   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
  lowerCAmelCase__								=			tokenized_datasets.rename_column('''label'''      ,   '''labels'''							)

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

  def collate_fn(__lowerCAmelCase							):
    # On TPU it's best to pad everything to the same length or training will be very slow.
    lowerCAmelCase__								=			1_28 if accelerator.distributed_type == DistributedType.TPU else None
    # When using mixed precision we want round multiples of 8/16
    if accelerator.mixed_precision == "fp8":
      lowerCAmelCase__								=			16
    elif accelerator.mixed_precision != "no":
      lowerCAmelCase__								=			8
    else:
      lowerCAmelCase__								=			None

    return tokenizer.pad(
        __lowerCAmelCase      ,   padding='''longest'''      ,   max_length=__lowerCAmelCase      ,   pad_to_multiple_of=__lowerCAmelCase      ,   return_tensors='''pt'''      ,   )

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

  set_seed(__lowerCAmelCase							)

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

  # New Code #
  # For FSDP feature, it is highly recommended and efficient to prepare the model before creating optimizer
  lowerCAmelCase__								=			accelerator.prepare(__lowerCAmelCase							)
  accelerator.print(__lowerCAmelCase							)

  # 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
  lowerCAmelCase__								=			torch.optim.AdamW(params=model.parameters()      ,   lr=__lowerCAmelCase      ,   weight_decay=2E-4							)

  # Instantiate scheduler
  lowerCAmelCase__								=			get_linear_schedule_with_warmup(
      optimizer=__lowerCAmelCase      ,   num_warmup_steps=10      ,   num_training_steps=(len(__lowerCAmelCase							) * 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.
  lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__								=			accelerator.prepare(
      __lowerCAmelCase      ,   __lowerCAmelCase      ,   __lowerCAmelCase      ,   __lowerCAmelCase							)

  lowerCAmelCase__								=			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							)
      lowerCAmelCase__								=			os.path.basename(args.resume_from_checkpoint							)
    else:
      # Get the most recent checkpoint
      lowerCAmelCase__								=			[f.name for f in os.scandir(os.getcwd()							) if f.is_dir()]
      dirs.sort(key=os.path.getctime							)
      lowerCAmelCase__								=			dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
    # Extract `epoch_{i}` or `step_{i}`
    lowerCAmelCase__								=			os.path.splitext(__lowerCAmelCase							)[0]

    if "epoch" in training_difference:
      num_epochs -= int(training_difference.replace('''epoch_'''      ,   ''''''							)							)
      lowerCAmelCase__								=			None
    else:
      lowerCAmelCase__								=			int(training_difference.replace('''step_'''      ,   ''''''							)							)
      num_epochs -= resume_step // len(__lowerCAmelCase							)
      # If resuming by step, we also need to know exactly how far into the DataLoader we went
      lowerCAmelCase__								=			(num_epochs * len(__lowerCAmelCase							)) - resume_step

    # Now we train the model
  for epoch in range(__lowerCAmelCase							):
    # New Code #
    # context manager to track the peak memory usage during the training epoch
    with TorchTracemalloc() as tracemalloc:
      model.train()
      if args.with_tracking:
        lowerCAmelCase__								=			0
      for step, batch in enumerate(__lowerCAmelCase							):
        # 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							)
        lowerCAmelCase__								=			model(**__lowerCAmelCase							)
        lowerCAmelCase__								=			outputs.loss
        lowerCAmelCase__								=			loss / gradient_accumulation_steps
        # We keep track of the loss at each epoch
        if args.with_tracking:
          total_loss += loss.detach().float()
        accelerator.backward(__lowerCAmelCase							)
        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(__lowerCAmelCase      ,   __lowerCAmelCase							):
          lowerCAmelCase__								=			F"""step_{overall_step}"""
          if overall_step % checkpointing_steps == 0:
            if args.output_dir is not None:
              lowerCAmelCase__								=			os.path.join(args.output_dir      ,   __lowerCAmelCase							)
            accelerator.save_state(__lowerCAmelCase							)
        # 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=__lowerCAmelCase      ,   )

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

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

      if checkpointing_steps == "epoch":
        lowerCAmelCase__								=			F"""epoch_{epoch}"""
        if args.output_dir is not None:
          lowerCAmelCase__								=			os.path.join(args.output_dir      ,   __lowerCAmelCase							)
        accelerator.save_state(__lowerCAmelCase							)
        # 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=__lowerCAmelCase      ,   )

  if args.with_tracking:
    accelerator.end_training()
def 		a_       (   ):
  lowerCAmelCase__								=			argparse.ArgumentParser(description='''Simple example of training script.'''							)
  parser.add_argument(
      '''--mixed_precision'''      ,   type=__lowerCAmelCase      ,   default=__lowerCAmelCase      ,   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=__lowerCAmelCase      ,   default=__lowerCAmelCase      ,   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=__lowerCAmelCase      ,   default=__lowerCAmelCase      ,   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=__lowerCAmelCase      ,   default='''.'''      ,   help='''Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'''      ,   )
  parser.add_argument(
      '''--logging_dir'''      ,   type=__lowerCAmelCase      ,   default='''logs'''      ,   help='''Location on where to store experiment tracking logs`'''      ,   )
  parser.add_argument(
      '''--model_name_or_path'''      ,   type=__lowerCAmelCase      ,   help='''Path to pretrained model or model identifier from huggingface.co/models.'''      ,   required=__lowerCAmelCase      ,   )
  lowerCAmelCase__								=			parser.parse_args()
  lowerCAmelCase__								=			{'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
  training_function(__lowerCAmelCase      ,   __lowerCAmelCase							)


if __name__ == "__main__":
      main()







import os
from dataclasses import replace

import jax
import wandb
from bigbird_flax import Args, DataCollator, FlaxBigBirdForNaturalQuestions, Trainer, build_tx, train_step, val_step
from datasets import load_dataset
from flax import jax_utils

from transformers import BigBirdTokenizerFast


if __name__ == "__main__":
      print("""#################### AVAILABLE DEVICES ####################""")
      print(jax.devices())
      print("""###########################################################""")

      # setup for wandb sweep
      __magic_name__     :     Optional[Any]  							=					Args()
      __magic_name__     :     List[Any]  							=					wandb.init(project="""bigbird-natural-questions""", config=args.__dict__)
      __magic_name__     :     Any  							=					dict(logger.config)
      del wandb_args["batch_size"]
      __magic_name__     :     int  							=					replace(args, **wandb_args)
      __magic_name__     :     Tuple  							=					args.base_dir + """-""" + wandb.run.id
      __magic_name__     :     Union[str, Any]  							=					replace(args, base_dir=base_dir)
      print(args)

      __magic_name__     :     Union[str, Any]  							=					load_dataset("""json""", data_files=args.tr_data_path)["""train"""]
      __magic_name__     :     List[Any]  							=					load_dataset("""json""", data_files=args.val_data_path)["""train"""]

      # drop extra batch for now
      __magic_name__     :     List[str]  							=					range(len(tr_dataset) - len(tr_dataset) % args.batch_size)
      __magic_name__     :     Dict  							=					tr_dataset.shuffle().select(indices)
      __magic_name__     :     Any  							=					range(len(val_dataset) - len(val_dataset) % args.batch_size)
      __magic_name__     :     Dict  							=					val_dataset.shuffle().select(indices)

      if os.environ.get("""TRAIN_ON_SMALL""", """false""") == "true":
            __magic_name__     :     Any  							=					tr_dataset.shuffle().select(range(8_00_00))
            __magic_name__     :     Optional[int]  							=					val_dataset.shuffle().select(range(80_00))

      print(tr_dataset)
      print(val_dataset)

      __magic_name__     :     Union[str, Any]  							=					FlaxBigBirdForNaturalQuestions.from_pretrained(
          args.model_id, block_size=args.block_size, num_random_blocks=args.num_random_blocks
      )
      __magic_name__     :     Tuple  							=					BigBirdTokenizerFast.from_pretrained(args.model_id)
      __magic_name__     :     Tuple  							=					DataCollator(pad_id=tokenizer.pad_token_id, max_length=40_96)

      __magic_name__     :     int  							=					{
          """lr""": args.lr,
          """init_lr""": args.init_lr,
          """warmup_steps""": args.warmup_steps,
          """num_train_steps""": args.max_epochs * (len(tr_dataset) // args.batch_size),
          """weight_decay""": args.weight_decay,
      }
      __magic_name__							,   __magic_name__     :     Optional[int]  							=					build_tx(**tx_args)

      __magic_name__     :     List[str]  							=					Trainer(
          args=args,
          data_collator=data_collator,
          model_save_fn=model.save_pretrained,
          train_step_fn=train_step,
          val_step_fn=val_step,
          logger=logger,
          scheduler_fn=lr,
      )

      __magic_name__     :     Union[str, Any]  							=					None
      __magic_name__     :     Dict  							=					trainer.create_state(model, tx, num_train_steps=tx_args["""num_train_steps"""], ckpt_dir=ckpt_dir)
      try:
            trainer.train(state, tr_dataset, val_dataset)
      except KeyboardInterrupt:
            print("""Oooops; TRAINING STOPPED UNFORTUNATELY""")

      print("""SAVING WEIGHTS IN `final-weights`""")
      __magic_name__     :     List[Any]  							=					jax_utils.unreplicate(state.params)
      model.save_pretrained(os.path.join(args.base_dir, """final-weights"""), params=params)

