




"""simple docstring"""






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			:					Any      ,       snake_case			:					List[str]      ,       snake_case			:					List[str]      ,       snake_case			:					Optional[int]    ):
					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[int]      ,       snake_case			:					Union[str, Any]      ,       snake_case			:					str      ,       snake_case			:					Optional[Any]      ,       snake_case			:					Tuple=True    ):
					model.train()
					_lowerCAmelCase:Union[str, Any]									=  model(snake_case    )
					_lowerCAmelCase:Union[str, 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			:					int      ,       snake_case			:					Optional[int]=False    ):
					set_seed(42    )
					_lowerCAmelCase:Union[str, Any]									=  RegressionModel()
					_lowerCAmelCase:Dict									=  deepcopy(snake_case    )
					_lowerCAmelCase:Optional[int]									=  RegressionDataset(length=80    )
					_lowerCAmelCase:Any									=  DataLoader(snake_case      ,       batch_size=16    )
					model.to(accelerator.device    )
					if sched:
										_lowerCAmelCase:Optional[Any]									=  AdamW(params=model.parameters()      ,       lr=1e-3    )
										_lowerCAmelCase:Tuple									=  AdamW(params=ddp_model.parameters()      ,       lr=1e-3    )
										_lowerCAmelCase:List[str]									=  LambdaLR(snake_case      ,       lr_lambda=lambda snake_case    : epoch**0.65    )
										_lowerCAmelCase:List[str]									=  LambdaLR(snake_case      ,       lr_lambda=lambda snake_case    : epoch**0.65    )
					# Make a copy of `model`
					if sched:
										_lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase:Any									=  accelerator.prepare(snake_case      ,       snake_case      ,       snake_case      ,       snake_case    )
					else:
										_lowerCAmelCase							,   _lowerCAmelCase:Optional[int]									=  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			:					Optional[Any]    ):
					# Test when on a single CPU or GPU that the context manager does nothing
					_lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase:List[str]									=  get_training_setup(snake_case    )
					# Use a single batch
					_lowerCAmelCase							,   _lowerCAmelCase:Dict									=  next(iter(snake_case    )    ).values()
					for iteration in range(3    ):
										# Gather the distributed inputs and targs for the base model
										_lowerCAmelCase							,   _lowerCAmelCase:Optional[Any]									=  accelerator.gather((ddp_input, ddp_target)    )
										_lowerCAmelCase							,   _lowerCAmelCase:Any									=  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    )
										_lowerCAmelCase:List[str]									=  ddp_input[torch.randperm(len(snake_case    )    )]





def 							UpperCAmelCase		(		snake_case			:					Dict    ):
					# Test on distributed setup that context manager behaves properly
					_lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase:List[Any]									=  get_training_setup(snake_case    )
					# Use a single batch
					_lowerCAmelCase							,   _lowerCAmelCase:Any									=  next(iter(snake_case    )    ).values()
					for iteration in range(3    ):
										# Gather the distributed inputs and targs for the base model
										_lowerCAmelCase							,   _lowerCAmelCase:Tuple									=  accelerator.gather((ddp_input, ddp_target)    )
										_lowerCAmelCase							,   _lowerCAmelCase:Tuple									=  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    )
										_lowerCAmelCase:str									=  ddp_input[torch.randperm(len(snake_case    )    )]





def 							UpperCAmelCase		(		snake_case			:					Optional[Any]=False      ,       snake_case			:					str=False    ):
					_lowerCAmelCase:str									=  Accelerator(
					    split_batches=snake_case      ,       dispatch_batches=snake_case      ,       gradient_accumulation_steps=2    )
					# Test that context manager behaves properly
					_lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase:List[Any]									=  get_training_setup(snake_case    )
					for iteration, batch in enumerate(snake_case    ):
										_lowerCAmelCase							,   _lowerCAmelCase:Optional[int]									=  batch.values()
										# Gather the distributed inputs and targs for the base model
										_lowerCAmelCase							,   _lowerCAmelCase:int									=  accelerator.gather((ddp_input, ddp_target)    )
										_lowerCAmelCase							,   _lowerCAmelCase:Tuple									=  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    )
										_lowerCAmelCase:List[str]									=  ddp_input[torch.randperm(len(snake_case    )    )]
					GradientState._reset_state()





