


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 :				Union[str, Any]						    =       16
UpperCamelCase :				Dict						    =       32
def 							UpperCamelCase_  (				__snake_case			:						Union[str, Any]     )->      str:




						'''simple docstring'''
						return int(x / 2**20     )

class     __snake_case							:
					def __enter__(	self					:	int     )       ->				Union[str, Any]:



											'''simple docstring'''





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

											__magic_name__ 		=			self.cpu_mem_used()
											__magic_name__ 		=			True
											__magic_name__ 		=			threading.Thread(target=self.peak_monitor_func     )
											__magic_name__ 		=			True
											peak_monitor_thread.start()
											return self
					def   a (	self					:	str     )       ->				List[Any]:



											'''simple docstring'''





											return self.process.memory_info().rss
					def   a (	self					:	Optional[int]     )       ->				int:



											'''simple docstring'''





											__magic_name__ 		=			-1

											while True:
																	__magic_name__ 		=			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					:	Tuple  ,			*_lowerCAmelCase					:	Tuple     )       ->				List[Any]:



											'''simple docstring'''





											__magic_name__ 		=			False

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

											__magic_name__ 		=			self.cpu_mem_used()
											__magic_name__ 		=			bamb(self.cpu_end - self.cpu_begin     )
											__magic_name__ 		=			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 :				int						    =       mocked_dataloaders  # noqa: F811
