


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      :			Dict      				=   16
_SCREAMING_SNAKE_CASE      :			Dict      				=   32
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_ = 16    ):



					"""simple docstring"""






					snake_case					=    AutoTokenizer.from_pretrained('''bert-base-cased'''    )
					snake_case					=    load_dataset('''glue'''							,'''mrpc'''    )

					def tokenize_function(UpperCamelCase_    ):
										# max_length=None => use the model max length (it's actually the default)
										snake_case					=    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():
										snake_case					=    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
					snake_case					=    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.
										snake_case					=    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":
															snake_case					=    16
										elif accelerator.mixed_precision != "no":
															snake_case					=    8
										else:
															snake_case					=    None

										return tokenizer.pad(
										    UpperCamelCase_							,padding='''longest'''							,max_length=UpperCamelCase_							,pad_to_multiple_of=UpperCamelCase_							,return_tensors='''pt'''							,)

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

					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      :			List[str]      				=   mocked_dataloaders  # noqa: F811
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_    ):



					"""simple docstring"""






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

					snake_case					=    evaluate.load('''glue'''							,'''mrpc'''    )

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

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

					# Instantiate optimizer
					snake_case					=    AdamW(params=model.parameters()							,lr=UpperCamelCase_    )

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

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

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

										snake_case					=    metric.compute()
										# Use accelerator.print to print only on the main process.
										accelerator.print(F'''epoch {epoch}:'''							,UpperCamelCase_    )
def       UpperCAmelCase__						():



					"""simple docstring"""






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


if __name__ == "__main__":
	main()








import unittest

from parameterized import parameterized

from transformers import AutoTokenizer, GPTNeoXConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import (
     GPTNeoXForCausalLM,
     GPTNeoXForQuestionAnswering,
     GPTNeoXForSequenceClassification,
     GPTNeoXForTokenClassification,
     GPTNeoXModel,
 )

