

"""simple docstring"""





import argparse
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.local_sgd import LocalSGD


########################################################################
# This is a fully working simple example to use Accelerate
# with LocalSGD, which is a method to synchronize model
# parameters every K batches. It is different, but complementary
# to gradient accumulation.
#
# 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)
#
# 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
#
########################################################################


_a							=							16
_a							=							32
def   _A  (  UpperCamelCase_					:	Accelerator,      UpperCamelCase_					:	int = 16)     ->						Union[str, Any]:
   '''simple docstring'''


   __lowercase								=						AutoTokenizer.from_pretrained("bert-base-cased")
   __lowercase								=						load_dataset("glue",      "mrpc")

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

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

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

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

   return train_dataloader, eval_dataloader


# For testing only
if os.environ.get('TESTING_MOCKED_DATALOADERS', None) == "1":
   from accelerate.test_utils.training import mocked_dataloaders

   _a							=							mocked_dataloaders  # noqa: F811
def   _A  (  UpperCamelCase_					:	Dict,      UpperCamelCase_					:	str)     ->						Optional[Any]:
   '''simple docstring'''


   if os.environ.get("TESTING_MOCKED_DATALOADERS",      __a) == "1":
      __lowercase								=						2
   # New Code #
   __lowercase								=						int(args.gradient_accumulation_steps)
   __lowercase								=						int(args.local_sgd_steps)
   # Initialize accelerator
   __lowercase								=						Accelerator(
       cpu=args.cpu,      mixed_precision=args.mixed_precision,      gradient_accumulation_steps=__a)
   if accelerator.distributed_type not in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
      raise NotImplementedError("LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)")
   # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
   __lowercase								=						config['lr']
   __lowercase								=						int(config["num_epochs"])
   __lowercase								=						int(config["seed"])
   __lowercase								=						int(config["batch_size"])

   __lowercase								=						evaluate.load("glue",      "mrpc")

   set_seed(__a)
   __lowercase								=						get_dataloaders(__a,      __a)
   # Instantiate the model (we build the model here so that the seed also control new weights initialization)
   __lowercase								=						AutoModelForSequenceClassification.from_pretrained("bert-base-cased",      return_dict=__a)

   # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
   # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
   # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
   __lowercase								=						model.to(accelerator.device)

   # Instantiate optimizer
   __lowercase								=						AdamW(params=model.parameters(),      lr=__a)

   # Instantiate scheduler
   __lowercase								=						get_linear_schedule_with_warmup(
       optimizer=__a,      num_warmup_steps=100,      num_training_steps=(len(__a) * num_epochs),      )

   # Prepare everything
   # 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.
   __lowercase								=						accelerator.prepare(
       __a,      __a,      __a,      __a,      __a)

   # Now we train the model
   for epoch in range(__a):
      model.train()
      with LocalSGD(
          accelerator=__a,      model=__a,      local_sgd_steps=__a,      enabled=local_sgd_steps is not None) as local_sgd:
         for step, batch in enumerate(__a):
            # We could avoid this line since we set the accelerator with `device_placement=True`.
            batch.to(accelerator.device)
            # New code #
            # We use the new `accumulate` context manager to perform gradient accumulation
            # We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
            with accelerator.accumulate(__a):
               __lowercase								=						model(**__a)
               __lowercase								=						output.loss
               accelerator.backward(__a)
               optimizer.step()
               lr_scheduler.step()
               optimizer.zero_grad()
               # LocalSGD-specific line
               local_sgd.step()

      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():
            __lowercase								=						model(**__a)
         __lowercase								=						outputs.logits.argmax(dim=-1)
         __lowercase								=						accelerator.gather_for_metrics((predictions, batch["labels"]))
         metric.add_batch(
             predictions=__a,      references=__a,      )

      __lowercase								=						metric.compute()
      # Use accelerator.print to print only on the main process.
      accelerator.print(F"""epoch {epoch}:""",      __a)
