




import argparse
import os

# New Code #
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.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to ensure out-of-memory errors never
# interrupt training, and builds off the `nlp_example.py` script.
#
# 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)
#
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# 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	:						Any			=     16
_UpperCAmelCase	:						Optional[int]			=     32



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase = 16   ):
  lowercase    :int      	=						AutoTokenizer.from_pretrained("bert-base-cased"   )
  lowercase    :int      	=						load_dataset("glue",			"mrpc"   )

  def tokenize_function(lowerCamelCase   ):
    # max_length=None => use the model max length (it's actually the default)
    lowercase    :Optional[int]      	=						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():
    lowercase    :List[Any]      	=						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
  lowercase    :Optional[int]      	=						tokenized_datasets.rename_column("label",			"labels"   )

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

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

  # Instantiate dataloaders.
  lowercase    :Dict      	=						DataLoader(
      tokenized_datasets["train"],			shuffle=lowerCamelCase,			collate_fn=lowerCamelCase,			batch_size=lowerCamelCase   )
  lowercase    :Tuple      	=						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

     _UpperCAmelCase	:						List[str]			=     mocked_dataloaders  # noqa: F811



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase   ):
  # For testing only
  if os.environ.get("TESTING_MOCKED_DATALOADERS",			lowerCamelCase   ) == "1":
    lowercase    :int      	=						2
  # Initialize accelerator
  lowercase    :int      	=						Accelerator(cpu=args.cpu,			mixed_precision=args.mixed_precision   )
  # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
  lowercase    :int      	=						config["lr"]
  lowercase    :int      	=						int(config["num_epochs"]   )
  lowercase    :List[Any]      	=						int(config["seed"]   )
  lowercase    :Tuple      	=						int(config["batch_size"]   )

  lowercase    :Any      	=						evaluate.load("glue",			"mrpc"   )

  # New Code #
  # We now can define an inner training loop function. It should take a batch size as the only parameter,
  # and build the dataloaders in there.
  # It also gets our decorator
  @find_executable_batch_size(starting_batch_size=lowerCamelCase   )
  def inner_training_loop(lowerCamelCase   ):
    # And now just move everything below under this function
    # We need to bring in the Accelerator object from earlier
    nonlocal accelerator
    # And reset all of its attributes that could hold onto any memory:
    accelerator.free_memory()

    # Then we can declare the model, optimizer, and everything else:
    set_seed(lowerCamelCase   )

    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
    lowercase    :int      	=						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).
    lowercase    :str      	=						model.to(accelerator.device   )

    # Instantiate optimizer
    lowercase    :List[Any]      	=						AdamW(params=model.parameters(),			lr=lowerCamelCase   )
    lowercase	,    lowercase    :Any      	=						get_dataloaders(lowerCamelCase,			lowerCamelCase   )

    # Instantiate scheduler
    lowercase    :Any      	=						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.
    lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :str      	=						accelerator.prepare(
        lowerCamelCase,			lowerCamelCase,			lowerCamelCase,			lowerCamelCase,			lowerCamelCase   )

    # Now we train the model
    for epoch in range(lowerCamelCase   ):
      model.train()
      for step, batch in enumerate(lowerCamelCase   ):
        # We could avoid this line since we set the accelerator with `device_placement=True`.
        batch.to(accelerator.device   )
        lowercase    :List[str]      	=						model(**lowerCamelCase   )
        lowercase    :Tuple      	=						outputs.loss
        accelerator.backward(lowerCamelCase   )
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

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

      lowercase    :Union[str, Any]      	=						metric.compute()
      # Use accelerator.print to print only on the main process.
      accelerator.print(F"epoch {epoch}:",			lowerCamelCase   )

    # New Code #
    # And call it at the end with no arguments
    # Note: You could also refactor this outside of your training loop function
  inner_training_loop()



