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


_UpperCAmelCase    			=					16
_UpperCAmelCase    			=					32


def 		_lowerCamelCase       (   _a  ):
  """simple docstring"""




  return int(x / 2**2_0  )

class 		__magic_name__				:


  """simple docstring"""





  def __enter__(					self		):
    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     _UpperCAmelCase  (					self		):
    return self.process.memory_info().rss


  def     _UpperCAmelCase  (					self		):
    _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	,		*a__		):
    _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

  _UpperCAmelCase    			=					mocked_dataloaders  # noqa: F811


def 		_lowerCamelCase       (   _a			,   _a  ):
  """simple docstring"""




  if os.environ.get('''TESTING_MOCKED_DATALOADERS'''			,   _a  ) == "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=_a			,   ranka_only=_a  )			,   optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=_a			,   ranka_only=_a  )			,   )

  # 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=_a			,   )
  else:
    _lowerCamelCase      =						Accelerator(fsdp_plugin=_a  )
  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(_a  )
    accelerator.init_trackers('''fsdp_glue_no_trainer'''			,   _a  )

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

  def tokenize_function(_a  ):
    # max_length=None => use the model max length (it's actually the default)
    _lowerCamelCase      =						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():
    _lowerCamelCase      =						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
  _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(_a  ):
    # On TPU it's best to pad everything to the same length or training will be very slow.
    _lowerCamelCase      =						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":
      _lowerCamelCase      =						1_6
    elif accelerator.mixed_precision != "no":
      _lowerCamelCase      =						8
    else:
      _lowerCamelCase      =						None

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

  # Instantiate dataloaders.
  _lowerCamelCase      =						DataLoader(
      tokenized_datasets['''train''']			,   shuffle=_a			,   collate_fn=_a			,   batch_size=_a  )
  _lowerCamelCase      =						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)
  _lowerCamelCase      =						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
  _lowerCamelCase      =						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
  _lowerCamelCase      =						torch.optim.AdamW(params=model.parameters()			,   lr=_a			,   weight_decay=2E-4  )

  # Instantiate scheduler
  _lowerCamelCase      =						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.
  _lowerCamelCase    ,    _lowerCamelCase    ,    _lowerCamelCase    ,    _lowerCamelCase      =						accelerator.prepare(
      _a			,   _a			,   _a			,   _a  )

  _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(_a  )[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(_a  )
      # If resuming by step, we also need to know exactly how far into the DataLoader we went
      _lowerCamelCase      =						(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:
        _lowerCamelCase      =						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  )
        _lowerCamelCase      =						model(**_a  )
        _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(_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  ):
          _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			,   _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():
          _lowerCamelCase      =						model(**_a  )
        _lowerCamelCase      =						outputs.logits.argmax(dim=-1  )
        _lowerCamelCase    ,    _lowerCamelCase      =						accelerator.gather_for_metrics((predictions, batch['''labels'''])  )
        metric.add_batch(
            predictions=_a			,   references=_a			,   )

      _lowerCamelCase      =						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":
        _lowerCamelCase      =						F'''epoch_{epoch}'''
        if args.output_dir is not None:
          _lowerCamelCase      =						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 		_lowerCamelCase       (   ):
  """simple docstring"""




  _lowerCamelCase      =						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			,   )
  _lowerCamelCase      =						parser.parse_args()
  _lowerCamelCase      =						{'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 4_2, '''batch_size''': 1_6}
  training_function(_a			,   _a  )


if __name__ == "__main__":
  main()






from collections import OrderedDict
from typing import Any, List, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


_UpperCAmelCase    			=					logging.get_logger(__name__)

_UpperCAmelCase    			=					{
    "gpt2": "https://huggingface.co/gpt2/resolve/main/config.json",
    "gpt2-medium": "https://huggingface.co/gpt2-medium/resolve/main/config.json",
    "gpt2-large": "https://huggingface.co/gpt2-large/resolve/main/config.json",
    "gpt2-xl": "https://huggingface.co/gpt2-xl/resolve/main/config.json",
    "distilgpt2": "https://huggingface.co/distilgpt2/resolve/main/config.json",
}

class 		__magic_name__				( lowercase_					):


  """simple docstring"""





  _UpperCamelCase						 = "gpt2"
  _UpperCamelCase						 = ["past_key_values"]
  _UpperCamelCase						 = {
      "hidden_size": "n_embd",
      "max_position_embeddings": "n_positions",
      "num_attention_heads": "n_head",
      "num_hidden_layers": "n_layer",
  }


  def __init__(					self	,		a__=5_02_57	,		a__=10_24	,		a__=7_68	,		a__=12	,		a__=12	,		a__=None	,		a__="gelu_new"	,		a__=0.1	,		a__=0.1	,		a__=0.1	,		a__=1E-5	,		a__=0.02	,		a__="cls_index"	,		a__=True	,		a__=None	,		a__=True	,		a__=0.1	,		a__=True	,		a__=True	,		a__=5_02_56	,		a__=5_02_56	,		a__=False	,		a__=False	,		**a__	,		):
    _lowerCamelCase      =						vocab_size
    _lowerCamelCase      =						n_positions
    _lowerCamelCase      =						n_embd
    _lowerCamelCase      =						n_layer
    _lowerCamelCase      =						n_head
    _lowerCamelCase      =						n_inner
    _lowerCamelCase      =						activation_function
    _lowerCamelCase      =						resid_pdrop
    _lowerCamelCase      =						embd_pdrop
    _lowerCamelCase      =						attn_pdrop
    _lowerCamelCase      =						layer_norm_epsilon
    _lowerCamelCase      =						initializer_range
    _lowerCamelCase      =						summary_type
    _lowerCamelCase      =						summary_use_proj
    _lowerCamelCase      =						summary_activation
    _lowerCamelCase      =						summary_first_dropout
    _lowerCamelCase      =						summary_proj_to_labels
    _lowerCamelCase      =						scale_attn_weights
    _lowerCamelCase      =						use_cache
    _lowerCamelCase      =						scale_attn_by_inverse_layer_idx
    _lowerCamelCase      =						reorder_and_upcast_attn

    _lowerCamelCase      =						bos_token_id
    _lowerCamelCase      =						eos_token_id

    super().__init__(bos_token_id=a__	,		eos_token_id=a__	,		**a__		)







class 		__magic_name__				( lowercase_					):


  """simple docstring"""





  def __init__(					self	,		a__	,		a__ = "default"	,		a__ = None	,		a__ = False	,		):
    super().__init__(a__	,		task=a__	,		patching_specs=a__	,		use_past=a__		)
    if not getattr(self._config	,		'''pad_token_id'''	,		a__		):
      # TODO: how to do that better?
      _lowerCamelCase      =						0


  @property
  def     _UpperCAmelCase  (					self		):
    _lowerCamelCase      =						OrderedDict({'''input_ids''': {0: '''batch''', 1: '''sequence'''}}		)
    if self.use_past:
      self.fill_with_past_key_values_(a__	,		direction='''inputs'''		)
      _lowerCamelCase      =						{0: '''batch''', 1: '''past_sequence + sequence'''}
    else:
      _lowerCamelCase      =						{0: '''batch''', 1: '''sequence'''}

    return common_inputs


  @property
  def     _UpperCAmelCase  (					self		):
    return self._config.n_layer


  @property
  def     _UpperCAmelCase  (					self		):
    return self._config.n_head


  def     _UpperCAmelCase  (					self	,		a__	,		a__ = -1	,		a__ = -1	,		a__ = False	,		a__ = None	,		):
    _lowerCamelCase      =						super(a__	,		self		).generate_dummy_inputs(
        a__	,		batch_size=a__	,		seq_length=a__	,		is_pair=a__	,		framework=a__		)

    # We need to order the input in the way they appears in the forward()
    _lowerCamelCase      =						OrderedDict({'''input_ids''': common_inputs['''input_ids''']}		)

    # Need to add the past_keys
    if self.use_past:
      if not is_torch_available():
        raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''		)
      else:
        import torch

        _lowerCamelCase    ,    _lowerCamelCase      =						common_inputs['''input_ids'''].shape
        # Not using the same length for past_key_values
        _lowerCamelCase      =						seqlen + 2
        _lowerCamelCase      =						(
            batch,
            self.num_attention_heads,
            past_key_values_length,
            self._config.hidden_size // self.num_attention_heads,
        )
        _lowerCamelCase      =						[
            (torch.zeros(a__		), torch.zeros(a__		)) for _ in range(self.num_layers		)
        ]

    _lowerCamelCase      =						common_inputs['''attention_mask''']
    if self.use_past:
      _lowerCamelCase      =						ordered_inputs['''attention_mask'''].dtype
      _lowerCamelCase      =						torch.cat(
          [ordered_inputs['''attention_mask'''], torch.ones(a__	,		a__	,		dtype=a__		)]	,		dim=1		)

    return ordered_inputs


  @property
  def     _UpperCAmelCase  (					self		):
    return 13




