






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)

__snake_case												=  '''bert-base-cased'''
__snake_case												=  '''fp16'''
__snake_case												=  '''bf16'''
__snake_case												=  [FPaa, BFaa]







@require_fsdp
@require_cuda
class 			UpperCAmelCase_						(					UpperCamelCase_ ):



      """simple docstring"""
      def     UpperCAmelCase						(				self )				->      List[str]:
          super().setUp()

          UpperCamelCase       :int								=   dict(
              ACCELERATE_USE_FSDP='''true'''    ,				MASTER_ADDR='''localhost'''    ,				MASTER_PORT='''10999'''    ,				RANK='''0'''    ,				LOCAL_RANK='''0'''    ,				WORLD_SIZE='''1'''    ,				)
      def     UpperCAmelCase						(				self )				->      Optional[Any]:
          from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

          for i, strategy in enumerate(A_ ):
              UpperCamelCase       :Optional[int]								=   self.dist_env.copy()
              UpperCamelCase       :Optional[Any]								=   F'''{i + 1}'''
              UpperCamelCase       :List[Any]								=   strategy
              with mockenv_context(**A_ ):
                  UpperCamelCase       :List[Any]								=   FullyShardedDataParallelPlugin()
                  self.assertEqual(fsdp_plugin.sharding_strategy    ,				ShardingStrategy(i + 1 ) )
      def     UpperCAmelCase						(				self )				->      str:
          from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

          for i, prefetch_policy in enumerate(A_ ):
              UpperCamelCase       :Tuple								=   self.dist_env.copy()
              UpperCamelCase       :str								=   prefetch_policy
              with mockenv_context(**A_ ):
                  UpperCamelCase       :List[str]								=   FullyShardedDataParallelPlugin()
                  if prefetch_policy == "NO_PREFETCH":
                      self.assertIsNone(fsdp_plugin.backward_prefetch )
                  else:
                      self.assertEqual(fsdp_plugin.backward_prefetch    ,				BackwardPrefetch(i + 1 ) )
      def     UpperCAmelCase						(				self )				->      Optional[int]:
          from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

          for i, state_dict_type in enumerate(A_ ):
              UpperCamelCase       :Optional[int]								=   self.dist_env.copy()
              UpperCamelCase       :str								=   state_dict_type
              with mockenv_context(**A_ ):
                  UpperCamelCase       :Optional[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     UpperCAmelCase						(				self )				->      Dict:
          UpperCamelCase       :str								=   AutoModel.from_pretrained(A_ )
          for policy in FSDP_AUTO_WRAP_POLICY:
              UpperCamelCase       :Tuple								=   self.dist_env.copy()
              UpperCamelCase       :List[str]								=   policy
              if policy == "TRANSFORMER_BASED_WRAP":
                  UpperCamelCase       :List[str]								=   '''BertLayer'''
              elif policy == "SIZE_BASED_WRAP":
                  UpperCamelCase       :Optional[int]								=   '''2000'''
              with mockenv_context(**A_ ):
                  UpperCamelCase       :Tuple								=   FullyShardedDataParallelPlugin()
                  fsdp_plugin.set_auto_wrap_policy(A_ )
                  if policy == "NO_WRAP":
                      self.assertIsNone(fsdp_plugin.auto_wrap_policy )
                  else:
                      self.assertIsNotNone(fsdp_plugin.auto_wrap_policy )

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

          UpperCamelCase       :Union[str, Any]								=   self.dist_env.copy()
          UpperCamelCase       :List[Any]								=   '''SIZE_BASED_WRAP'''
          UpperCamelCase       :List[Any]								=   '''0'''
          with mockenv_context(**A_ ):
              UpperCamelCase       :str								=   FullyShardedDataParallelPlugin()
              fsdp_plugin.set_auto_wrap_policy(A_ )
              self.assertIsNone(fsdp_plugin.auto_wrap_policy )
      def     UpperCAmelCase						(				self )				->      int:
          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       :Dict								=   self.dist_env.copy()
              UpperCamelCase       :List[str]								=   mp_dtype
              with mockenv_context(**A_ ):
                  UpperCamelCase       :List[Any]								=   Accelerator()
                  if mp_dtype == "fp16":
                      UpperCamelCase       :Optional[Any]								=   torch.floataa
                  elif mp_dtype == "bf16":
                      UpperCamelCase       :Any								=   torch.bfloataa
                  UpperCamelCase       :Optional[Any]								=   MixedPrecision(param_dtype=A_    ,				reduce_dtype=A_    ,				buffer_dtype=A_ )
                  self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy    ,				A_ )
                  if mp_dtype == FPaa:
                      self.assertTrue(isinstance(accelerator.scaler    ,				A_ ) )
                  elif mp_dtype == BFaa:
                      self.assertIsNone(accelerator.scaler )
                  AcceleratorState._reset_state(A_ )




      def     UpperCAmelCase						(				self )				->      List[Any]:
          from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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







