



"""simple docstring"""




import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu






class      UpperCamelCase    (						nn.Module     ):



	def __init__(       self		):
							super().__init__()
							A__							=       nn.Linear(3			,	4		)
							A__							=       nn.BatchNormad(4		)
							A__							=       nn.Linear(4			,	5		)



	def 			__A							(       self			,	UpperCAmelCase__		):
							return self.lineara(self.batchnorm(self.lineara(UpperCAmelCase__		)		)		)






class      UpperCamelCase    (						_UpperCAmelCase     ):



	def 			__A							(       self			,	UpperCAmelCase__			,	*UpperCAmelCase__			,	**UpperCAmelCase__		):
							return (args[0] + 1,) + args[1:], kwargs






class      UpperCamelCase    (						_UpperCAmelCase     ):



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							return output + 1






class      UpperCamelCase    (						unittest.TestCase     ):



	def 			__A							(       self		):
							A__							=       ModelForTest()
							A__							=       ModelHook()

							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)
							self.assertEqual(test_model._hf_hook			,	UpperCAmelCase__		)
							self.assertTrue(hasattr(UpperCAmelCase__			,	"_old_forward"		)		)

							# Check adding the hook did not change the name or the signature
							self.assertEqual(test_model.forward.__name__			,	"forward"		)
							self.assertListEqual(list(inspect.signature(test_model.forward		).parameters		)			,	["x"]		)

							remove_hook_from_module(UpperCAmelCase__		)
							self.assertFalse(hasattr(UpperCAmelCase__			,	"_hf_hook"		)		)
							self.assertFalse(hasattr(UpperCAmelCase__			,	"_old_forward"		)		)



	def 			__A							(       self		):
							A__							=       ModelForTest()
							A__							=       ModelHook()

							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__			,	append=UpperCAmelCase__		)

							self.assertEqual(isinstance(test_model._hf_hook			,	UpperCAmelCase__		)			,	UpperCAmelCase__		)
							self.assertEqual(len(test_model._hf_hook.hooks		)			,	2		)
							self.assertTrue(hasattr(UpperCAmelCase__			,	"_old_forward"		)		)

							# Check adding the hook did not change the name or the signature
							self.assertEqual(test_model.forward.__name__			,	"forward"		)
							self.assertListEqual(list(inspect.signature(test_model.forward		).parameters		)			,	["x"]		)

							remove_hook_from_module(UpperCAmelCase__		)
							self.assertFalse(hasattr(UpperCAmelCase__			,	"_hf_hook"		)		)
							self.assertFalse(hasattr(UpperCAmelCase__			,	"_old_forward"		)		)



	def 			__A							(       self		):
							A__							=       ModelForTest()
							A__							=       torch.randn(2			,	3		)
							A__							=       test_model(x + 1		)
							A__							=       test_model(x + 2		)

							A__							=       PreForwardHook()
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)
							A__							=       test_model(UpperCAmelCase__		)
							self.assertTrue(torch.allclose(UpperCAmelCase__			,	UpperCAmelCase__			,	atol=1e-5		)		)

							# Attaching a hook to a model when it already has one replaces, does not chain
							A__							=       PreForwardHook()
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)
							A__							=       test_model(UpperCAmelCase__		)
							self.assertTrue(torch.allclose(UpperCAmelCase__			,	UpperCAmelCase__			,	atol=1e-5		)		)

							# You need to use the sequential hook to chain two or more hooks
							A__							=       SequentialHook(PreForwardHook()			,	PreForwardHook()		)
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)

							A__							=       test_model(UpperCAmelCase__		)
							assert torch.allclose(UpperCAmelCase__			,	UpperCAmelCase__			,	atol=1e-5		)



	def 			__A							(       self		):
							A__							=       ModelForTest()
							A__							=       torch.randn(2			,	3		)
							A__							=       test_model(UpperCAmelCase__		)

							A__							=       PostForwardHook()
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)
							A__							=       test_model(UpperCAmelCase__		)
							self.assertTrue(torch.allclose(UpperCAmelCase__			,	output + 1			,	atol=1e-5		)		)

							# Attaching a hook to a model when it already has one replaces, does not chain
							A__							=       PostForwardHook()
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)
							A__							=       test_model(UpperCAmelCase__		)
							self.assertTrue(torch.allclose(UpperCAmelCase__			,	output + 1			,	atol=1e-5		)		)

							# You need to use the sequential hook to chain two or more hooks
							A__							=       SequentialHook(PostForwardHook()			,	PostForwardHook()		)
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)

							A__							=       test_model(UpperCAmelCase__		)
							assert torch.allclose(UpperCAmelCase__			,	output + 2			,	atol=1e-5		)



	def 			__A							(       self		):
							A__							=       ModelForTest()
							A__							=       torch.randn(2			,	3		)
							A__							=       test_model(UpperCAmelCase__		)

							A__							=       PostForwardHook()
							add_hook_to_module(UpperCAmelCase__			,	UpperCAmelCase__		)
							A__							=       test_model(UpperCAmelCase__		)
							self.assertTrue(torch.allclose(UpperCAmelCase__			,	output + 1		)		)
							self.assertTrue(outputa.requires_grad		)

							A__							=       True
							A__							=       test_model(UpperCAmelCase__		)
							self.assertFalse(outputa.requires_grad		)



	@require_multi_gpu
	def 			__A							(       self		):
							A__							=       ModelForTest()
							# Everything is on CPU
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)

							# This will move each submodule on different devices
							add_hook_to_module(model.lineara			,	AlignDevicesHook(execution_device=0		)		)
							add_hook_to_module(model.batchnorm			,	AlignDevicesHook(execution_device=0		)		)
							add_hook_to_module(model.lineara			,	AlignDevicesHook(execution_device=1		)		)

							self.assertEqual(model.lineara.weight.device			,	torch.device(0		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device(0		)		)
							self.assertEqual(model.batchnorm.running_mean.device			,	torch.device(0		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device(1		)		)

							# We can still make a forward pass. The input does not need to be on any particular device
							A__							=       torch.randn(2			,	3		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	torch.device(1		)		)

							# We can add a general hook to put back output on same device as input.
							add_hook_to_module(UpperCAmelCase__			,	AlignDevicesHook(io_same_device=UpperCAmelCase__		)		)
							A__							=       torch.randn(2			,	3		).to(0		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	torch.device(0		)		)



	def 			__A							(       self		):
							A__							=       ModelForTest()

							# Everything is on CPU
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)

							# This will move each submodule on different devices
							A__							=       {"execution_device": 0 if torch.cuda.is_available() else "cpu", "offload": True}

							add_hook_to_module(model.lineara			,	AlignDevicesHook(**UpperCAmelCase__		)		)
							add_hook_to_module(model.batchnorm			,	AlignDevicesHook(**UpperCAmelCase__		)		)
							add_hook_to_module(model.lineara			,	AlignDevicesHook(**UpperCAmelCase__		)		)

							# Parameters have been offloaded, so on the meta device
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							# Buffers are not included in the offload by default, so are on the execution device
							A__							=       torch.device(hook_kwargs["execution_device"]		)
							self.assertEqual(model.batchnorm.running_mean.device			,	UpperCAmelCase__		)

							A__							=       torch.randn(2			,	3		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	UpperCAmelCase__		)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_module(model.lineara		)
							remove_hook_from_module(model.batchnorm		)
							remove_hook_from_module(model.lineara		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)

							# Now test with buffers included in the offload
							A__							=       {
							    "execution_device": 0 if torch.cuda.is_available() else "cpu",
							    "offload": True,
							    "offload_buffers": True,
							}

							add_hook_to_module(model.lineara			,	AlignDevicesHook(**UpperCAmelCase__		)		)
							add_hook_to_module(model.batchnorm			,	AlignDevicesHook(**UpperCAmelCase__		)		)
							add_hook_to_module(model.lineara			,	AlignDevicesHook(**UpperCAmelCase__		)		)

							# Parameters have been offloaded, so on the meta device, buffers included
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.running_mean.device			,	torch.device("meta"		)		)

							A__							=       torch.randn(2			,	3		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	UpperCAmelCase__		)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_module(model.lineara		)
							remove_hook_from_module(model.batchnorm		)
							remove_hook_from_module(model.lineara		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)



	def 			__A							(       self		):
							A__							=       ModelForTest()

							# Everything is on CPU
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)

							# This will move each submodule on different devices
							A__							=       0 if torch.cuda.is_available() else "cpu"
							attach_align_device_hook(UpperCAmelCase__			,	execution_device=UpperCAmelCase__			,	offload=UpperCAmelCase__		)

							# Parameters have been offloaded, so on the meta device
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							# Buffers are not included in the offload by default, so are on the execution device
							A__							=       torch.device(UpperCAmelCase__		)
							self.assertEqual(model.batchnorm.running_mean.device			,	UpperCAmelCase__		)

							A__							=       torch.randn(2			,	3		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	UpperCAmelCase__		)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_submodules(UpperCAmelCase__		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)

							# Now test with buffers included in the offload
							attach_align_device_hook(UpperCAmelCase__			,	execution_device=UpperCAmelCase__			,	offload=UpperCAmelCase__			,	offload_buffers=UpperCAmelCase__		)

							# Parameters have been offloaded, so on the meta device, buffers included
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.running_mean.device			,	torch.device("meta"		)		)

							A__							=       torch.randn(2			,	3		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	UpperCAmelCase__		)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_submodules(UpperCAmelCase__		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)



	def 			__A							(       self		):
							A__							=       ModelForTest()

							# Everything is on CPU
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)

							# This will move each submodule on different devices
							A__							=       0 if torch.cuda.is_available() else "cpu"
							attach_align_device_hook(
							    UpperCAmelCase__			,	execution_device=UpperCAmelCase__			,	offload=UpperCAmelCase__			,	weights_map=model.state_dict()		)

							# Parameters have been offloaded, so on the meta device
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							# Buffers are not included in the offload by default, so are on the execution device
							A__							=       torch.device(UpperCAmelCase__		)
							self.assertEqual(model.batchnorm.running_mean.device			,	UpperCAmelCase__		)

							A__							=       torch.randn(2			,	3		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	UpperCAmelCase__		)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_submodules(UpperCAmelCase__		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)

							# Now test with buffers included in the offload
							attach_align_device_hook(
							    UpperCAmelCase__			,	execution_device=UpperCAmelCase__			,	offload=UpperCAmelCase__			,	weights_map=model.state_dict()			,	offload_buffers=UpperCAmelCase__			,	)

							# Parameters have been offloaded, so on the meta device, buffers included
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("meta"		)		)
							self.assertEqual(model.batchnorm.running_mean.device			,	torch.device("meta"		)		)

							A__							=       torch.randn(2			,	3		)
							A__							=       model(UpperCAmelCase__		)
							self.assertEqual(output.device			,	UpperCAmelCase__		)

							# Removing hooks loads back the weights in the model.
							remove_hook_from_submodules(UpperCAmelCase__		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.batchnorm.weight.device			,	torch.device("cpu"		)		)
							self.assertEqual(model.lineara.weight.device			,	torch.device("cpu"		)		)










