


import argparse
import os
import sys
from unittest.mock import patch

import pytorch_lightning as pl
import timeout_decorator
import torch
from distillation import SummarizationDistiller, distill_main
from finetune import SummarizationModule, main

from transformers import MarianMTModel
from transformers.file_utils import cached_path
from transformers.testing_utils import TestCasePlus, require_torch_gpu, slow
from utils import load_json


lowerCAmelCase__      		= 'sshleifer/mar_enro_6_3_student'







class      lowerCAmelCase__       (				a):




     '''simple docstring'''


     def  _lowerCamelCase (						self)			->				Any:
            super().setUp()

            _A      :       List[str]				    = cached_path(
                "https://cdn-datasets.huggingface.co/translation/wmt_en_ro-tr40k-va0.5k-te0.5k.tar.gz" ,						extract_compressed_file=__lowerCamelCase ,						)
            _A      :       Union[str, Any]				    = F"{data_cached}/wmt_en_ro-tr40k-va0.5k-te0.5k"


     @slow
     @require_torch_gpu
     def  _lowerCamelCase (						self)			->				List[Any]:
            MarianMTModel.from_pretrained(__lowerCamelCase)


     @slow
     @require_torch_gpu
     def  _lowerCamelCase (						self)			->				Any:
            _A      :       str				    = {
                """$MAX_LEN""": 6_4,
                """$BS""": 6_4,
                """$GAS""": 1,
                """$ENRO_DIR""": self.data_dir,
                """facebook/mbart-large-cc25""": MARIAN_MODEL,
                # "val_check_interval=0.25": "val_check_interval=1.0",
                """--learning_rate=3e-5""": """--learning_rate 3e-4""",
                """--num_train_epochs 6""": """--num_train_epochs 1""",
            }

            # Clean up bash script
            _A      :       List[Any]				    = (self.test_file_dir / """train_mbart_cc25_enro.sh""").open().read().split("finetune.py")[1].strip()
            _A      :       Optional[Any]				    = bash_script.replace("\\\n" ,						"").strip().replace("\"$@\"" ,						"")
            for k, v in env_vars_to_replace.items():
                   _A      :       List[str]				    = bash_script.replace(__lowerCamelCase ,						str(__lowerCamelCase))
            _A      :       List[Any]				    = self.get_auto_remove_tmp_dir()

            # bash_script = bash_script.replace("--fp16 ", "")
            _A      :       str				    = F"\n            --output_dir {output_dir}\n            --tokenizer_name Helsinki-NLP/opus-mt-en-ro\n            --sortish_sampler\n            --do_predict\n            --gpus 1\n            --freeze_encoder\n            --n_train 40000\n            --n_val 500\n            --n_test 500\n            --fp16_opt_level O1\n            --num_sanity_val_steps 0\n            --eval_beams 2\n        ".split()
            # XXX: args.gpus > 1 : handle multi_gpu in the future

            _A      :       str				    = ["""finetune.py"""] + bash_script.split() + args
            with patch.object(__lowerCamelCase ,						"argv" ,						__lowerCamelCase):
                   _A      :       str				    = argparse.ArgumentParser()
                   _A      :       List[str]				    = pl.Trainer.add_argparse_args(__lowerCamelCase)
                   _A      :       List[Any]				    = SummarizationModule.add_model_specific_args(__lowerCamelCase ,						os.getcwd())
                   _A      :       List[Any]				    = parser.parse_args()
                   _A      :       Union[str, Any]				    = main(__lowerCamelCase)

            # Check metrics
            _A      :       Tuple				    = load_json(model.metrics_save_path)
            _A      :       str				    = metrics["""val"""][0]
            _A      :       Optional[Any]				    = metrics["""val"""][-1]
            self.assertEqual(len(metrics["val"]) ,						(args.max_epochs / args.val_check_interval))
            assert isinstance(last_step_stats[F"val_avg_{model.val_metric}"] ,						__lowerCamelCase)

            self.assertGreater(last_step_stats["val_avg_gen_time"] ,						0.0_1)
            # model hanging on generate. Maybe bad config was saved. (XXX: old comment/assert?)
            self.assertLessEqual(last_step_stats["val_avg_gen_time"] ,						1.0)

            # test learning requirements:

            # 1. BLEU improves over the course of training by more than 2 pts
            self.assertGreater(last_step_stats["val_avg_bleu"] - first_step_stats["val_avg_bleu"] ,						2)

            # 2. BLEU finishes above 17
            self.assertGreater(last_step_stats["val_avg_bleu"] ,						1_7)

            # 3. test BLEU and val BLEU within ~1.1 pt.
            self.assertLess(abs(metrics["val"][-1]["val_avg_bleu"] - metrics["test"][-1]["test_avg_bleu"]) ,						1.1)

            # check lightning ckpt can be loaded and has a reasonable statedict
            _A      :       List[Any]				    = os.listdir(__lowerCamelCase)
            _A      :       List[Any]				    = [x for x in contents if x.endswith(".ckpt")][0]
            _A      :       int				    = os.path.join(args.output_dir ,						__lowerCamelCase)
            _A      :       str				    = torch.load(__lowerCamelCase ,						map_location="cpu")
            _A      :       int				    = """model.model.decoder.layers.0.encoder_attn_layer_norm.weight"""
            assert expected_key in ckpt["state_dict"]
            assert ckpt["state_dict"]["model.model.decoder.layers.0.encoder_attn_layer_norm.weight"].dtype == torch.floataa

            # TODO: turn on args.do_predict when PL bug fixed.
            if args.do_predict:
                   _A      :       Union[str, Any]				    = {os.path.basename(__lowerCamelCase) for p in contents}
                   assert "test_generations.txt" in contents
                   assert "test_results.txt" in contents
                   # assert len(metrics["val"]) ==  desired_n_evals
                   assert len(metrics["test"]) == 1







