


'''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)

UpperCAmelCase_			:  Dict										=   'bert-base-cased'
UpperCAmelCase_			:  Optional[int]										=   'fp16'
UpperCAmelCase_			:  Union[str, Any]										=   'bf16'
UpperCAmelCase_			:  List[Any]										=   [FPaa, BFaa]

@require_fsdp
@require_cuda
class     _lowerCamelCase       (     snake_case_					):




 '''simple docstring'''


 def 			snake_case__ (   self    ):

  """simple docstring"""







  super().setUp()

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


 def 			snake_case__ (   self    ):

  """simple docstring"""







  from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

  for i, strategy in enumerate(__lowercase    ):
   __A	:		Optional[Any]	  =      self.dist_env.copy()
   __A	:		Any	  =      F"""{i + 1}"""
   __A	:		Union[str, Any]	  =      strategy
   with mockenv_context(**__lowercase    ):
    __A	:		Any	  =      FullyShardedDataParallelPlugin()
    self.assertEqual(fsdp_plugin.sharding_strategy   ,   ShardingStrategy(i + 1    )    )


 def 			snake_case__ (   self    ):

  """simple docstring"""







  from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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


 def 			snake_case__ (   self    ):

  """simple docstring"""







  from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

  for i, state_dict_type in enumerate(__lowercase    ):
   __A	:		List[Any]	  =      self.dist_env.copy()
   __A	:		Tuple	  =      state_dict_type
   with mockenv_context(**__lowercase    ):
    __A	:		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 			snake_case__ (   self    ):

  """simple docstring"""







  __A	:		str	  =      AutoModel.from_pretrained(__lowercase    )
  for policy in FSDP_AUTO_WRAP_POLICY:
   __A	:		List[Any]	  =      self.dist_env.copy()
   __A	:		Tuple	  =      policy
   if policy == "TRANSFORMER_BASED_WRAP":
    __A	:		Dict	  =      'BertLayer'
   elif policy == "SIZE_BASED_WRAP":
    __A	:		Tuple	  =      '2000'
   with mockenv_context(**__lowercase    ):
    __A	:		Optional[int]	  =      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    )

  __A	:		Tuple	  =      self.dist_env.copy()
  __A	:		Tuple	  =      'TRANSFORMER_BASED_WRAP'
  __A	:		Optional[Any]	  =      'T5Layer'
  with mockenv_context(**__lowercase    ):
   __A	:		int	  =      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    )    )

  __A	:		Union[str, Any]	  =      self.dist_env.copy()
  __A	:		str	  =      'SIZE_BASED_WRAP'
  __A	:		int	  =      '0'
  with mockenv_context(**__lowercase    ):
   __A	:		Any	  =      FullyShardedDataParallelPlugin()
   fsdp_plugin.set_auto_wrap_policy(__lowercase    )
   self.assertIsNone(fsdp_plugin.auto_wrap_policy    )


 def 			snake_case__ (   self    ):

  """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:
   __A	:		Dict	  =      self.dist_env.copy()
   __A	:		Optional[int]	  =      mp_dtype
   with mockenv_context(**__lowercase    ):
    __A	:		Optional[Any]	  =      Accelerator()
    if mp_dtype == "fp16":
     __A	:		List[str]	  =      torch.floataa
    elif mp_dtype == "bf16":
     __A	:		List[Any]	  =      torch.bfloataa
    __A	:		List[str]	  =      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 			snake_case__ (   self    ):

  """simple docstring"""







  from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

  for flag in [True, False]:
   __A	:		int	  =      self.dist_env.copy()
   __A	:		Dict	  =      str(__lowercase    ).lower()
   with mockenv_context(**__lowercase    ):
    __A	:		Optional[Any]	  =      FullyShardedDataParallelPlugin()
    self.assertEqual(fsdp_plugin.cpu_offload   ,   CPUOffload(offload_params=__lowercase    )    )


@require_fsdp
@require_multi_gpu
@slow
class     _lowerCamelCase       (     snake_case_					):




 '''simple docstring'''


 def 			snake_case__ (   self    ):

  """simple docstring"""







  super().setUp()
  __A	:		Dict	  =      0.8_2
  __A	:		int	  =      [
      'fsdp_shard_grad_op_transformer_based_wrap',
      'fsdp_full_shard_transformer_based_wrap',
  ]
  __A	:		Tuple	  =      {
      'multi_gpu_fp16': 3_200,
      'fsdp_shard_grad_op_transformer_based_wrap_fp16': 2_000,
      'fsdp_full_shard_transformer_based_wrap_fp16': 1_900,
      # 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
  }
  __A	:		Union[str, Any]	  =      160
  __A	:		List[Any]	  =      160

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


 def 			snake_case__ (   self    ):

  """simple docstring"""







  __A	:		Tuple	  =      os.path.join(self.test_scripts_folder   ,   'test_performance.py'    )
  __A	:		Optional[Any]	  =      ['accelerate', 'launch', '--num_processes=2', '--num_machines=1', '--machine_rank=0', '--use_fsdp']
  for config in self.performance_configs:
   __A	:		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 			snake_case__ (   self    ):

  """simple docstring"""







  __A	:		Tuple	  =      os.path.join(self.test_scripts_folder   ,   'test_checkpointing.py'    )
  __A	:		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    ):
   __A	:		str	  =      cmd.copy()
   cmd_config.append(F"""--fsdp_sharding_strategy={i+1}"""    )
   if strategy != "FULL_SHARD":
    continue
   __A	:		Union[str, Any]	  =      len(__lowercase    )
   for state_dict_type in FSDP_STATE_DICT_TYPE:
    __A	:		Tuple	  =      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()    )

    __A	:		Union[str, Any]	  =      cmd_config[:-1]
    __A	:		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(__lowercase   ,   env=os.environ.copy()    )







 def 			snake_case__ (   self    ):

  """simple docstring"""







  __A	:		Optional[Any]	  =      os.path.join(self.test_scripts_folder   ,   'test_peak_memory_usage.py'    )
  __A	:		Union[str, Any]	  =      [
      '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():
   __A	:		str	  =      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()    )










'''simple docstring'''
import unittest
from typing import Tuple

import torch

from diffusers.utils import floats_tensor, randn_tensor, torch_all_close, torch_device
from diffusers.utils.testing_utils import require_torch

@require_torch
class     _lowerCamelCase       :




 '''simple docstring'''


 @property
 def 			snake_case__ (   self    ):

  """simple docstring"""







  return self.get_dummy_input()


 @property
 def 			snake_case__ (   self    ):

  """simple docstring"""







  if self.block_type == "down":
   return (4, 32, 16, 16)
  elif self.block_type == "mid":
   return (4, 32, 32, 32)
  elif self.block_type == "up":
   return (4, 32, 64, 64)

  raise ValueError(F"""'{self.block_type}' is not a supported block_type. Set it to 'up', 'mid', or 'down'."""    )


 def 			snake_case__ (   self   ,   __lowercase=True   ,   __lowercase=False   ,   __lowercase=False   ,   __lowercase=False   ,   ):

  """simple docstring"""







  __A	:		List[str]	  =      4
  __A	:		List[Any]	  =      32
  __A	:		Optional[Any]	  =      (32, 32)

  __A	:		List[str]	  =      torch.manual_seed(0    )
  __A	:		Union[str, Any]	  =      torch.device(__lowercase    )
  __A	:		Optional[int]	  =      (batch_size, num_channels) + sizes
  __A	:		str	  =      randn_tensor(__lowercase   ,   generator=__lowercase   ,   device=__lowercase    )
  __A	:		int	  =      {'hidden_states': hidden_states}

  if include_temb:
   __A	:		Optional[Any]	  =      128
   __A	:		Any	  =      randn_tensor((batch_size, temb_channels)   ,   generator=__lowercase   ,   device=__lowercase    )

  if include_res_hidden_states_tuple:
   __A	:		Any	  =      torch.manual_seed(1    )
   __A	:		List[Any]	  =      (randn_tensor(__lowercase   ,   generator=__lowercase   ,   device=__lowercase    ),)

  if include_encoder_hidden_states:
   __A	:		str	  =      floats_tensor((batch_size, 32, 32)    ).to(__lowercase    )

  if include_skip_sample:
   __A	:		List[Any]	  =      randn_tensor(((batch_size, 3) + sizes)   ,   generator=__lowercase   ,   device=__lowercase    )

  return dummy_input


 def 			snake_case__ (   self    ):

  """simple docstring"""







  __A	:		Tuple	  =      {
      'in_channels': 32,
      'out_channels': 32,
      'temb_channels': 128,
  }
  if self.block_type == "up":
   __A	:		int	  =      32

  if self.block_type == "mid":
   init_dict.pop('out_channels'    )

  __A	:		Optional[int]	  =      self.dummy_input
  return init_dict, inputs_dict


 def 			snake_case__ (   self   ,   __lowercase    ):

  """simple docstring"""







  __A  ,__A	:		List[Any]	  =      self.prepare_init_args_and_inputs_for_common()
  __A	:		int	  =      self.block_class(**__lowercase    )
  unet_block.to(__lowercase    )
  unet_block.eval()

  with torch.no_grad():
   __A	:		List[Any]	  =      unet_block(**__lowercase    )

  if isinstance(__lowercase   ,   __lowercase    ):
   __A	:		Tuple	  =      output[0]

  self.assertEqual(output.shape   ,   self.output_shape    )

  __A	:		Any	  =      output[0, -1, -3:, -3:]
  __A	:		Dict	  =      torch.tensor(__lowercase    ).to(__lowercase    )
  assert torch_all_close(output_slice.flatten()   ,   __lowercase   ,   atol=5E-3    )







 @unittest.skipIf(torch_device == 'mps'   ,   'Training is not supported in mps'    )
 def 			snake_case__ (   self    ):

  """simple docstring"""







  __A  ,__A	:		str	  =      self.prepare_init_args_and_inputs_for_common()
  __A	:		Optional[Any]	  =      self.block_class(**__lowercase    )
  model.to(__lowercase    )
  model.train()
  __A	:		Any	  =      model(**__lowercase    )

  if isinstance(__lowercase   ,   __lowercase    ):
   __A	:		Any	  =      output[0]

  __A	:		int	  =      torch.device(__lowercase    )
  __A	:		Dict	  =      randn_tensor(output.shape   ,   device=__lowercase    )
  __A	:		Dict	  =      torch.nn.functional.mse_loss(__lowercase   ,   __lowercase    )
  loss.backward()