from .dependency_versions_table import deps
from .utils.versions import require_version, require_version_core


# define which module versions we always want to check at run time
# (usually the ones defined in `install_requires` in setup.py)
#
# order specific notes:
# - tqdm must be checked before tokenizers

UpperCAmelCase_							:      List[Any]      =  [
    "python",
    "tqdm",
    "regex",
    "requests",
    "packaging",
    "filelock",
    "numpy",
    "tokenizers",
    "huggingface-hub",
    "safetensors",
    "accelerate",
    "pyyaml",
]

for pkg in pkgs_to_check_at_runtime:
						if pkg in deps:
												if pkg == "tokenizers":
																		# must be loaded here, or else tqdm check may fail
																		from .utils import is_tokenizers_available

																		if not is_tokenizers_available():
																								continue  # not required, check version only if installed
												elif pkg == "accelerate":
																		# must be loaded here, or else tqdm check may fail
																		from .utils import is_accelerate_available

																		# Maybe switch to is_torch_available in the future here so that Accelerate is hard dep of
																		# Transformers with PyTorch
																		if not is_accelerate_available():
																								continue  # not required, check version only if installed

												require_version_core(deps[pkg])
						else:
												raise ValueError(F'''can\'t find {pkg} in {deps.keys()}, check dependency_versions_table.py''')
def  UpperCamelCase						(       _A			:				List[Any]			,							_A			:				int=None    )->       Optional[int]:





						"""simple docstring"""





						require_version(deps[pkg]			,							_A    )





