






'''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_				: Tuple     	=							'''bert-base-cased'''
lowerCamelCase_				: Any     	=							'''fp16'''
lowerCamelCase_				: List[Any]     	=							'''bf16'''
lowerCamelCase_				: Optional[int]     	=							[FPaa, BFaa]
@require_fsdp
@require_cuda
class    _SCREAMING_SNAKE_CASE							(		SCREAMING_SNAKE_CASE	):


			'''simple docstring'''





			def 				A      (   self							:						Optional[int]			)							->    List[Any]:

									'''simple docstring'''







									super().setUp()

									UpperCamelCase__         = dict(
									    ACCELERATE_USE_FSDP="""true""" ,      MASTER_ADDR="""localhost""" ,      MASTER_PORT="""10999""" ,      RANK="""0""" ,      LOCAL_RANK="""0""" ,      WORLD_SIZE="""1""" ,      )






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

									'''simple docstring'''







									from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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






			def 				A      (   self							:						int			)							->    Dict:

									'''simple docstring'''







									from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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






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

									'''simple docstring'''







									from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

									for i, state_dict_type in enumerate(lowercase			):
															UpperCamelCase__         = self.dist_env.copy()
															UpperCamelCase__         = state_dict_type
															with mockenv_context(**lowercase			):
																					UpperCamelCase__         = 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 				A      (   self							:						Optional[int]			)							->    Optional[Any]:

									'''simple docstring'''







									UpperCamelCase__         = AutoModel.from_pretrained(lowercase			)
									for policy in FSDP_AUTO_WRAP_POLICY:
															UpperCamelCase__         = self.dist_env.copy()
															UpperCamelCase__         = policy
															if policy == "TRANSFORMER_BASED_WRAP":
																					UpperCamelCase__         = """BertLayer"""
															elif policy == "SIZE_BASED_WRAP":
																					UpperCamelCase__         = """2000"""
															with mockenv_context(**lowercase			):
																					UpperCamelCase__         = 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			)

									UpperCamelCase__         = self.dist_env.copy()
									UpperCamelCase__         = """TRANSFORMER_BASED_WRAP"""
									UpperCamelCase__         = """T5Layer"""
									with mockenv_context(**lowercase			):
															UpperCamelCase__         = 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			)			)

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






			def 				A      (   self							:						Tuple			)							->    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:
															UpperCamelCase__         = self.dist_env.copy()
															UpperCamelCase__         = mp_dtype
															with mockenv_context(**lowercase			):
																					UpperCamelCase__         = Accelerator()
																					if mp_dtype == "fp16":
																											UpperCamelCase__         = torch.floataa
																					elif mp_dtype == "bf16":
																											UpperCamelCase__         = torch.bfloataa
																					UpperCamelCase__         = 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 				A      (   self							:						Any			)							->    List[str]:

									'''simple docstring'''







									from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

									for flag in [True, False]:
															UpperCamelCase__         = self.dist_env.copy()
															UpperCamelCase__         = str(lowercase			).lower()
															with mockenv_context(**lowercase			):
																					UpperCamelCase__         = FullyShardedDataParallelPlugin()
																					self.assertEqual(fsdp_plugin.cpu_offload ,      CPUOffload(offload_params=lowercase			)			)
@require_fsdp
@require_multi_gpu
@slow
class    _SCREAMING_SNAKE_CASE							(		SCREAMING_SNAKE_CASE	):


			'''simple docstring'''





			def 				A      (   self							:						Optional[int]			)							->    str:

									'''simple docstring'''







									super().setUp()
									UpperCamelCase__         = 0.8_2
									UpperCamelCase__         = [
									    """fsdp_shard_grad_op_transformer_based_wrap""",
									    """fsdp_full_shard_transformer_based_wrap""",
									]
									UpperCamelCase__         = {
									    """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
									}
									UpperCamelCase__         = 1_6_0
									UpperCamelCase__         = 1_6_0

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






			def 				A      (   self							:						List[Any]			)							->    Dict:

									'''simple docstring'''







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

									'''simple docstring'''







									UpperCamelCase__         = os.path.join(self.test_scripts_folder ,      """test_checkpointing.py"""			)
									UpperCamelCase__         = [
									    """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			):
															UpperCamelCase__         = cmd.copy()
															cmd_config.append(f"--fsdp_sharding_strategy={i+1}"			)
															if strategy != "FULL_SHARD":
																					continue
															UpperCamelCase__         = len(lowercase			)
															for state_dict_type in FSDP_STATE_DICT_TYPE:
																					UpperCamelCase__         = 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()			)

																					UpperCamelCase__         = cmd_config[:-1]
																					UpperCamelCase__         = 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 				A      (   self							:						Tuple			)							->    Union[str, Any]:

									'''simple docstring'''







									UpperCamelCase__         = os.path.join(self.test_scripts_folder ,      """test_peak_memory_usage.py"""			)
									UpperCamelCase__         = [
									    """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():
															UpperCamelCase__         = 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'''


# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation

import warnings

from .state import AcceleratorState, GradientState


warnings.filterwarnings('''ignore''', category=UserWarning, module='''torch.optim.lr_scheduler''')
class    _SCREAMING_SNAKE_CASE							:


			'''simple docstring'''





			def __init__(   self							:						Optional[int] ,      lowercase							:						Optional[int] ,      lowercase							:						Optional[Any] ,      lowercase							:						bool = True ,      lowercase							:						bool = False			)							->    Optional[int]:

									'''simple docstring'''







									UpperCamelCase__         = scheduler
									UpperCamelCase__         = optimizers if isinstance(lowercase ,      (list, tuple)			) else [optimizers]
									UpperCamelCase__         = split_batches
									UpperCamelCase__         = step_with_optimizer
									UpperCamelCase__         = GradientState()






			def 				A      (   self							:						int ,      *lowercase							:						int ,      **lowercase							:						str			)							->    str:

									'''simple docstring'''







									if not self.step_with_optimizer:
															# No link between scheduler and optimizer -> just step
															self.scheduler.step(*lowercase ,      **lowercase			)
															return

									# Otherwise, first make sure the optimizer was stepped.
									if not self.gradient_state.sync_gradients:
															if self.gradient_state.adjust_scheduler:
																					self.scheduler._step_count += 1
															return

									for opt in self.optimizers:
															if opt.step_was_skipped:
																					return
									if self.split_batches:
															# Split batches -> the training dataloader batch size is not changed so one step per training step
															self.scheduler.step(*lowercase ,      **lowercase			)
									else:
															# Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
															# num_processes steps per training step
															UpperCamelCase__         = AcceleratorState().num_processes
															for _ in range(lowercase			):
																					# Special case when using OneCycle and `drop_last` was not used
																					if hasattr(self.scheduler ,      """total_steps"""			):
																											if self.scheduler._step_count <= self.scheduler.total_steps:
																																	self.scheduler.step(*lowercase ,      **lowercase			)
																					else:
																											self.scheduler.step(*lowercase ,      **lowercase			)






			def 				A      (   self							:						List[str]			)							->    Optional[Any]:

									'''simple docstring'''







									return self.scheduler.get_last_lr()






			def 				A      (   self							:						Any			)							->    int:

									'''simple docstring'''







									return self.scheduler.state_dict()






			def 				A      (   self							:						Any ,      lowercase							:						int			)							->    Optional[int]:

									'''simple docstring'''







									self.scheduler.load_state_dict(lowercase			)






			def 				A      (   self							:						str			)							->    int:

									'''simple docstring'''







									return self.scheduler.get_lr()







			def 				A      (   self							:						Union[str, Any] ,      *lowercase							:						List[Any] ,      **lowercase							:						str			)							->    str:

									'''simple docstring'''







									return self.scheduler.print_lr(*lowercase ,      **lowercase			)
