






"""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      =      16
__SCREAMING_SNAKE_CASE      =      32


def 					UpperCAmelCase			(				a__   ):





						'''simple docstring'''






						return int(x / 2**20   )




class 	__UpperCamelCase    :
					def __enter__(       self       :  List[str]      )  ->					str:
											gc.collect()
											torch.cuda.empty_cache()
											torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
											lowerCAmelCase      :Union[str, Any]				  = torch.cuda.memory_allocated()
											lowerCAmelCase      :Tuple				  = psutil.Process()

											lowerCAmelCase      :Optional[Any]				  = self.cpu_mem_used()
											lowerCAmelCase      :Optional[int]				  = True
											lowerCAmelCase      :Optional[Any]				  = threading.Thread(target=self.peak_monitor_func      )
											lowerCAmelCase      :Any				  = True
											peak_monitor_thread.start()
											return self
					def 					UpperCAmelCase__   (       self       :  Union[str, Any]      )  ->					Dict:
											return self.process.memory_info().rss
					def 					UpperCAmelCase__   (       self       :  Tuple      )  ->					int:
											lowerCAmelCase      :str				  = -1

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

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

											lowerCAmelCase      :Any				  = self.cpu_mem_used()
											lowerCAmelCase      :Union[str, Any]				  = bamb(self.cpu_end - self.cpu_begin      )
											lowerCAmelCase      :List[str]				  = 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      =      mocked_dataloaders  # noqa: F811


def 					UpperCAmelCase			(				a__    ,			a__   ):





						'''simple docstring'''






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

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

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

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

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

						def tokenize_function(a__   ):
												# max_length=None => use the model max length (it's actually the default)
												lowerCAmelCase      :int				  = 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      :Tuple				  = 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      :Dict				  = tokenized_datasets.rename_column('label'    ,			'labels'   )

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

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

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

						# Instantiate scheduler
						lowerCAmelCase      :Optional[int]				  = get_linear_schedule_with_warmup(
						    optimizer=a__    ,			num_warmup_steps=10    ,			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      :int				  = accelerator.prepare(
						    a__    ,			a__    ,			a__    ,			a__   )

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

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

																		lowerCAmelCase      :Any				  = 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      :int				  = F"""epoch_{epoch}"""
																								if args.output_dir is not None:
																														lowerCAmelCase      :Any				  = 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 					UpperCAmelCase			(				):





						'''simple docstring'''






						lowerCAmelCase      :Optional[Any]				  = 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      :Optional[int]				  = parser.parse_args()
						lowerCAmelCase      :Optional[Any]				  = {'lr': 2e-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
						training_function(a__    ,			a__   )


if __name__ == "__main__":
		main()







"""simple docstring"""


import unittest
import warnings
from dataclasses import dataclass

from transformers.convert_slow_tokenizer import SpmConverter
from transformers.testing_utils import get_tests_dir




@dataclass
class 	__UpperCamelCase    :
					lowercase_						:							str




class 	__UpperCamelCase    (    unittest.TestCase	):
					def 					UpperCAmelCase__   (       self       :  Union[str, Any]      )  ->					Optional[Any]:
											lowerCAmelCase      :List[str]				  = get_tests_dir('fixtures/test_sentencepiece.model'      )
											lowerCAmelCase      :Tuple				  = get_tests_dir('fixtures/test_sentencepiece_with_bytefallback.model'      )

											lowerCAmelCase      :List[str]				  = FakeOriginalTokenizer(vocab_file=UpperCAmelCase      )

											with warnings.catch_warnings(record=UpperCAmelCase      ) as w:
																	lowerCAmelCase      :Dict				  = SpmConverter(UpperCAmelCase      )
											self.assertEqual(len(UpperCAmelCase      )    ,   0      )

											lowerCAmelCase      :int				  = FakeOriginalTokenizer(vocab_file=UpperCAmelCase      )

											with warnings.catch_warnings(record=UpperCAmelCase      ) as w:
																	lowerCAmelCase      :List[str]				  = SpmConverter(UpperCAmelCase      )
											self.assertEqual(len(UpperCAmelCase      )    ,   1      )

											self.assertIn(
											    'The sentencepiece tokenizer that you are converting to a fast tokenizer uses the byte fallback option'
											    ' which is not implemented in the fast tokenizers.'    ,   str(w[0].message      )    ,   )