



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)

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



@require_fsdp
@require_cuda
class 				__snake_case						(     UpperCamelCase_						):




							def    UpperCAmelCase__			(       self     :       str):
									super().setUp()

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




							def    UpperCAmelCase__			(       self     :       Any):
									from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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




							def    UpperCAmelCase__			(       self     :       Any):
									from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

									for i, prefetch_policy in enumerate(A_):
											lowerCAmelCase_			:     Tuple									= self.dist_env.copy()
											lowerCAmelCase_			:     str									= prefetch_policy
											with mockenv_context(**A_):
													lowerCAmelCase_			:     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     :       int):
									from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

									for i, state_dict_type in enumerate(A_):
											lowerCAmelCase_			:     Optional[int]									= self.dist_env.copy()
											lowerCAmelCase_			:     str									= state_dict_type
											with mockenv_context(**A_):
													lowerCAmelCase_			:     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     :       int):
									lowerCAmelCase_			:     str									= AutoModel.from_pretrained(A_)
									for policy in FSDP_AUTO_WRAP_POLICY:
											lowerCAmelCase_			:     Tuple									= self.dist_env.copy()
											lowerCAmelCase_			:     List[str]									= policy
											if policy == "TRANSFORMER_BASED_WRAP":
													lowerCAmelCase_			:     List[str]									= '''BertLayer'''
											elif policy == "SIZE_BASED_WRAP":
													lowerCAmelCase_			:     Optional[int]									= '''2000'''
											with mockenv_context(**A_):
													lowerCAmelCase_			:     Tuple									= 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)

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

									lowerCAmelCase_			:     Union[str, Any]									= self.dist_env.copy()
									lowerCAmelCase_			:     List[Any]									= '''SIZE_BASED_WRAP'''
									lowerCAmelCase_			:     List[Any]									= '''0'''
									with mockenv_context(**A_):
											lowerCAmelCase_			:     str									= FullyShardedDataParallelPlugin()
											fsdp_plugin.set_auto_wrap_policy(A_)
											self.assertIsNone(fsdp_plugin.auto_wrap_policy)




							def    UpperCAmelCase__			(       self     :       Optional[Any]):
									from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
									from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

									for mp_dtype in dtypes:
											lowerCAmelCase_			:     Dict									= self.dist_env.copy()
											lowerCAmelCase_			:     List[str]									= mp_dtype
											with mockenv_context(**A_):
													lowerCAmelCase_			:     List[Any]									= Accelerator()
													if mp_dtype == "fp16":
															lowerCAmelCase_			:     Optional[Any]									= torch.floataa
													elif mp_dtype == "bf16":
															lowerCAmelCase_			:     Any									= torch.bfloataa
													lowerCAmelCase_			:     Optional[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    UpperCAmelCase__			(       self     :       Any):
									from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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





@require_fsdp
@require_multi_gpu
@slow
class 				__snake_case						(     UpperCamelCase_						):




							def    UpperCAmelCase__			(       self     :       List[Any]):
									super().setUp()
									lowerCAmelCase_			:     Optional[Any]									= 0.82
									lowerCAmelCase_			:     Union[str, Any]									= [
									    '''fsdp_shard_grad_op_transformer_based_wrap''',
									    '''fsdp_full_shard_transformer_based_wrap''',
									]
									lowerCAmelCase_			:     Union[str, Any]									= {
									    '''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
									}
									lowerCAmelCase_			:     Union[str, Any]									= 1_6_0
									lowerCAmelCase_			:     List[Any]									= 1_6_0

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




							def    UpperCAmelCase__			(       self     :       Optional[int]):
									lowerCAmelCase_			:     Tuple									= os.path.join(self.test_scripts_folder				,				'''test_performance.py''')
									lowerCAmelCase_			:     Tuple									= ['''accelerate''', '''launch''', '''--num_processes=2''', '''--num_machines=1''', '''--machine_rank=0''', '''--use_fsdp''']
									for config in self.performance_configs:
											lowerCAmelCase_			:     Union[str, Any]									= 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    UpperCAmelCase__			(       self     :       Tuple):
									lowerCAmelCase_			:     Union[str, Any]									= os.path.join(self.test_scripts_folder				,				'''test_checkpointing.py''')
									lowerCAmelCase_			:     List[str]									= [
									    '''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_):
											lowerCAmelCase_			:     List[Any]									= cmd.copy()
											cmd_config.append(F"""--fsdp_sharding_strategy={i+1}""")
											if strategy != "FULL_SHARD":
													continue
											lowerCAmelCase_			:     int									= len(A_)
											for state_dict_type in FSDP_STATE_DICT_TYPE:
													lowerCAmelCase_			:     List[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())

													lowerCAmelCase_			:     int									= cmd_config[:-1]
													lowerCAmelCase_			:     List[str]									= 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    UpperCAmelCase__			(       self     :       Dict):
									lowerCAmelCase_			:     int									= os.path.join(self.test_scripts_folder				,				'''test_peak_memory_usage.py''')
									lowerCAmelCase_			:     Dict									= [
									    '''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():
											lowerCAmelCase_			:     Tuple									= 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())










import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
						from sagemaker import Session, TrainingJobAnalytics
						from sagemaker.huggingface import HuggingFace



@pytest.mark.skipif(
    literal_eval(os.getenv('''TEST_SAGEMAKER'''						,'''False'''						)						) is not True						,reason='''Skipping test because should only be run when releasing minor transformers version'''						,)
@pytest.mark.usefixtures('''sm_env'''						)
@parameterized_class(
    [
        {
            '''framework''': '''pytorch''',
            '''script''': '''run_glue.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.g4dn.xlarge''',
            '''results''': {'''train_runtime''': 650, '''eval_accuracy''': 0.6, '''eval_loss''': 0.9},
        },
        {
            '''framework''': '''tensorflow''',
            '''script''': '''run_tf.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.g4dn.xlarge''',
            '''results''': {'''train_runtime''': 600, '''eval_accuracy''': 0.3, '''eval_loss''': 0.9},
        },
    ]						)
class 				__snake_case						(     unittest.TestCase						):




							def    UpperCAmelCase__			(       self     :       int):
									if self.framework == "pytorch":
											subprocess.run(
											    F"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()				,				encoding='''utf-8'''				,				check=A_				,				)
									assert hasattr(self				,				'''env''')




							def    UpperCAmelCase__			(       self     :       Union[str, Any]				,				A_     :       str=1):
									# creates estimator
									return HuggingFace(
									    entry_point=self.script				,				source_dir=self.env.test_path				,				role=self.env.role				,				image_uri=self.env.image_uri				,				base_job_name=F"""{self.env.base_job_name}-single"""				,				instance_count=A_				,				instance_type=self.instance_type				,				debugger_hook_config=A_				,				hyperparameters={**self.env.hyperparameters, '''model_name_or_path''': self.model_name_or_path}				,				metric_definitions=self.env.metric_definitions				,				py_version='''py36'''				,				)




							def    UpperCAmelCase__			(       self     :       List[str]				,				A_     :       Optional[Any]):
									TrainingJobAnalytics(A_).export_csv(F"""{self.env.test_path}/{job_name}_metrics.csv""")







							def    UpperCAmelCase__			(       self     :       int):
									# create estimator
									lowerCAmelCase_			:     List[Any]									= self.create_estimator()

									# run training
									estimator.fit()

									# result dataframe
									lowerCAmelCase_			:     Union[str, Any]									= TrainingJobAnalytics(estimator.latest_training_job.name).dataframe()

									# extract kpis
									lowerCAmelCase_			:     Optional[Any]									= list(result_metrics_df[result_metrics_df.metric_name == '''eval_accuracy''']['''value'''])
									lowerCAmelCase_			:     Optional[int]									= list(result_metrics_df[result_metrics_df.metric_name == '''eval_loss''']['''value'''])
									# get train time from SageMaker job, this includes starting, preprocessing, stopping
									lowerCAmelCase_			:     Dict									= (
									    Session().describe_training_job(estimator.latest_training_job.name).get('''TrainingTimeInSeconds'''				,				9_9_9_9_9_9)
									)

									# assert kpis
									assert train_runtime <= self.results["train_runtime"]
									assert all(t >= self.results['''eval_accuracy'''] for t in eval_accuracy)
									assert all(t <= self.results['''eval_loss'''] for t in eval_loss)

									# dump tests result into json file to share in PR
									with open(F"""{estimator.latest_training_job.name}.json"""				,				'''w''') as outfile:
											json.dump({'''train_time''': train_runtime, '''eval_accuracy''': eval_accuracy, '''eval_loss''': eval_loss}				,				A_)





