




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(42)

__SCREAMING_SNAKE_CASE 						="""bert-base-cased"""
__SCREAMING_SNAKE_CASE 						="""fp16"""
__SCREAMING_SNAKE_CASE 						="""bf16"""
__SCREAMING_SNAKE_CASE 						=[FPaa, BFaa]






@require_fsdp
@require_cuda
class  __magic_name__						(  __lowerCamelCase):
    '''simple docstring'''







    def     _A							(							self:   Dict		):
           super().setUp()

           SCREAMING_SNAKE_CASE_   	=    dict(
               ACCELERATE_USE_FSDP='''true'''		,					MASTER_ADDR='''localhost'''		,					MASTER_PORT='''10999'''		,					RANK='''0'''		,					LOCAL_RANK='''0'''		,					WORLD_SIZE='''1'''		,					)


    def     _A							(							self:   Dict		):
           from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

           for i, strategy in enumerate(UpperCAmelCase_		):
                  SCREAMING_SNAKE_CASE_   	=    self.dist_env.copy()
                  SCREAMING_SNAKE_CASE_   	=    f"{i + 1}"
                  SCREAMING_SNAKE_CASE_   	=    strategy
                  with mockenv_context(**UpperCAmelCase_		):
                         SCREAMING_SNAKE_CASE_   	=    FullyShardedDataParallelPlugin()
                         self.assertEqual(fsdp_plugin.sharding_strategy		,					ShardingStrategy(i + 1		)		)


    def     _A							(							self:   Optional[int]		):
           from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

           for i, prefetch_policy in enumerate(UpperCAmelCase_		):
                  SCREAMING_SNAKE_CASE_   	=    self.dist_env.copy()
                  SCREAMING_SNAKE_CASE_   	=    prefetch_policy
                  with mockenv_context(**UpperCAmelCase_		):
                         SCREAMING_SNAKE_CASE_   	=    FullyShardedDataParallelPlugin()
                         if prefetch_policy == "NO_PREFETCH":
                                self.assertIsNone(fsdp_plugin.backward_prefetch		)
                         else:
                                self.assertEqual(fsdp_plugin.backward_prefetch		,					BackwardPrefetch(i + 1		)		)


    def     _A							(							self:   Any		):
           from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

           for i, state_dict_type in enumerate(UpperCAmelCase_		):
                  SCREAMING_SNAKE_CASE_   	=    self.dist_env.copy()
                  SCREAMING_SNAKE_CASE_   	=    state_dict_type
                  with mockenv_context(**UpperCAmelCase_		):
                         SCREAMING_SNAKE_CASE_   	=    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     _A							(							self:   Union[str, Any]		):
           SCREAMING_SNAKE_CASE_   	=    AutoModel.from_pretrained(UpperCAmelCase_		)
           for policy in FSDP_AUTO_WRAP_POLICY:
                  SCREAMING_SNAKE_CASE_   	=    self.dist_env.copy()
                  SCREAMING_SNAKE_CASE_   	=    policy
                  if policy == "TRANSFORMER_BASED_WRAP":
                         SCREAMING_SNAKE_CASE_   	=    'BertLayer'
                  elif policy == "SIZE_BASED_WRAP":
                         SCREAMING_SNAKE_CASE_   	=    '2000'
                  with mockenv_context(**UpperCAmelCase_		):
                         SCREAMING_SNAKE_CASE_   	=    FullyShardedDataParallelPlugin()
                         fsdp_plugin.set_auto_wrap_policy(UpperCAmelCase_		)
                         if policy == "NO_WRAP":
                                self.assertIsNone(fsdp_plugin.auto_wrap_policy		)
                         else:
                                self.assertIsNotNone(fsdp_plugin.auto_wrap_policy		)

           SCREAMING_SNAKE_CASE_   	=    self.dist_env.copy()
           SCREAMING_SNAKE_CASE_   	=    'TRANSFORMER_BASED_WRAP'
           SCREAMING_SNAKE_CASE_   	=    'T5Layer'
           with mockenv_context(**UpperCAmelCase_		):
                  SCREAMING_SNAKE_CASE_   	=    FullyShardedDataParallelPlugin()
                  with self.assertRaises(UpperCAmelCase_		) as cm:
                         fsdp_plugin.set_auto_wrap_policy(UpperCAmelCase_		)
                  self.assertTrue('''Could not find the transformer layer class to wrap in the model.''' in str(cm.exception		)		)

           SCREAMING_SNAKE_CASE_   	=    self.dist_env.copy()
           SCREAMING_SNAKE_CASE_   	=    'SIZE_BASED_WRAP'
           SCREAMING_SNAKE_CASE_   	=    '0'
           with mockenv_context(**UpperCAmelCase_		):
                  SCREAMING_SNAKE_CASE_   	=    FullyShardedDataParallelPlugin()
                  fsdp_plugin.set_auto_wrap_policy(UpperCAmelCase_		)
                  self.assertIsNone(fsdp_plugin.auto_wrap_policy		)


    def     _A							(							self:   Any		):
           from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
           from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

           for mp_dtype in dtypes:
                  SCREAMING_SNAKE_CASE_   	=    self.dist_env.copy()
                  SCREAMING_SNAKE_CASE_   	=    mp_dtype
                  with mockenv_context(**UpperCAmelCase_		):
                         SCREAMING_SNAKE_CASE_   	=    Accelerator()
                         if mp_dtype == "fp16":
                                SCREAMING_SNAKE_CASE_   	=    torch.floataa
                         elif mp_dtype == "bf16":
                                SCREAMING_SNAKE_CASE_   	=    torch.bfloataa
                         SCREAMING_SNAKE_CASE_   	=    MixedPrecision(param_dtype=UpperCAmelCase_		,					reduce_dtype=UpperCAmelCase_		,					buffer_dtype=UpperCAmelCase_		)
                         self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy		,					UpperCAmelCase_		)
                         if mp_dtype == FPaa:
                                self.assertTrue(isinstance(accelerator.scaler		,					UpperCAmelCase_		)		)
                         elif mp_dtype == BFaa:
                                self.assertIsNone(accelerator.scaler		)
                         AcceleratorState._reset_state(UpperCAmelCase_		)





    def     _A							(							self:   Optional[Any]		):
           from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

           for flag in [True, False]:
                  SCREAMING_SNAKE_CASE_   	=    self.dist_env.copy()
                  SCREAMING_SNAKE_CASE_   	=    str(UpperCAmelCase_		).lower()
                  with mockenv_context(**UpperCAmelCase_		):
                         SCREAMING_SNAKE_CASE_   	=    FullyShardedDataParallelPlugin()
                         self.assertEqual(fsdp_plugin.cpu_offload		,					CPUOffload(offload_params=UpperCAmelCase_		)		)