class      lowerCAmelCase__       (				a):




     '''simple docstring'''


     @timeout_decorator.timeout(6_0_0)
     @slow
     @require_torch_gpu
     def  _lowerCamelCase (						self)			->				int:
            _A      :       str				    = F"{self.test_file_dir_str}/test_data/wmt_en_ro"
            _A      :       Tuple				    = {
                """--fp16_opt_level=O1""": """""",
                """$MAX_LEN""": 1_2_8,
                """$BS""": 1_6,
                """$GAS""": 1,
                """$ENRO_DIR""": data_dir,
                """$m""": """sshleifer/student_marian_en_ro_6_1""",
                """val_check_interval=0.25""": """val_check_interval=1.0""",
            }

            # Clean up bash script
            _A      :       Dict				    = (
                (self.test_file_dir / """distil_marian_no_teacher.sh""").open().read().split("distillation.py")[1].strip()
            )
            _A      :       Tuple				    = bash_script.replace("\\\n" ,						"").strip().replace("\"$@\"" ,						"")
            _A      :       Any				    = bash_script.replace("--fp16 " ,						" ")

            for k, v in env_vars_to_replace.items():
                   _A      :       Dict				    = bash_script.replace(__lowerCamelCase ,						str(__lowerCamelCase))
            _A      :       str				    = self.get_auto_remove_tmp_dir()
            _A      :       int				    = bash_script.replace("--fp16" ,						"")
            _A      :       Dict				    = 6
            _A      :       Any				    = (
                ["""distillation.py"""]
                + bash_script.split()
                + [
                    F"--output_dir={output_dir}",
                    """--gpus=1""",
                    """--learning_rate=1e-3""",
                    F"--num_train_epochs={epochs}",
                    """--warmup_steps=10""",
                    """--val_check_interval=1.0""",
                    """--do_predict""",
                ]
            )
            with patch.object(__lowerCamelCase ,						"argv" ,						__lowerCamelCase):
                   _A      :       List[str]				    = argparse.ArgumentParser()
                   _A      :       Optional[Any]				    = pl.Trainer.add_argparse_args(__lowerCamelCase)
                   _A      :       int				    = SummarizationDistiller.add_model_specific_args(__lowerCamelCase ,						os.getcwd())
                   _A      :       List[str]				    = parser.parse_args()
                   # assert args.gpus == gpus THIS BREAKS for multi_gpu

                   _A      :       List[str]				    = distill_main(__lowerCamelCase)

            # Check metrics
            _A      :       Dict				    = load_json(model.metrics_save_path)
            _A      :       List[str]				    = metrics["""val"""][0]
            _A      :       Union[str, Any]				    = metrics["""val"""][-1]
            assert len(metrics["val"]) >= (args.max_epochs / args.val_check_interval)  # +1 accounts for val_sanity_check

            assert last_step_stats["val_avg_gen_time"] >= 0.0_1

            assert first_step_stats["val_avg_bleu"] < last_step_stats["val_avg_bleu"]  # model learned nothing
            assert 1.0 >= last_step_stats["val_avg_gen_time"]  # model hanging on generate. Maybe bad config was saved.
            assert isinstance(last_step_stats[F"val_avg_{model.val_metric}"] ,						__lowerCamelCase)

            # check lightning ckpt can be loaded and has a reasonable statedict
            _A      :       str				    = os.listdir(__lowerCamelCase)
            _A      :       int				    = [x for x in contents if x.endswith(".ckpt")][0]
            _A      :       Any				    = os.path.join(args.output_dir ,						__lowerCamelCase)
            _A      :       Any				    = torch.load(__lowerCamelCase ,						map_location="cpu")
            _A      :       str				    = """model.model.decoder.layers.0.encoder_attn_layer_norm.weight"""
            assert expected_key in ckpt["state_dict"]
            assert ckpt["state_dict"]["model.model.decoder.layers.0.encoder_attn_layer_norm.weight"].dtype == torch.floataa

            # TODO: turn on args.do_predict when PL bug fixed.
            if args.do_predict:
                   _A      :       Tuple				    = {os.path.basename(__lowerCamelCase) for p in contents}
                   assert "test_generations.txt" in contents
                   assert "test_results.txt" in contents
                   # assert len(metrics["val"]) ==  desired_n_evals
                   assert len(metrics["test"]) == 1




