
"""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__        =		"""bert-base-cased"""
lowerCamelCase__        =		"""fp16"""
lowerCamelCase__        =		"""bf16"""
lowerCamelCase__        =		[FPaa, BFaa]




@require_fsdp
@require_cuda
class 			A__	(     _lowerCamelCase):

		def 							__lowerCamelCase							(				self			):
			super().setUp()

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

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

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

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

			for i, prefetch_policy in enumerate(_SCREAMING_SNAKE_CASE			):
				__lowerCAmelCase      :						Dict						      =				self.dist_env.copy()
				__lowerCAmelCase      :						Optional[int]						      =				prefetch_policy
				with mockenv_context(**_SCREAMING_SNAKE_CASE			):
					__lowerCAmelCase      :						Union[str, Any]						      =				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			):
			from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

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

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

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

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







		def 							__lowerCamelCase							(				self			):
			from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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




@require_fsdp
@require_multi_gpu
@slow
class 			A__	(     _lowerCamelCase):

		def 							__lowerCamelCase							(				self			):
			super().setUp()
			__lowerCAmelCase      :						Dict						      =				0.82
			__lowerCAmelCase      :						List[Any]						      =				[
			    'fsdp_shard_grad_op_transformer_based_wrap',
			    'fsdp_full_shard_transformer_based_wrap',
			]
			__lowerCAmelCase      :						List[Any]						      =				{
			    '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      :						Dict						      =				1_60
			__lowerCAmelCase      :						int						      =				1_60

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

		def 							__lowerCamelCase							(				self			):
			__lowerCAmelCase      :						List[str]						      =				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      :						Tuple						      =				cmd.copy()
				for i, strategy in enumerate(_SCREAMING_SNAKE_CASE			):
					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(_SCREAMING_SNAKE_CASE	,		env=os.environ.copy()			)

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

					__lowerCAmelCase      :						Optional[int]						      =				cmd_config[:-1]
					__lowerCAmelCase      :						List[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(_SCREAMING_SNAKE_CASE	,		env=os.environ.copy()			)







		def 							__lowerCamelCase							(				self			):
			__lowerCAmelCase      :						Tuple						      =				os.path.join(self.test_scripts_folder	,		'test_peak_memory_usage.py'			)
			__lowerCAmelCase      :						List[str]						      =				[
			    '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      :						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(_SCREAMING_SNAKE_CASE			):
						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(_SCREAMING_SNAKE_CASE	,		env=os.environ.copy()			)

"""simple docstring"""







import copy
from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import ClassLabel, Features, Value
from .base import TaskTemplate




@dataclass(frozen=_lowerCamelCase)
class 			A__	(     _lowerCamelCase):
		# `task` is not a ClassVar since we want it to be part of the `asdict` output for JSON serialization
		A_						:	str           =     field(default='text-classification'	,					metadata={'include_in_asdict_even_if_is_default': True})
		A_						:	ClassVar[Features]           =     Features({'text': Value('string')})
		A_						:	ClassVar[Features]           =     Features({'labels': ClassLabel})
		A_						:	str           =     "text"
		A_						:	str           =     "labels"

		def 							__lowerCamelCase							(				self	,		_SCREAMING_SNAKE_CASE			):
			if self.label_column not in features:
				raise ValueError(f"Column {self.label_column} is not present in features."			)
			if not isinstance(features[self.label_column]	,		_SCREAMING_SNAKE_CASE			):
				raise ValueError(f"Column {self.label_column} is not a ClassLabel."			)
			__lowerCAmelCase      :						Any						      =				copy.deepcopy(self			)
			__lowerCAmelCase      :						Dict						      =				self.label_schema.copy()
			__lowerCAmelCase      :						List[Any]						      =				features[self.label_column]
			__lowerCAmelCase      :						Dict						      =				label_schema
			return task_template







		@property
		def 							__lowerCamelCase							(				self			):
			return {
			    self.text_column: "text",
			    self.label_column: "labels",
			}