



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)

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



@require_fsdp
@require_cuda
class 	lowerCamelCase__      (							__snake_case):




		'''simple docstring'''



		def 		lowerCAmelCase__				(self			)					->      Optional[int]:

						"""simple docstring"""






						super().setUp()

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



		def 		lowerCAmelCase__				(self			)					->      Dict:

						"""simple docstring"""






						from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

						for i, strategy in enumerate(lowerCamelCase_			):
										lowerCAmelCase__					:    Tuple												=					self.dist_env.copy()
										lowerCAmelCase__					:    int												=					f"""{i + 1}"""
										lowerCAmelCase__					:    Optional[int]												=					strategy
										with mockenv_context(**lowerCamelCase_			):
														lowerCAmelCase__					:    List[str]												=					FullyShardedDataParallelPlugin()
														self.assertEqual(fsdp_plugin.sharding_strategy		,ShardingStrategy(i + 1			)			)



		def 		lowerCAmelCase__				(self			)					->      Dict:

						"""simple docstring"""






						from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

						for i, prefetch_policy in enumerate(lowerCamelCase_			):
										lowerCAmelCase__					:    int												=					self.dist_env.copy()
										lowerCAmelCase__					:    str												=					prefetch_policy
										with mockenv_context(**lowerCamelCase_			):
														lowerCAmelCase__					:    Tuple												=					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			)					->      Any:

						"""simple docstring"""






						from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

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

						"""simple docstring"""






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

						lowerCAmelCase__					:    Dict												=					self.dist_env.copy()
						lowerCAmelCase__					:    List[str]												=					"""TRANSFORMER_BASED_WRAP"""
						lowerCAmelCase__					:    str												=					"""T5Layer"""
						with mockenv_context(**lowerCamelCase_			):
										lowerCAmelCase__					:    Tuple												=					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			)			)

						lowerCAmelCase__					:    Optional[int]												=					self.dist_env.copy()
						lowerCAmelCase__					:    int												=					"""SIZE_BASED_WRAP"""
						lowerCAmelCase__					:    List[Any]												=					"""0"""
						with mockenv_context(**lowerCamelCase_			):
										lowerCAmelCase__					:    Optional[Any]												=					FullyShardedDataParallelPlugin()
										fsdp_plugin.set_auto_wrap_policy(lowerCamelCase_			)
										self.assertIsNone(fsdp_plugin.auto_wrap_policy			)



		def 		lowerCAmelCase__				(self			)					->      str:

						"""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:
										lowerCAmelCase__					:    List[str]												=					self.dist_env.copy()
										lowerCAmelCase__					:    List[str]												=					mp_dtype
										with mockenv_context(**lowerCamelCase_			):
														lowerCAmelCase__					:    int												=					Accelerator()
														if mp_dtype == "fp16":
																		lowerCAmelCase__					:    List[Any]												=					torch.floataa
														elif mp_dtype == "bf16":
																		lowerCAmelCase__					:    Optional[Any]												=					torch.bfloataa
														lowerCAmelCase__					:    Dict												=					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 		lowerCAmelCase__				(self			)					->      Union[str, Any]:

						"""simple docstring"""






						from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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



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




		'''simple docstring'''



		def 		lowerCAmelCase__				(self			)					->      Optional[int]:

						"""simple docstring"""






						super().setUp()
						lowerCAmelCase__					:    List[str]												=					0.82
						lowerCAmelCase__					:    List[str]												=					[
						    """fsdp_shard_grad_op_transformer_based_wrap""",
						    """fsdp_full_shard_transformer_based_wrap""",
						]
						lowerCAmelCase__					:    List[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
						}
						lowerCAmelCase__					:    List[Any]												=					1_60
						lowerCAmelCase__					:    str												=					1_60

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



		def 		lowerCAmelCase__				(self			)					->      Dict:

						"""simple docstring"""






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

						"""simple docstring"""






						lowerCAmelCase__					:    Dict												=					os.path.join(self.test_scripts_folder		,'''test_checkpointing.py'''			)
						lowerCAmelCase__					:    Tuple												=					[
						    """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_			):
										lowerCAmelCase__					:    Any												=					cmd.copy()
										cmd_config.append(f"""--fsdp_sharding_strategy={i+1}"""			)
										if strategy != "FULL_SHARD":
														continue
										lowerCAmelCase__					:    List[Any]												=					len(lowerCamelCase_			)
										for state_dict_type in FSDP_STATE_DICT_TYPE:
														lowerCAmelCase__					:    Union[str, Any]												=					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()			)

														lowerCAmelCase__					:    Union[str, Any]												=					cmd_config[:-1]
														lowerCAmelCase__					:    Optional[int]												=					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 		lowerCAmelCase__				(self			)					->      List[str]:

						"""simple docstring"""






						lowerCAmelCase__					:    Union[str, Any]												=					os.path.join(self.test_scripts_folder		,'''test_peak_memory_usage.py'''			)
						lowerCAmelCase__					:    List[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():
										lowerCAmelCase__					:    Union[str, 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(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()			)






import unittest

import numpy as np

from transformers.testing_utils import is_flaky, require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
							import torch

if is_vision_available():
							from PIL import Image

							from transformers import DonutImageProcessor



class 	lowerCamelCase__      (							unittest.TestCase):




		'''simple docstring'''



		def __init__(self		,__lowerCamelCase		,__lowerCamelCase=7		,__lowerCamelCase=3		,__lowerCamelCase=18		,__lowerCamelCase=30		,__lowerCamelCase=4_00		,__lowerCamelCase=True		,__lowerCamelCase=None		,__lowerCamelCase=True		,__lowerCamelCase=False		,__lowerCamelCase=True		,__lowerCamelCase=True		,__lowerCamelCase=[0.5, 0.5, 0.5]		,__lowerCamelCase=[0.5, 0.5, 0.5]		,)					->      Union[str, Any]:

						"""simple docstring"""






						lowerCAmelCase__					:    List[Any]												=					parent
						lowerCAmelCase__					:    Optional[Any]												=					batch_size
						lowerCAmelCase__					:    str												=					num_channels
						lowerCAmelCase__					:    List[str]												=					image_size
						lowerCAmelCase__					:    List[Any]												=					min_resolution
						lowerCAmelCase__					:    Union[str, Any]												=					max_resolution
						lowerCAmelCase__					:    Union[str, Any]												=					do_resize
						lowerCAmelCase__					:    str												=					size if size is not None else {'''height''': 18, '''width''': 20}
						lowerCAmelCase__					:    List[str]												=					do_thumbnail
						lowerCAmelCase__					:    str												=					do_align_axis
						lowerCAmelCase__					:    Optional[Any]												=					do_pad
						lowerCAmelCase__					:    Tuple												=					do_normalize
						lowerCAmelCase__					:    List[Any]												=					image_mean
						lowerCAmelCase__					:    List[str]												=					image_std



		def 		lowerCAmelCase__				(self			)					->      Optional[Any]:

						"""simple docstring"""






						return {
						    "do_resize": self.do_resize,
						    "size": self.size,
						    "do_thumbnail": self.do_thumbnail,
						    "do_align_long_axis": self.do_align_axis,
						    "do_pad": self.do_pad,
						    "do_normalize": self.do_normalize,
						    "image_mean": self.image_mean,
						    "image_std": self.image_std,
						}



@require_torch
@require_vision
class 	lowerCamelCase__      (							lowerCamelCase__				,    unittest.TestCase):




		'''simple docstring'''



		snake_case_     =DonutImageProcessor if is_vision_available() else None



		def 		lowerCAmelCase__				(self			)					->      List[Any]:

						"""simple docstring"""






						lowerCAmelCase__					:    Union[str, Any]												=					DonutImageProcessingTester(self			)



		@property
		def 		lowerCAmelCase__				(self			)					->      Any:

						"""simple docstring"""






						return self.image_processor_tester.prepare_image_processor_dict()



		def 		lowerCAmelCase__				(self			)					->      int:

						"""simple docstring"""






						lowerCAmelCase__					:    Optional[Any]												=					self.image_processing_class(**self.image_processor_dict			)
						self.assertTrue(hasattr(__lowerCamelCase		,'''do_resize'''			)			)
						self.assertTrue(hasattr(__lowerCamelCase		,'''size'''			)			)
						self.assertTrue(hasattr(__lowerCamelCase		,'''do_thumbnail'''			)			)
						self.assertTrue(hasattr(__lowerCamelCase		,'''do_align_long_axis'''			)			)
						self.assertTrue(hasattr(__lowerCamelCase		,'''do_pad'''			)			)
						self.assertTrue(hasattr(__lowerCamelCase		,'''do_normalize'''			)			)
						self.assertTrue(hasattr(__lowerCamelCase		,'''image_mean'''			)			)
						self.assertTrue(hasattr(__lowerCamelCase		,'''image_std'''			)			)



		def 		lowerCAmelCase__				(self			)					->      Tuple:

						"""simple docstring"""






						lowerCAmelCase__					:    Optional[int]												=					self.image_processing_class.from_dict(self.image_processor_dict			)
						self.assertEqual(image_processor.size		,{'''height''': 18, '''width''': 20}			)

						lowerCAmelCase__					:    int												=					self.image_processing_class.from_dict(self.image_processor_dict		,size=42			)
						self.assertEqual(image_processor.size		,{'''height''': 42, '''width''': 42}			)

						# Previous config had dimensions in (width, height) order
						lowerCAmelCase__					:    Union[str, Any]												=					self.image_processing_class.from_dict(self.image_processor_dict		,size=(42, 84)			)
						self.assertEqual(image_processor.size		,{'''height''': 84, '''width''': 42}			)



		def 		lowerCAmelCase__				(self			)					->      Optional[Any]:

						"""simple docstring"""






						pass



		@is_flaky()
		def 		lowerCAmelCase__				(self			)					->      Optional[Any]:

						"""simple docstring"""






						lowerCAmelCase__					:    str												=					self.image_processing_class(**self.image_processor_dict			)
						# create random PIL images
						lowerCAmelCase__					:    Dict												=					prepare_image_inputs(self.image_processor_tester		,equal_resolution=__lowerCamelCase			)
						for image in image_inputs:
										self.assertIsInstance(__lowerCamelCase		,Image.Image			)

						# Test not batched input
						lowerCAmelCase__					:    Optional[Any]												=					image_processing(image_inputs[0]		,return_tensors='''pt'''			).pixel_values
						self.assertEqual(
						    encoded_images.shape		,(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.size['''height'''],
						        self.image_processor_tester.size['''width'''],
						    )		,)

						# Test batched
						lowerCAmelCase__					:    Optional[int]												=					image_processing(__lowerCamelCase		,return_tensors='''pt'''			).pixel_values
						self.assertEqual(
						    encoded_images.shape		,(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.size['''height'''],
						        self.image_processor_tester.size['''width'''],
						    )		,)



		@is_flaky()
		def 		lowerCAmelCase__				(self			)					->      Any:

						"""simple docstring"""






						lowerCAmelCase__					:    Optional[int]												=					self.image_processing_class(**self.image_processor_dict			)
						# create random numpy tensors
						lowerCAmelCase__					:    Any												=					prepare_image_inputs(self.image_processor_tester		,equal_resolution=__lowerCamelCase		,numpify=__lowerCamelCase			)
						for image in image_inputs:
										self.assertIsInstance(__lowerCamelCase		,np.ndarray			)

						# Test not batched input
						lowerCAmelCase__					:    Union[str, Any]												=					image_processing(image_inputs[0]		,return_tensors='''pt'''			).pixel_values
						self.assertEqual(
						    encoded_images.shape		,(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.size['''height'''],
						        self.image_processor_tester.size['''width'''],
						    )		,)

						# Test batched
						lowerCAmelCase__					:    int												=					image_processing(__lowerCamelCase		,return_tensors='''pt'''			).pixel_values
						self.assertEqual(
						    encoded_images.shape		,(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.size['''height'''],
						        self.image_processor_tester.size['''width'''],
						    )		,)



		@is_flaky()
		def 		lowerCAmelCase__				(self			)					->      Optional[Any]:

						"""simple docstring"""






						lowerCAmelCase__					:    Any												=					self.image_processing_class(**self.image_processor_dict			)
						# create random PyTorch tensors
						lowerCAmelCase__					:    int												=					prepare_image_inputs(self.image_processor_tester		,equal_resolution=__lowerCamelCase		,torchify=__lowerCamelCase			)
						for image in image_inputs:
										self.assertIsInstance(__lowerCamelCase		,torch.Tensor			)

						# Test not batched input
						lowerCAmelCase__					:    Any												=					image_processing(image_inputs[0]		,return_tensors='''pt'''			).pixel_values
						self.assertEqual(
						    encoded_images.shape		,(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.size['''height'''],
						        self.image_processor_tester.size['''width'''],
						    )		,)

						# Test batched
						lowerCAmelCase__					:    Dict												=					image_processing(__lowerCamelCase		,return_tensors='''pt'''			).pixel_values
						self.assertEqual(
						    encoded_images.shape		,(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.size['''height'''],
						        self.image_processor_tester.size['''width'''],
						    )		,)

