

"""simple docstring"""






import inspect
import os

import torch
from transformers import AutoModel
from transformers.testing_utils import mockenv_context
from transformers.trainer_utils import set_seed

import accelerate
from accelerate.accelerator import Accelerator
from accelerate.state import AcceleratorState
from accelerate.test_utils.testing import (
    AccelerateTestCase,
    TempDirTestCase,
    execute_subprocess_async,
    require_cuda,
    require_fsdp,
    require_multi_gpu,
    slow,
)
from accelerate.utils.constants import (
    FSDP_AUTO_WRAP_POLICY,
    FSDP_BACKWARD_PREFETCH,
    FSDP_SHARDING_STRATEGY,
    FSDP_STATE_DICT_TYPE,
)
from accelerate.utils.dataclasses import FullyShardedDataParallelPlugin
from accelerate.utils.other import patch_environment


set_seed(4_2)

snake_case_ : Union[str, Any]   =							"""bert-base-cased"""
snake_case_ : Optional[Any]   =							"""fp16"""
snake_case_ : List[Any]   =							"""bf16"""
snake_case_ : Any   =							[FPaa, BFaa]







@require_fsdp
@require_cuda
class    snake_case__	(			__a  ):







			def       __lowerCAmelCase					(      self      :			Dict	):





								'''simple docstring'''



								super().setUp()

								UpperCAmelCase      : List[Any]   =					dict(
								    ACCELERATE_USE_FSDP="true"      ,	MASTER_ADDR="localhost"      ,	MASTER_PORT="10999"      ,	RANK="0"      ,	LOCAL_RANK="0"      ,	WORLD_SIZE="1"      ,	)







			def       __lowerCAmelCase					(      self      :			List[Any]	):





								'''simple docstring'''



								from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

								for i, strategy in enumerate(_lowercase	):
													UpperCAmelCase      : Optional[int]   =					self.dist_env.copy()
													UpperCAmelCase      : List[str]   =					f"""{i + 1}"""
													UpperCAmelCase      : Dict   =					strategy
													with mockenv_context(**_lowercase	):
																		UpperCAmelCase      : List[str]   =					FullyShardedDataParallelPlugin()
																		self.assertEqual(fsdp_plugin.sharding_strategy      ,	ShardingStrategy(i + 1	)	)







			def       __lowerCAmelCase					(      self      :			Union[str, Any]	):





								'''simple docstring'''



								from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

								for i, prefetch_policy in enumerate(_lowercase	):
													UpperCAmelCase      : Any   =					self.dist_env.copy()
													UpperCAmelCase      : Any   =					prefetch_policy
													with mockenv_context(**_lowercase	):
																		UpperCAmelCase      : Union[str, Any]   =					FullyShardedDataParallelPlugin()
																		if prefetch_policy == "NO_PREFETCH":
																							self.assertIsNone(fsdp_plugin.backward_prefetch	)
																		else:
																							self.assertEqual(fsdp_plugin.backward_prefetch      ,	BackwardPrefetch(i + 1	)	)







			def       __lowerCAmelCase					(      self      :			Dict	):





								'''simple docstring'''



								from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

								for i, state_dict_type in enumerate(_lowercase	):
													UpperCAmelCase      : Dict   =					self.dist_env.copy()
													UpperCAmelCase      : str   =					state_dict_type
													with mockenv_context(**_lowercase	):
																		UpperCAmelCase      : Any   =					FullyShardedDataParallelPlugin()
																		self.assertEqual(fsdp_plugin.state_dict_type      ,	StateDictType(i + 1	)	)
																		if state_dict_type == "FULL_STATE_DICT":
																							self.assertTrue(fsdp_plugin.state_dict_config.offload_to_cpu	)
																							self.assertTrue(fsdp_plugin.state_dict_config.ranka_only	)







			def       __lowerCAmelCase					(      self      :			Optional[int]	):





								'''simple docstring'''



								UpperCAmelCase      : List[Any]   =					AutoModel.from_pretrained(_lowercase	)
								for policy in FSDP_AUTO_WRAP_POLICY:
													UpperCAmelCase      : str   =					self.dist_env.copy()
													UpperCAmelCase      : str   =					policy
													if policy == "TRANSFORMER_BASED_WRAP":
																		UpperCAmelCase      : Optional[Any]   =					'''BertLayer'''
													elif policy == "SIZE_BASED_WRAP":
																		UpperCAmelCase      : Tuple   =					'''2000'''
													with mockenv_context(**_lowercase	):
																		UpperCAmelCase      : Union[str, Any]   =					FullyShardedDataParallelPlugin()
																		fsdp_plugin.set_auto_wrap_policy(_lowercase	)
																		if policy == "NO_WRAP":
																							self.assertIsNone(fsdp_plugin.auto_wrap_policy	)
																		else:
																							self.assertIsNotNone(fsdp_plugin.auto_wrap_policy	)

								UpperCAmelCase      : Optional[int]   =					self.dist_env.copy()
								UpperCAmelCase      : int   =					'''TRANSFORMER_BASED_WRAP'''
								UpperCAmelCase      : List[Any]   =					'''T5Layer'''
								with mockenv_context(**_lowercase	):
													UpperCAmelCase      : str   =					FullyShardedDataParallelPlugin()
													with self.assertRaises(_lowercase	) as cm:
																		fsdp_plugin.set_auto_wrap_policy(_lowercase	)
													self.assertTrue("Could not find the transformer layer class to wrap in the model." in str(cm.exception	)	)

								UpperCAmelCase      : str   =					self.dist_env.copy()
								UpperCAmelCase      : Tuple   =					'''SIZE_BASED_WRAP'''
								UpperCAmelCase      : Optional[int]   =					'''0'''
								with mockenv_context(**_lowercase	):
													UpperCAmelCase      : List[str]   =					FullyShardedDataParallelPlugin()
													fsdp_plugin.set_auto_wrap_policy(_lowercase	)
													self.assertIsNone(fsdp_plugin.auto_wrap_policy	)







			def       __lowerCAmelCase					(      self      :			List[Any]	):





								'''simple docstring'''



								from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
								from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

								for mp_dtype in dtypes:
													UpperCAmelCase      : int   =					self.dist_env.copy()
													UpperCAmelCase      : Dict   =					mp_dtype
													with mockenv_context(**_lowercase	):
																		UpperCAmelCase      : Union[str, Any]   =					Accelerator()
																		if mp_dtype == "fp16":
																							UpperCAmelCase      : Optional[Any]   =					torch.floataa
																		elif mp_dtype == "bf16":
																							UpperCAmelCase      : int   =					torch.bfloataa
																		UpperCAmelCase      : Any   =					MixedPrecision(param_dtype=_lowercase      ,	reduce_dtype=_lowercase      ,	buffer_dtype=_lowercase	)
																		self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy      ,	_lowercase	)
																		if mp_dtype == FPaa:
																							self.assertTrue(isinstance(accelerator.scaler      ,	_lowercase	)	)
																		elif mp_dtype == BFaa:
																							self.assertIsNone(accelerator.scaler	)
																		AcceleratorState._reset_state(_lowercase	)







			def       __lowerCAmelCase					(      self      :			Union[str, Any]	):





								'''simple docstring'''



								from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

								for flag in [True, False]:
													UpperCAmelCase      : Union[str, Any]   =					self.dist_env.copy()
													UpperCAmelCase      : Dict   =					str(_lowercase	).lower()
													with mockenv_context(**_lowercase	):
																		UpperCAmelCase      : List[Any]   =					FullyShardedDataParallelPlugin()
																		self.assertEqual(fsdp_plugin.cpu_offload      ,	CPUOffload(offload_params=_lowercase	)	)







