



from copy import deepcopy

import torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader

from accelerate.accelerator import Accelerator
from accelerate.state import GradientState
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import DistributedType, is_torch_version, set_seed
def 						UpperCamelCase      (  snake_case__			:       Optional[int]			,	snake_case__			:       Dict			,	snake_case__			:       Union[str, Any]			,	snake_case__			:       Optional[int]      )      ->	Dict:
   for param, grad_param in zip(model_a.parameters()			,	model_b.parameters()      ):
      if not param.requires_grad:
         continue
      if not did_step:
         # Grads should not be in sync
         assert (
             torch.allclose(param.grad			,	grad_param.grad      ) is False
         ), F"""Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})"""
      else:
         # Grads should be in sync
         assert (
             torch.allclose(param.grad			,	grad_param.grad      ) is True
         ), F"""Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})"""
def 						UpperCamelCase      (  snake_case__			:       Optional[Any]			,	snake_case__			:       Tuple			,	snake_case__			:       Dict			,	snake_case__			:       Tuple			,	snake_case__			:       Union[str, Any]=True      )      ->	Tuple:
   model.train()
   UpperCamelCase       :    Union[str, Any]		      =      model(snake_case__      )
   UpperCamelCase       :    Optional[Any]		      =      F.mse_loss(snake_case__			,	target.to(output.device      )      )
   if not do_backward:
      loss /= accelerator.gradient_accumulation_steps
      loss.backward()
   else:
      accelerator.backward(snake_case__      )
def 						UpperCamelCase      (  snake_case__			:       Tuple			,	snake_case__			:       Any=False      )      ->	List[str]:
   set_seed(42      )
   UpperCamelCase       :    Tuple		      =      RegressionModel()
   UpperCamelCase       :    Dict		      =      deepcopy(snake_case__      )
   UpperCamelCase       :    Tuple		      =      RegressionDataset(length=80      )
   UpperCamelCase       :    Dict		      =      DataLoader(snake_case__			,	batch_size=16      )
   model.to(accelerator.device      )
   if sched:
      UpperCamelCase       :    str		      =      AdamW(params=model.parameters()			,	lr=1E-3      )
      UpperCamelCase       :    Dict		      =      AdamW(params=ddp_model.parameters()			,	lr=1E-3      )
      UpperCamelCase       :    Optional[int]		      =      LambdaLR(snake_case__			,	lr_lambda=lambda snake_case__      : epoch**0.65      )
      UpperCamelCase       :    int		      =      LambdaLR(snake_case__			,	lr_lambda=lambda snake_case__      : epoch**0.65      )
   # Make a copy of `model`
   if sched:
      UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    Any		      =      accelerator.prepare(snake_case__			,	snake_case__			,	snake_case__			,	snake_case__      )
   else:
      UpperCamelCase      ,	UpperCamelCase       :    List[str]		      =      accelerator.prepare(snake_case__			,	snake_case__      )
   if sched:
      return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
   return model, ddp_model, dataloader
def 						UpperCamelCase      (  snake_case__			:       str      )      ->	Optional[int]:
   # Test when on a single CPU or GPU that the context manager does nothing
   UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    Optional[Any]		      =      get_training_setup(snake_case__      )
   # Use a single batch
   UpperCamelCase      ,	UpperCamelCase       :    Any		      =      next(iter(snake_case__      )      ).values()
   for iteration in range(3      ):
      # Gather the distributed inputs and targs for the base model
      UpperCamelCase      ,	UpperCamelCase       :    Tuple		      =      accelerator.gather((ddp_input, ddp_target)      )
      UpperCamelCase      ,	UpperCamelCase       :    List[str]		      =      input.to(accelerator.device      ), target.to(accelerator.device      )
      # Perform our initial ground truth step in non "DDP"
      step_model(snake_case__			,	snake_case__			,	snake_case__			,	snake_case__      )
      # Do "gradient accumulation" (noop)
      if iteration % 2 == 0:
         # Accumulate grads locally
         with accelerator.no_sync(snake_case__      ):
            step_model(snake_case__			,	snake_case__			,	snake_case__			,	snake_case__      )
      else:
         # Sync grads
         step_model(snake_case__			,	snake_case__			,	snake_case__			,	snake_case__      )

      # Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
      check_model_parameters(snake_case__			,	snake_case__			,	snake_case__			,	snake_case__      )
      for param, ddp_param in zip(model.parameters()			,	ddp_model.parameters()      ):
         if not param.requires_grad:
            continue
         assert torch.allclose(
             param.grad			,	ddp_param.grad      ), F"""Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"""

      # Shuffle ddp_input on each iteration
      torch.manual_seed(1337 + iteration      )
      UpperCamelCase       :    str		      =      ddp_input[torch.randperm(len(snake_case__      )      )]
