


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




							"""simple docstring"""


							def __init__(							self	:  str    )     ->   Dict:
														super().__init__()
														__snake_case		:      str										=       nn.Linear(3   ,     4    )
														__snake_case		:      Tuple										=       nn.BatchNormad(4    )
														__snake_case		:      List[str]										=       nn.Linear(4   ,     5    )






							def    __snake_case (							self	:  Tuple   ,     lowerCamelCase	:  Union[str, Any]    )     ->   Optional[Any]:
														return self.lineara(self.batchnorm(self.lineara(lowerCamelCase    )    )    )

class     a		(_lowerCAmelCase				):




							"""simple docstring"""


							def    __snake_case (							self	:  Union[str, Any]   ,     lowerCamelCase	:  str   ,     *lowerCamelCase	:  Any   ,     **lowerCamelCase	:  Dict    )     ->   List[Any]:
														return (args[0] + 1,) + args[1:], kwargs

class     a		(_lowerCAmelCase				):




							"""simple docstring"""


							def    __snake_case (							self	:  Tuple   ,     lowerCamelCase	:  Any   ,     lowerCamelCase	:  str    )     ->   Optional[Any]:
														return output + 1

class     a		(unittest.TestCase				):




							"""simple docstring"""


							def    __snake_case (							self	:  Optional[int]    )     ->   Tuple:
														__snake_case		:      Tuple										=       ModelForTest()
														__snake_case		:      Optional[Any]										=       ModelHook()

														add_hook_to_module(lowerCamelCase   ,     lowerCamelCase    )
														self.assertEqual(test_model._hf_hook   ,     lowerCamelCase    )
														self.assertTrue(hasattr(lowerCamelCase   ,     "_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(lowerCamelCase    )
														self.assertFalse(hasattr(lowerCamelCase   ,     "_hf_hook"    )    )
														self.assertFalse(hasattr(lowerCamelCase   ,     "_old_forward"    )    )

							def    __snake_case (							self	:  int    )     ->   Dict:
														__snake_case		:      Optional[int]										=       ModelForTest()
														__snake_case		:      List[str]										=       ModelHook()

														add_hook_to_module(lowerCamelCase   ,     lowerCamelCase    )
														add_hook_to_module(lowerCamelCase   ,     lowerCamelCase   ,     append=lowerCamelCase    )

														self.assertEqual(isinstance(test_model._hf_hook   ,     lowerCamelCase    )   ,     lowerCamelCase    )
														self.assertEqual(len(test_model._hf_hook.hooks    )   ,     2    )
														self.assertTrue(hasattr(lowerCamelCase   ,     "_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(lowerCamelCase    )
														self.assertFalse(hasattr(lowerCamelCase   ,     "_hf_hook"    )    )
														self.assertFalse(hasattr(lowerCamelCase   ,     "_old_forward"    )    )

							def    __snake_case (							self	:  int    )     ->   List[Any]:
														__snake_case		:      Any										=       ModelForTest()
														__snake_case		:      Any										=       torch.randn(2   ,     3    )
														__snake_case		:      Tuple										=       test_model(x + 1    )
														__snake_case		:      Dict										=       test_model(x + 2    )

														__snake_case		:      List[Any]										=       PreForwardHook()
														add_hook_to_module(lowerCamelCase   ,     lowerCamelCase    )
														__snake_case		:      Any										=       test_model(lowerCamelCase    )
														self.assertTrue(torch.allclose(lowerCamelCase   ,     lowerCamelCase   ,     atol=1E-5    )    )

														# Attaching a hook to a model when it already has one replaces, does not chain
														__snake_case		:      int										=       PreForwardHook()
														add_hook_to_module(lowerCamelCase   ,     lowerCamelCase    )
														__snake_case		:      List[Any]										=       test_model(lowerCamelCase    )
														self.assertTrue(torch.allclose(lowerCamelCase   ,     lowerCamelCase   ,     atol=1E-5    )    )

														# You need to use the sequential hook to chain two or more hooks
														__snake_case		:      Optional[int]										=       SequentialHook(PreForwardHook()   ,     PreForwardHook()    )
														add_hook_to_module(lowerCamelCase   ,     lowerCamelCase    )

														__snake_case		:      Optional[Any]										=       test_model(lowerCamelCase    )
														assert torch.allclose(lowerCamelCase   ,     lowerCamelCase   ,     atol=1E-5    )

							def    __snake_case (							self	:  Dict    )     ->   str:
														__snake_case		:      Union[str, Any]										=       ModelForTest()
														__snake_case		:      Optional[int]										=       torch.randn(2   ,     3    )
														__snake_case		:      Any										=       test_model(lowerCamelCase    )

														__snake_case		:      List[str]										=       PostForwardHook()
														add_hook_to_module(lowerCamelCase   ,     lowerCamelCase    )
														__snake_case		:      Tuple										=       test_model(lowerCamelCase    )
														self.assertTrue(torch.allclose(lowerCamelCase   ,     output + 1   ,     atol=1E-5    )    )

														# Attaching a hook to a model when it already has one replaces, does not chain
														__snake_case		:      List[Any]										=       PostForwardHook()
														add_hook_to_module(lowerCamelCase   ,     lowerCamelCase    )
														__snake_case		:      Union[str, Any]										=       test_model(lowerCamelCase    )
														self.assertTrue(torch.allclose(lowerCamelCase   ,     output + 1   ,     atol=1E-5    )    )

														# You need to use the sequential hook to chain two or more hooks
														__snake_case		:      Any										=       SequentialHook(PostForwardHook()   ,     PostForwardHook()    )
														add_hook_to_module(lowerCamelCase   ,     lowerCamelCase    )

														__snake_case		:      Tuple										=       test_model(lowerCamelCase    )
														assert torch.allclose(lowerCamelCase   ,     output + 2   ,     atol=1E-5    )

							def    __snake_case (							self	:  List[Any]    )     ->   int:
														__snake_case		:      int										=       ModelForTest()
														__snake_case		:      Optional[Any]										=       torch.randn(2   ,     3    )
														__snake_case		:      Tuple										=       test_model(lowerCamelCase    )

														__snake_case		:      List[Any]										=       PostForwardHook()
														add_hook_to_module(lowerCamelCase   ,     lowerCamelCase    )
														__snake_case		:      Any										=       test_model(lowerCamelCase    )
														self.assertTrue(torch.allclose(lowerCamelCase   ,     output + 1    )    )
														self.assertTrue(outputa.requires_grad    )

														__snake_case		:      str										=       True
														__snake_case		:      str										=       test_model(lowerCamelCase    )
														self.assertFalse(outputa.requires_grad    )

							@require_multi_gpu
							def    __snake_case (							self	:  str    )     ->   int:
														__snake_case		:      Optional[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
														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
														__snake_case		:      int										=       torch.randn(2   ,     3    )
														__snake_case		:      str										=       model(lowerCamelCase    )
														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(lowerCamelCase   ,     AlignDevicesHook(io_same_device=lowerCamelCase    )    )
														__snake_case		:      Optional[int]										=       torch.randn(2   ,     3    ).to(0    )
														__snake_case		:      Tuple										=       model(lowerCamelCase    )
														self.assertEqual(output.device   ,     torch.device(0    )    )

							def    __snake_case (							self	:  Dict    )     ->   str:
														__snake_case		:      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
														__snake_case		:      List[Any]										=       {"execution_device": 0 if torch.cuda.is_available() else "cpu", "offload": True}

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

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

														__snake_case		:      Optional[Any]										=       torch.randn(2   ,     3    )
														__snake_case		:      int										=       model(lowerCamelCase    )
														self.assertEqual(output.device   ,     lowerCamelCase    )

														# 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
														__snake_case		:      Optional[Any]										=       {
														    "execution_device": 0 if torch.cuda.is_available() else "cpu",
														    "offload": True,
														    "offload_buffers": True,
														}

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

														# 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"    )    )

														__snake_case		:      Optional[int]										=       torch.randn(2   ,     3    )
														__snake_case		:      Tuple										=       model(lowerCamelCase    )
														self.assertEqual(output.device   ,     lowerCamelCase    )

														# 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    __snake_case (							self	:  str    )     ->   List[Any]:
														__snake_case		:      Optional[Any]										=       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
														__snake_case		:      Union[str, Any]										=       0 if torch.cuda.is_available() else "cpu"
														attach_align_device_hook(lowerCamelCase   ,     execution_device=lowerCamelCase   ,     offload=lowerCamelCase    )

														# 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
														__snake_case		:      str										=       torch.device(lowerCamelCase    )
														self.assertEqual(model.batchnorm.running_mean.device   ,     lowerCamelCase    )

														__snake_case		:      List[Any]										=       torch.randn(2   ,     3    )
														__snake_case		:      int										=       model(lowerCamelCase    )
														self.assertEqual(output.device   ,     lowerCamelCase    )

														# Removing hooks loads back the weights in the model.
														remove_hook_from_submodules(lowerCamelCase    )
														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(lowerCamelCase   ,     execution_device=lowerCamelCase   ,     offload=lowerCamelCase   ,     offload_buffers=lowerCamelCase    )

														# 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"    )    )

														__snake_case		:      List[str]										=       torch.randn(2   ,     3    )
														__snake_case		:      Any										=       model(lowerCamelCase    )
														self.assertEqual(output.device   ,     lowerCamelCase    )

														# Removing hooks loads back the weights in the model.
														remove_hook_from_submodules(lowerCamelCase    )
														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    __snake_case (							self	:  Tuple    )     ->   List[Any]:
														__snake_case		:      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
														__snake_case		:      Optional[Any]										=       0 if torch.cuda.is_available() else "cpu"
														attach_align_device_hook(
														    lowerCamelCase   ,     execution_device=lowerCamelCase   ,     offload=lowerCamelCase   ,     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
														__snake_case		:      List[Any]										=       torch.device(lowerCamelCase    )
														self.assertEqual(model.batchnorm.running_mean.device   ,     lowerCamelCase    )

														__snake_case		:      Tuple										=       torch.randn(2   ,     3    )
														__snake_case		:      List[Any]										=       model(lowerCamelCase    )
														self.assertEqual(output.device   ,     lowerCamelCase    )

														# Removing hooks loads back the weights in the model.
														remove_hook_from_submodules(lowerCamelCase    )
														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(
														    lowerCamelCase   ,     execution_device=lowerCamelCase   ,     offload=lowerCamelCase   ,     weights_map=model.state_dict()   ,     offload_buffers=lowerCamelCase   ,     )

														# 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"    )    )

														__snake_case		:      List[str]										=       torch.randn(2   ,     3    )
														__snake_case		:      Dict										=       model(lowerCamelCase    )
														self.assertEqual(output.device   ,     lowerCamelCase    )

														# Removing hooks loads back the weights in the model.
														remove_hook_from_submodules(lowerCamelCase    )
														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"    )    )





# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

import torch

from ...utils import is_npu_available, is_xpu_available
from .config_args import ClusterConfig, default_json_config_file
from .config_utils import SubcommandHelpFormatter


_snake_case					:		Optional[Any]      =    "Create a default config file for Accelerate with only a few flags set."







def     lowerCAmelCase_					( __lowerCamelCase="no"							,     __lowerCamelCase = default_json_config_file							,     __lowerCamelCase = False     ):
							__snake_case		:      int										=       Path(__lowerCamelCase     )
							path.parent.mkdir(parents=__lowerCamelCase							,     exist_ok=__lowerCamelCase     )
							if path.exists():
														print(
														    F'Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`.'     )
														return False
							__snake_case		:      Any										=       mixed_precision.lower()
							if mixed_precision not in ["no", "fp16", "bf16", "fp8"]:
														raise ValueError(
														    F'`mixed_precision` should be one of \'no\', \'fp16\', \'bf16\', or \'fp8\'. Received {mixed_precision}'     )
							__snake_case		:      Optional[int]										=       {
							    "compute_environment": "LOCAL_MACHINE",
							    "mixed_precision": mixed_precision,
							}
							if torch.cuda.is_available():
														__snake_case		:      Dict										=       torch.cuda.device_count()
														__snake_case		:      Tuple										=       num_gpus
														__snake_case		:      List[str]										=       False
														if num_gpus > 1:
																					__snake_case		:      Optional[int]										=       "MULTI_GPU"
														else:
																					__snake_case		:      Dict										=       "NO"
							elif is_xpu_available() and use_xpu:
														__snake_case		:      List[str]										=       torch.xpu.device_count()
														__snake_case		:      str										=       num_xpus
														__snake_case		:      int										=       False
														if num_xpus > 1:
																					__snake_case		:      Optional[int]										=       "MULTI_XPU"
														else:
																					__snake_case		:      str										=       "NO"
							elif is_npu_available():
														__snake_case		:      Any										=       torch.npu.device_count()
														__snake_case		:      str										=       num_npus
														__snake_case		:      str										=       False
														if num_npus > 1:
																					__snake_case		:      Optional[int]										=       "MULTI_NPU"
														else:
																					__snake_case		:      int										=       "NO"
							else:
														__snake_case		:      List[Any]										=       0
														__snake_case		:      Dict										=       True
														__snake_case		:      Tuple										=       1
														__snake_case		:      Tuple										=       "NO"
							__snake_case		:      str										=       ClusterConfig(**__lowerCamelCase     )
							config.to_json_file(__lowerCamelCase     )
							return path







def     lowerCAmelCase_					( __lowerCamelCase							,     __lowerCamelCase     ):
							__snake_case		:      Optional[Any]										=       parser.add_parser("default"							,     parents=__lowerCamelCase							,     help=__lowerCamelCase							,     formatter_class=__lowerCamelCase     )
							parser.add_argument(
							    "--config_file"							,     default=__lowerCamelCase							,     help=(
							        "The path to use to store the config file. Will default to a file named default_config.yaml in the cache "
							        "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
							        "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
							        "with 'huggingface'."
							    )							,     dest="save_location"							,     )

							parser.add_argument(
							    "--mixed_precision"							,     choices=["no", "fp16", "bf16"]							,     type=__lowerCamelCase							,     help="Whether or not to use mixed precision training. "
							    "Choose between FP16 and BF16 (bfloat16) training. "
							    "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later."							,     default="no"							,     )
							parser.set_defaults(func=__lowerCamelCase     )
							return parser







def     lowerCAmelCase_					( __lowerCamelCase     ):
							__snake_case		:      List[Any]										=       write_basic_config(args.mixed_precision							,     args.save_location     )
							if config_file:
														print(F'accelerate configuration saved at {config_file}'     )