@require_fsdp
@require_multi_gpu
@slow
class    snake_case__	(			__a  ):







			def       __lowerCAmelCase					(      self      :			Any	):





								'''simple docstring'''



								super().setUp()
								UpperCAmelCase      : List[Any]   =					0.8_2
								UpperCAmelCase      : Optional[Any]   =					[
								    '''fsdp_shard_grad_op_transformer_based_wrap''',
								    '''fsdp_full_shard_transformer_based_wrap''',
								]
								UpperCAmelCase      : Dict   =					{
								    '''multi_gpu_fp16''': 32_00,
								    '''fsdp_shard_grad_op_transformer_based_wrap_fp16''': 20_00,
								    '''fsdp_full_shard_transformer_based_wrap_fp16''': 19_00,
								    # Disabling below test as it overwhelms the RAM memory usage
								    # on CI self-hosted runner leading to tests getting killed.
								    # "fsdp_full_shard_cpu_offload_transformer_based_wrap_fp32": 1500,  # fp16 was leading to indefinite hang
								}
								UpperCAmelCase      : int   =					1_60
								UpperCAmelCase      : Any   =					1_60

								UpperCAmelCase      : Optional[Any]   =					inspect.getfile(accelerate.test_utils	)
								UpperCAmelCase      : List[Any]   =					os.path.sep.join(mod_file.split(os.path.sep	)[:-1] + ["scripts", "external_deps"]	)







			def       __lowerCAmelCase					(      self      :			Optional[Any]	):





								'''simple docstring'''



								UpperCAmelCase      : Any   =					os.path.join(self.test_scripts_folder      ,	"test_performance.py"	)
								UpperCAmelCase      : Optional[int]   =					['''accelerate''', '''launch''', '''--num_processes=2''', '''--num_machines=1''', '''--machine_rank=0''', '''--use_fsdp''']
								for config in self.performance_configs:
													UpperCAmelCase      : List[str]   =					cmd.copy()
													for i, strategy in enumerate(_lowercase	):
																		if strategy.lower() in config:
																							cmd_config.append(f"""--fsdp_sharding_strategy={i+1}"""	)
																							break

													if "fp32" in config:
																		cmd_config.append("--mixed_precision=no"	)
													else:
																		cmd_config.append("--mixed_precision=fp16"	)

													if "cpu_offload" in config:
																		cmd_config.append("--fsdp_offload_params=True"	)

													for policy in FSDP_AUTO_WRAP_POLICY:
																		if policy.lower() in config:
																							cmd_config.append(f"""--fsdp_auto_wrap_policy={policy}"""	)
																							break

													if policy == "TRANSFORMER_BASED_WRAP":
																		cmd_config.append("--fsdp_transformer_layer_cls_to_wrap=BertLayer"	)
													elif policy == "SIZE_BASED_WRAP":
																		cmd_config.append("--fsdp_min_num_params=2000"	)

													cmd_config.extend(
													    [
													        self.test_file_path,
													        f"""--output_dir={self.tmpdir}""",
													        f"""--performance_lower_bound={self.performance_lower_bound}""",
													    ]	)
													with patch_environment(omp_num_threads=1	):
																		execute_subprocess_async(_lowercase      ,	env=os.environ.copy()	)







			def       __lowerCAmelCase					(      self      :			Union[str, Any]	):





								'''simple docstring'''



								UpperCAmelCase      : Dict   =					os.path.join(self.test_scripts_folder      ,	"test_checkpointing.py"	)
								UpperCAmelCase      : int   =					[
								    '''accelerate''',
								    '''launch''',
								    '''--num_processes=2''',
								    '''--num_machines=1''',
								    '''--machine_rank=0''',
								    '''--use_fsdp''',
								    '''--mixed_precision=fp16''',
								    '''--fsdp_transformer_layer_cls_to_wrap=BertLayer''',
								]

								for i, strategy in enumerate(_lowercase	):
													UpperCAmelCase      : Any   =					cmd.copy()
													cmd_config.append(f"""--fsdp_sharding_strategy={i+1}"""	)
													if strategy != "FULL_SHARD":
																		continue
													UpperCAmelCase      : Union[str, Any]   =					len(_lowercase	)
													for state_dict_type in FSDP_STATE_DICT_TYPE:
																		UpperCAmelCase      : Optional[Any]   =					cmd_config[:state_dict_config_index]
																		cmd_config.append(f"""--fsdp_state_dict_type={state_dict_type}"""	)
																		cmd_config.extend(
																		    [
																		        self.test_file_path,
																		        f"""--output_dir={self.tmpdir}""",
																		        "--partial_train_epoch=1",
																		    ]	)
																		with patch_environment(omp_num_threads=1	):
																							execute_subprocess_async(_lowercase      ,	env=os.environ.copy()	)

																		UpperCAmelCase      : List[Any]   =					cmd_config[:-1]
																		UpperCAmelCase      : Optional[Any]   =					os.path.join(self.tmpdir      ,	"epoch_0"	)
																		cmd_config.extend(
																		    [
																		        f"""--resume_from_checkpoint={resume_from_checkpoint}""",
																		    ]	)
																		with patch_environment(omp_num_threads=1	):
																							execute_subprocess_async(_lowercase      ,	env=os.environ.copy()	)







			def       __lowerCAmelCase					(      self      :			Optional[Any]	):





								'''simple docstring'''



								UpperCAmelCase      : Dict   =					os.path.join(self.test_scripts_folder      ,	"test_peak_memory_usage.py"	)
								UpperCAmelCase      : List[str]   =					[
								    '''accelerate''',
								    '''launch''',
								    '''--num_processes=2''',
								    '''--num_machines=1''',
								    '''--machine_rank=0''',
								]
								for spec, peak_mem_upper_bound in self.peak_memory_usage_upper_bound.items():
													UpperCAmelCase      : int   =					cmd.copy()
													if "fp16" in spec:
																		cmd_config.extend(["--mixed_precision=fp16"]	)
													else:
																		cmd_config.extend(["--mixed_precision=no"]	)

													if "multi_gpu" in spec:
																		continue
													else:
																		cmd_config.extend(["--use_fsdp"]	)
																		for i, strategy in enumerate(_lowercase	):
																							if strategy.lower() in spec:
																												cmd_config.append(f"""--fsdp_sharding_strategy={i+1}"""	)
																												break

																		if "cpu_offload" in spec:
																							cmd_config.append("--fsdp_offload_params=True"	)

																		for policy in FSDP_AUTO_WRAP_POLICY:
																							if policy.lower() in spec:
																												cmd_config.append(f"""--fsdp_auto_wrap_policy={policy}"""	)
																												break

																		if policy == "TRANSFORMER_BASED_WRAP":
																							cmd_config.append("--fsdp_transformer_layer_cls_to_wrap=BertLayer"	)
																		elif policy == "SIZE_BASED_WRAP":
																							cmd_config.append("--fsdp_min_num_params=2000"	)

													cmd_config.extend(
													    [
													        self.test_file_path,
													        f"""--output_dir={self.tmpdir}""",
													        f"""--peak_memory_upper_bound={peak_mem_upper_bound}""",
													        f"""--n_train={self.n_train}""",
													        f"""--n_val={self.n_val}""",
													    ]	)
													with patch_environment(omp_num_threads=1	):
																		execute_subprocess_async(_lowercase      ,	env=os.environ.copy()	)



