
'''simple docstring'''
import unittest

from transformers import SqueezeBertConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import (
						    SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
						    SqueezeBertForMaskedLM,
						    SqueezeBertForMultipleChoice,
						    SqueezeBertForQuestionAnswering,
						    SqueezeBertForSequenceClassification,
						    SqueezeBertForTokenClassification,
						    SqueezeBertModel,
						)


class       _lowerCAmelCase   (   __snake_case				):







							'''simple docstring'''

							def __init__(self  ,  UpperCAmelCase  ,  UpperCAmelCase=13  ,  UpperCAmelCase=7  ,  UpperCAmelCase=True  ,  UpperCAmelCase=True  ,  UpperCAmelCase=False  ,  UpperCAmelCase=True  ,  UpperCAmelCase=99  ,  UpperCAmelCase=32  ,  UpperCAmelCase=5  ,  UpperCAmelCase=4  ,  UpperCAmelCase=64  ,  UpperCAmelCase="gelu"  ,  UpperCAmelCase=0.1  ,  UpperCAmelCase=0.1  ,  UpperCAmelCase=512  ,  UpperCAmelCase=16  ,  UpperCAmelCase=2  ,  UpperCAmelCase=0.02  ,  UpperCAmelCase=3  ,  UpperCAmelCase=4  ,  UpperCAmelCase=None  ,  UpperCAmelCase=2  ,  UpperCAmelCase=2  ,  UpperCAmelCase=2  ,  UpperCAmelCase=2  ,  UpperCAmelCase=4  ,  UpperCAmelCase=1  ,  )					->       Any:
														_snake_case       =  parent
														_snake_case       =  batch_size
														_snake_case       =  seq_length
														_snake_case       =  is_training
														_snake_case       =  use_input_mask
														_snake_case       =  use_token_type_ids
														_snake_case       =  use_labels
														_snake_case       =  vocab_size
														_snake_case       =  hidden_size
														_snake_case       =  num_hidden_layers
														_snake_case       =  num_attention_heads
														_snake_case       =  intermediate_size
														_snake_case       =  hidden_act
														_snake_case       =  hidden_dropout_prob
														_snake_case       =  attention_probs_dropout_prob
														_snake_case       =  max_position_embeddings
														_snake_case       =  type_vocab_size
														_snake_case       =  type_sequence_label_size
														_snake_case       =  initializer_range
														_snake_case       =  num_labels
														_snake_case       =  num_choices
														_snake_case       =  scope
														_snake_case       =  q_groups
														_snake_case       =  k_groups
														_snake_case       =  v_groups
														_snake_case       =  post_attention_groups
														_snake_case       =  intermediate_groups
														_snake_case       =  output_groups

							def 	lowercase							(self  )					->       Any:
														_snake_case       =  ids_tensor([self.batch_size, self.seq_length]  ,  self.vocab_size  )

														_snake_case       =  None
														if self.use_input_mask:
																					_snake_case       =  random_attention_mask([self.batch_size, self.seq_length]  )

														_snake_case       =  None
														_snake_case       =  None
														_snake_case       =  None
														if self.use_labels:
																					_snake_case       =  ids_tensor([self.batch_size]  ,  self.type_sequence_label_size  )
																					_snake_case       =  ids_tensor([self.batch_size, self.seq_length]  ,  self.num_labels  )
																					_snake_case       =  ids_tensor([self.batch_size]  ,  self.num_choices  )

														_snake_case       =  self.get_config()

														return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels

							def 	lowercase							(self  )					->       Union[str, Any]:
														return SqueezeBertConfig(
														    embedding_size=self.hidden_size  ,  vocab_size=self.vocab_size  ,  hidden_size=self.hidden_size  ,  num_hidden_layers=self.num_hidden_layers  ,  num_attention_heads=self.num_attention_heads  ,  intermediate_size=self.intermediate_size  ,  hidden_act=self.hidden_act  ,  attention_probs_dropout_prob=self.hidden_dropout_prob  ,  attention_dropout=self.attention_probs_dropout_prob  ,  max_position_embeddings=self.max_position_embeddings  ,  initializer_range=self.initializer_range  ,  q_groups=self.q_groups  ,  k_groups=self.k_groups  ,  v_groups=self.v_groups  ,  post_attention_groups=self.post_attention_groups  ,  intermediate_groups=self.intermediate_groups  ,  output_groups=self.output_groups  ,  )

							def 	lowercase							(self  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  )					->       Optional[int]:
														_snake_case       =  SqueezeBertModel(config=UpperCAmelCase  )
														model.to(UpperCAmelCase  )
														model.eval()
														_snake_case       =  model(UpperCAmelCase  ,  UpperCAmelCase  )
														_snake_case       =  model(UpperCAmelCase  )
														self.parent.assertEqual(result.last_hidden_state.shape  ,  (self.batch_size, self.seq_length, self.hidden_size)  )

							def 	lowercase							(self  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  )					->       Optional[Any]:
														_snake_case       =  SqueezeBertForMaskedLM(config=UpperCAmelCase  )
														model.to(UpperCAmelCase  )
														model.eval()
														_snake_case       =  model(UpperCAmelCase  ,  attention_mask=UpperCAmelCase  ,  labels=UpperCAmelCase  )
														self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.vocab_size)  )

							def 	lowercase							(self  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  )					->       List[Any]:
														_snake_case       =  SqueezeBertForQuestionAnswering(config=UpperCAmelCase  )
														model.to(UpperCAmelCase  )
														model.eval()
														_snake_case       =  model(
														    UpperCAmelCase  ,  attention_mask=UpperCAmelCase  ,  start_positions=UpperCAmelCase  ,  end_positions=UpperCAmelCase  )
														self.parent.assertEqual(result.start_logits.shape  ,  (self.batch_size, self.seq_length)  )
														self.parent.assertEqual(result.end_logits.shape  ,  (self.batch_size, self.seq_length)  )

							def 	lowercase							(self  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  )					->       Any:
														_snake_case       =  self.num_labels
														_snake_case       =  SqueezeBertForSequenceClassification(UpperCAmelCase  )
														model.to(UpperCAmelCase  )
														model.eval()
														_snake_case       =  model(UpperCAmelCase  ,  attention_mask=UpperCAmelCase  ,  labels=UpperCAmelCase  )
														self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_labels)  )

							def 	lowercase							(self  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  )					->       Optional[Any]:
														_snake_case       =  self.num_labels
														_snake_case       =  SqueezeBertForTokenClassification(config=UpperCAmelCase  )
														model.to(UpperCAmelCase  )
														model.eval()

														_snake_case       =  model(UpperCAmelCase  ,  attention_mask=UpperCAmelCase  ,  labels=UpperCAmelCase  )
														self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.num_labels)  )

							def 	lowercase							(self  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  )					->       Union[str, Any]:
														_snake_case       =  self.num_choices
														_snake_case       =  SqueezeBertForMultipleChoice(config=UpperCAmelCase  )
														model.to(UpperCAmelCase  )
														model.eval()
														_snake_case       =  input_ids.unsqueeze(1  ).expand(-1  ,  self.num_choices  ,  -1  ).contiguous()
														_snake_case       =  input_mask.unsqueeze(1  ).expand(-1  ,  self.num_choices  ,  -1  ).contiguous()
														_snake_case       =  model(
														    UpperCAmelCase  ,  attention_mask=UpperCAmelCase  ,  labels=UpperCAmelCase  ,  )
														self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_choices)  )





							def 	lowercase							(self  )					->       int:
														_snake_case       =  self.prepare_config_and_inputs()
														((_snake_case),							(_snake_case),							(_snake_case),							(_snake_case),							(_snake_case),							(_snake_case))       =  config_and_inputs
														_snake_case       =  {"""input_ids""": input_ids, """attention_mask""": input_mask}
														return config, inputs_dict