class   A__ :







    """simple docstring"""

    def __init__(	self      , __snake_case      , __snake_case=1_3      , __snake_case=7      , __snake_case=True      , __snake_case=True      , __snake_case=True      , __snake_case=True      , __snake_case=9_9      , __snake_case=6_4      , __snake_case=5      , __snake_case=4      , __snake_case=3_7      , __snake_case="gelu"      , __snake_case=0.1      , __snake_case=0.1      , __snake_case=5_1_2      , __snake_case=1_6      , __snake_case=2      , __snake_case=0.02      , __snake_case=3      , __snake_case=4      , __snake_case=None      , ):
         snake_case					=    parent
         snake_case					=    batch_size
         snake_case					=    seq_length
         snake_case					=    is_training
         snake_case					=    use_input_mask
         snake_case					=    use_token_type_ids
         snake_case					=    use_labels
         snake_case					=    vocab_size
         snake_case					=    hidden_size
         snake_case					=    num_hidden_layers
         snake_case					=    num_attention_heads
         snake_case					=    intermediate_size
         snake_case					=    hidden_act
         snake_case					=    hidden_dropout_prob
         snake_case					=    attention_probs_dropout_prob
         snake_case					=    max_position_embeddings
         snake_case					=    type_vocab_size
         snake_case					=    type_sequence_label_size
         snake_case					=    initializer_range
         snake_case					=    num_labels
         snake_case					=    num_choices
         snake_case					=    scope
         snake_case					=    vocab_size - 1
    def 		a_      (	self   ):
         snake_case					=    ids_tensor([self.batch_size, self.seq_length]      , self.vocab_size   )

         snake_case					=    None
         if self.use_input_mask:
              snake_case					=    random_attention_mask([self.batch_size, self.seq_length]   )

         snake_case					=    None
         if self.use_labels:
              snake_case					=    ids_tensor([self.batch_size, self.seq_length]      , self.num_labels   )

         snake_case					=    self.get_config()

         return config, input_ids, input_mask, token_labels
    def 		a_      (	self   ):
         return GPTNeoXConfig(
             vocab_size=self.vocab_size      , hidden_size=self.hidden_size      , num_hidden_layers=self.num_hidden_layers      , num_attention_heads=self.num_attention_heads      , intermediate_size=self.intermediate_size      , hidden_act=self.hidden_act      , hidden_dropout_prob=self.hidden_dropout_prob      , attention_probs_dropout_prob=self.attention_probs_dropout_prob      , max_position_embeddings=self.max_position_embeddings      , type_vocab_size=self.type_vocab_size      , is_decoder=__snake_case      , initializer_range=self.initializer_range      , pad_token_id=self.pad_token_id      , )
    def 		a_      (	self   ):
         snake_case			,  snake_case			,  snake_case			,  snake_case					=    self.prepare_config_and_inputs()

         snake_case					=    True

         return config, input_ids, input_mask, token_labels
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    GPTNeoXModel(config=__snake_case   )
         model.to(__snake_case   )
         model.eval()
         snake_case					=    model(__snake_case      , attention_mask=__snake_case   )
         snake_case					=    model(__snake_case   )
         self.parent.assertEqual(result.last_hidden_state.shape      , (self.batch_size, self.seq_length, self.hidden_size)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    True
         snake_case					=    GPTNeoXModel(__snake_case   )
         model.to(__snake_case   )
         model.eval()
         snake_case					=    model(__snake_case      , attention_mask=__snake_case   )
         self.parent.assertEqual(result.last_hidden_state.shape      , (self.batch_size, self.seq_length, self.hidden_size)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    GPTNeoXForCausalLM(config=__snake_case   )
         model.to(__snake_case   )
         model.eval()
         snake_case					=    model(__snake_case      , attention_mask=__snake_case      , labels=__snake_case   )
         self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.seq_length, self.vocab_size)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    self.num_labels
         snake_case					=    GPTNeoXForQuestionAnswering(__snake_case   )
         model.to(__snake_case   )
         model.eval()
         snake_case					=    model(__snake_case      , attention_mask=__snake_case   )
         self.parent.assertEqual(result.start_logits.shape      , (self.batch_size, self.seq_length)   )
         self.parent.assertEqual(result.end_logits.shape      , (self.batch_size, self.seq_length)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    self.num_labels
         snake_case					=    GPTNeoXForSequenceClassification(__snake_case   )
         model.to(__snake_case   )
         model.eval()
         snake_case					=    ids_tensor([self.batch_size]      , self.type_sequence_label_size   )
         snake_case					=    model(__snake_case      , attention_mask=__snake_case      , labels=__snake_case   )
         self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.num_labels)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    self.num_labels
         snake_case					=    GPTNeoXForTokenClassification(__snake_case   )
         model.to(__snake_case   )
         model.eval()
         snake_case					=    model(__snake_case      , attention_mask=__snake_case      , labels=__snake_case   )
         self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.seq_length, self.num_labels)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    True
         snake_case					=    GPTNeoXForCausalLM(config=__snake_case   )
         model.to(__snake_case   )
         model.eval()

         # first forward pass
         snake_case					=    model(__snake_case      , attention_mask=__snake_case      , use_cache=__snake_case   )
         snake_case					=    outputs.past_key_values

         # create hypothetical multiple next token and extent to next_input_ids
         snake_case					=    ids_tensor((self.batch_size, 3)      , config.vocab_size   )
         snake_case					=    ids_tensor((self.batch_size, 3)      , vocab_size=2   )

         # append to next input_ids and
         snake_case					=    torch.cat([input_ids, next_tokens]      , dim=-1   )
         snake_case					=    torch.cat([input_mask, next_mask]      , dim=-1   )

         snake_case					=    model(__snake_case      , attention_mask=__snake_case      , output_hidden_states=__snake_case   )
         snake_case					=    output_from_no_past['''hidden_states'''][0]
         snake_case					=    model(
             __snake_case      , attention_mask=__snake_case      , past_key_values=__snake_case      , output_hidden_states=__snake_case      , )['''hidden_states'''][0]

         # select random slice
         snake_case					=    ids_tensor((1,)      , output_from_past.shape[-1]   ).item()
         snake_case					=    output_from_no_past[:, -3:, random_slice_idx].detach()
         snake_case					=    output_from_past[:, :, random_slice_idx].detach()

         self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]   )

         # test that outputs are equal for slice
         self.parent.assertTrue(torch.allclose(__snake_case      , __snake_case      , atol=1E-3   )   )




    def 		a_      (	self   ):
         snake_case					=    self.prepare_config_and_inputs()
         snake_case			,  snake_case			,  snake_case			,  snake_case					=    config_and_inputs
         snake_case					=    {'''input_ids''': input_ids, '''attention_mask''': input_mask}
         return config, inputs_dict






