





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

lowerCamelCase_     :						Optional[int]						=					'bert-base-cased'
lowerCamelCase_     :						List[Any]						=					'fp16'
lowerCamelCase_     :						int						=					'bf16'
lowerCamelCase_     :						Any						=					[FPaa, BFaa]

@require_fsdp
@require_cuda
class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def 					lowerCamelCase_			(     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 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

        for i, prefetch_policy in enumerate(snake_case_  ):
          A_  :       Any										=  self.dist_env.copy()
          A_  :       Optional[int]										=  prefetch_policy
          with mockenv_context(**snake_case_  ):
            A_  :       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 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

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

        """simple docstring"""

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

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

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

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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


@require_fsdp
@require_multi_gpu
@slow
class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        super().setUp()
        A_  :       int										=  0.82
        A_  :       List[Any]										=  [
            'fsdp_shard_grad_op_transformer_based_wrap',
            'fsdp_full_shard_transformer_based_wrap',
        ]
        A_  :       int										=  {
            'multi_gpu_fp16': 3_2_0_0,
            'fsdp_shard_grad_op_transformer_based_wrap_fp16': 2_0_0_0,
            'fsdp_full_shard_transformer_based_wrap_fp16': 1_9_0_0,
            # 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_  :       Tuple										=  1_6_0
        A_  :       Optional[int]										=  1_6_0

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

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

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

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






"""simple docstring"""




from typing import List, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase_     :						Union[str, Any]						=					logging.get_logger(__name__)

lowerCamelCase_     :						Optional[Any]						=					{
    'huggingface/informer-tourism-monthly': (
        'https://huggingface.co/huggingface/informer-tourism-monthly/resolve/main/config.json'
    ),
    # See all Informer models at https://huggingface.co/models?filter=informer
}

class  _UpperCAmelCase  (   UpperCAmelCase__		):



      '''simple docstring'''
      lowercase_      : Tuple	    =   """informer"""
      lowercase_      : str	    =   {
          """hidden_size""": """d_model""",
          """num_attention_heads""": """encoder_attention_heads""",
          """num_hidden_layers""": """encoder_layers""",
      }

      def __init__(     self					,      snake_case_ = None					,      snake_case_ = None					,      snake_case_ = "student_t"					,      snake_case_ = "nll"					,      snake_case_ = 1					,      snake_case_ = None					,      snake_case_ = "mean"					,      snake_case_ = 0					,      snake_case_ = 0					,      snake_case_ = 0					,      snake_case_ = 0					,      snake_case_ = None					,      snake_case_ = None					,      snake_case_ = 6_4					,      snake_case_ = 3_2					,      snake_case_ = 3_2					,      snake_case_ = 2					,      snake_case_ = 2					,      snake_case_ = 2					,      snake_case_ = 2					,      snake_case_ = True					,      snake_case_ = "gelu"					,      snake_case_ = 0.05					,      snake_case_ = 0.1					,      snake_case_ = 0.1					,      snake_case_ = 0.1					,      snake_case_ = 0.1					,      snake_case_ = 1_0_0					,      snake_case_ = 0.02					,      snake_case_=True					,      snake_case_ = "prob"					,      snake_case_ = 5					,      snake_case_ = True					,      **snake_case_					,      ):

        """simple docstring"""

        A_  :       str										=  prediction_length
        A_  :       List[Any]										=  context_length or prediction_length
        A_  :       str										=  distribution_output
        A_  :       Dict										=  loss
        A_  :       Any										=  input_size
        A_  :       Union[str, Any]										=  num_time_features
        A_  :       Optional[Any]										=  lags_sequence if lags_sequence is not None else [1, 2, 3, 4, 5, 6, 7]
        A_  :       List[Any]										=  scaling
        A_  :       Tuple										=  num_dynamic_real_features
        A_  :       Any										=  num_static_real_features
        A_  :       str										=  num_static_categorical_features

        # set cardinality
        if cardinality and num_static_categorical_features > 0:
          if len(snake_case_  ) != num_static_categorical_features:
            raise ValueError(
                'The cardinality should be a list of the same length as `num_static_categorical_features`'  )
          A_  :       Optional[int]										=  cardinality
        else:
          A_  :       Optional[Any]										=  [0]

        # set embedding_dimension
        if embedding_dimension and num_static_categorical_features > 0:
          if len(snake_case_  ) != num_static_categorical_features:
            raise ValueError(
                'The embedding dimension should be a list of the same length as `num_static_categorical_features`'  )
          A_  :       Any										=  embedding_dimension
        else:
          A_  :       Optional[Any]										=  [min(5_0					,      (cat + 1) // 2  ) for cat in self.cardinality]

        A_  :       int										=  num_parallel_samples

        # Transformer architecture configuration
        A_  :       str										=  input_size * len(self.lags_sequence  ) + self._number_of_features
        A_  :       List[Any]										=  d_model
        A_  :       Dict										=  encoder_attention_heads
        A_  :       Dict										=  decoder_attention_heads
        A_  :       List[Any]										=  encoder_ffn_dim
        A_  :       Union[str, Any]										=  decoder_ffn_dim
        A_  :       int										=  encoder_layers
        A_  :       Any										=  decoder_layers

        A_  :       List[Any]										=  dropout
        A_  :       str										=  attention_dropout
        A_  :       Tuple										=  activation_dropout
        A_  :       List[str]										=  encoder_layerdrop
        A_  :       List[str]										=  decoder_layerdrop

        A_  :       str										=  activation_function
        A_  :       Optional[int]										=  init_std

        A_  :       List[Any]										=  use_cache

        # Informer
        A_  :       Tuple										=  attention_type
        A_  :       List[Any]										=  sampling_factor
        A_  :       Optional[int]										=  distil

        super().__init__(is_encoder_decoder=snake_case_					,      **snake_case_  )

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        return (
            sum(self.embedding_dimension  )
            + self.num_dynamic_real_features
            + self.num_time_features
            + self.num_static_real_features
            + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
        )