




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 							_A       (						lowerCAmelCase_					: List[str]       ,    lowerCAmelCase_					: str       ,    lowerCAmelCase_					: Optional[Any]       ,    lowerCAmelCase_					: Optional[Any]      ):
     """simple docstring"""

     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 							_A       (						lowerCAmelCase_					: Any       ,    lowerCAmelCase_					: Union[str, Any]       ,    lowerCAmelCase_					: List[Any]       ,    lowerCAmelCase_					: List[str]       ,    lowerCAmelCase_					: Tuple=True      ):
     """simple docstring"""

     model.train()
     lowerCAmelCase__          =			model(lowerCAmelCase_      )
     lowerCAmelCase__          =			F.mse_loss(lowerCAmelCase_       ,    target.to(output.device      )      )
     if not do_backward:
          loss /= accelerator.gradient_accumulation_steps
          loss.backward()
     else:
          accelerator.backward(lowerCAmelCase_      )



def 							_A       (						lowerCAmelCase_					: Optional[int]       ,    lowerCAmelCase_					: Optional[Any]=False      ):
     """simple docstring"""

     set_seed(42      )
     lowerCAmelCase__          =			RegressionModel()
     lowerCAmelCase__          =			deepcopy(lowerCAmelCase_      )
     lowerCAmelCase__          =			RegressionDataset(length=80      )
     lowerCAmelCase__          =			DataLoader(lowerCAmelCase_       ,    batch_size=16      )
     model.to(accelerator.device      )
     if sched:
          lowerCAmelCase__          =			AdamW(params=model.parameters()       ,    lr=1E-3      )
          lowerCAmelCase__          =			AdamW(params=ddp_model.parameters()       ,    lr=1E-3      )
          lowerCAmelCase__          =			LambdaLR(lowerCAmelCase_       ,    lr_lambda=lambda lowerCAmelCase_      : epoch**0.65      )
          lowerCAmelCase__          =			LambdaLR(lowerCAmelCase_       ,    lr_lambda=lambda lowerCAmelCase_      : epoch**0.65      )
     # Make a copy of `model`
     if sched:
          lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__          =			accelerator.prepare(lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_      )
     else:
          lowerCAmelCase__     ,		lowerCAmelCase__          =			accelerator.prepare(lowerCAmelCase_       ,    lowerCAmelCase_      )
     if sched:
          return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
     return model, ddp_model, dataloader



def 							_A       (						lowerCAmelCase_					: Tuple      ):
     """simple docstring"""

     lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__          =			get_training_setup(lowerCAmelCase_      )
     # Use a single batch
     lowerCAmelCase__     ,		lowerCAmelCase__          =			next(iter(lowerCAmelCase_      )      ).values()
     for iteration in range(3      ):
          # Gather the distributed inputs and targs for the base model
          lowerCAmelCase__     ,		lowerCAmelCase__          =			accelerator.gather((ddp_input, ddp_target)      )
          lowerCAmelCase__     ,		lowerCAmelCase__          =			input.to(accelerator.device      ), target.to(accelerator.device      )
          # Perform our initial ground truth step in non "DDP"
          step_model(lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_      )
          # Do "gradient accumulation" (noop)
          if iteration % 2 == 0:
               # Accumulate grads locally
               with accelerator.no_sync(lowerCAmelCase_      ):
                    step_model(lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_      )
          else:
               # Sync grads
               step_model(lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_      )

          # Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
          check_model_parameters(lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_      )
          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      )
          lowerCAmelCase__          =			ddp_input[torch.randperm(len(lowerCAmelCase_      )      )]



def 							_A       (						lowerCAmelCase_					: Tuple      ):
     """simple docstring"""

     lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__          =			get_training_setup(lowerCAmelCase_      )
     # Use a single batch
     lowerCAmelCase__     ,		lowerCAmelCase__          =			next(iter(lowerCAmelCase_      )      ).values()
     for iteration in range(3      ):
          # Gather the distributed inputs and targs for the base model
          lowerCAmelCase__     ,		lowerCAmelCase__          =			accelerator.gather((ddp_input, ddp_target)      )
          lowerCAmelCase__     ,		lowerCAmelCase__          =			input.to(accelerator.device      ), target.to(accelerator.device      )
          # Perform our initial ground truth step in non "DDP"
          step_model(lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_      )
          # Do "gradient accumulation" (noop)
          if iteration % 2 == 0:
               # Accumulate grads locally
               with accelerator.no_sync(lowerCAmelCase_      ):
                    step_model(lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_      )
          else:
               # Sync grads
               step_model(lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_      )

          # 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      )
          lowerCAmelCase__          =			ddp_input[torch.randperm(len(lowerCAmelCase_      )      )]



