
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_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							):
	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_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE=True							):
	model.train()
	lowercase       =    model(__SCREAMING_SNAKE_CASE							)
	lowercase       =    F.mse_loss(__SCREAMING_SNAKE_CASE ,							target.to(output.device							)							)
	if not do_backward:
		loss /= accelerator.gradient_accumulation_steps
		loss.backward()
	else:
		accelerator.backward(__SCREAMING_SNAKE_CASE							)






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






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE							):
	# Test when on a single CPU or GPU that the context manager does nothing
	lowercase  ,					lowercase  ,					lowercase       =    get_training_setup(__SCREAMING_SNAKE_CASE							)
	# Use a single batch
	lowercase  ,					lowercase       =    next(iter(__SCREAMING_SNAKE_CASE							)							).values()
	for iteration in range(3							):
		# Gather the distributed inputs and targs for the base model
		lowercase  ,					lowercase       =    accelerator.gather((ddp_input, ddp_target)							)
		lowercase  ,					lowercase       =    input.to(accelerator.device							), target.to(accelerator.device							)
		# Perform our initial ground truth step in non "DDP"
		step_model(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
		# Do "gradient accumulation" (noop)
		if iteration % 2 == 0:
			# Accumulate grads locally
			with accelerator.no_sync(__SCREAMING_SNAKE_CASE							):
				step_model(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
		else:
			# Sync grads
			step_model(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)

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






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE							):
	# Test on distributed setup that context manager behaves properly
	lowercase  ,					lowercase  ,					lowercase       =    get_training_setup(__SCREAMING_SNAKE_CASE							)
	# Use a single batch
	lowercase  ,					lowercase       =    next(iter(__SCREAMING_SNAKE_CASE							)							).values()
	for iteration in range(3							):
		# Gather the distributed inputs and targs for the base model
		lowercase  ,					lowercase       =    accelerator.gather((ddp_input, ddp_target)							)
		lowercase  ,					lowercase       =    input.to(accelerator.device							), target.to(accelerator.device							)
		# Perform our initial ground truth step in non "DDP"
		step_model(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
		# Do "gradient accumulation" (noop)
		if iteration % 2 == 0:
			# Accumulate grads locally
			with accelerator.no_sync(__SCREAMING_SNAKE_CASE							):
				step_model(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
		else:
			# Sync grads
			step_model(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_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							)
		lowercase       =    ddp_input[torch.randperm(len(__SCREAMING_SNAKE_CASE							)							)]






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE=False ,							__SCREAMING_SNAKE_CASE=False							):
	lowercase       =    Accelerator(
	    split_batches=__SCREAMING_SNAKE_CASE ,							dispatch_batches=__SCREAMING_SNAKE_CASE ,							gradient_accumulation_steps=2							)
	# Test that context manager behaves properly
	lowercase  ,					lowercase  ,					lowercase       =    get_training_setup(__SCREAMING_SNAKE_CASE							)
	for iteration, batch in enumerate(__SCREAMING_SNAKE_CASE							):
		lowercase  ,					lowercase       =    batch.values()
		# Gather the distributed inputs and targs for the base model
		lowercase  ,					lowercase       =    accelerator.gather((ddp_input, ddp_target)							)
		lowercase  ,					lowercase       =    input.to(accelerator.device							), target.to(accelerator.device							)
		# Perform our initial ground truth step in non "DDP"
		step_model(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
		# Do "gradient accumulation" (noop)
		with accelerator.accumulate(__SCREAMING_SNAKE_CASE							):
			step_model(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_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(__SCREAMING_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							)
		lowercase       =    ddp_input[torch.randperm(len(__SCREAMING_SNAKE_CASE							)							)]
	GradientState._reset_state()






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

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






def 							UpperCAmelCase_		(							):
	lowercase       =    Accelerator()

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






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE							):
	# For xla_spawn (TPUs)
	main()


if __name__ == "__main__":
	main()



import warnings

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase	      = logging.get_logger(__name__)

UpperCAmelCase	      = {
    '''xlnet-base-cased''': '''https://huggingface.co/xlnet-base-cased/resolve/main/config.json''',
    '''xlnet-large-cased''': '''https://huggingface.co/xlnet-large-cased/resolve/main/config.json''',
}







class        A_   (			__lowerCamelCase   ):




	'''simple docstring'''


	_UpperCamelCase     :     Tuple							 =       """xlnet"""
	_UpperCamelCase     :     Optional[Any]							 =       ["""mems"""]
	_UpperCamelCase     :     Tuple							 =       {
	    """n_token""": """vocab_size""",  # Backward compatibility
	    """hidden_size""": """d_model""",
	    """num_attention_heads""": """n_head""",
	    """num_hidden_layers""": """n_layer""",
	}





	def __init__(       self		,    snake_case=3_2000		,    snake_case=1024		,    snake_case=24		,    snake_case=16		,    snake_case=4096		,    snake_case="gelu"		,    snake_case=True		,    snake_case="bi"		,    snake_case=0.02		,    snake_case=1E-12		,    snake_case=0.1		,    snake_case=512		,    snake_case=None		,    snake_case=True		,    snake_case=False		,    snake_case=False		,    snake_case=-1		,    snake_case=False		,    snake_case="last"		,    snake_case=True		,    snake_case="tanh"		,    snake_case=0.1		,    snake_case=5		,    snake_case=5		,    snake_case=5		,    snake_case=1		,    snake_case=2		,    **snake_case		,    ):
		lowercase       =    vocab_size
		lowercase       =    d_model
		lowercase       =    n_layer
		lowercase       =    n_head
		if d_model % n_head != 0:
			raise ValueError(F'''\'d_model % n_head\' ({d_model % n_head}) should be equal to 0'''       )
		if "d_head" in kwargs:
			if kwargs["d_head"] != d_model // n_head:
				raise ValueError(
				    F'''`d_head` ({kwargs['d_head']}) should be equal to `d_model // n_head` ({d_model // n_head})'''       )
		lowercase       =    d_model // n_head
		lowercase       =    ff_activation
		lowercase       =    d_inner
		lowercase       =    untie_r
		lowercase       =    attn_type

		lowercase       =    initializer_range
		lowercase       =    layer_norm_eps

		lowercase       =    dropout
		lowercase       =    mem_len
		lowercase       =    reuse_len
		lowercase       =    bi_data
		lowercase       =    clamp_len
		lowercase       =    same_length

		lowercase       =    summary_type
		lowercase       =    summary_use_proj
		lowercase       =    summary_activation
		lowercase       =    summary_last_dropout
		lowercase       =    start_n_top
		lowercase       =    end_n_top

		lowercase       =    bos_token_id
		lowercase       =    pad_token_id
		lowercase       =    eos_token_id

		if "use_cache" in kwargs:
			warnings.warn(
			    'The `use_cache` argument is deprecated and will be removed in a future version, use `use_mems_eval`'
			    ' instead.'		,    snake_case		,    )
			lowercase       =    kwargs['use_cache']

		lowercase       =    use_mems_eval
		lowercase       =    use_mems_train
		super().__init__(pad_token_id=snake_case		,    bos_token_id=snake_case		,    eos_token_id=snake_case		,    **snake_case       )





	@property
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		logger.info(F'''The model {self.model_type} is one of the few models that has no sequence length limit.'''       )
		return -1





	@max_position_embeddings.setter
	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case       ):
		# Message copied from Transformer-XL documentation
		raise NotImplementedError(
		    F'''The model {self.model_type} is one of the few models that has no sequence length limit.'''       )