@require_fsdp
@require_multi_gpu
@slow
class  __magic_name__						(  __lowerCamelCase):
    '''simple docstring'''







    def     _A							(							self:   Dict		):
           super().setUp()
           SCREAMING_SNAKE_CASE_   	=    0.82
           SCREAMING_SNAKE_CASE_   	=    [
               'fsdp_shard_grad_op_transformer_based_wrap',
               'fsdp_full_shard_transformer_based_wrap',
           ]
           SCREAMING_SNAKE_CASE_   	=    {
               '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
           }
           SCREAMING_SNAKE_CASE_   	=    1_60
           SCREAMING_SNAKE_CASE_   	=    1_60

           SCREAMING_SNAKE_CASE_   	=    inspect.getfile(accelerate.test_utils		)
           SCREAMING_SNAKE_CASE_   	=    os.path.sep.join(mod_file.split(os.path.sep		)[:-1] + ['''scripts''', '''external_deps''']		)


    def     _A							(							self:   List[str]		):
           SCREAMING_SNAKE_CASE_   	=    os.path.join(self.test_scripts_folder		,					'''test_performance.py'''		)
           SCREAMING_SNAKE_CASE_   	=    ['accelerate', 'launch', '--num_processes=2', '--num_machines=1', '--machine_rank=0', '--use_fsdp']
           for config in self.performance_configs:
                  SCREAMING_SNAKE_CASE_   	=    cmd.copy()
                  for i, strategy in enumerate(UpperCAmelCase_		):
                         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(UpperCAmelCase_		,					env=os.environ.copy()		)


    def     _A							(							self:   Any		):
           SCREAMING_SNAKE_CASE_   	=    os.path.join(self.test_scripts_folder		,					'''test_checkpointing.py'''		)
           SCREAMING_SNAKE_CASE_   	=    [
               '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(UpperCAmelCase_		):
                  SCREAMING_SNAKE_CASE_   	=    cmd.copy()
                  cmd_config.append(f"--fsdp_sharding_strategy={i+1}"		)
                  if strategy != "FULL_SHARD":
                         continue
                  SCREAMING_SNAKE_CASE_   	=    len(UpperCAmelCase_		)
                  for state_dict_type in FSDP_STATE_DICT_TYPE:
                         SCREAMING_SNAKE_CASE_   	=    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(UpperCAmelCase_		,					env=os.environ.copy()		)

                         SCREAMING_SNAKE_CASE_   	=    cmd_config[:-1]
                         SCREAMING_SNAKE_CASE_   	=    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(UpperCAmelCase_		,					env=os.environ.copy()		)





    def     _A							(							self:   Optional[Any]		):
           SCREAMING_SNAKE_CASE_   	=    os.path.join(self.test_scripts_folder		,					'''test_peak_memory_usage.py'''		)
           SCREAMING_SNAKE_CASE_   	=    [
               '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():
                  SCREAMING_SNAKE_CASE_   	=    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(UpperCAmelCase_		):
                                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(UpperCAmelCase_		,					env=os.environ.copy()		)








import unittest

from diffusers import FlaxAutoencoderKL
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax

from .test_modeling_common_flax import FlaxModelTesterMixin


if is_flax_available():
			import jax






@require_flax
class  __magic_name__						(  __UpperCAmelCase    ,  unittest.TestCase):
				'''simple docstring'''







				SCREAMING_SNAKE_CASE__					:							str	      =					FlaxAutoencoderKL


				@property
				def     _A							(							self:   Any		):
											SCREAMING_SNAKE_CASE_   	=    4
											SCREAMING_SNAKE_CASE_   	=    3
											SCREAMING_SNAKE_CASE_   	=    (32, 32)

											SCREAMING_SNAKE_CASE_   	=    jax.random.PRNGKey(0		)
											SCREAMING_SNAKE_CASE_   	=    jax.random.uniform(_lowerCamelCase		,					((batch_size, num_channels) + sizes)		)

											return {"sample": image, "prng_key": prng_key}





				def     _A							(							self:   List[str]		):
											SCREAMING_SNAKE_CASE_   	=    {
											    '''block_out_channels''': [32, 64],
											    '''in_channels''': 3,
											    '''out_channels''': 3,
											    '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
											    '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
											    '''latent_channels''': 4,
											}
											SCREAMING_SNAKE_CASE_   	=    self.dummy_input
											return init_dict, inputs_dict