def 							_A       (						lowerCAmelCase_					: int=False       ,    lowerCAmelCase_					: str=False      ):
     """simple docstring"""

     lowerCAmelCase__          =			Accelerator(
         split_batches=lowerCAmelCase_       ,    dispatch_batches=lowerCAmelCase_       ,    gradient_accumulation_steps=2      )
     # Test that context manager behaves properly
     lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__          =			get_training_setup(lowerCAmelCase_      )
     for iteration, batch in enumerate(lowerCAmelCase_      ):
          lowerCAmelCase__     ,		lowerCAmelCase__          =			batch.values()
          # Gather the distributed inputs and targs for the base model
          lowerCAmelCase__     ,		lowerCAmelCase__          =			accelerator.gather((ddp_input, ddp_target)      )
          lowerCAmelCase__     ,		lowerCAmelCase__          =			input.to(accelerator.device      ), target.to(accelerator.device      )
          # Perform our initial ground truth step in non "DDP"
          step_model(lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_      )
          # Do "gradient accumulation" (noop)
          with accelerator.accumulate(lowerCAmelCase_      ):
               step_model(lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_      )

          # 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(lowerCAmelCase_      ) - 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      )
          lowerCAmelCase__          =			ddp_input[torch.randperm(len(lowerCAmelCase_      )      )]
     GradientState._reset_state()



def 							_A       (						lowerCAmelCase_					: int=False       ,    lowerCAmelCase_					: List[Any]=False      ):
     """simple docstring"""

     lowerCAmelCase__          =			Accelerator(
         split_batches=lowerCAmelCase_       ,    dispatch_batches=lowerCAmelCase_       ,    gradient_accumulation_steps=2      )
     # Test that context manager behaves properly
     lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__          =			get_training_setup(lowerCAmelCase_       ,    lowerCAmelCase_      )
     for iteration, batch in enumerate(lowerCAmelCase_      ):
          lowerCAmelCase__     ,		lowerCAmelCase__          =			batch.values()
          # Gather the distributed inputs and targs for the base model
          lowerCAmelCase__     ,		lowerCAmelCase__          =			accelerator.gather((ddp_input, ddp_target)      )
          lowerCAmelCase__     ,		lowerCAmelCase__          =			input.to(accelerator.device      ), target.to(accelerator.device      )
          # Perform our initial ground truth step in non "DDP"
          model.train()
          ddp_model.train()
          step_model(lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_      )
          opt.step()

          if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(lowerCAmelCase_      )):
               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(lowerCAmelCase_      ):
               step_model(lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_      )
               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'
          lowerCAmelCase__          =			(((iteration + 1) % 2) == 0) or ((iteration + 1) == len(lowerCAmelCase_      ))
          if accelerator.num_processes > 1:
               check_model_parameters(lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_       ,    lowerCAmelCase_      )
          # Shuffle ddp_input on each iteration
          torch.manual_seed(1337 + iteration      )
     GradientState._reset_state()



