




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

_UpperCAmelCase	:						Optional[int]			=     "bert-base-cased"
_UpperCAmelCase	:						List[Any]			=     "fp16"
_UpperCAmelCase	:						List[str]			=     "bf16"
_UpperCAmelCase	:						Optional[Any]			=     [FPaa, BFaa]




@require_fsdp
@require_cuda
class    __lowerCAmelCase       (							lowerCAmelCase):






       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         super().setUp()

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






       def 					SCREAMING_SNAKE_CASE    (   self:					str				):
         from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

         for i, strategy in enumerate(_lowerCAmelCase				):
           lowercase    :Any      	=						self.dist_env.copy()
           lowercase    :str      	=						F"{i + 1}"
           lowercase    :Any      	=						strategy
           with mockenv_context(**_lowerCAmelCase				):
             lowercase    :str      	=						FullyShardedDataParallelPlugin()
             self.assertEqual(fsdp_plugin.sharding_strategy       ,					ShardingStrategy(i + 1				)				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Union[str, Any]				):
         from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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






       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

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

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

         lowercase    :str      	=						self.dist_env.copy()
         lowercase    :Dict      	=						"SIZE_BASED_WRAP"
         lowercase    :Dict      	=						"0"
         with mockenv_context(**_lowerCAmelCase				):
           lowercase    :str      	=						FullyShardedDataParallelPlugin()
           fsdp_plugin.set_auto_wrap_policy(_lowerCAmelCase				)
           self.assertIsNone(fsdp_plugin.auto_wrap_policy				)






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






       def 					SCREAMING_SNAKE_CASE    (   self:					Union[str, Any]				):
         from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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







@require_fsdp
@require_multi_gpu
@slow
class    __lowerCAmelCase       (							lowerCAmelCase):






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[Any]				):
         super().setUp()
         lowercase    :Dict      	=						0.82
         lowercase    :str      	=						[
             "fsdp_shard_grad_op_transformer_based_wrap",
             "fsdp_full_shard_transformer_based_wrap",
         ]
         lowercase    :str      	=						{
             "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
         }
         lowercase    :Dict      	=						1_60
         lowercase    :str      	=						1_60

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






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






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

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






       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         lowercase    :Optional[int]      	=						os.path.join(self.test_scripts_folder       ,					"test_peak_memory_usage.py"				)
         lowercase    :Optional[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():
           lowercase    :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(_lowerCAmelCase				):
               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(_lowerCAmelCase       ,					env=os.environ.copy()				)









_UpperCAmelCase	:						Union[str, Any]			=     "0.21.0"

from .accelerator import Accelerator
from .big_modeling import (
    cpu_offload,
    cpu_offload_with_hook,
    disk_offload,
    dispatch_model,
    init_empty_weights,
    init_on_device,
    load_checkpoint_and_dispatch,
)
from .data_loader import skip_first_batches
from .launchers import debug_launcher, notebook_launcher
from .state import PartialState
from .utils import (
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    FullyShardedDataParallelPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    find_executable_batch_size,
    infer_auto_device_map,
    is_rich_available,
    load_checkpoint_in_model,
    synchronize_rng_states,
)


if is_rich_available():
     from .utils import rich