import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()




class 			__magic_name__	(		__a   ,						__a   ,						unittest.TestCase     ):





   """simple docstring"""




   lowerCAmelCase   :     Dict	      = VQModel
   lowerCAmelCase   :     List[str]	      = '''sample'''



   @property
   def 	lowerCAmelCase    (						self       :	List[Any]		,	_lowercase       :	Tuple=(32, 32)				):



    """simple docstring"""



    _UpperCamelCase:  List[Any]    						=			4
    _UpperCamelCase:  Tuple    						=			3

    _UpperCamelCase:  List[str]    						=			floats_tensor((batch_size, num_channels) + sizes				).to(_lowercase				)

    return {"sample": image}



   @property
   def 	lowerCAmelCase    (						self       :	List[Any]				):



    """simple docstring"""



    return (3, 32, 32)



   @property
   def 	lowerCAmelCase    (						self       :	str				):



    """simple docstring"""



    return (3, 32, 32)



   def 	lowerCAmelCase    (						self       :	Optional[int]				):



    """simple docstring"""



    _UpperCamelCase:  Any    						=			{
        '''block_out_channels''': [32, 64],
        '''in_channels''': 3,
        '''out_channels''': 3,
        '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
        '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
        '''latent_channels''': 3,
    }
    _UpperCamelCase:  Optional[Any]    						=			self.dummy_input
    return init_dict, inputs_dict



   def 	lowerCAmelCase    (						self       :	List[Any]				):



    """simple docstring"""



    pass



   def 	lowerCAmelCase    (						self       :	List[Any]				):



    """simple docstring"""



    pass



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    _UpperCamelCase						,				_UpperCamelCase:  int    						=			VQModel.from_pretrained('''fusing/vqgan-dummy'''		,	output_loading_info=_lowercase				)
    self.assertIsNotNone(_lowercase				)
    self.assertEqual(len(loading_info['''missing_keys''']				)		,	0				)

    model.to(_lowercase				)
    _UpperCamelCase:  Tuple    						=			model(**self.dummy_input				)

    assert image is not None, "Make sure output is not None"



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			VQModel.from_pretrained('''fusing/vqgan-dummy'''				)
    model.to(_lowercase				).eval()

    torch.manual_seed(0				)
    if torch.cuda.is_available():
     torch.cuda.manual_seed_all(0				)

    _UpperCamelCase:  int    						=			torch.randn(1		,	model.config.in_channels		,	model.config.sample_size		,	model.config.sample_size				)
    _UpperCamelCase:  Tuple    						=			image.to(_lowercase				)
    with torch.no_grad():
     _UpperCamelCase:  Optional[int]    						=			model(_lowercase				).sample

    _UpperCamelCase:  List[Any]    						=			output[0, -1, -3:, -3:].flatten().cpu()
    # fmt: off
    _UpperCamelCase:  Any    						=			torch.tensor([-0.0153, -0.4044, -0.1880, -0.5161, -0.2418, -0.4072, -0.1612, -0.0633, -0.0143]				)
    # fmt: on
    self.assertTrue(torch.allclose(_lowercase		,	_lowercase		,	atol=1E-3				)				)