def       UpperCAmelCase__  (     ):
  lowercase    :Optional[Any]      	=						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."   )
  lowercase    :Union[str, Any]      	=						parser.parse_args()
  lowercase    :Union[str, Any]      	=						{"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
  training_function(lowerCamelCase,			lowerCamelCase   )


if __name__ == "__main__":
     main()









import json
import os
import tempfile
from unittest.mock import patch

import torch
from torch.utils.data import DataLoader, TensorDataset

from accelerate import DistributedType, infer_auto_device_map, init_empty_weights
from accelerate.accelerator import Accelerator
from accelerate.state import GradientState, PartialState
from accelerate.test_utils import require_bnb, require_multi_gpu, slow
from accelerate.test_utils.testing import AccelerateTestCase, require_cuda
from accelerate.utils import patch_environment



def       UpperCAmelCase__  (     ):
  lowercase    :List[str]      	=						torch.nn.Linear(2,			4   )
  lowercase    :List[Any]      	=						torch.optim.AdamW(model.parameters(),			lr=1.0   )
  lowercase    :int      	=						torch.optim.lr_scheduler.OneCycleLR(lowerCamelCase,			max_lr=0.01,			steps_per_epoch=2,			epochs=1   )
  lowercase    :Optional[int]      	=						DataLoader(TensorDataset(torch.tensor([1, 2, 3]   )   )   )
  lowercase    :List[str]      	=						DataLoader(TensorDataset(torch.tensor([4, 5, 6]   )   )   )

  return model, optimizer, scheduler, train_dl, valid_dl



def       UpperCAmelCase__  (     lowerCamelCase   ):
  return (model.weight.abs().sum() + model.bias.abs().sum()).item()



def       UpperCAmelCase__  (     lowerCamelCase   ):
  lowercase    :Dict      	=						torch.nn.Linear(*tuple(model.weight.T.shape   )   ).state_dict()
  model.load_state_dict(lowerCamelCase   )







class    __lowerCAmelCase       (							lowerCAmelCase):






       @require_cuda
       def 					SCREAMING_SNAKE_CASE    (   self:					str				):
         lowercase    :List[str]      	=						Accelerator()
         assert PartialState._shared_state["_cpu"] is False
         assert PartialState._shared_state["device"].type == "cuda"
         with self.assertRaises(_lowerCAmelCase				):
           lowercase    :Any      	=						Accelerator(cpu=_lowerCAmelCase				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Tuple				):
         lowercase    :Any      	=						Accelerator()
         lowercase    :Dict      	=						GradientState()
         assert state.num_steps == 1
         lowercase    :Tuple      	=						4
         assert state.num_steps == 4

         assert state.sync_gradients is True
         lowercase    :List[Any]      	=						False
         assert state.sync_gradients is False
         GradientState._reset_state()






       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         lowercase    :str      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :Any      	=						create_components()

         (
             (
             lowercase
         )	,    (
             lowercase
         )	,    (
             lowercase
         )	,    (
             lowercase
         )	,    (
             lowercase
         )	,    
         )    :Tuple      	=						accelerator.prepare(_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)

         self.assertTrue(prepared_model in accelerator._models				)
         self.assertTrue(prepared_optimizer in accelerator._optimizers				)
         self.assertTrue(prepared_scheduler in accelerator._schedulers				)
         self.assertTrue(prepared_train_dl in accelerator._dataloaders				)
         self.assertTrue(prepared_valid_dl in accelerator._dataloaders				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         lowercase    :Dict      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :List[str]      	=						create_components()
         accelerator.prepare(_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)
         accelerator.free_memory()

         self.assertTrue(len(accelerator._models				) == 0				)
         self.assertTrue(len(accelerator._optimizers				) == 0				)
         self.assertTrue(len(accelerator._schedulers				) == 0				)
         self.assertTrue(len(accelerator._dataloaders				) == 0				)






       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         PartialState._reset_state()

         # Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
         def noop(*_lowerCAmelCase:					List[str]       ,					**_lowerCAmelCase:					Optional[int]				):
           pass

         with patch("torch.cuda.set_device"       ,					_lowerCAmelCase				), patch_environment(ACCELERATE_TORCH_DEVICE="cuda:64"				):
           lowercase    :List[Any]      	=						Accelerator()
           self.assertEqual(str(accelerator.state.device				)       ,					"cuda:64"				)






       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         lowercase    :Tuple      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :Optional[Any]      	=						create_components()
         accelerator.prepare(_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)

         lowercase    :Tuple      	=						get_signature(_lowerCAmelCase				)

         with tempfile.TemporaryDirectory() as tmpdirname:
           accelerator.save_state(_lowerCAmelCase				)

           # make sure random weights don't match
           load_random_weights(_lowerCAmelCase				)
           self.assertTrue(abs(model_signature - get_signature(_lowerCAmelCase				)				) > 1e-3				)

           # make sure loaded weights match
           accelerator.load_state(_lowerCAmelCase				)
           self.assertTrue(abs(model_signature - get_signature(_lowerCAmelCase				)				) < 1e-3				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]				):
         lowercase    :Dict      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :int      	=						create_components()
         accelerator.prepare(_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)

         lowercase    :int      	=						get_signature(_lowerCAmelCase				)

         # saving hook
         def save_config(_lowerCAmelCase:					Union[str, Any]       ,					_lowerCAmelCase:					str       ,					_lowerCAmelCase:					List[Any]				):
           lowercase    :Dict      	=						{"class_name": models[0].__class__.__name__}

           with open(os.path.join(_lowerCAmelCase       ,					"data.json"				)       ,					"w"				) as f:
             json.dump(_lowerCAmelCase       ,					_lowerCAmelCase				)

        # loading hook
         def load_config(_lowerCAmelCase:					Union[str, Any]       ,					_lowerCAmelCase:					Union[str, Any]				):
           with open(os.path.join(_lowerCAmelCase       ,					"data.json"				)       ,					"r"				) as f:
             lowercase    :int      	=						json.load(_lowerCAmelCase				)

           lowercase    :Optional[int]      	=						config["class_name"]

         lowercase    :Optional[Any]      	=						accelerator.register_save_state_pre_hook(_lowerCAmelCase				)
         lowercase    :Tuple      	=						accelerator.register_load_state_pre_hook(_lowerCAmelCase				)

         with tempfile.TemporaryDirectory() as tmpdirname:
           accelerator.save_state(_lowerCAmelCase				)

           # make sure random weights don't match with hooks
           load_random_weights(_lowerCAmelCase				)
           self.assertTrue(abs(model_signature - get_signature(_lowerCAmelCase				)				) > 1e-3				)

           # random class name to verify correct one is loaded
           lowercase    :Optional[int]      	=						"random"

           # make sure loaded weights match with hooks
           accelerator.load_state(_lowerCAmelCase				)
           self.assertTrue(abs(model_signature - get_signature(_lowerCAmelCase				)				) < 1e-3				)

           # mode.class_name is loaded from config
           self.assertTrue(model.class_name == model.__class__.__name__				)

         # remove hooks
         save_hook.remove()
         load_hook.remove()

         with tempfile.TemporaryDirectory() as tmpdirname:
           accelerator.save_state(_lowerCAmelCase				)

           # make sure random weights don't match with hooks removed
           load_random_weights(_lowerCAmelCase				)
           self.assertTrue(abs(model_signature - get_signature(_lowerCAmelCase				)				) > 1e-3				)

           # random class name to verify correct one is loaded
           lowercase    :str      	=						"random"

           # make sure loaded weights match with hooks removed
           accelerator.load_state(_lowerCAmelCase				)
           self.assertTrue(abs(model_signature - get_signature(_lowerCAmelCase				)				) < 1e-3				)

           # mode.class_name is NOT loaded from config
           self.assertTrue(model.class_name != model.__class__.__name__				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]				):
         lowercase    :List[str]      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :List[Any]      	=						create_components()
         lowercase    :List[str]      	=						None

         # This should work
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :str      	=						accelerator.prepare(
             _lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)
         self.assertTrue(dummy_obj is None				)






       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]				):
         lowercase    :Union[str, Any]      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :Union[str, Any]      	=						create_components()
         lowercase    :Union[str, Any]      	=						[1, 2, 3]

         # This should work
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :List[str]      	=						accelerator.prepare(
             _lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Dummy object should have `_is_accelerate_prepared` set to `True`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Model is missing `_is_accelerator_prepared` or is set to `False`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Optimizer is missing `_is_accelerator_prepared` or is set to `False`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Scheduler is missing `_is_accelerator_prepared` or is set to `False`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`"       ,					)






       @slow
       @require_bnb
       def 					SCREAMING_SNAKE_CASE    (   self:					Union[str, Any]				):
         from transformers import AutoModelForCausalLM

         lowercase    :Any      	=						AutoModelForCausalLM.from_pretrained(
             "EleutherAI/gpt-neo-125m"       ,					load_in_abit=_lowerCAmelCase       ,					device_map={"": 0}       ,					)
         lowercase    :int      	=						Accelerator()

         # This should work
         lowercase    :Optional[int]      	=						accelerator.prepare(_lowerCAmelCase				)






       @slow
       @require_bnb
       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]				):
         from transformers import AutoModelForCausalLM

         lowercase    :Optional[Any]      	=						Accelerator()

         with init_empty_weights():
           lowercase    :Any      	=						AutoModelForCausalLM.from_pretrained(
               "EleutherAI/gpt-neo-125m"       ,					)
           model.tie_weights()
           lowercase    :Dict      	=						infer_auto_device_map(_lowerCAmelCase				)
           lowercase    :int      	=						"cpu"

         lowercase    :str      	=						AutoModelForCausalLM.from_pretrained(
             "EleutherAI/gpt-neo-125m"       ,					device_map=_lowerCAmelCase       ,					load_in_abit=_lowerCAmelCase       ,					llm_inta_enable_fpaa_cpu_offload=_lowerCAmelCase				)

         # This should not work and get value error
         with self.assertRaises(_lowerCAmelCase				):
           lowercase    :Dict      	=						accelerator.prepare(_lowerCAmelCase				)






       @slow
       @require_bnb
       @require_multi_gpu
       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         from transformers import AutoModelForCausalLM

         lowercase    :List[str]      	=						{"distributed_type": DistributedType.MULTI_GPU}

         with init_empty_weights():
           lowercase    :str      	=						AutoModelForCausalLM.from_pretrained(
               "EleutherAI/gpt-neo-125m"       ,					)
           model.tie_weights()
           lowercase    :Optional[Any]      	=						infer_auto_device_map(_lowerCAmelCase				)
           lowercase    :Any      	=						1

         lowercase    :str      	=						AutoModelForCausalLM.from_pretrained(
             "EleutherAI/gpt-neo-125m"       ,					load_in_abit=_lowerCAmelCase       ,					device_map=_lowerCAmelCase       ,					)
         lowercase    :List[str]      	=						Accelerator()

         # This should not work and get value error
         with self.assertRaises(_lowerCAmelCase				):
           lowercase    :Optional[Any]      	=						accelerator.prepare(_lowerCAmelCase				)

         PartialState._reset_state()






       @slow
       @require_bnb
       @require_multi_gpu
       def 					SCREAMING_SNAKE_CASE    (   self:					List[Any]				):
         from transformers import AutoModelForCausalLM

         with init_empty_weights():
           lowercase    :List[str]      	=						AutoModelForCausalLM.from_pretrained(
               "EleutherAI/gpt-neo-125m"       ,					)
           lowercase    :List[Any]      	=						infer_auto_device_map(_lowerCAmelCase				)
           lowercase    :Optional[int]      	=						1

         lowercase    :Union[str, Any]      	=						AutoModelForCausalLM.from_pretrained(
             "EleutherAI/gpt-neo-125m"       ,					load_in_abit=_lowerCAmelCase       ,					device_map=_lowerCAmelCase       ,					)
         lowercase    :Union[str, Any]      	=						Accelerator()

         # This should work
         lowercase    :List[Any]      	=						accelerator.prepare(_lowerCAmelCase				)






       @require_cuda
       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         lowercase    :Optional[int]      	=						torch.nn.Linear(10       ,					10				)
         lowercase    :Optional[int]      	=						torch.optim.SGD(model.parameters()       ,					lr=0.01				)
         lowercase    :List[Any]      	=						Accelerator(cpu=_lowerCAmelCase				)
         lowercase    :List[str]      	=						accelerator.prepare(_lowerCAmelCase				)



