

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


SCREAMING_SNAKE_CASE_       :						int					=    1_6
SCREAMING_SNAKE_CASE_       :						List[str]					=    3_2





def 			_snake_case  (     UpperCAmelCase_	:       Union[str, Any]      ):
     return int(x / 2**20      )



class        a    :



    """simple docstring"""



    def __enter__(  self:			Tuple							):






         """simple docstring"""


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

         A__				      =	self.cpu_mem_used()
         A__				      =	True
         A__				      =	threading.Thread(target=self.peak_monitor_func							)
         A__				      =	True
         peak_monitor_thread.start()
         return self

    def 					UpperCamelCase						(  self:			Union[str, Any]							):






         """simple docstring"""


         return self.process.memory_info().rss

    def 					UpperCamelCase						(  self:			Union[str, Any]							):






         """simple docstring"""


         A__				      =	-1

         while True:
              A__				      =	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:			Optional[Any]     ,    *UpperCamelCase:			List[str]							):






         """simple docstring"""


         A__				      =	False

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

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

 SCREAMING_SNAKE_CASE_       :						int					=    mocked_dataloaders  # noqa: F811





def 			_snake_case  (     UpperCAmelCase_	:       Optional[int]       , UpperCAmelCase_	:       int      ):
     # For testing only
     if os.environ.get("""TESTING_MOCKED_DATALOADERS"""       , UpperCAmelCase_      ) == "1":
          A__				      =	2

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

     # Initialize accelerator
     if args.with_tracking:
          A__				      =	Accelerator(
              cpu=args.cpu       , mixed_precision=args.mixed_precision       , log_with="""wandb"""       , project_dir=args.logging_dir       , fsdp_plugin=UpperCAmelCase_       , )
     else:
          A__				      =	Accelerator(fsdp_plugin=UpperCAmelCase_      )
     accelerator.print(accelerator.distributed_type      )

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

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

     A__				      =	AutoTokenizer.from_pretrained(args.model_name_or_path      )
     A__				      =	load_dataset("""glue"""       , """mrpc"""      )
     A__				      =	evaluate.load("""glue"""       , """mrpc"""      )

     def tokenize_function(UpperCAmelCase_	:       Union[str, Any]      ):
          # max_length=None => use the model max length (it's actually the default)
          A__				      =	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
     # starting with the main process first:
     with accelerator.main_process_first():
          A__				      =	datasets.map(
              UpperCAmelCase_       , batched=UpperCAmelCase_       , 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
     A__				      =	tokenized_datasets.rename_column("""label"""       , """labels"""      )

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

     def collate_fn(UpperCAmelCase_	:       Dict      ):
          # On TPU it's best to pad everything to the same length or training will be very slow.
          A__				      =	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":
               A__				      =	16
          elif accelerator.mixed_precision != "no":
               A__				      =	8
          else:
               A__				      =	None

          return tokenizer.pad(
              UpperCAmelCase_       , padding="""longest"""       , max_length=UpperCAmelCase_       , pad_to_multiple_of=UpperCAmelCase_       , return_tensors="""pt"""       , )

     # Instantiate dataloaders.
     A__				      =	DataLoader(
         tokenized_datasets["""train"""]       , shuffle=UpperCAmelCase_       , collate_fn=UpperCAmelCase_       , batch_size=UpperCAmelCase_      )
     A__				      =	DataLoader(
         tokenized_datasets["""validation"""]       , shuffle=UpperCAmelCase_       , collate_fn=UpperCAmelCase_       , batch_size=UpperCAmelCase_      )

     set_seed(UpperCAmelCase_      )

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

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

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

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

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

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

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

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

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

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

     if args.with_tracking:
          accelerator.end_training()





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


if __name__ == "__main__":
 main()









"""simple docstring"""



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


SCREAMING_SNAKE_CASE_       :						Union[str, Any]					=    {'configuration_gptj': ['GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP', 'GPTJConfig', 'GPTJOnnxConfig']}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 SCREAMING_SNAKE_CASE_       :						Optional[int]					=    [
     'GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST',
     'GPTJForCausalLM',
     'GPTJForQuestionAnswering',
     'GPTJForSequenceClassification',
     'GPTJModel',
     'GPTJPreTrainedModel',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 SCREAMING_SNAKE_CASE_       :						List[str]					=    [
     'TFGPTJForCausalLM',
     'TFGPTJForQuestionAnswering',
     'TFGPTJForSequenceClassification',
     'TFGPTJModel',
     'TFGPTJPreTrainedModel',
 ]

try:
 if not is_flax_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 SCREAMING_SNAKE_CASE_       :						Tuple					=    [
     'FlaxGPTJForCausalLM',
     'FlaxGPTJModel',
     'FlaxGPTJPreTrainedModel',
 ]


if TYPE_CHECKING:
 from .configuration_gptj import GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTJConfig, GPTJOnnxConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_gptj import (
      GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST,
      GPTJForCausalLM,
      GPTJForQuestionAnswering,
      GPTJForSequenceClassification,
      GPTJModel,
      GPTJPreTrainedModel,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_gptj import (
      TFGPTJForCausalLM,
      TFGPTJForQuestionAnswering,
      TFGPTJForSequenceClassification,
      TFGPTJModel,
      TFGPTJPreTrainedModel,
  )

 try:
  if not is_flax_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel, FlaxGPTJPreTrainedModel

else:
 import sys

 SCREAMING_SNAKE_CASE_       :						List[str]					=    _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)






