





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

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






@require_fsdp
@require_cuda
class  _SCREAMING_SNAKE_CASE(       A  ):




       def    _UpperCamelCase					(						self							)      ->				Optional[int]:





             """simple docstring"""






             super().setUp()

             __SCREAMING_SNAKE_CASE       :List[str]		    =  dict(
                 ACCELERATE_USE_FSDP='''true'''		,MASTER_ADDR='''localhost'''		,MASTER_PORT='''10999'''		,RANK='''0'''		,LOCAL_RANK='''0'''		,WORLD_SIZE='''1'''		,)




       def    _UpperCamelCase					(						self							)      ->				int:





             """simple docstring"""






             from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

             for i, strategy in enumerate(SCREAMING_SNAKE_CASE__							):
                   __SCREAMING_SNAKE_CASE       :List[Any]		    =  self.dist_env.copy()
                   __SCREAMING_SNAKE_CASE       :Dict		    =  f'''{i + 1}'''
                   __SCREAMING_SNAKE_CASE       :Any		    =  strategy
                   with mockenv_context(**SCREAMING_SNAKE_CASE__							):
                         __SCREAMING_SNAKE_CASE       :str		    =  FullyShardedDataParallelPlugin()
                         self.assertEqual(fsdp_plugin.sharding_strategy		,ShardingStrategy(i + 1							)							)




       def    _UpperCamelCase					(						self							)      ->				Union[str, Any]:





             """simple docstring"""






             from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

             for i, prefetch_policy in enumerate(SCREAMING_SNAKE_CASE__							):
                   __SCREAMING_SNAKE_CASE       :Optional[int]		    =  self.dist_env.copy()
                   __SCREAMING_SNAKE_CASE       :str		    =  prefetch_policy
                   with mockenv_context(**SCREAMING_SNAKE_CASE__							):
                         __SCREAMING_SNAKE_CASE       :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							)      ->				List[str]:





             """simple docstring"""






             from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

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





             """simple docstring"""






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

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

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




       def    _UpperCamelCase					(						self							)      ->				Dict:





             """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:
                   __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  self.dist_env.copy()
                   __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  mp_dtype
                   with mockenv_context(**SCREAMING_SNAKE_CASE__							):
                         __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  Accelerator()
                         if mp_dtype == "fp16":
                               __SCREAMING_SNAKE_CASE       :Tuple		    =  torch.floataa
                         elif mp_dtype == "bf16":
                               __SCREAMING_SNAKE_CASE       :List[Any]		    =  torch.bfloataa
                         __SCREAMING_SNAKE_CASE       :List[Any]		    =  MixedPrecision(param_dtype=SCREAMING_SNAKE_CASE__		,reduce_dtype=SCREAMING_SNAKE_CASE__		,buffer_dtype=SCREAMING_SNAKE_CASE__							)
                         self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy		,SCREAMING_SNAKE_CASE__							)
                         if mp_dtype == FPaa:
                               self.assertTrue(isinstance(accelerator.scaler		,SCREAMING_SNAKE_CASE__							)							)
                         elif mp_dtype == BFaa:
                               self.assertIsNone(accelerator.scaler							)
                         AcceleratorState._reset_state(SCREAMING_SNAKE_CASE__							)




       def    _UpperCamelCase					(						self							)      ->				str:





             """simple docstring"""






             from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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






@require_fsdp
@require_multi_gpu
@slow
class  _SCREAMING_SNAKE_CASE(       A  ):




       def    _UpperCamelCase					(						self							)      ->				Tuple:





             """simple docstring"""






             super().setUp()
             __SCREAMING_SNAKE_CASE       :Dict		    =  0.8_2
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  [
                 '''fsdp_shard_grad_op_transformer_based_wrap''',
                 '''fsdp_full_shard_transformer_based_wrap''',
             ]
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  {
                 '''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       :Optional[int]		    =  1_60
             __SCREAMING_SNAKE_CASE       :int		    =  1_60

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




       def    _UpperCamelCase					(						self							)      ->				Optional[int]:





             """simple docstring"""






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




       def    _UpperCamelCase					(						self							)      ->				Any:





             """simple docstring"""






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

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




       def    _UpperCamelCase					(						self							)      ->				Optional[int]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :int		    =  os.path.join(self.test_scripts_folder		,'''test_peak_memory_usage.py'''							)
             __SCREAMING_SNAKE_CASE       :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():
                   __SCREAMING_SNAKE_CASE       :Dict		    =  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(SCREAMING_SNAKE_CASE__							):
                               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(SCREAMING_SNAKE_CASE__		,env=os.environ.copy()							)






"""simple docstring"""
from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase_       =		logging.get_logger(__name__)

lowerCamelCase_       =		{
    "google/vivit-b-16x2-kinetics400": (
        "https://huggingface.co/google/vivit-b-16x2-kinetics400/resolve/main/config.json"
    ),
    # See all Vivit models at https://huggingface.co/models?filter=vivit
}






class  _SCREAMING_SNAKE_CASE(       A  ):
       SCREAMING_SNAKE_CASE_       :				Optional[Any]     						=  '''vivit'''




       def __init__(						self		,SCREAMING_SNAKE_CASE__=2_24		,SCREAMING_SNAKE_CASE__=32		,SCREAMING_SNAKE_CASE__=[2, 16, 16]		,SCREAMING_SNAKE_CASE__=3		,SCREAMING_SNAKE_CASE__=7_68		,SCREAMING_SNAKE_CASE__=12		,SCREAMING_SNAKE_CASE__=12		,SCREAMING_SNAKE_CASE__=30_72		,SCREAMING_SNAKE_CASE__="gelu_fast"		,SCREAMING_SNAKE_CASE__=0.0		,SCREAMING_SNAKE_CASE__=0.0		,SCREAMING_SNAKE_CASE__=0.0_2		,SCREAMING_SNAKE_CASE__=1E-06		,SCREAMING_SNAKE_CASE__=True		,**SCREAMING_SNAKE_CASE__		,)      ->				str:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  hidden_size
             __SCREAMING_SNAKE_CASE       :List[str]		    =  num_hidden_layers
             __SCREAMING_SNAKE_CASE       :List[str]		    =  num_attention_heads
             __SCREAMING_SNAKE_CASE       :Any		    =  intermediate_size
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  hidden_act
             __SCREAMING_SNAKE_CASE       :int		    =  hidden_dropout_prob
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  attention_probs_dropout_prob
             __SCREAMING_SNAKE_CASE       :Dict		    =  initializer_range
             __SCREAMING_SNAKE_CASE       :Dict		    =  layer_norm_eps

             __SCREAMING_SNAKE_CASE       :Any		    =  image_size
             __SCREAMING_SNAKE_CASE       :Any		    =  num_frames
             __SCREAMING_SNAKE_CASE       :Any		    =  tubelet_size
             __SCREAMING_SNAKE_CASE       :Tuple		    =  num_channels
             __SCREAMING_SNAKE_CASE       :List[str]		    =  qkv_bias

             super().__init__(**SCREAMING_SNAKE_CASE__							)