def 						UpperCamelCase      (  snake_case__			:       int      )      ->	Tuple:
   # Test on distributed setup that context manager behaves properly
   UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    Any		      =      get_training_setup(snake_case__      )
   # Use a single batch
   UpperCamelCase      ,	UpperCamelCase       :    Dict		      =      next(iter(snake_case__      )      ).values()
   for iteration in range(3      ):
      # Gather the distributed inputs and targs for the base model
      UpperCamelCase      ,	UpperCamelCase       :    int		      =      accelerator.gather((ddp_input, ddp_target)      )
      UpperCamelCase      ,	UpperCamelCase       :    Optional[int]		      =      input.to(accelerator.device      ), target.to(accelerator.device      )
      # Perform our initial ground truth step in non "DDP"
      step_model(snake_case__			,	snake_case__			,	snake_case__			,	snake_case__      )
      # Do "gradient accumulation" (noop)
      if iteration % 2 == 0:
         # Accumulate grads locally
         with accelerator.no_sync(snake_case__      ):
            step_model(snake_case__			,	snake_case__			,	snake_case__			,	snake_case__      )
      else:
         # Sync grads
         step_model(snake_case__			,	snake_case__			,	snake_case__			,	snake_case__      )

      # DDP model and model should only be in sync when not (iteration % 2 == 0)
      for param, ddp_param in zip(model.parameters()			,	ddp_model.parameters()      ):
         if not param.requires_grad:
            continue
         if iteration % 2 == 0:
            # Grads should not be in sync
            assert (
                torch.allclose(param.grad			,	ddp_param.grad      ) is False
            ), F"""Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"""
         else:
            # Grads should be in sync
            assert (
                torch.allclose(param.grad			,	ddp_param.grad      ) is True
            ), F"""Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"""

        # Shuffle ddp_input on each iteration
      torch.manual_seed(1337 + iteration      )
      UpperCamelCase       :    str		      =      ddp_input[torch.randperm(len(snake_case__      )      )]
def 						UpperCamelCase      (  snake_case__			:       Optional[int]=False			,	snake_case__			:       Dict=False      )      ->	Tuple:
   UpperCamelCase       :    Tuple		      =      Accelerator(
       split_batches=snake_case__			,	dispatch_batches=snake_case__			,	gradient_accumulation_steps=2      )
   # Test that context manager behaves properly
   UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    Union[str, Any]		      =      get_training_setup(snake_case__      )
   for iteration, batch in enumerate(snake_case__      ):
      UpperCamelCase      ,	UpperCamelCase       :    Optional[int]		      =      batch.values()
      # Gather the distributed inputs and targs for the base model
      UpperCamelCase      ,	UpperCamelCase       :    str		      =      accelerator.gather((ddp_input, ddp_target)      )
      UpperCamelCase      ,	UpperCamelCase       :    str		      =      input.to(accelerator.device      ), target.to(accelerator.device      )
      # Perform our initial ground truth step in non "DDP"
      step_model(snake_case__			,	snake_case__			,	snake_case__			,	snake_case__			,	snake_case__      )
      # Do "gradient accumulation" (noop)
      with accelerator.accumulate(snake_case__      ):
         step_model(snake_case__			,	snake_case__			,	snake_case__			,	snake_case__      )

      # DDP model and model should only be in sync when not (iteration % 2 == 0)
      for param, ddp_param in zip(model.parameters()			,	ddp_model.parameters()      ):
         if not param.requires_grad:
            continue
         if ((iteration + 1) % 2 == 0) or (iteration == len(snake_case__      ) - 1):
            # Grads should be in sync
            assert (
                torch.allclose(param.grad			,	ddp_param.grad      ) is True
            ), F"""Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"""
         else:
            # Grads should not be in sync
            assert (
                torch.allclose(param.grad			,	ddp_param.grad      ) is False
            ), F"""Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"""

        # Shuffle ddp_input on each iteration
      torch.manual_seed(1337 + iteration      )
      UpperCamelCase       :    Any		      =      ddp_input[torch.randperm(len(snake_case__      )      )]
   GradientState._reset_state()