def 							UpperCamelCase_  (				__snake_case			:						Any	,			__snake_case			:						Dict     )->      int:




						'''simple docstring'''
						# For testing only
						if os.environ.get('TESTING_MOCKED_DATALOADERS'	,			__snake_case     ) == "1":
												__magic_name__ 		=			2

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

						# Initialize accelerator
						if args.with_tracking:
												__magic_name__ 		=			Accelerator(
												    cpu=args.cpu	,			mixed_precision=args.mixed_precision	,			log_with='wandb'	,			project_dir=args.logging_dir	,			fsdp_plugin=__snake_case	,			)
						else:
												__magic_name__ 		=			Accelerator(fsdp_plugin=__snake_case     )
						accelerator.print(accelerator.distributed_type     )

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

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

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

						def tokenize_function(__snake_case			:						List[Any]     ):
												# max_length=None => use the model max length (it's actually the default)
												__magic_name__ 		=			tokenizer(examples['sentence1']	,			examples['sentence2']	,			truncation=__snake_case	,			max_length=__snake_case     )
												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():
												__magic_name__ 		=			datasets.map(
												    __snake_case	,			batched=__snake_case	,			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
						__magic_name__ 		=			tokenized_datasets.rename_column('label'	,			'labels'     )

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

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

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

						# Instantiate dataloaders.
						__magic_name__ 		=			DataLoader(
						    tokenized_datasets['train']	,			shuffle=__snake_case	,			collate_fn=__snake_case	,			batch_size=__snake_case     )
						__magic_name__ 		=			DataLoader(
						    tokenized_datasets['validation']	,			shuffle=__snake_case	,			collate_fn=__snake_case	,			batch_size=__snake_case     )

						set_seed(__snake_case     )

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

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

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

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

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

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

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

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

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

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

						if args.with_tracking:
												accelerator.end_training()
def 							UpperCamelCase_  (				)->      List[Any]:




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


if __name__ == "__main__":
					main()





import os
import tempfile
import unittest
from typing import List

from transformers.models.esm.tokenization_esm import VOCAB_FILES_NAMES, EsmTokenizer
from transformers.testing_utils import require_tokenizers
from transformers.tokenization_utils import PreTrainedTokenizer
from transformers.tokenization_utils_base import PreTrainedTokenizerBase

@require_tokenizers
class     __snake_case							(      unittest.TestCase):
					SCREAMING_SNAKE_CASE_		:					Tuple 				= EsmTokenizer
					def   a (	self					:	List[Any]     )       ->				Dict:



											'''simple docstring'''





											super().setUp()
											__magic_name__ 		=			tempfile.mkdtemp()
											# fmt: off
											__magic_name__ 		=			["<cls>", "<pad>", "<eos>", "<unk>", "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K", "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z", "O", ".", "-", "<null_1>", "<mask>"]  # noqa: E501
											# fmt: on
											__magic_name__ 		=			os.path.join(self.tmpdirname  ,			VOCAB_FILES_NAMES['vocab_file']     )
											with open(self.vocab_file  ,			'w'  ,			encoding='utf-8'     ) as vocab_writer:
																	vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]     )     )
					def   a (	self					:	Optional[int]  ,			**_lowerCAmelCase					:	int     )       ->				List[PreTrainedTokenizerBase]:



											'''simple docstring'''





											return [self.get_tokenizer(**_lowerCAmelCase     )]
					def   a (	self					:	List[str]  ,			**_lowerCAmelCase					:	int     )       ->				PreTrainedTokenizer:



											'''simple docstring'''





											return self.tokenizer_class.from_pretrained(self.tmpdirname  ,			**_lowerCAmelCase     )
					def   a (	self					:	int     )       ->				int:



											'''simple docstring'''





											__magic_name__ 		=			self.tokenizer_class(self.vocab_file     )

											__magic_name__ 		=			tokenizer.tokenize('LAGVS'     )
											self.assertListEqual(_lowerCAmelCase  ,			['L', 'A', 'G', 'V', 'S']     )
											self.assertListEqual(tokenizer.convert_tokens_to_ids(_lowerCAmelCase     )  ,			[4, 5, 6, 7, 8]     )
					def   a (	self					:	str     )       ->				str:



											'''simple docstring'''





											__magic_name__ 		=			self.tokenizer_class(self.vocab_file     )

											__magic_name__ 		=			'LAGVS'
											self.assertListEqual(tokenizer.encode(_lowerCAmelCase     )  ,			[0, 4, 5, 6, 7, 8, 2]     )
					def   a (	self					:	int     )       ->				List[Any]:



											'''simple docstring'''





											__magic_name__ 		=			self.tokenizer_class(self.vocab_file     )

											__magic_name__ 		=			['LAGVS', 'WCB']
											__magic_name__ 		=			tokenizer(_lowerCAmelCase  ,			padding=_lowerCAmelCase     )['input_ids']

											self.assertListEqual(_lowerCAmelCase  ,			[[0, 4, 5, 6, 7, 8, 2], [0, 2_2, 2_3, 2_5, 2]]     )
					def   a (	self					:	Dict     )       ->				List[Any]:



											'''simple docstring'''





											__magic_name__ 		=			self.tokenizer_class(self.vocab_file     )

											__magic_name__ 		=			['LAGVS', 'WCB']
											__magic_name__ 		=			tokenizer(_lowerCAmelCase  ,			padding=_lowerCAmelCase     )['input_ids']

											self.assertListEqual(_lowerCAmelCase  ,			[[0, 4, 5, 6, 7, 8, 2], [0, 2_2, 2_3, 2_5, 2, 1, 1]]     )







					def   a (	self					:	Optional[int]     )       ->				str:



											'''simple docstring'''





											__magic_name__ 		=			self.get_tokenizers(fast=_lowerCAmelCase     )
											for tokenizer in tokenizers:
																	with self.subTest(F'''{tokenizer.__class__.__name__}'''     ):
																							__magic_name__ 		=			'<unk>'
																							__magic_name__ 		=			'<mask>'

																							__magic_name__ 		=			tokenizer.tokenize(_lowerCAmelCase     )
																							__magic_name__ 		=			tokenizer.tokenize(_lowerCAmelCase     )

																							self.assertEqual(len(_lowerCAmelCase     )  ,			1     )
																							self.assertEqual(len(_lowerCAmelCase     )  ,			1     )
																							self.assertEqual(token_a[0]  ,			_lowerCAmelCase     )
																							self.assertEqual(token_a[0]  ,			_lowerCAmelCase     )