def 							_A       (						):
     """simple docstring"""

     lowerCAmelCase__          =			Accelerator()

     lowerCAmelCase__          =			RegressionDataset(length=80      )
     lowerCAmelCase__          =			DataLoader(lowerCAmelCase_       ,    batch_size=16      )
     lowerCAmelCase__          =			RegressionDataset(length=96      )
     lowerCAmelCase__          =			DataLoader(lowerCAmelCase_       ,    batch_size=16      )
     lowerCAmelCase__     ,		lowerCAmelCase__          =			accelerator.prepare(lowerCAmelCase_       ,    lowerCAmelCase_      )
     assert accelerator.gradient_state.active_dataloader is None
     for iteration, _ in enumerate(lowerCAmelCase_      ):
          assert id(accelerator.gradient_state.active_dataloader      ) == id(lowerCAmelCase_      )
          if iteration < len(lowerCAmelCase_      ) - 1:
               assert not accelerator.gradient_state.end_of_dataloader
               if iteration == 1:
                    for batch_num, _ in enumerate(lowerCAmelCase_      ):
                         assert id(accelerator.gradient_state.active_dataloader      ) == id(lowerCAmelCase_      )
                         if batch_num < len(lowerCAmelCase_      ) - 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 							_A       (						):
     """simple docstring"""

     lowerCAmelCase__          =			Accelerator()
     lowerCAmelCase__          =			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(lowerCAmelCase_      )
     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(lowerCAmelCase_      )
     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(lowerCAmelCase_       ,    lowerCAmelCase_      )

    # 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(lowerCAmelCase_       ,    lowerCAmelCase_      )



def 							_A       (						lowerCAmelCase_					: Union[str, Any]      ):
     """simple docstring"""

     main()


if __name__ == "__main__":
     main()






import time
from contextlib import contextmanager
from pathlib import Path

import pytest
import requests
from huggingface_hub.hf_api import HfApi, HfFolder


UpperCamelCase								=					'__DUMMY_TRANSFORMERS_USER__'
UpperCamelCase								=					'Dummy User'
UpperCamelCase								=					'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'

UpperCamelCase								=					'https://hub-ci.huggingface.co'
UpperCamelCase								=					CI_HUB_ENDPOINT + '/datasets/{repo_id}/resolve/{revision}/{path}'
UpperCamelCase								=					CI_HUB_ENDPOINT + '/{repo_id}/resolve/{revision}/{filename}'
UpperCamelCase								=					Path('~/.huggingface/hub_ci_token').expanduser()



@pytest.fixture
def 							_A       (						lowerCAmelCase_					: Dict      ):
     """simple docstring"""

     monkeypatch.setattr(
         "huggingface_hub.file_download.HUGGINGFACE_CO_URL_TEMPLATE"       ,    lowerCAmelCase_      )



@pytest.fixture
def 							_A       (						lowerCAmelCase_					: int      ):
     """simple docstring"""

     monkeypatch.setattr("datasets.config.HF_ENDPOINT"       ,    lowerCAmelCase_      )
     monkeypatch.setattr("datasets.config.HUB_DATASETS_URL"       ,    lowerCAmelCase_      )



@pytest.fixture
def 							_A       (						lowerCAmelCase_					: str      ):
     """simple docstring"""

     monkeypatch.setattr("huggingface_hub.hf_api.HfFolder.path_token"       ,    lowerCAmelCase_      )



@pytest.fixture
def 							_A       (						lowerCAmelCase_					: int       ,    lowerCAmelCase_					: List[Any]      ):
     """simple docstring"""

     HfFolder.save_token(lowerCAmelCase_      )
     yield
     HfFolder.delete_token()



@pytest.fixture(scope="session"      )
def 							_A       (						):
     """simple docstring"""

     return HfApi(endpoint=lowerCAmelCase_      )



@pytest.fixture(scope="session"      )
def 							_A       (						lowerCAmelCase_					: HfApi      ):
     """simple docstring"""

     lowerCAmelCase__          =			HfFolder.get_token()
     HfFolder.save_token(lowerCAmelCase_      )
     yield CI_HUB_USER_TOKEN
     if previous_token is not None:
          HfFolder.save_token(lowerCAmelCase_      )



@pytest.fixture
def 							_A       (						lowerCAmelCase_					: Tuple      ):
     """simple docstring"""

     def _cleanup_repo(lowerCAmelCase_					: Optional[Any]      ):
          hf_api.delete_repo(lowerCAmelCase_       ,    token=lowerCAmelCase_       ,    repo_type="dataset"      )

     return _cleanup_repo



@pytest.fixture
def 							_A       (						lowerCAmelCase_					: Union[str, Any]      ):
     """simple docstring"""

     @contextmanager
     def _temporary_repo(lowerCAmelCase_					: str      ):
          try:
               yield repo_id
          finally:
               cleanup_repo(lowerCAmelCase_      )

     return _temporary_repo



