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
#
########################################################################


_SCREAMING_SNAKE_CASE      						=	1_6
_SCREAMING_SNAKE_CASE      						=	3_2





def        lowercase(     UpperCamelCase_				, UpperCamelCase_ = 16    )						->     Tuple:







 '''simple docstring'''
 UpperCamelCase          =	AutoTokenizer.from_pretrained("""bert-base-cased"""    )
 UpperCamelCase          =	load_dataset("""glue"""				, """mrpc"""    )

 def tokenize_function(UpperCamelCase_    ):
  # 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"""    )

 def collate_fn(UpperCamelCase_    ):
  # On TPU it's best to pad everything to the same length or training will be very slow.
  UpperCamelCase          =	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":
   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__    )

 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

     _SCREAMING_SNAKE_CASE      						=	mocked_dataloaders  # noqa: F811





def        lowercase(     UpperCamelCase_				, UpperCamelCase_    )						->     Any:







 '''simple docstring'''




 if os.environ.get("""TESTING_MOCKED_DATALOADERS"""				, lowerCAmelCase__    ) == "1":
  UpperCamelCase          =	2
 # New Code #
 UpperCamelCase          =	int(args.gradient_accumulation_steps    )
 UpperCamelCase          =	int(args.local_sgd_steps    )
 # Initialize accelerator
 UpperCamelCase          =	Accelerator(
     cpu=args.cpu				, mixed_precision=args.mixed_precision				, gradient_accumulation_steps=lowerCAmelCase__    )
 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
 UpperCamelCase          =	config['lr']
 UpperCamelCase          =	int(config["""num_epochs"""]    )
 UpperCamelCase          =	int(config["""seed"""]    )
 UpperCamelCase          =	int(config["""batch_size"""]    )

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

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

 # 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).
 UpperCamelCase          =	model.to(accelerator.device    )

 # Instantiate optimizer
 UpperCamelCase          =	AdamW(params=model.parameters()				, lr=lowerCAmelCase__    )

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

 # Now we train the model
 for epoch in range(lowerCAmelCase__    ):
  model.train()
  with LocalSGD(
      accelerator=lowerCAmelCase__				, model=lowerCAmelCase__				, local_sgd_steps=lowerCAmelCase__				, enabled=local_sgd_steps is not None    ) as local_sgd:
   for step, batch in enumerate(lowerCAmelCase__    ):
    # 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(lowerCAmelCase__    ):
     UpperCamelCase          =	model(**lowerCAmelCase__    )
     UpperCamelCase          =	output.loss
     accelerator.backward(lowerCAmelCase__    )
     optimizer.step()
     lr_scheduler.step()
     optimizer.zero_grad()
     # LocalSGD-specific line
     local_sgd.step()

  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          =	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__    )





