




'''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      :      List[str]   							=		16
_snake_case      :      Tuple   							=		32
def 		snake_case_    (UpperCamelCase       :  Union[str, Any]      ):





							'''simple docstring'''


							return int(x / 2**20      )

class 			A    :







							def __enter__(  self       :  str					)			->							Dict:


														"""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    __lowerCAmelCase			(  self       :  str					)			->							Tuple:


														"""simple docstring"""


														return self.process.memory_info().rss







							def    __lowerCAmelCase			(  self       :  int					)			->							Optional[int]:


														"""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       :  Union[str, Any]							,			*lowerCAmelCase_       :  List[Any]					)			->							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

				_snake_case      :      Union[str, Any]   							=		mocked_dataloaders  # noqa: F811
def 		snake_case_    (UpperCamelCase       :  str ,						UpperCamelCase       :  Union[str, Any]      ):





							'''simple docstring'''


							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       :  Optional[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       :  int      ):
														# 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_    ():





							'''simple docstring'''


							_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'''


import unittest
from pathlib import Path

from transformers import is_vision_available, load_tool
from transformers.testing_utils import get_tests_dir

from .test_tools_common import ToolTesterMixin


if is_vision_available():
				from PIL import Image

class 			A    (		unittest.TestCase     ,_a    ):







							def    __lowerCAmelCase			(  self       :  Optional[int]					)			->							Optional[int]:


														"""simple docstring"""


														_a	     =				load_tool('''image-question-answering'''					)
														self.tool.setup()
														_a	     =				load_tool('''image-question-answering'''							,			remote=lowerCAmelCase_					)







							def    __lowerCAmelCase			(  self       :  Tuple					)			->							Optional[int]:


														"""simple docstring"""


														_a	     =				Image.open(Path(get_tests_dir('''fixtures/tests_samples/COCO'''					)					) / '''000000039769.png'''					)
														_a	     =				self.tool(lowerCAmelCase_							,			'''How many cats are sleeping on the couch?'''					)
														self.assertEqual(lowerCAmelCase_							,			'''2'''					)







							def    __lowerCAmelCase			(  self       :  List[str]					)			->							Union[str, Any]:


														"""simple docstring"""


														_a	     =				Image.open(Path(get_tests_dir('''fixtures/tests_samples/COCO'''					)					) / '''000000039769.png'''					)
														_a	     =				self.remote_tool(lowerCAmelCase_							,			'''How many cats are sleeping on the couch?'''					)
														self.assertEqual(lowerCAmelCase_							,			'''2'''					)







							def    __lowerCAmelCase			(  self       :  Dict					)			->							Tuple:


														"""simple docstring"""


														_a	     =				Image.open(Path(get_tests_dir('''fixtures/tests_samples/COCO'''					)					) / '''000000039769.png'''					)
														_a	     =				self.tool(image=lowerCAmelCase_							,			question='''How many cats are sleeping on the couch?'''					)
														self.assertEqual(lowerCAmelCase_							,			'''2'''					)







							def    __lowerCAmelCase			(  self       :  Tuple					)			->							Dict:


														"""simple docstring"""


														_a	     =				Image.open(Path(get_tests_dir('''fixtures/tests_samples/COCO'''					)					) / '''000000039769.png'''					)
														_a	     =				self.remote_tool(image=lowerCAmelCase_							,			question='''How many cats are sleeping on the couch?'''					)
														self.assertEqual(lowerCAmelCase_							,			'''2'''					)