def   _A  (  )     ->						List[str]:
   '''simple docstring'''


   __lowercase								=						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.",      )
   # New Code #
   parser.add_argument(
       "--gradient_accumulation_steps",      type=__a,      default=1,      help="The number of minibatches to be ran before gradients are accumulated.",      )
   parser.add_argument(
       "--local_sgd_steps",      type=__a,      default=8,      help="Number of local SGD steps or None to disable local SGD")
   parser.add_argument("--cpu",      action="store_true",      help="If passed, will train on the CPU.")
   __lowercase								=						parser.parse_args()
   __lowercase								=						{'lr': 2E-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
   training_function(__a,      __a)


if __name__ == "__main__":
   main()







"""simple docstring"""





import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


_a							=							5_00_00
_a							=							50_00

_a				,						_a							=							os.path.split(__file__)
_a							=							os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))




@get_duration
def   _A  (  UpperCamelCase_					:	datasets.Dataset,      UpperCamelCase_					:	List[str])     ->						List[str]:
   '''simple docstring'''


   for i in range(UpperCamelCase_):
      __lowercase								=						dataset[i]




@get_duration
def   _A  (  UpperCamelCase_					:	datasets.Dataset,      UpperCamelCase_					:	List[Any],      UpperCamelCase_					:	int)     ->						Dict:
   '''simple docstring'''


   for i in range(0,      len(UpperCamelCase_),      UpperCamelCase_):
      __lowercase								=						dataset[i : i + batch_size]




@get_duration
def   _A  (  UpperCamelCase_					:	datasets.Dataset,      UpperCamelCase_					:	Any,      UpperCamelCase_					:	Optional[int])     ->						List[str]:
   '''simple docstring'''


   with dataset.formatted_as(type=UpperCamelCase_):
      for i in range(UpperCamelCase_):
         __lowercase								=						dataset[i]




@get_duration
def   _A  (  UpperCamelCase_					:	datasets.Dataset,      UpperCamelCase_					:	str,      UpperCamelCase_					:	Union[str, Any],      UpperCamelCase_					:	Union[str, Any])     ->						Dict:
   '''simple docstring'''


   with dataset.formatted_as(type=UpperCamelCase_):
      for i in range(0,      UpperCamelCase_,      UpperCamelCase_):
         __lowercase								=						dataset[i : i + batch_size]




def   _A  (  )     ->						List[str]:
   '''simple docstring'''


   __lowercase								=						{"num examples": SPEED_TEST_N_EXAMPLES}
   __lowercase								=						[
       (read, {"length": SMALL_TEST}),
       (read, {"length": SPEED_TEST_N_EXAMPLES}),
       (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 10}),
       (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 100}),
       (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 1000}),
       (read_formatted, {"type": "numpy", "length": SMALL_TEST}),
       (read_formatted, {"type": "pandas", "length": SMALL_TEST}),
       (read_formatted, {"type": "torch", "length": SMALL_TEST}),
       (read_formatted, {"type": "tensorflow", "length": SMALL_TEST}),
       (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 10}),
       (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 1000}),
   ]

   __lowercase								=						[
       (read, {"length": SMALL_TEST}),
       (read, {"length": SPEED_TEST_N_EXAMPLES}),
       (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 10}),
       (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 100}),
       (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 1000}),
       (read_formatted, {"type": "numpy", "length": SMALL_TEST}),
       (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 10}),
       (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 1000}),
   ]
   with tempfile.TemporaryDirectory() as tmp_dir:
      print("generating dataset")
      __lowercase								=						datasets.Features(
          {"list": datasets.Sequence(datasets.Value("float32")), "numbers": datasets.Value("float32")})
      __lowercase								=						generate_example_dataset(
          os.path.join(UpperCamelCase_,      "dataset.arrow"),      UpperCamelCase_,      num_examples=UpperCamelCase_,      seq_shapes={"list": (100,)},      )
      print("first set of iterations")
      for func, kwargs in functions:
         print(func.__name__,      str(UpperCamelCase_))
         __lowercase								=						func(UpperCamelCase_,      **UpperCamelCase_)

      print("shuffling dataset")
      __lowercase								=						dataset.shuffle()
      print("Second set of iterations (after shuffling")
      for func, kwargs in functions_shuffled:
         print("shuffled ",      func.__name__,      str(UpperCamelCase_))
         __lowercase								=						func(
             UpperCamelCase_,      **UpperCamelCase_)

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


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