def 							UpperCAmelCase		(		snake_case			:					List[Any]=False      ,       snake_case			:					Union[str, Any]=False    ):
					_lowerCAmelCase:List[str]									=  Accelerator(
					    split_batches=snake_case      ,       dispatch_batches=snake_case      ,       gradient_accumulation_steps=2    )
					# Test that context manager behaves properly
					_lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase:int									=  get_training_setup(snake_case      ,       snake_case    )
					for iteration, batch in enumerate(snake_case    ):
										_lowerCAmelCase							,   _lowerCAmelCase:Optional[Any]									=  batch.values()
										# Gather the distributed inputs and targs for the base model
										_lowerCAmelCase							,   _lowerCAmelCase:Any									=  accelerator.gather((ddp_input, ddp_target)    )
										_lowerCAmelCase							,   _lowerCAmelCase:str									=  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'
										_lowerCAmelCase:List[str]									=  (((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		(		):
					_lowerCAmelCase:List[Any]									=  Accelerator()

					_lowerCAmelCase:int									=  RegressionDataset(length=80    )
					_lowerCAmelCase:Union[str, Any]									=  DataLoader(snake_case      ,       batch_size=16    )
					_lowerCAmelCase:List[str]									=  RegressionDataset(length=96    )
					_lowerCAmelCase:Optional[int]									=  DataLoader(snake_case      ,       batch_size=16    )
					_lowerCAmelCase							,   _lowerCAmelCase:List[str]									=  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		(		):
					_lowerCAmelCase:List[str]									=  Accelerator()
					_lowerCAmelCase:List[str]									=  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			:					Any    ):
					# For xla_spawn (TPUs)
					main()


if __name__ == "__main__":
				main()












"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase__						=			{
    '''configuration_convbert''': ['''CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ConvBertConfig''', '''ConvBertOnnxConfig'''],
    '''tokenization_convbert''': ['''ConvBertTokenizer'''],
}

try:
				if not is_tokenizers_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCamelCase__						=			['''ConvBertTokenizerFast''']

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCamelCase__						=			[
				    '''CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
				    '''ConvBertForMaskedLM''',
				    '''ConvBertForMultipleChoice''',
				    '''ConvBertForQuestionAnswering''',
				    '''ConvBertForSequenceClassification''',
				    '''ConvBertForTokenClassification''',
				    '''ConvBertLayer''',
				    '''ConvBertModel''',
				    '''ConvBertPreTrainedModel''',
				    '''load_tf_weights_in_convbert''',
				]


try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCamelCase__						=			[
				    '''TF_CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
				    '''TFConvBertForMaskedLM''',
				    '''TFConvBertForMultipleChoice''',
				    '''TFConvBertForQuestionAnswering''',
				    '''TFConvBertForSequenceClassification''',
				    '''TFConvBertForTokenClassification''',
				    '''TFConvBertLayer''',
				    '''TFConvBertModel''',
				    '''TFConvBertPreTrainedModel''',
				]


if TYPE_CHECKING:
				from .configuration_convbert import CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, ConvBertConfig, ConvBertOnnxConfig
				from .tokenization_convbert import ConvBertTokenizer

				try:
								if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tokenization_convbert_fast import ConvBertTokenizerFast

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_convbert import (
								    CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
								    ConvBertForMaskedLM,
								    ConvBertForMultipleChoice,
								    ConvBertForQuestionAnswering,
								    ConvBertForSequenceClassification,
								    ConvBertForTokenClassification,
								    ConvBertLayer,
								    ConvBertModel,
								    ConvBertPreTrainedModel,
								    load_tf_weights_in_convbert,
								)

				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tf_convbert import (
								    TF_CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
								    TFConvBertForMaskedLM,
								    TFConvBertForMultipleChoice,
								    TFConvBertForQuestionAnswering,
								    TFConvBertForSequenceClassification,
								    TFConvBertForTokenClassification,
								    TFConvBertLayer,
								    TFConvBertModel,
								    TFConvBertPreTrainedModel,
								)


else:
				import sys

				UpperCamelCase__						=			_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