@require_torch
class       _lowerCAmelCase   (   __snake_case					,							__snake_case					,							unittest.TestCase				):







							'''simple docstring'''

							lowerCAmelCase_ 						=						(
							    (
							        SqueezeBertModel,
							        SqueezeBertForMaskedLM,
							        SqueezeBertForMultipleChoice,
							        SqueezeBertForQuestionAnswering,
							        SqueezeBertForSequenceClassification,
							        SqueezeBertForTokenClassification,
							    )
							    if is_torch_available()
							    else None
							)
							lowerCAmelCase_ 						=						(
							    {
							        "feature-extraction": SqueezeBertModel,
							        "fill-mask": SqueezeBertForMaskedLM,
							        "question-answering": SqueezeBertForQuestionAnswering,
							        "text-classification": SqueezeBertForSequenceClassification,
							        "token-classification": SqueezeBertForTokenClassification,
							        "zero-shot": SqueezeBertForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							lowerCAmelCase_ 						=						False
							lowerCAmelCase_ 						=						True
							lowerCAmelCase_ 						=						False

							def 	lowercase							(self  )					->       int:
														_snake_case       =  SqueezeBertModelTester(self  )
														_snake_case       =  ConfigTester(self  ,  config_class=UpperCAmelCase  ,  dim=37  )

							def 	lowercase							(self  )					->       Tuple:
														self.config_tester.run_common_tests()

							def 	lowercase							(self  )					->       Union[str, Any]:
														_snake_case       =  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_squeezebert_model(*UpperCAmelCase  )

							def 	lowercase							(self  )					->       Optional[Any]:
														_snake_case       =  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_squeezebert_for_masked_lm(*UpperCAmelCase  )

							def 	lowercase							(self  )					->       Dict:
														_snake_case       =  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_squeezebert_for_question_answering(*UpperCAmelCase  )

							def 	lowercase							(self  )					->       Tuple:
														_snake_case       =  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_squeezebert_for_sequence_classification(*UpperCAmelCase  )

							def 	lowercase							(self  )					->       Union[str, Any]:
														_snake_case       =  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_squeezebert_for_token_classification(*UpperCAmelCase  )

							def 	lowercase							(self  )					->       List[str]:
														_snake_case       =  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_squeezebert_for_multiple_choice(*UpperCAmelCase  )





							@slow
							def 	lowercase							(self  )					->       Any:
														for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					_snake_case       =  SqueezeBertModel.from_pretrained(UpperCAmelCase  )
																					self.assertIsNotNone(UpperCAmelCase  )


@require_sentencepiece
@require_tokenizers
@require_torch
class       _lowerCAmelCase   (   unittest.TestCase				):







							'''simple docstring'''

							@slow
							def 	lowercase							(self  )					->       Optional[Any]:
														_snake_case       =  SqueezeBertForSequenceClassification.from_pretrained("""squeezebert/squeezebert-mnli"""  )

														_snake_case       =  torch.tensor([[1, 29414, 232, 328, 740, 1140, 12695, 69, 13, 1588, 2]]  )
														_snake_case       =  model(UpperCAmelCase  )[0]
														_snake_case       =  torch.Size((1, 3)  )
														self.assertEqual(output.shape  ,  UpperCAmelCase  )
														_snake_case       =  torch.tensor([[0.6401, -0.0349, -0.6041]]  )
														self.assertTrue(torch.allclose(UpperCAmelCase  ,  UpperCAmelCase  ,  atol=1e-4  )  )

'''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       _lowerCAmelCase   (   __snake_case				):







							'''simple docstring'''

							def 	lowercase							(self  )					->       Optional[Any]:
														super().setUp()

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

							def 	lowercase							(self  )					->       Optional[int]:
														from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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

							def 	lowercase							(self  )					->       List[str]:
														from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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

							def 	lowercase							(self  )					->       int:
														from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

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

														_snake_case       =  self.dist_env.copy()
														_snake_case       =  """TRANSFORMER_BASED_WRAP"""
														_snake_case       =  """T5Layer"""
														with mockenv_context(**UpperCAmelCase  ):
																					_snake_case       =  FullyShardedDataParallelPlugin()
																					with self.assertRaises(UpperCAmelCase  ) as cm:
																												fsdp_plugin.set_auto_wrap_policy(UpperCAmelCase  )
																					self.assertTrue("""Could not find the transformer layer class to wrap in the model.""" in str(cm.exception  )  )

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

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





							def 	lowercase							(self  )					->       Optional[int]:
														from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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


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







							'''simple docstring'''

							def 	lowercase							(self  )					->       List[Any]:
														super().setUp()
														_snake_case       =  0.82
														_snake_case       =  [
														    """fsdp_shard_grad_op_transformer_based_wrap""",
														    """fsdp_full_shard_transformer_based_wrap""",
														]
														_snake_case       =  {
														    """multi_gpu_fp16""": 3200,
														    """fsdp_shard_grad_op_transformer_based_wrap_fp16""": 2000,
														    """fsdp_full_shard_transformer_based_wrap_fp16""": 1900,
														    # 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
														}
														_snake_case       =  160
														_snake_case       =  160

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

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

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

																												_snake_case       =  cmd_config[:-1]
																												_snake_case       =  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(UpperCAmelCase  ,  env=os.environ.copy()  )





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