@require_fsdp
@require_multi_gpu
@slow
class 			UpperCAmelCase_						(					UpperCamelCase_ ):



      """simple docstring"""
      def     UpperCAmelCase						(				self )				->      List[Any]:
          super().setUp()
          UpperCamelCase       :Optional[Any]								=   0.82
          UpperCamelCase       :Union[str, Any]								=   [
              '''fsdp_shard_grad_op_transformer_based_wrap''',
              '''fsdp_full_shard_transformer_based_wrap''',
          ]
          UpperCamelCase       :Union[str, Any]								=   {
              '''multi_gpu_fp16''': 3200,
              '''fsdp_shard_grad_op_transformer_based_wrap_fp16''': 2000,
              '''fsdp_full_shard_transformer_based_wrap_fp16''': 1900,
              # 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       :Union[str, Any]								=   160
          UpperCamelCase       :List[Any]								=   160

          UpperCamelCase       :Dict								=   inspect.getfile(accelerate.test_utils )
          UpperCamelCase       :str								=   os.path.sep.join(mod_file.split(os.path.sep )[:-1] + ['''scripts''', '''external_deps'''] )
      def     UpperCAmelCase						(				self )				->      Tuple:
          UpperCamelCase       :Tuple								=   os.path.join(self.test_scripts_folder    ,				'''test_performance.py''' )
          UpperCamelCase       :Tuple								=   ['''accelerate''', '''launch''', '''--num_processes=2''', '''--num_machines=1''', '''--machine_rank=0''', '''--use_fsdp''']
          for config in self.performance_configs:
              UpperCamelCase       :Union[str, Any]								=   cmd.copy()
              for i, strategy in enumerate(A_ ):
                  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(A_    ,				env=os.environ.copy() )
      def     UpperCAmelCase						(				self )				->      Optional[Any]:
          UpperCamelCase       :Union[str, Any]								=   os.path.join(self.test_scripts_folder    ,				'''test_checkpointing.py''' )
          UpperCamelCase       :List[str]								=   [
              '''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(A_ ):
              UpperCamelCase       :List[Any]								=   cmd.copy()
              cmd_config.append(F'''--fsdp_sharding_strategy={i+1}''' )
              if strategy != "FULL_SHARD":
                  continue
              UpperCamelCase       :int								=   len(A_ )
              for state_dict_type in FSDP_STATE_DICT_TYPE:
                  UpperCamelCase       :List[str]								=   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(A_    ,				env=os.environ.copy() )

                  UpperCamelCase       :int								=   cmd_config[:-1]
                  UpperCamelCase       :List[str]								=   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(A_    ,				env=os.environ.copy() )




      def     UpperCAmelCase						(				self )				->      Optional[int]:
          UpperCamelCase       :int								=   os.path.join(self.test_scripts_folder    ,				'''test_peak_memory_usage.py''' )
          UpperCamelCase       :Dict								=   [
              '''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       :Tuple								=   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(A_ ):
                      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(A_    ,				env=os.environ.copy() )










import gc
import unittest

from diffusers import FlaxStableDiffusionInpaintPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
						import jax
						import jax.numpy as jnp
						from flax.jax_utils import replicate
						from flax.training.common_utils import shard



@slow
@require_flax
class 				__snake_case						(     unittest.TestCase						):




							def    UpperCAmelCase__			(       self     :       Tuple):
									# clean up the VRAM after each test
									super().tearDown()
									gc.collect()







							def    UpperCAmelCase__			(       self     :       List[str]):
									lowerCAmelCase_			:     Any									= load_image(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
									    '''/sd2-inpaint/init_image.png''')
									lowerCAmelCase_			:     str									= load_image(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png''')

									lowerCAmelCase_			:     List[Any]									= '''xvjiarui/stable-diffusion-2-inpainting'''
									lowerCAmelCase_					,  lowerCAmelCase_			:     Optional[Any]									= FlaxStableDiffusionInpaintPipeline.from_pretrained(A_				,				safety_checker=A_)

									lowerCAmelCase_			:     List[str]									= '''Face of a yellow cat, high resolution, sitting on a park bench'''

									lowerCAmelCase_			:     List[Any]									= jax.random.PRNGKey(0)
									lowerCAmelCase_			:     str									= 5_0

									lowerCAmelCase_			:     List[Any]									= jax.device_count()
									lowerCAmelCase_			:     Union[str, Any]									= num_samples * [prompt]
									lowerCAmelCase_			:     str									= num_samples * [init_image]
									lowerCAmelCase_			:     Union[str, Any]									= num_samples * [mask_image]
									lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_			:     Optional[Any]									= pipeline.prepare_inputs(A_				,				A_				,				A_)

									# shard inputs and rng
									lowerCAmelCase_			:     str									= replicate(A_)
									lowerCAmelCase_			:     str									= jax.random.split(A_				,				jax.device_count())
									lowerCAmelCase_			:     List[Any]									= shard(A_)
									lowerCAmelCase_			:     str									= shard(A_)
									lowerCAmelCase_			:     Tuple									= shard(A_)

									lowerCAmelCase_			:     int									= pipeline(
									    A_				,				A_				,				A_				,				A_				,				A_				,				A_				,				jit=A_)

									lowerCAmelCase_			:     Optional[int]									= output.images.reshape(A_				,				5_1_2				,				5_1_2				,				3)

									lowerCAmelCase_			:     List[str]									= images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

									lowerCAmelCase_			:     int									= jnp.asarray(jax.device_get(image_slice.flatten()))
									lowerCAmelCase_			:     str									= jnp.array(
									    [0.361_1307, 0.3764_9736, 0.375_7408, 0.3821_3953, 0.3929_5167, 0.384_1631, 0.4155_4978, 0.413_7475, 0.421_7084])
									print(F"""output_slice: {output_slice}""")
									assert jnp.abs(output_slice - expected_slice).max() < 1e-2