def 						UpperCamelCase      (  snake_case__			:       Optional[int]=False			,	snake_case__			:       List[Any]=False      )      ->	Union[str, Any]:
   UpperCamelCase       :    List[Any]		      =      Accelerator(
       split_batches=snake_case__			,	dispatch_batches=snake_case__			,	gradient_accumulation_steps=2      )
   # Test that context manager behaves properly
   UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    str		      =      get_training_setup(snake_case__			,	snake_case__      )
   for iteration, batch in enumerate(snake_case__      ):
      UpperCamelCase      ,	UpperCamelCase       :    List[str]		      =      batch.values()
      # Gather the distributed inputs and targs for the base model
      UpperCamelCase      ,	UpperCamelCase       :    str		      =      accelerator.gather((ddp_input, ddp_target)      )
      UpperCamelCase      ,	UpperCamelCase       :    Dict		      =      input.to(accelerator.device      ), target.to(accelerator.device      )
      # Perform our initial ground truth step in non "DDP"
      model.train()
      ddp_model.train()
      step_model(snake_case__			,	snake_case__			,	snake_case__			,	snake_case__			,	snake_case__      )
      opt.step()

      if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(snake_case__      )):
         if split_batches:
            sched.step()
         else:
            for _ in range(accelerator.num_processes      ):
               sched.step()
      opt.zero_grad()
      # Perform gradient accumulation under wrapper
      with accelerator.accumulate(snake_case__      ):
         step_model(snake_case__			,	snake_case__			,	snake_case__			,	snake_case__      )
         ddp_opt.step()
         ddp_sched.step()
         ddp_opt.zero_grad()

      # Learning rates should be the same
      assert (
          opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
      ), F"""Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]["lr"]}\nDDP opt: {ddp_opt.param_groups[0]["lr"]}\n"""
      UpperCamelCase       :    int		      =      (((iteration + 1) % 2) == 0) or ((iteration + 1) == len(snake_case__      ))
      if accelerator.num_processes > 1:
         check_model_parameters(snake_case__			,	snake_case__			,	snake_case__			,	snake_case__      )
      # Shuffle ddp_input on each iteration
      torch.manual_seed(1337 + iteration      )
   GradientState._reset_state()
def 						UpperCamelCase      (  )      ->	Optional[int]:
   UpperCamelCase       :    Any		      =      Accelerator()

   UpperCamelCase       :    int		      =      RegressionDataset(length=80      )
   UpperCamelCase       :    Dict		      =      DataLoader(snake_case__			,	batch_size=16      )
   UpperCamelCase       :    Optional[Any]		      =      RegressionDataset(length=96      )
   UpperCamelCase       :    List[str]		      =      DataLoader(snake_case__			,	batch_size=16      )
   UpperCamelCase      ,	UpperCamelCase       :    List[Any]		      =      accelerator.prepare(snake_case__			,	snake_case__      )
   assert accelerator.gradient_state.active_dataloader is None
   for iteration, _ in enumerate(snake_case__      ):
      assert id(accelerator.gradient_state.active_dataloader      ) == id(snake_case__      )
      if iteration < len(snake_case__      ) - 1:
         assert not accelerator.gradient_state.end_of_dataloader
         if iteration == 1:
            for batch_num, _ in enumerate(snake_case__      ):
               assert id(accelerator.gradient_state.active_dataloader      ) == id(snake_case__      )
               if batch_num < len(snake_case__      ) - 1:
                  assert not accelerator.gradient_state.end_of_dataloader
               else:
                  assert accelerator.gradient_state.end_of_dataloader
      else:
         assert accelerator.gradient_state.end_of_dataloader
   assert accelerator.gradient_state.active_dataloader is None
def 						UpperCamelCase      (  )      ->	int:
   UpperCamelCase       :    str		      =      Accelerator()
   UpperCamelCase       :    List[Any]		      =      accelerator.state
   if state.local_process_index == 0:
      print('**Test `accumulate` gradient accumulation with dataloader break**'      )
   test_dataloader_break()
   if state.distributed_type == DistributedType.NO:
      if state.local_process_index == 0:
         print('**Test NOOP `no_sync` context manager**'      )
      test_noop_sync(snake_case__      )
   if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
      if state.local_process_index == 0:
         print('**Test Distributed `no_sync` context manager**'      )
      test_distributed_sync(snake_case__      )
   if state.distributed_type == DistributedType.MULTI_GPU:
      for split_batch in [True, False]:
         for dispatch_batches in [True, False]:
            if state.local_process_index == 0:
               print(
                   '**Test `accumulate` gradient accumulation, '			,	F"""`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**"""			,	)
            test_gradient_accumulation(snake_case__			,	snake_case__      )

    # Currently will break on torch 2.0 +, need to investigate why
   if is_torch_version('<'			,	'2.0'      ) or state.distributed_type == DistributedType.NO:
      if state.local_process_index == 0:
         print(
             '**Test `accumulate` gradient accumulation with optimizer and scheduler, '			,	'`split_batches=False`, `dispatch_batches=False`**'			,	)
      test_gradient_accumulation_with_opt_and_scheduler()
      if state.distributed_type == DistributedType.MULTI_GPU:
         for split_batch in [True, False]:
            for dispatch_batches in [True, False]:
               if not split_batch and not dispatch_batches:
                  continue
               if state.local_process_index == 0:
                  print(
                      '**Test `accumulate` gradient accumulation with optimizer and scheduler, '			,	F"""`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**"""			,	)
               test_gradient_accumulation_with_opt_and_scheduler(snake_case__			,	snake_case__      )
