
'''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 UpperCAmelCase__	( lowercase__	):





  """simple docstring"""

  def   __lowercase		(		self      :     Dict       ):


     '''simple docstring'''






     super().setUp()

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


  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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


  def   __lowercase		(		self      :     str       ):


     '''simple docstring'''






     from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

     for i, prefetch_policy in enumerate(_a       ):
        _a				:	Tuple			=   self.dist_env.copy()
        _a				:	List[str]			=   prefetch_policy
        with mockenv_context(**_a       ):
           _a				:	Optional[int]			=   FullyShardedDataParallelPlugin()
           if prefetch_policy == "NO_PREFETCH":
              self.assertIsNone(fsdp_plugin.backward_prefetch       )
           else:
              self.assertEqual(fsdp_plugin.backward_prefetch							,BackwardPrefetch(i + 1       )       )


  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

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


     '''simple docstring'''






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

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

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


  def   __lowercase		(		self      :     int       ):


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


     '''simple docstring'''






     from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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







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





  """simple docstring"""

  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     super().setUp()
     _a				:	int			=   0.82
     _a				:	Union[str, Any]			=   [
         'fsdp_shard_grad_op_transformer_based_wrap',
         'fsdp_full_shard_transformer_based_wrap',
     ]
     _a				:	Dict			=   {
         '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
     }
     _a				:	Tuple			=   160
     _a				:	Dict			=   160

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


  def   __lowercase		(		self      :     Union[str, Any]       ):


     '''simple docstring'''






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


     '''simple docstring'''






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

           _a				:	Optional[int]			=   cmd_config[:-1]
           _a				:	Tuple			=   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   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     _a				:	int			=   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				:	Any			=   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()       )





'''simple docstring'''







import doctest
import logging
import os
import unittest
from pathlib import Path
from typing import List, Union

import transformers
from transformers.testing_utils import require_tf, require_torch, slow


__lowerCAmelCase         =    logging.getLogger()







@unittest.skip('''Temporarily disable the doc tests.'''	)
@require_torch
@require_tf
@slow
class UpperCAmelCase__	( unittest.TestCase	):





  """simple docstring"""

  def   __lowercase		(		self      :     str							,_a      :     Path							,_a      :     Union[str, None] = None							,_a      :     Union[List[str], None] = None							,_a      :     Union[str, List[str], None] = None							,_a      :     bool = True							,):


     '''simple docstring'''






     _a				:	Optional[int]			=   [file for file in os.listdir(_a       ) if os.path.isfile(os.path.join(_a							,_a       )       )]

     if identifier is not None:
        _a				:	List[str]			=   [file for file in files if identifier in file]

     if n_identifier is not None:
        if isinstance(_a							,_a       ):
           for n_ in n_identifier:
              _a				:	Tuple			=   [file for file in files if n_ not in file]
        else:
           _a				:	Optional[Any]			=   [file for file in files if n_identifier not in file]

     _a				:	List[str]			=   ignore_files or []
     ignore_files.append('__init__.py'       )
     _a				:	Tuple			=   [file for file in files if file not in ignore_files]

     for file in files:
        # Open all files
        print('Testing'							,_a       )

        if only_modules:
           _a				:	Any			=   file.split('.'       )[0]
           try:
              _a				:	List[str]			=   getattr(_a							,_a       )
              _a				:	int			=   doctest.DocTestSuite(_a       )
              _a				:	Any			=   unittest.TextTestRunner().run(_a       )
              self.assertIs(len(result.failures       )							,0       )
           except AttributeError:
              logger.info(F"""{module_identifier} is not a module."""       )
        else:
           _a				:	Union[str, Any]			=   doctest.testfile(str('..' / directory / file       )							,optionflags=doctest.ELLIPSIS       )
           self.assertIs(result.failed							,0       )


  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     _a				:	int			=   Path('src/transformers'       )
     _a				:	List[Any]			=   'modeling'
     _a				:	Optional[Any]			=   [
         'modeling_ctrl.py',
         'modeling_tf_ctrl.py',
     ]
     self.analyze_directory(_a							,identifier=_a							,ignore_files=_a       )


  def   __lowercase		(		self      :     Optional[int]       ):


     '''simple docstring'''






     _a				:	Optional[Any]			=   Path('src/transformers'       )
     _a				:	Optional[Any]			=   'tokenization'
     self.analyze_directory(_a							,identifier=_a       )


  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     _a				:	Dict			=   Path('src/transformers'       )
     _a				:	str			=   'configuration'
     self.analyze_directory(_a							,identifier=_a       )


  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     _a				:	Tuple			=   Path('src/transformers'       )
     _a				:	List[Any]			=   ['configuration', 'modeling', 'tokenization']
     self.analyze_directory(_a							,n_identifier=_a       )


  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     _a				:	List[Any]			=   Path('docs/source'       )
     _a				:	List[str]			=   ['favicon.ico']
     self.analyze_directory(_a							,ignore_files=_a							,only_modules=_a       )