def        lowercase(     )						->     Any:







 '''simple docstring'''
 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."""				, )
 # New Code #
 parser.add_argument(
     """--gradient_accumulation_steps"""				, type=lowerCAmelCase__				, default=1				, help="""The number of minibatches to be ran before gradients are accumulated."""				, )
 parser.add_argument(
     """--local_sgd_steps"""				, type=lowerCAmelCase__				, 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."""    )
 UpperCamelCase          =	parser.parse_args()
 UpperCamelCase          =	{'lr': 2E-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
 training_function(lowerCAmelCase__				, lowerCAmelCase__    )


if __name__ == "__main__":
     main()





"""simple docstring"""







import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
     import tensorflow as tf
     from tensorflow.python.eager import context
     from tensorflow.python.framework import ops

     from transformers import GradientAccumulator, create_optimizer







@require_tf
class       UpperCamelCase__   (	unittest.TestCase  ):




      """simple docstring"""




      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      str    ,			SCREAMING_SNAKE_CASE_				:      List[Any]    ,			SCREAMING_SNAKE_CASE_				:      List[str]       ):
        self.assertEqual(len(SCREAMING_SNAKE_CASE_       )    ,			len(SCREAMING_SNAKE_CASE_       )       )
        for a, b in zip(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       ):
          self.assertAlmostEqual(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			delta=SCREAMING_SNAKE_CASE_       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      int       ):
        lowerCAmelCase_   :				int								= GradientAccumulator()
        accumulator([tf.constant([1.0, 2.0]       )]       )
        accumulator([tf.constant([-2.0, 1.0]       )]       )
        accumulator([tf.constant([-1.0, 2.0]       )]       )
        with self.assertRaises(SCREAMING_SNAKE_CASE_       ):
          accumulator([tf.constant([1.0, 1.0]       ), tf.constant([2.0, 2.0]       )]       )
        self.assertEqual(accumulator.step    ,			3       )
        self.assertEqual(len(accumulator.gradients       )    ,			1       )
        self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()    ,			[-2.0, 5.0]    ,			tol=1E-2       )
        accumulator.reset()
        self.assertEqual(accumulator.step    ,			0       )
        self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()    ,			[0.0, 0.0]    ,			tol=1E-2       )





      def      SCREAMING_SNAKE_CASE__ (      self				:      str       ):
        lowerCAmelCase_   :				Optional[int]								= None
        ops.enable_eager_execution_internal()
        lowerCAmelCase_   :				str								= tf.config.list_physical_devices('CPU'       )
        if len(SCREAMING_SNAKE_CASE_       ) == 1:
          tf.config.set_logical_device_configuration(
              physical_devices[0]    ,			[tf.config.LogicalDeviceConfiguration(), tf.config.LogicalDeviceConfiguration()]       )
        lowerCAmelCase_   :				Dict								= tf.config.list_logical_devices(device_type='CPU'       )
        lowerCAmelCase_   :				Optional[int]								= tf.distribute.MirroredStrategy(devices=devices[:2]       )

        with strategy.scope():
          lowerCAmelCase_   :				Union[str, Any]								= GradientAccumulator()
          lowerCAmelCase_   :				int								= tf.Variable([4.0, 3.0]       )
          lowerCAmelCase_						,lowerCAmelCase_   :				Optional[int]								= create_optimizer(5E-5    ,			1_0    ,			5       )
          lowerCAmelCase_   :				Union[str, Any]								= tf.Variable([0.0, 0.0]    ,			trainable=SCREAMING_SNAKE_CASE_       )

        def accumulate_on_replica(SCREAMING_SNAKE_CASE_				:      Optional[int]       ):
          accumulator([gradient]       )

        def apply_on_replica():
          optimizer.apply_gradients(list(zip(accumulator.gradients    ,			[variable]       )       )       )

        @tf.function
        def accumulate(SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]       ):
          with strategy.scope():
            lowerCAmelCase_   :				Tuple								= strategy.experimental_local_results(SCREAMING_SNAKE_CASE_       )
            local_variables[0].assign(SCREAMING_SNAKE_CASE_       )
            local_variables[1].assign(SCREAMING_SNAKE_CASE_       )
            strategy.run(SCREAMING_SNAKE_CASE_    ,			args=(gradient_placeholder,)       )

        @tf.function
        def apply_grad():
          with strategy.scope():
            strategy.run(SCREAMING_SNAKE_CASE_       )

        def _check_local_values(SCREAMING_SNAKE_CASE_				:      Union[str, Any]    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]       ):
          lowerCAmelCase_   :				List[Any]								= strategy.experimental_local_results(accumulator._gradients[0]       )
          self.assertListAlmostEqual(values[0].value()    ,			SCREAMING_SNAKE_CASE_    ,			tol=1E-2       )
          self.assertListAlmostEqual(values[1].value()    ,			SCREAMING_SNAKE_CASE_    ,			tol=1E-2       )

        accumulate([1.0, 2.0]    ,			[-1.0, 1.0]       )
        accumulate([3.0, -1.0]    ,			[-1.0, -1.0]       )
        accumulate([-2.0, 2.0]    ,			[3.0, -2.0]       )
        self.assertEqual(accumulator.step    ,			3       )
        _check_local_values([2.0, 3.0]    ,			[1.0, -2.0]       )
        apply_grad()
        self.assertListAlmostEqual(variable.value()    ,			[4.0, 3.0]    ,			tol=1E-2       )
        accumulator.reset()
        self.assertEqual(accumulator.step    ,			0       )
        _check_local_values([0.0, 0.0]    ,			[0.0, 0.0]       )