"""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      lowerCamelCase__   (     nn.Module				):







					"""simple docstring"""


					def __init__(       self    : List[Any]  ):
												'''simple docstring'''






												super().__init__()
												__UpperCAmelCase   :						Optional[Any]      							=					nn.Linear(3 ,       4  )
												__UpperCAmelCase   :						Optional[int]      							=					nn.BatchNormad(4  )
												__UpperCAmelCase   :						int      							=					nn.Linear(4 ,       5  )






					def lowerCamelCase__							(       self    : List[str] ,       UpperCamelCase    : Dict  ):
												'''simple docstring'''






												return self.lineara(self.batchnorm(self.lineara(UpperCamelCase  )  )  )






class      lowerCamelCase__   (     A				):







					"""simple docstring"""


					def lowerCamelCase__							(       self    : str ,       UpperCamelCase    : Any ,       *UpperCamelCase    : List[Any] ,       **UpperCamelCase    : Dict  ):
												'''simple docstring'''






												return (args[0] + 1,) + args[1:], kwargs






class      lowerCamelCase__   (     A				):







					"""simple docstring"""


					def lowerCamelCase__							(       self    : Any ,       UpperCamelCase    : Dict ,       UpperCamelCase    : List[Any]  ):
												'''simple docstring'''






												return output + 1






class      lowerCamelCase__   (     unittest.TestCase				):







					"""simple docstring"""


					def lowerCamelCase__							(       self    : str  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Dict      							=					ModelForTest()
												__UpperCAmelCase   :						List[str]      							=					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 lowerCamelCase__							(       self    : str  ):
												'''simple docstring'''






												__UpperCAmelCase   :						int      							=					ModelForTest()
												__UpperCAmelCase   :						Optional[Any]      							=					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 lowerCamelCase__							(       self    : Any  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Union[str, Any]      							=					ModelForTest()
												__UpperCAmelCase   :						str      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						List[str]      							=					test_model(x + 1  )
												__UpperCAmelCase   :						Optional[int]      							=					test_model(x + 2  )

												__UpperCAmelCase   :						Optional[Any]      							=					PreForwardHook()
												add_hook_to_module(UpperCamelCase ,       UpperCamelCase  )
												__UpperCAmelCase   :						Optional[Any]      							=					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
												__UpperCAmelCase   :						Optional[Any]      							=					PreForwardHook()
												add_hook_to_module(UpperCamelCase ,       UpperCamelCase  )
												__UpperCAmelCase   :						Optional[int]      							=					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
												__UpperCAmelCase   :						Optional[Any]      							=					SequentialHook(PreForwardHook() ,       PreForwardHook()  )
												add_hook_to_module(UpperCamelCase ,       UpperCamelCase  )

												__UpperCAmelCase   :						List[str]      							=					test_model(UpperCamelCase  )
												assert torch.allclose(UpperCamelCase ,       UpperCamelCase ,       atol=1e-5  )






					def lowerCamelCase__							(       self    : Dict  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Optional[int]      							=					ModelForTest()
												__UpperCAmelCase   :						str      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						int      							=					test_model(UpperCamelCase  )

												__UpperCAmelCase   :						Dict      							=					PostForwardHook()
												add_hook_to_module(UpperCamelCase ,       UpperCamelCase  )
												__UpperCAmelCase   :						Any      							=					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
												__UpperCAmelCase   :						List[str]      							=					PostForwardHook()
												add_hook_to_module(UpperCamelCase ,       UpperCamelCase  )
												__UpperCAmelCase   :						Optional[Any]      							=					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
												__UpperCAmelCase   :						Optional[int]      							=					SequentialHook(PostForwardHook() ,       PostForwardHook()  )
												add_hook_to_module(UpperCamelCase ,       UpperCamelCase  )

												__UpperCAmelCase   :						int      							=					test_model(UpperCamelCase  )
												assert torch.allclose(UpperCamelCase ,       output + 2 ,       atol=1e-5  )






					def lowerCamelCase__							(       self    : Optional[int]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						int      							=					ModelForTest()
												__UpperCAmelCase   :						str      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						Optional[Any]      							=					test_model(UpperCamelCase  )

												__UpperCAmelCase   :						Union[str, Any]      							=					PostForwardHook()
												add_hook_to_module(UpperCamelCase ,       UpperCamelCase  )
												__UpperCAmelCase   :						Optional[int]      							=					test_model(UpperCamelCase  )
												self.assertTrue(torch.allclose(UpperCamelCase ,       output + 1  )  )
												self.assertTrue(outputa.requires_grad  )

												__UpperCAmelCase   :						List[str]      							=					True
												__UpperCAmelCase   :						Optional[int]      							=					test_model(UpperCamelCase  )
												self.assertFalse(outputa.requires_grad  )






					@require_multi_gpu
					def lowerCamelCase__							(       self    : Optional[Any]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						str      							=					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
												__UpperCAmelCase   :						List[str]      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						Any      							=					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  )  )
												__UpperCAmelCase   :						int      							=					torch.randn(2 ,       3  ).to(0  )
												__UpperCAmelCase   :						Any      							=					model(UpperCamelCase  )
												self.assertEqual(output.device ,       torch.device(0  )  )






					def lowerCamelCase__							(       self    : Optional[int]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Tuple      							=					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
												__UpperCAmelCase   :						Union[str, Any]      							=					{"""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
												__UpperCAmelCase   :						Any      							=					torch.device(hook_kwargs["""execution_device"""]  )
												self.assertEqual(model.batchnorm.running_mean.device ,       UpperCamelCase  )

												__UpperCAmelCase   :						Union[str, Any]      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						Optional[int]      							=					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
												__UpperCAmelCase   :						Union[str, Any]      							=					{
												    """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"""  )  )

												__UpperCAmelCase   :						str      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						Dict      							=					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 lowerCamelCase__							(       self    : Union[str, Any]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						int      							=					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
												__UpperCAmelCase   :						Optional[int]      							=					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
												__UpperCAmelCase   :						Any      							=					torch.device(UpperCamelCase  )
												self.assertEqual(model.batchnorm.running_mean.device ,       UpperCamelCase  )

												__UpperCAmelCase   :						str      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						Any      							=					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"""  )  )

												__UpperCAmelCase   :						List[str]      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						Any      							=					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 lowerCamelCase__							(       self    : Any  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Dict      							=					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
												__UpperCAmelCase   :						Any      							=					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
												__UpperCAmelCase   :						str      							=					torch.device(UpperCamelCase  )
												self.assertEqual(model.batchnorm.running_mean.device ,       UpperCamelCase  )

												__UpperCAmelCase   :						Dict      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						Optional[Any]      							=					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"""  )  )

												__UpperCAmelCase   :						int      							=					torch.randn(2 ,       3  )
												__UpperCAmelCase   :						List[str]      							=					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"""  )  )
