

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


lowerCAmelCase__							=			16
lowerCAmelCase__							=			32
def   _lowerCAmelCase(						__A      ):
  return int(x / 2**20      )






class     __magic_name__     :


 def __enter__(    self		:			Optional[int]     )     ->   List[Any]:
   gc.collect()
   torch.cuda.empty_cache()
   torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
   UpperCAmelCase          =	torch.cuda.memory_allocated()
   UpperCAmelCase          =	psutil.Process()

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


 def _UpperCamelCase						(    self		:			Any     )     ->   Any:
   return self.process.memory_info().rss


 def _UpperCamelCase						(    self		:			Tuple     )     ->   List[Any]:
   UpperCAmelCase          =	-1

   while True:
     UpperCAmelCase          =	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		:			Union[str, Any]      ,    *lowerCAmelCase__		:			int     )     ->   List[Any]:
   UpperCAmelCase          =	False

   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     )

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

      lowerCAmelCase__							=			mocked_dataloaders  # noqa: F811
def   _lowerCAmelCase(						__A ,							__A      ):





  if os.environ.get("TESTING_MOCKED_DATALOADERS" ,							_UpperCAmelCase      ) == "1":
    UpperCAmelCase          =	2

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

  if hasattr(args.checkpointing_steps ,							"isdigit"      ):
    if args.checkpointing_steps == "epoch":
      UpperCAmelCase          =	args.checkpointing_steps
    elif args.checkpointing_steps.isdigit():
      UpperCAmelCase          =	int(args.checkpointing_steps      )
    else:
      raise ValueError(
          F"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."      )
  else:
    UpperCAmelCase          =	None
  # 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"]      )

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

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

  def tokenize_function(__A      ):
    # 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
  # starting with the main process first:
  with accelerator.main_process_first():
    UpperCAmelCase          =	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
  UpperCAmelCase          =	tokenized_datasets.rename_column("label" ,							"labels"      )

  # If the batch size is too big we use gradient accumulation
  UpperCAmelCase          =	1
  if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
    UpperCAmelCase          =	batch_size // MAX_GPU_BATCH_SIZE
    UpperCAmelCase          =	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.
    UpperCAmelCase          =	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          =	16
    elif accelerator.mixed_precision != "no":
      UpperCAmelCase          =	8
    else:
      UpperCAmelCase          =	None

    return tokenizer.pad(
        _UpperCAmelCase ,							padding="longest" ,							max_length=_UpperCAmelCase ,							pad_to_multiple_of=_UpperCAmelCase ,							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      )

  set_seed(_UpperCAmelCase      )

  # Instantiate the model (we build the model here so that the seed also control new weights initialization)
  UpperCAmelCase          =	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
  UpperCAmelCase          =	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
  UpperCAmelCase          =	torch.optim.AdamW(params=model.parameters() ,							lr=_UpperCAmelCase ,							weight_decay=2E-4      )

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

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

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

      UpperCAmelCase          =	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":
        UpperCAmelCase          =	F"epoch_{epoch}"
        if args.output_dir is not None:
          UpperCAmelCase          =	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   _lowerCAmelCase(						):
  UpperCAmelCase          =	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 ,							)
  UpperCAmelCase          =	parser.parse_args()
  UpperCAmelCase          =	{'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
  training_function(_UpperCAmelCase ,							_UpperCAmelCase      )


if __name__ == "__main__":
      main()








import json
import os
import unittest

from transformers import HerbertTokenizer, HerbertTokenizerFast
from transformers.models.herbert.tokenization_herbert import VOCAB_FILES_NAMES
from transformers.testing_utils import get_tests_dir, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin

@require_tokenizers
class 						UpperCAmelCase__     (     __UpperCamelCase		,unittest.TestCase  ):




			'''simple docstring'''






			UpperCamelCase     		=			HerbertTokenizer
			UpperCamelCase     		=			HerbertTokenizerFast
			UpperCamelCase     		=			True




			def 					snake_case__	(  self				:					Tuple					):
					'''simple docstring'''





					super().setUp()

					# Use a simpler test file without japanese/chinese characters
					with open(F'{get_tests_dir()}/fixtures/sample_text_no_unicode.txt'			,			encoding='''utf-8'''					) as f_data:
							__UpperCAmelCase       :	List[str]						=      f_data.read().replace('''\n\n'''			,			'''\n'''					).strip()

					__UpperCAmelCase       :	Dict						=      [
					    '''<s>''',
					    '''</s>''',
					    '''l''',
					    '''o''',
					    '''w''',
					    '''e''',
					    '''r''',
					    '''s''',
					    '''t''',
					    '''i''',
					    '''d''',
					    '''n''',
					    '''w</w>''',
					    '''r</w>''',
					    '''t</w>''',
					    '''lo''',
					    '''low''',
					    '''er</w>''',
					    '''low</w>''',
					    '''lowest</w>''',
					    '''newer</w>''',
					    '''wider</w>''',
					    ''',</w>''',
					    '''<unk>''',
					]
					__UpperCAmelCase       :	Tuple						=      dict(zip(a_			,			range(len(a_					)					)					)					)
					__UpperCAmelCase       :	Optional[int]						=      ['''l o 123''', '''lo w 1456''', '''e r</w> 1789''', '''''']

					__UpperCAmelCase       :	str						=      os.path.join(self.tmpdirname			,			VOCAB_FILES_NAMES['''vocab_file''']					)
					__UpperCAmelCase       :	List[str]						=      os.path.join(self.tmpdirname			,			VOCAB_FILES_NAMES['''merges_file''']					)
					with open(self.vocab_file			,			'''w'''					) as fp:
							fp.write(json.dumps(a_					)					)
					with open(self.merges_file			,			'''w'''					) as fp:
							fp.write('''\n'''.join(a_					)					)




			def 					snake_case__	(  self				:					Union[str, Any]			,			a_				:					Optional[Any]					):
					'''simple docstring'''





					__UpperCAmelCase       :	List[str]						=      '''lower newer'''
					__UpperCAmelCase       :	List[Any]						=      '''lower newer'''
					return input_text, output_text




			def 					snake_case__	(  self				:					int					):
					'''simple docstring'''





					__UpperCAmelCase       :	Any						=      self.tokenizer_class(vocab_file=self.vocab_file			,			merges_file=self.merges_file					)

					__UpperCAmelCase       :	Optional[int]						=      '''lower'''
					__UpperCAmelCase       :	Any						=      ['''low''', '''er</w>''']
					__UpperCAmelCase       :	Any						=      tokenizer.tokenize(a_					)
					self.assertListEqual(a_			,			a_					)

					__UpperCAmelCase       :	int						=      tokens + ['''<unk>''']
					__UpperCAmelCase       :	Optional[Any]						=      [16, 17, 23]
					self.assertListEqual(tokenizer.convert_tokens_to_ids(a_					)			,			a_					)




			def 					snake_case__	(  self				:					List[Any]					):
					'''simple docstring'''





					if not self.test_rust_tokenizer:
							return

					__UpperCAmelCase       :	str						=      self.get_tokenizer()
					__UpperCAmelCase       :	Optional[int]						=      self.get_rust_tokenizer()

					__UpperCAmelCase       :	Tuple						=      '''lower,newer'''

					__UpperCAmelCase       :	Optional[int]						=      tokenizer.tokenize(a_					)
					__UpperCAmelCase       :	Tuple						=      rust_tokenizer.tokenize(a_					)
					self.assertListEqual(a_			,			a_					)

					__UpperCAmelCase       :	Optional[Any]						=      tokenizer.encode(a_			,			add_special_tokens=a_					)
					__UpperCAmelCase       :	Union[str, Any]						=      rust_tokenizer.encode(a_			,			add_special_tokens=a_					)
					self.assertListEqual(a_			,			a_					)

					__UpperCAmelCase       :	Tuple						=      self.get_rust_tokenizer()
					__UpperCAmelCase       :	Union[str, Any]						=      tokenizer.encode(a_					)
					__UpperCAmelCase       :	str						=      rust_tokenizer.encode(a_					)
					self.assertListEqual(a_			,			a_					)




			@slow
			def 					snake_case__	(  self				:					Tuple					):
					'''simple docstring'''





					__UpperCAmelCase       :	List[Any]						=      self.tokenizer_class.from_pretrained('''allegro/herbert-base-cased'''					)

					__UpperCAmelCase       :	Tuple						=      tokenizer.encode('''konstruowanie sekwencji'''			,			add_special_tokens=a_					)
					__UpperCAmelCase       :	int						=      tokenizer.encode('''konstruowanie wielu sekwencji'''			,			add_special_tokens=a_					)

					__UpperCAmelCase       :	Union[str, Any]						=      tokenizer.build_inputs_with_special_tokens(a_					)
					__UpperCAmelCase       :	Any						=      tokenizer.build_inputs_with_special_tokens(a_			,			a_					)

					assert encoded_sentence == [0] + text + [2]
					assert encoded_pair == [0] + text + [2] + text_a + [2]