@pytest.fixture(scope="session"      )
def 							_A       (						lowerCAmelCase_					: HfApi       ,    lowerCAmelCase_					: Union[str, Any]       ,    lowerCAmelCase_					: Any      ):
     """simple docstring"""

     lowerCAmelCase__          =			F'repo_txt_data-{int(time.time() * 1_0E3      )}'
     lowerCAmelCase__          =			F'{CI_HUB_USER}/{repo_name}'
     hf_api.create_repo(lowerCAmelCase_       ,    token=lowerCAmelCase_       ,    repo_type="dataset"       ,    private=lowerCAmelCase_      )
     hf_api.upload_file(
         token=lowerCAmelCase_       ,    path_or_fileobj=str(lowerCAmelCase_      )       ,    path_in_repo="data/text_data.txt"       ,    repo_id=lowerCAmelCase_       ,    repo_type="dataset"       ,    )
     yield repo_id
     try:
          hf_api.delete_repo(lowerCAmelCase_       ,    token=lowerCAmelCase_       ,    repo_type="dataset"      )
     except (requests.exceptions.HTTPError, ValueError):  # catch http error and token invalid error
          pass



@pytest.fixture()
def 							_A       (						lowerCAmelCase_					: Optional[Any]       ,    lowerCAmelCase_					: Optional[Any]       ,    lowerCAmelCase_					: int      ):
     """simple docstring"""

     return hf_private_dataset_repo_txt_data_



@pytest.fixture(scope="session"      )
def 							_A       (						lowerCAmelCase_					: HfApi       ,    lowerCAmelCase_					: Tuple       ,    lowerCAmelCase_					: Tuple      ):
     """simple docstring"""

     lowerCAmelCase__          =			F'repo_zipped_txt_data-{int(time.time() * 1_0E3      )}'
     lowerCAmelCase__          =			F'{CI_HUB_USER}/{repo_name}'
     hf_api.create_repo(lowerCAmelCase_       ,    token=lowerCAmelCase_       ,    repo_type="dataset"       ,    private=lowerCAmelCase_      )
     hf_api.upload_file(
         token=lowerCAmelCase_       ,    path_or_fileobj=str(lowerCAmelCase_      )       ,    path_in_repo="data.zip"       ,    repo_id=lowerCAmelCase_       ,    repo_type="dataset"       ,    )
     yield repo_id
     try:
          hf_api.delete_repo(lowerCAmelCase_       ,    token=lowerCAmelCase_       ,    repo_type="dataset"      )
     except (requests.exceptions.HTTPError, ValueError):  # catch http error and token invalid error
          pass



@pytest.fixture()
def 							_A       (						lowerCAmelCase_					: int       ,    lowerCAmelCase_					: List[str]       ,    lowerCAmelCase_					: Tuple      ):
     """simple docstring"""

     return hf_private_dataset_repo_zipped_txt_data_



@pytest.fixture(scope="session"      )
def 							_A       (						lowerCAmelCase_					: HfApi       ,    lowerCAmelCase_					: List[str]       ,    lowerCAmelCase_					: int      ):
     """simple docstring"""

     lowerCAmelCase__          =			F'repo_zipped_img_data-{int(time.time() * 1_0E3      )}'
     lowerCAmelCase__          =			F'{CI_HUB_USER}/{repo_name}'
     hf_api.create_repo(lowerCAmelCase_       ,    token=lowerCAmelCase_       ,    repo_type="dataset"       ,    private=lowerCAmelCase_      )
     hf_api.upload_file(
         token=lowerCAmelCase_       ,    path_or_fileobj=str(lowerCAmelCase_      )       ,    path_in_repo="data.zip"       ,    repo_id=lowerCAmelCase_       ,    repo_type="dataset"       ,    )
     yield repo_id
     try:
          hf_api.delete_repo(lowerCAmelCase_       ,    token=lowerCAmelCase_       ,    repo_type="dataset"      )
     except (requests.exceptions.HTTPError, ValueError):  # catch http error and token invalid error
          pass






@pytest.fixture()
def 							_A       (						lowerCAmelCase_					: Optional[Any]       ,    lowerCAmelCase_					: Dict       ,    lowerCAmelCase_					: List[str]      ):
     """simple docstring"""

     return hf_private_dataset_repo_zipped_img_data_
