'''simple docstring'''


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


snake_case__     :			Union[str, Any]      			=  16
snake_case__     :			Tuple      			=  32


def 					_lowerCamelCase    (      lowerCamelCase_					:							Union[str, Any]   ):



     """simple docstring"""





     return int(x / 2**20   )






class      __SCREAMING_SNAKE_CASE						:







  '''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_							:   Dict   =      torch.cuda.memory_allocated()
       UpperCAmelCase_							:   str   =      psutil.Process()

       UpperCAmelCase_							:   Optional[Any]   =      self.cpu_mem_used()
       UpperCAmelCase_							:   int   =      True
       UpperCAmelCase_							:   Optional[Any]   =      threading.Thread(target=self.peak_monitor_func						)
       UpperCAmelCase_							:   Any   =      True
       peak_monitor_thread.start()
       return self


  def _UpperCamelCase     ( self						):




       '''simple docstring'''

       return self.process.memory_info().rss


  def _UpperCamelCase     ( self						):




       '''simple docstring'''

       UpperCAmelCase_							:   List[str]   =      -1

       while True:
            UpperCAmelCase_							:   Union[str, Any]   =      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      ,		*snake_case_						):




       '''simple docstring'''

       UpperCAmelCase_							:   Optional[int]   =      False

       gc.collect()
       torch.cuda.empty_cache()
       UpperCAmelCase_							:   List[Any]   =      torch.cuda.memory_allocated()
       UpperCAmelCase_							:   int   =      torch.cuda.max_memory_allocated()
       UpperCAmelCase_							:   Any   =      bamb(self.end - self.begin						)
       UpperCAmelCase_							:   List[str]   =      bamb(self.peak - self.begin						)

       UpperCAmelCase_							:   str   =      self.cpu_mem_used()
       UpperCAmelCase_							:   Tuple   =      bamb(self.cpu_end - self.cpu_begin						)
       UpperCAmelCase_							:   Optional[int]   =      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

 snake_case__     :			Tuple      			=  mocked_dataloaders  # noqa: F811


def 					_lowerCamelCase    (      lowerCamelCase_					:							Optional[int]	,     lowerCamelCase_					:							List[Any]   ):



     """simple docstring"""





     if os.environ.get('TESTING_MOCKED_DATALOADERS'	,     lowerCamelCase_   ) == "1":
          UpperCAmelCase_							:   Optional[Any]   =      2

     # New Code #
     # Pass the advanced FSDP settings not part of the accelerate config by creating fsdp_plugin
     UpperCAmelCase_							:   int   =      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:
          UpperCAmelCase_							:   Dict   =      Accelerator(
              cpu=args.cpu	,     mixed_precision=args.mixed_precision	,     log_with='wandb'	,     project_dir=args.logging_dir	,     fsdp_plugin=lowerCamelCase_	,     )
     else:
          UpperCAmelCase_							:   Tuple   =      Accelerator(fsdp_plugin=lowerCamelCase_   )
     accelerator.print(accelerator.distributed_type   )

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

     # We need to initialize the trackers we use, and also store our configuration
     if args.with_tracking:
          UpperCAmelCase_							:   List[str]   =      vars(lowerCamelCase_   )
          accelerator.init_trackers('fsdp_glue_no_trainer'	,     lowerCamelCase_   )

     UpperCAmelCase_							:   Optional[Any]   =      AutoTokenizer.from_pretrained(args.model_name_or_path   )
     UpperCAmelCase_							:   Optional[int]   =      load_dataset('glue'	,     'mrpc'   )
     UpperCAmelCase_							:   Any   =      evaluate.load('glue'	,     'mrpc'   )

     def tokenize_function(lowerCamelCase_					:							Optional[int]   ):
          # max_length=None => use the model max length (it's actually the default)
          UpperCAmelCase_							:   List[str]   =      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():
          UpperCAmelCase_							:   int   =      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
     UpperCAmelCase_							:   List[Any]   =      tokenized_datasets.rename_column('label'	,     'labels'   )

     # If the batch size is too big we use gradient accumulation
     UpperCAmelCase_							:   List[Any]   =      1
     if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
          UpperCAmelCase_							:   str   =      batch_size // MAX_GPU_BATCH_SIZE
          UpperCAmelCase_							:   Any   =      MAX_GPU_BATCH_SIZE

     def collate_fn(lowerCamelCase_					:							str   ):
          # On TPU it's best to pad everything to the same length or training will be very slow.
          UpperCAmelCase_							:   Dict   =      128 if accelerator.distributed_type == DistributedType.TPU else None
          # When using mixed precision we want round multiples of 8/16
          if accelerator.mixed_precision == "fp8":
               UpperCAmelCase_							:   List[Any]   =      16
          elif accelerator.mixed_precision != "no":
               UpperCAmelCase_							:   Any   =      8
          else:
               UpperCAmelCase_							:   Tuple   =      None

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

     # Instantiate dataloaders.
     UpperCAmelCase_							:   Tuple   =      DataLoader(
         tokenized_datasets['train']	,     shuffle=lowerCamelCase_	,     collate_fn=lowerCamelCase_	,     batch_size=lowerCamelCase_   )
     UpperCAmelCase_							:   int   =      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)
     UpperCAmelCase_							:   Tuple   =      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
     UpperCAmelCase_							:   Tuple   =      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
     UpperCAmelCase_							:   str   =      torch.optim.AdamW(params=model.parameters()	,     lr=lowerCamelCase_	,     weight_decay=2e-4   )

     # Instantiate scheduler
     UpperCAmelCase_							:   Optional[int]   =      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.
     UpperCAmelCase_   ,			UpperCAmelCase_   ,			UpperCAmelCase_   ,			UpperCAmelCase_							:   Optional[int]   =      accelerator.prepare(
         lowerCamelCase_	,     lowerCamelCase_	,     lowerCamelCase_	,     lowerCamelCase_   )

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

          if "epoch" in training_difference:
               num_epochs -= int(training_difference.replace('epoch_'	,     ''   )   )
               UpperCAmelCase_							:   List[Any]   =      None
          else:
               UpperCAmelCase_							:   Dict   =      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
               UpperCAmelCase_							:   str   =      (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:
                    UpperCAmelCase_							:   List[Any]   =      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   )
                    UpperCAmelCase_							:   str   =      model(**lowerCamelCase_   )
                    UpperCAmelCase_							:   str   =      outputs.loss
                    UpperCAmelCase_							:   List[Any]   =      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_   ):
                         UpperCAmelCase_							:   int   =      F'''step_{overall_step}'''
                         if overall_step % checkpointing_steps == 0:
                              if args.output_dir is not None:
                                   UpperCAmelCase_							:   Tuple   =      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():
                         UpperCAmelCase_							:   Any   =      model(**lowerCamelCase_   )
                    UpperCAmelCase_							:   Dict   =      outputs.logits.argmax(dim=-1   )
                    UpperCAmelCase_   ,			UpperCAmelCase_							:   Any   =      accelerator.gather_for_metrics((predictions, batch['labels'])   )
                    metric.add_batch(
                        predictions=lowerCamelCase_	,     references=lowerCamelCase_	,     )

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



     """simple docstring"""





     UpperCAmelCase_							:   Tuple   =      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_	,     )
     UpperCAmelCase_							:   List[Any]   =      parser.parse_args()
     UpperCAmelCase_							:   Tuple   =      {'lr': 2e-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
     training_function(lowerCamelCase_	,     lowerCamelCase_   )


if __name__ == "__main__":
 main()







'''simple docstring'''


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


snake_case__     :			Optional[Any]      			=  {'''configuration_nat''': ['''NAT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''NatConfig''']}


try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 snake_case__     :			Tuple      			=  [
     '''NAT_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''NatForImageClassification''',
     '''NatModel''',
     '''NatPreTrainedModel''',
     '''NatBackbone''',
 ]

if TYPE_CHECKING:
 from .configuration_nat import NAT_PRETRAINED_CONFIG_ARCHIVE_MAP, NatConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_nat import (
      NAT_PRETRAINED_MODEL_ARCHIVE_LIST,
      NatBackbone,
      NatForImageClassification,
      NatModel,
      NatPreTrainedModel,
  )

else:
 import sys

 snake_case__     :			int      			=  _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






