



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 :				Tuple							=16
UpperCAmelCase :				str							=32



def  _lowerCAmelCase     (_lowerCAmelCase):
					return int(x / 2**20)


class    _lowercase :
				'''simple docstring'''







				def __enter__(			self     ):

									'''simple docstring'''


									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    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									return self.process.memory_info().rss

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									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  ,			*snake_case__     ):

									'''simple docstring'''


									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

	UpperCAmelCase :				Union[str, Any]							=mocked_dataloaders  # noqa: F811



def  _lowerCAmelCase     (_lowerCAmelCase ,						_lowerCAmelCase):
					# For testing only
					if os.environ.get("TESTING_MOCKED_DATALOADERS" ,						_lowerCAmelCase) == "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=_lowerCAmelCase ,						ranka_only=_lowerCAmelCase) ,						optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=_lowerCAmelCase ,						ranka_only=_lowerCAmelCase) ,						)

					# 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=_lowerCAmelCase ,						)
					else:
										UpperCamelCase_			= Accelerator(fsdp_plugin=_lowerCAmelCase)
					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(_lowerCAmelCase)
										accelerator.init_trackers("fsdp_glue_no_trainer" ,						_lowerCAmelCase)

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

					def tokenize_function(_lowerCAmelCase):
										# max_length=None => use the model max length (it's actually the default)
										UpperCamelCase_			= tokenizer(examples["sentence1"] ,						examples["sentence2"] ,						truncation=_lowerCAmelCase ,						max_length=_lowerCAmelCase)
										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(
										    _lowerCAmelCase ,						batched=_lowerCAmelCase ,						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(_lowerCAmelCase):
										# On TPU it's best to pad everything to the same length or training will be very slow.
										UpperCamelCase_			= 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":
															UpperCamelCase_			= 16
										elif accelerator.mixed_precision != "no":
															UpperCamelCase_			= 8
										else:
															UpperCamelCase_			= None

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

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

					set_seed(_lowerCAmelCase)

					# 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=_lowerCAmelCase ,						low_cpu_mem_usage=_lowerCAmelCase)

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

					# 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=_lowerCAmelCase ,						weight_decay=2e-4)

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

					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(_lowerCAmelCase)[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(_lowerCAmelCase)
															# If resuming by step, we also need to know exactly how far into the DataLoader we went
															UpperCamelCase_			= (num_epochs * len(_lowerCAmelCase)) - resume_step

    # Now we train the model
					for epoch in range(_lowerCAmelCase):
										# 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(_lowerCAmelCase):
																				# 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(**_lowerCAmelCase)
																				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(_lowerCAmelCase)
																				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(_lowerCAmelCase ,						_lowerCAmelCase):
																									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 ,						_lowerCAmelCase)
																														accelerator.save_state(_lowerCAmelCase)
        # 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=_lowerCAmelCase ,						)

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

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

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

					if args.with_tracking:
										accelerator.end_training()



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


if __name__ == "__main__":
	main()





import json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


UpperCAmelCase :				Any							=1000_0000_0000
UpperCAmelCase :				Any							=1_0000

UpperCAmelCase			,			UpperCAmelCase :				Tuple							=os.path.split(__file__)
UpperCAmelCase :				Tuple							=os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))



def  _lowerCAmelCase     (_lowerCAmelCase ,						_lowerCAmelCase):
					UpperCamelCase_			= pa.Table.from_pydict({"col": [0] * chunk_size})
					UpperCamelCase_			= pa.concat_tables([table] * (num_examples // chunk_size))
					return datasets.Dataset(_lowerCAmelCase ,						fingerprint="table_100B")


@dataclass
class    _lowercase :
				'''simple docstring'''







				lowercase__     =    42
				lowercase__     =    42
				lowercase__     =    42
				lowercase__     =    42

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= np.random.default_rng(self.seed     )
									UpperCamelCase_			= rng.integers(low=self.low  ,			high=self.high  ,			size=self.size     ).tolist()

				def __iter__(			self     ):

									'''simple docstring'''


									return iter(self._sampled_values     )

				def __len__(			self     ):

									'''simple docstring'''


									return self.size




@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[0]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[-1]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[range(len(_lowerCAmelCase) // 2 ,						len(_lowerCAmelCase) // 2 + 10_24)]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[RandIter(0 ,						len(_lowerCAmelCase) ,						10_24 ,						seed=42)]



def  _lowerCAmelCase     ():
					UpperCamelCase_			= {"num examples": SPEED_TEST_N_EXAMPLES}
					UpperCamelCase_			= (get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
					print("generating dataset")
					UpperCamelCase_			= generate_100B_dataset(num_examples=_lowerCAmelCase ,						chunk_size=_lowerCAmelCase)
					print("Functions")
					for func in functions:
										print(func.__name__)
										UpperCamelCase_			= func(_lowerCAmelCase)

					with open(_lowerCAmelCase ,						"wb") as f:
										f.write(json.dumps(_lowerCAmelCase).encode("utf-8"))


if __name__ == "__main__":  # useful to run the profiler
	benchmark_table_100B()
