

import unittest

from transformers import BigBirdTokenizer, BigBirdTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, require_torch, slow
from transformers.utils import cached_property

from ...test_tokenization_common import TokenizerTesterMixin


_lowerCamelCase    :       List[str]        =							"▁"

_lowerCamelCase    :       Any        =							get_tests_dir("fixtures/test_sentencepiece.model")





@require_sentencepiece
@require_tokenizers
class  __UpperCAmelCase					(	lowerCamelCase__      ,  unittest.TestCase						):
   UpperCamelCase       					=					BigBirdTokenizer
   UpperCamelCase       					=					BigBirdTokenizerFast
   UpperCamelCase       					=					True
   UpperCamelCase       					=					True



   def 	__magic_name__   (							self							:					int						):
     super().setUp()

     UpperCAmelCase	:							str						=		self.tokenizer_class(__A,		keep_accents=__A						)
     tokenizer.save_pretrained(self.tmpdirname						)



   def 	__magic_name__   (							self							:					Optional[int]						):
     UpperCAmelCase	:							Tuple						=		'''<s>'''
     UpperCAmelCase	:							int						=		1

     self.assertEqual(self.get_tokenizer()._convert_token_to_id(__A						),		__A						)
     self.assertEqual(self.get_tokenizer()._convert_id_to_token(__A						),		__A						)



   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase	:							int						=		list(self.get_tokenizer().get_vocab().keys()						)

     self.assertEqual(vocab_keys[0],		'''<unk>'''						)
     self.assertEqual(vocab_keys[1],		'''<s>'''						)
     self.assertEqual(vocab_keys[-1],		'''[MASK]'''						)
     self.assertEqual(len(__A						),		1_0_0_4						)



   def 	__magic_name__   (							self							:					Dict						):
     self.assertEqual(self.get_tokenizer().vocab_size,		1_0_0_0						)



   def 	__magic_name__   (							self							:					Optional[int]						):
     if not self.test_rust_tokenizer:
       return

     UpperCAmelCase	:							Tuple						=		self.get_tokenizer()
     UpperCAmelCase	:							Any						=		self.get_rust_tokenizer()

     UpperCAmelCase	:							str						=		'''I was born in 92000, and this is falsé.'''

     UpperCAmelCase	:							Union[str, Any]						=		tokenizer.tokenize(__A						)
     UpperCAmelCase	:							Union[str, Any]						=		rust_tokenizer.tokenize(__A						)
     self.assertListEqual(__A,		__A						)

     UpperCAmelCase	:							Dict						=		tokenizer.encode(__A,		add_special_tokens=__A						)
     UpperCAmelCase	:							str						=		rust_tokenizer.encode(__A,		add_special_tokens=__A						)
     self.assertListEqual(__A,		__A						)

     UpperCAmelCase	:							List[Any]						=		self.get_rust_tokenizer()
     UpperCAmelCase	:							Optional[int]						=		tokenizer.encode(__A						)
     UpperCAmelCase	:							str						=		rust_tokenizer.encode(__A						)
     self.assertListEqual(__A,		__A						)



   def 	__magic_name__   (							self							:					List[Any]						):
     UpperCAmelCase	:							Dict						=		BigBirdTokenizer(__A,		keep_accents=__A						)

     UpperCAmelCase	:							Optional[Any]						=		tokenizer.tokenize('''This is a test'''						)
     self.assertListEqual(__A,		['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est''']						)

     self.assertListEqual(
         tokenizer.convert_tokens_to_ids(__A						),		[2_8_5, 4_6, 1_0, 1_7_0, 3_8_2],		)

     UpperCAmelCase	:							Any						=		tokenizer.tokenize('''I was born in 92000, and this is falsé.'''						)
     self.assertListEqual(
         __A,		[
             SPIECE_UNDERLINE + '''I''',
             SPIECE_UNDERLINE + '''was''',
             SPIECE_UNDERLINE + '''b''',
             '''or''',
             '''n''',
             SPIECE_UNDERLINE + '''in''',
             SPIECE_UNDERLINE + '''''',
             '''9''',
             '''2''',
             '''0''',
             '''0''',
             '''0''',
             ''',''',
             SPIECE_UNDERLINE + '''and''',
             SPIECE_UNDERLINE + '''this''',
             SPIECE_UNDERLINE + '''is''',
             SPIECE_UNDERLINE + '''f''',
             '''al''',
             '''s''',
             '''é''',
             '''.''',
         ],		)
     UpperCAmelCase	:							List[str]						=		tokenizer.convert_tokens_to_ids(__A						)
     self.assertListEqual(
         __A,		[8, 2_1, 8_4, 5_5, 2_4, 1_9, 7, 0, 6_0_2, 3_4_7, 3_4_7, 3_4_7, 3, 1_2, 6_6, 4_6, 7_2, 8_0, 6, 0, 4],		)

     UpperCAmelCase	:							Union[str, Any]						=		tokenizer.convert_ids_to_tokens(__A						)
     self.assertListEqual(
         __A,		[
             SPIECE_UNDERLINE + '''I''',
             SPIECE_UNDERLINE + '''was''',
             SPIECE_UNDERLINE + '''b''',
             '''or''',
             '''n''',
             SPIECE_UNDERLINE + '''in''',
             SPIECE_UNDERLINE + '''''',
             '''<unk>''',
             '''2''',
             '''0''',
             '''0''',
             '''0''',
             ''',''',
             SPIECE_UNDERLINE + '''and''',
             SPIECE_UNDERLINE + '''this''',
             SPIECE_UNDERLINE + '''is''',
             SPIECE_UNDERLINE + '''f''',
             '''al''',
             '''s''',
             '''<unk>''',
             '''.''',
         ],		)



   @cached_property
   def 	__magic_name__   (							self							:					List[str]						):
     return BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base'''						)



   @slow
   def 	__magic_name__   (							self							:					int						):
     UpperCAmelCase	:							Dict						=		'''Hello World!'''
     UpperCAmelCase	:							Dict						=		[6_5, 1_8_5_3_6, 2_2_6_0, 1_0_1, 6_6]

     self.assertListEqual(__A,		self.big_tokenizer.encode(__A						)						)



   @slow
   def 	__magic_name__   (							self							:					List[Any]						):
     UpperCAmelCase	:							List[str]						=		(
         '''This is a very long text with a lot of weird characters, such as: . , ~ ? ( ) " [ ] ! : - . Also we will'''
         ''' add words that should not exsist and be tokenized to <unk>, such as saoneuhaoesuth'''
     )
     # fmt: off
     UpperCAmelCase	:							Tuple						=		[6_5, 8_7_1, 4_1_9, 3_5_8, 9_4_6, 9_9_1, 2_5_2_1, 4_5_2, 3_5_8, 1_3_5_7, 3_8_7, 7_7_5_1, 3_5_3_6, 1_1_2, 9_8_5, 4_5_6, 1_2_6, 8_6_5, 9_3_8, 5_4_0_0, 5_7_3_4, 4_5_8, 1_3_6_8, 4_6_7, 7_8_6, 2_4_6_2, 5_2_4_6, 1_1_5_9, 6_3_3, 8_6_5, 4_5_1_9, 4_5_7, 5_8_2, 8_5_2, 2_5_5_7, 4_2_7, 9_1_6, 5_0_8, 4_0_5, 3_4_3_2_4, 4_9_7, 3_9_1, 4_0_8, 1_1_3_4_2, 1_2_4_4, 3_8_5, 1_0_0, 9_3_8, 9_8_5, 4_5_6, 5_7_4, 3_6_2, 1_2_5_9_7, 3_2_0_0, 3_1_2_9, 1_1_7_2, 6_6]  # noqa: E231
     # fmt: on
     self.assertListEqual(__A,		self.big_tokenizer.encode(__A						)						)



   @require_torch
   @slow
   def 	__magic_name__   (							self							:					List[str]						):
     import torch

     from transformers import BigBirdConfig, BigBirdModel

     # Build sequence
     UpperCAmelCase	:							int						=		list(self.big_tokenizer.get_vocab().keys()						)[:1_0]
     UpperCAmelCase	:							str						=		''' '''.join(__A						)
     UpperCAmelCase	:							Optional[int]						=		self.big_tokenizer.encode_plus(__A,		return_tensors='''pt''',		return_token_type_ids=__A						)
     UpperCAmelCase	:							Tuple						=		self.big_tokenizer.batch_encode_plus(
         [sequence + ''' ''' + sequence],		return_tensors='''pt''',		return_token_type_ids=__A						)

     UpperCAmelCase	:							Any						=		BigBirdConfig(attention_type='''original_full'''						)
     UpperCAmelCase	:							str						=		BigBirdModel(__A						)

     assert model.get_input_embeddings().weight.shape[0] >= self.big_tokenizer.vocab_size

     with torch.no_grad():
       model(**__A						)
       model(**__A						)



   @slow
   def 	__magic_name__   (							self							:					Tuple						):
     UpperCAmelCase	:							Tuple						=		BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base'''						)
     UpperCAmelCase	:							Dict						=		tokenizer.decode(tokenizer('''Paris is the [MASK].'''						).input_ids						)

     self.assertTrue(decoded_text == '''[CLS] Paris is the[MASK].[SEP]'''						)





   @slow
   def 	__magic_name__   (							self							:					Dict						):
     # fmt: off
     UpperCAmelCase	:							List[Any]						=		{'''input_ids''': [[6_5, 3_9_2_8_6, 4_5_8, 3_6_3_3_5, 2_0_0_1, 4_5_6, 1_3_0_7_3, 1_3_2_6_6, 4_5_5, 1_1_3, 7_7_4_6, 1_7_4_1, 1_1_1_5_7, 3_9_1, 1_3_0_7_3, 1_3_2_6_6, 4_5_5, 1_1_3, 3_9_6_7, 3_5_4_1_2, 1_1_3, 4_9_3_6, 1_0_9, 3_8_7_0, 2_3_7_7, 1_1_3, 3_0_0_8_4, 4_5_7_2_0, 4_5_8, 1_3_4, 1_7_4_9_6, 1_1_2, 5_0_3, 1_1_6_7_2, 1_1_3, 1_1_8, 1_1_2, 5_6_6_5, 1_3_3_4_7, 3_8_6_8_7, 1_1_2, 1_4_9_6, 3_1_3_8_9, 1_1_2, 3_2_6_8, 4_7_2_6_4, 1_3_4, 9_6_2, 1_1_2, 1_6_3_7_7, 8_0_3_5, 2_3_1_3_0, 4_3_0, 1_2_1_6_9, 1_5_5_1_8, 2_8_5_9_2, 4_5_8, 1_4_6, 4_1_6_9_7, 1_0_9, 3_9_1, 1_2_1_6_9, 1_5_5_1_8, 1_6_6_8_9, 4_5_8, 1_4_6, 4_1_3_5_8, 1_0_9, 4_5_2, 7_2_6, 4_0_3_4, 1_1_1, 7_6_3, 3_5_4_1_2, 5_0_8_2, 3_8_8, 1_9_0_3, 1_1_1, 9_0_5_1, 3_9_1, 2_8_7_0, 4_8_9_1_8, 1_9_0_0, 1_1_2_3, 5_5_0, 9_9_8, 1_1_2, 9_5_8_6, 1_5_9_8_5, 4_5_5, 3_9_1, 4_1_0, 2_2_9_5_5, 3_7_6_3_6, 1_1_4, 6_6], [6_5, 4_4_8, 1_7_4_9_6, 4_1_9, 3_6_6_3, 3_8_5, 7_6_3, 1_1_3, 2_7_5_3_3, 2_8_7_0, 3_2_8_3, 1_3_0_4_3, 1_6_3_9, 2_4_7_1_3, 5_2_3, 6_5_6, 2_4_0_1_3, 1_8_5_5_0, 2_5_2_1, 5_1_7, 2_7_0_1_4, 2_1_2_4_4, 4_2_0, 1_2_1_2, 1_4_6_5, 3_9_1, 9_2_7, 4_8_3_3, 3_8_8, 5_7_8, 1_1_7_8_6, 1_1_4, 6_6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [6_5, 4_8_4, 2_1_6_9, 7_6_8_7, 2_1_9_3_2, 1_8_1_4_6, 7_2_6, 3_6_3, 1_7_0_3_2, 3_3_9_1, 1_1_4, 6_6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], '''attention_mask''': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
     # fmt: on

     self.tokenizer_integration_test_util(
         expected_encoding=__A,		model_name='''google/bigbird-roberta-base''',		revision='''215c99f1600e06f83acce68422f2035b2b5c3510''',		)








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)

_lowerCamelCase    :       List[str]        =							"bert-base-cased"
_lowerCamelCase    :       str        =							"fp16"
_lowerCamelCase    :       Optional[int]        =							"bf16"
_lowerCamelCase    :       List[str]        =							[FPaa, BFaa]





@require_fsdp
@require_cuda
class  __UpperCAmelCase					(	lowerCamelCase__						):



   def 	__magic_name__   (							self							:					Tuple						):
     super().setUp()

     UpperCAmelCase	:							Optional[Any]						=			dict(
         ACCELERATE_USE_FSDP='''true''',		MASTER_ADDR='''localhost''',		MASTER_PORT='''10999''',		RANK='''0''',		LOCAL_RANK='''0''',		WORLD_SIZE='''1''',		)



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

     for i, strategy in enumerate(__A						):
       UpperCAmelCase	:							List[Any]						=			self.dist_env.copy()
       UpperCAmelCase	:							Union[str, Any]						=			F'''{i + 1}'''
       UpperCAmelCase	:							Union[str, Any]						=			strategy
       with mockenv_context(**__A						):
         UpperCAmelCase	:							Union[str, Any]						=			FullyShardedDataParallelPlugin()
         self.assertEqual(fsdp_plugin.sharding_strategy,		ShardingStrategy(i + 1						)						)



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

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



   def 	__magic_name__   (							self							:					Any						):
     from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

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

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

     UpperCAmelCase	:							List[Any]						=			self.dist_env.copy()
     UpperCAmelCase	:							str						=			'''SIZE_BASED_WRAP'''
     UpperCAmelCase	:							str						=			'''0'''
     with mockenv_context(**__A						):
       UpperCAmelCase	:							Optional[int]						=			FullyShardedDataParallelPlugin()
       fsdp_plugin.set_auto_wrap_policy(__A						)
       self.assertIsNone(fsdp_plugin.auto_wrap_policy						)



   def 	__magic_name__   (							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:
       UpperCAmelCase	:							List[Any]						=			self.dist_env.copy()
       UpperCAmelCase	:							int						=			mp_dtype
       with mockenv_context(**__A						):
         UpperCAmelCase	:							int						=			Accelerator()
         if mp_dtype == "fp16":
           UpperCAmelCase	:							Any						=			torch.floataa
         elif mp_dtype == "bf16":
           UpperCAmelCase	:							Any						=			torch.bfloataa
         UpperCAmelCase	:							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 	__magic_name__   (							self							:					Optional[int]						):
     from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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






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



   def 	__magic_name__   (							self							:					List[Any]						):
     super().setUp()
     UpperCAmelCase	:							int						=			0.8_2
     UpperCAmelCase	:							List[str]						=			[
         '''fsdp_shard_grad_op_transformer_based_wrap''',
         '''fsdp_full_shard_transformer_based_wrap''',
     ]
     UpperCAmelCase	:							int						=			{
         '''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	:							List[str]						=			1_6_0
     UpperCAmelCase	:							Optional[int]						=			1_6_0

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



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

         UpperCAmelCase	:							Tuple						=			cmd_config[:-1]
         UpperCAmelCase	:							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 	__magic_name__   (							self							:					Any						):
     UpperCAmelCase	:							Optional[Any]						=			os.path.join(self.test_scripts_folder,		'''test_peak_memory_usage.py'''						)
     UpperCAmelCase	:							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():
       UpperCAmelCase	:							List[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(__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()						)