@require_torch
class   A__ (   snake_case__	, snake_case__	, snake_case__	, unittest.TestCase     ):







    """simple docstring"""

    __magic_name__        =		(
        (
            GPTNeoXModel,
            GPTNeoXForCausalLM,
            GPTNeoXForQuestionAnswering,
            GPTNeoXForSequenceClassification,
            GPTNeoXForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    __magic_name__        =		(GPTNeoXForCausalLM,) if is_torch_available() else ()
    __magic_name__        =		(
        {
            'feature-extraction': GPTNeoXModel,
            'question-answering': GPTNeoXForQuestionAnswering,
            'text-classification': GPTNeoXForSequenceClassification,
            'text-generation': GPTNeoXForCausalLM,
            'token-classification': GPTNeoXForTokenClassification,
            'zero-shot': GPTNeoXForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    __magic_name__        =		False
    __magic_name__        =		False
    __magic_name__        =		False
    __magic_name__        =		False
    def 		a_      (	self   ):
         snake_case					=    GPTNeoXModelTester(self   )
         snake_case					=    ConfigTester(self      , config_class=__snake_case      , hidden_size=6_4      , num_attention_heads=8   )
    def 		a_      (	self   ):
         self.config_tester.run_common_tests()
    def 		a_      (	self   ):
         snake_case			,  snake_case			,  snake_case			,  snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(__snake_case      , __snake_case      , __snake_case   )
    def 		a_      (	self   ):
         snake_case			,  snake_case			,  snake_case			,  snake_case					=    self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_model_as_decoder(__snake_case      , __snake_case      , __snake_case   )
    def 		a_      (	self   ):
         # This regression test was failing with PyTorch < 1.3
         snake_case			,  snake_case			,  snake_case			,  snake_case					=    self.model_tester.prepare_config_and_inputs_for_decoder()

         snake_case					=    None

         self.model_tester.create_and_check_model_as_decoder(__snake_case      , __snake_case      , __snake_case   )
    def 		a_      (	self   ):
         snake_case			,  snake_case			,  snake_case			,  snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_decoder_model_past_large_inputs(__snake_case      , __snake_case      , __snake_case   )
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_causal_lm(*__snake_case   )
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*__snake_case   )
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*__snake_case   )
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*__snake_case   )
    @unittest.skip(reason='''Feed forward chunking is not implemented'''   )
    def 		a_      (	self   ):
         pass




    @parameterized.expand([('''linear''',), ('''dynamic''',)]   )
    def 		a_      (	self      , __snake_case   ):
         snake_case			,  snake_case					=    self.model_tester.prepare_config_and_inputs_for_common()
         snake_case					=    ids_tensor([1, 1_0]      , config.vocab_size   )
         snake_case					=    ids_tensor([1, int(config.max_position_embeddings * 1.5   )]      , config.vocab_size   )

         set_seed(4_2   )  # Fixed seed at init time so the two models get the same random weights
         snake_case					=    GPTNeoXModel(__snake_case   )
         original_model.to(__snake_case   )
         original_model.eval()
         snake_case					=    original_model(__snake_case   ).last_hidden_state
         snake_case					=    original_model(__snake_case   ).last_hidden_state

         set_seed(4_2   )  # Fixed seed at init time so the two models get the same random weights
         snake_case					=    {'''type''': scaling_type, '''factor''': 10.0}
         snake_case					=    GPTNeoXModel(__snake_case   )
         scaled_model.to(__snake_case   )
         scaled_model.eval()
         snake_case					=    scaled_model(__snake_case   ).last_hidden_state
         snake_case					=    scaled_model(__snake_case   ).last_hidden_state

         # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
         # maximum sequence length, so the outputs for the short input should match.
         if scaling_type == "dynamic":
              self.assertTrue(torch.allclose(__snake_case      , __snake_case      , atol=1E-5   )   )
         else:
              self.assertFalse(torch.allclose(__snake_case      , __snake_case      , atol=1E-5   )   )

         # The output should be different for long inputs
         self.assertFalse(torch.allclose(__snake_case      , __snake_case      , atol=1E-5   )   )






@require_torch
class   A__ (   unittest.TestCase     ):







    """simple docstring"""

    @slow
    def 		a_      (	self   ):
         snake_case					=    AutoTokenizer.from_pretrained('''EleutherAI/pythia-410m-deduped'''   )
         for checkpointing in [True, False]:
              snake_case					=    GPTNeoXForCausalLM.from_pretrained('''EleutherAI/pythia-410m-deduped'''   )

              if checkpointing:
                   model.gradient_checkpointing_enable()
              else:
                   model.gradient_checkpointing_disable()
              model.to(__snake_case   )

              snake_case					=    tokenizer('''My favorite food is'''      , return_tensors='''pt'''   ).to(__snake_case   )
              # The hub repo. is updated on 2023-04-04, resulting in poor outputs.
              # See: https://github.com/huggingface/transformers/pull/24193
              snake_case					=    '''My favorite food is a good old-fashioned, old-fashioned, old-fashioned.\n\nI\'m not sure'''

              snake_case					=    model.generate(**__snake_case      , do_sample=__snake_case      , max_new_tokens=2_0   )
              snake_case					=    tokenizer.batch_decode(__snake_case   )[0]

              self.assertEqual(__snake_case      , __snake_case   )