def 						UpperCamelCase      (  snake_case__			:       Optional[Any]      )      ->	Optional[int]:
   # For xla_spawn (TPUs)
   main()


if __name__ == "__main__":
    main()






'''simple docstring'''




import subprocess
import sys

from transformers import BertConfig, BertModel, BertTokenizer, pipeline
from transformers.testing_utils import TestCasePlus, require_torch



class 							lowercase_						(		a			):

  '''simple docstring'''





  @require_torch
  def     snake_case_      (  self  )				->					Optional[Any]:


    """simple docstring"""
    # this test is a bit tricky since TRANSFORMERS_OFFLINE can only be changed before
    # `transformers` is loaded, and it's too late for inside pytest - so we are changing it
    # while running an external program

    # python one-liner segments

    # this must be loaded before socket.socket is monkey-patched
    UpperCAmelCase    				=						'\nfrom transformers import BertConfig, BertModel, BertTokenizer, pipeline\n        '

    UpperCAmelCase    				=						'\nmname = "hf-internal-testing/tiny-random-bert"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nBertTokenizer.from_pretrained(mname)\npipe = pipeline(task="fill-mask", model=mname)\nprint("success")\n        '

    UpperCAmelCase    				=						'\nimport socket\ndef offline_socket(*args, **kwargs): raise RuntimeError("Offline mode is enabled, we shouldn\'t access internet")\nsocket.socket = offline_socket\n        '

    # Force fetching the files so that we can use the cache
    UpperCAmelCase    				=						'hf-internal-testing/tiny-random-bert'
    BertConfig.from_pretrained(a_  )
    BertModel.from_pretrained(a_  )
    BertTokenizer.from_pretrained(a_  )
    pipeline(task='fill-mask' ,			model=a_  )

    # baseline - just load from_pretrained with normal network
    UpperCAmelCase    				=						[sys.executable, '-c', '\n'.join([load, run, mock]  )]

    # should succeed
    UpperCAmelCase    				=						self.get_env()
    # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
    UpperCAmelCase    				=						'1'
    UpperCAmelCase    				=						subprocess.run(a_ ,			env=a_ ,			check=a_ ,			capture_output=a_  )
    self.assertEqual(result.returncode ,			0 ,			result.stderr  )
    self.assertIn('success' ,			result.stdout.decode()  )
  @require_torch
  def     snake_case_      (  self  )				->					Union[str, Any]:


    """simple docstring"""
    # python one-liner segments
    # this must be loaded before socket.socket is monkey-patched
    UpperCAmelCase    				=						'\nfrom transformers import BertConfig, BertModel, BertTokenizer, pipeline\n        '

    UpperCAmelCase    				=						'\nmname = "hf-internal-testing/tiny-random-bert"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nBertTokenizer.from_pretrained(mname)\npipe = pipeline(task="fill-mask", model=mname)\nprint("success")\n        '

    UpperCAmelCase    				=						'\nimport socket\ndef offline_socket(*args, **kwargs): raise socket.error("Faking flaky internet")\nsocket.socket = offline_socket\n        '

    # Force fetching the files so that we can use the cache
    UpperCAmelCase    				=						'hf-internal-testing/tiny-random-bert'
    BertConfig.from_pretrained(a_  )
    BertModel.from_pretrained(a_  )
    BertTokenizer.from_pretrained(a_  )
    pipeline(task='fill-mask' ,			model=a_  )

    # baseline - just load from_pretrained with normal network
    UpperCAmelCase    				=						[sys.executable, '-c', '\n'.join([load, run, mock]  )]

    # should succeed
    UpperCAmelCase    				=						self.get_env()
    UpperCAmelCase    				=						subprocess.run(a_ ,			env=a_ ,			check=a_ ,			capture_output=a_  )
    self.assertEqual(result.returncode ,			0 ,			result.stderr  )
    self.assertIn('success' ,			result.stdout.decode()  )
  @require_torch
  def     snake_case_      (  self  )				->					List[str]:


    """simple docstring"""
    # this test is a bit tricky since TRANSFORMERS_OFFLINE can only be changed before
    # `transformers` is loaded, and it's too late for inside pytest - so we are changing it
    # while running an external program

    # python one-liner segments

    # this must be loaded before socket.socket is monkey-patched
    UpperCAmelCase    				=						'\nfrom transformers import BertConfig, BertModel, BertTokenizer\n        '

    UpperCAmelCase    				=						'\nmname = "hf-internal-testing/tiny-random-bert-sharded"\nBertConfig.from_pretrained(mname)\nBertModel.from_pretrained(mname)\nprint("success")\n        '

    UpperCAmelCase    				=						'\nimport socket\ndef offline_socket(*args, **kwargs): raise ValueError("Offline mode is enabled")\nsocket.socket = offline_socket\n        '

    # baseline - just load from_pretrained with normal network
    UpperCAmelCase    				=						[sys.executable, '-c', '\n'.join([load, run]  )]

    # should succeed
    UpperCAmelCase    				=						self.get_env()
    UpperCAmelCase    				=						subprocess.run(a_ ,			env=a_ ,			check=a_ ,			capture_output=a_  )
    self.assertEqual(result.returncode ,			0 ,			result.stderr  )
    self.assertIn('success' ,			result.stdout.decode()  )

    # next emulate no network
    UpperCAmelCase    				=						[sys.executable, '-c', '\n'.join([load, mock, run]  )]

    # Doesn't fail anymore since the model is in the cache due to other tests, so commenting this.
    # env["TRANSFORMERS_OFFLINE"] = "0"
    # result = subprocess.run(cmd, env=env, check=False, capture_output=True)
    # self.assertEqual(result.returncode, 1, result.stderr)

    # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
    UpperCAmelCase    				=						'1'
    UpperCAmelCase    				=						subprocess.run(a_ ,			env=a_ ,			check=a_ ,			capture_output=a_  )
    self.assertEqual(result.returncode ,			0 ,			result.stderr  )
    self.assertIn('success' ,			result.stdout.decode()  )
  @require_torch
  def     snake_case_      (  self  )				->					Optional[Any]:


    """simple docstring"""
    UpperCAmelCase    				=						'\nfrom transformers import pipeline\n        '
    UpperCAmelCase    				=						'\nmname = "hf-internal-testing/tiny-random-bert"\npipe = pipeline(model=mname)\n        '

    UpperCAmelCase    				=						'\nimport socket\ndef offline_socket(*args, **kwargs): raise socket.error("Offline mode is enabled")\nsocket.socket = offline_socket\n        '
    UpperCAmelCase    				=						self.get_env()
    UpperCAmelCase    				=						'1'
    UpperCAmelCase    				=						[sys.executable, '-c', '\n'.join([load, mock, run]  )]
    UpperCAmelCase    				=						subprocess.run(a_ ,			env=a_ ,			check=a_ ,			capture_output=a_  )
    self.assertEqual(result.returncode ,			1 ,			result.stderr  )
    self.assertIn(
        'You cannot infer task automatically within `pipeline` when using offline mode' ,			result.stderr.decode().replace('\n' ,			''  ) ,			)





  @require_torch
  def     snake_case_      (  self  )				->					str:


    """simple docstring"""
    UpperCAmelCase    				=						'\nfrom transformers import AutoModel\n        '
    UpperCAmelCase    				=						'\nmname = "hf-internal-testing/test_dynamic_model"\nAutoModel.from_pretrained(mname, trust_remote_code=True)\nprint("success")\n        '

    # baseline - just load from_pretrained with normal network
    UpperCAmelCase    				=						[sys.executable, '-c', '\n'.join([load, run]  )]

    # should succeed
    UpperCAmelCase    				=						self.get_env()
    UpperCAmelCase    				=						subprocess.run(a_ ,			env=a_ ,			check=a_ ,			capture_output=a_  )
    self.assertEqual(result.returncode ,			0 ,			result.stderr  )
    self.assertIn('success' ,			result.stdout.decode()  )

    # should succeed as TRANSFORMERS_OFFLINE=1 tells it to use local files
    UpperCAmelCase    				=						'1'
    UpperCAmelCase    				=						subprocess.run(a_ ,			env=a_ ,			check=a_ ,			capture_output=a_  )
    self.assertEqual(result.returncode ,			0 ,			result.stderr  )
    self.assertIn('success' ,			result.stdout.decode()  )


