




"""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)

lowercase__			=		"""bert-base-cased"""
lowercase__			=		"""fp16"""
lowercase__			=		"""bf16"""
lowercase__			=		[FPaa, BFaa]






@require_fsdp
@require_cuda
class lowerCAmelCase__      (	lowercase		):






			'''simple docstring'''



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

							_lowerCamelCase	:     int       =       dict(
							    ACCELERATE_USE_FSDP='true'    , MASTER_ADDR='localhost'    , MASTER_PORT='10999'    , RANK='0'    , LOCAL_RANK='0'    , WORLD_SIZE='1'    , )
			def    A_       (    self    ):
							from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

							for i, strategy in enumerate(lowercase    ):
											_lowerCamelCase	:     Tuple       =       self.dist_env.copy()
											_lowerCamelCase	:     int       =       F'''{i + 1}'''
											_lowerCamelCase	:     int       =       strategy
											with mockenv_context(**lowercase    ):
															_lowerCamelCase	:     Optional[Any]       =       FullyShardedDataParallelPlugin()
															self.assertEqual(fsdp_plugin.sharding_strategy    , ShardingStrategy(i + 1    )    )
			def    A_       (    self    ):
							from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

							for i, prefetch_policy in enumerate(lowercase    ):
											_lowerCamelCase	:     Tuple       =       self.dist_env.copy()
											_lowerCamelCase	:     str       =       prefetch_policy
											with mockenv_context(**lowercase    ):
															_lowerCamelCase	:     Optional[Any]       =       FullyShardedDataParallelPlugin()
															if prefetch_policy == "NO_PREFETCH":
																			self.assertIsNone(fsdp_plugin.backward_prefetch    )
															else:
																			self.assertEqual(fsdp_plugin.backward_prefetch    , BackwardPrefetch(i + 1    )    )
			def    A_       (    self    ):
							from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

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

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

							_lowerCamelCase	:     int       =       self.dist_env.copy()
							_lowerCamelCase	:     Optional[int]       =       'SIZE_BASED_WRAP'
							_lowerCamelCase	:     Dict       =       '0'
							with mockenv_context(**lowercase    ):
											_lowerCamelCase	:     List[str]       =       FullyShardedDataParallelPlugin()
											fsdp_plugin.set_auto_wrap_policy(lowercase    )
											self.assertIsNone(fsdp_plugin.auto_wrap_policy    )
			def    A_       (    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	:     Any       =       mp_dtype
											with mockenv_context(**lowercase    ):
															_lowerCamelCase	:     Optional[Any]       =       Accelerator()
															if mp_dtype == "fp16":
																			_lowerCamelCase	:     Optional[int]       =       torch.floataa
															elif mp_dtype == "bf16":
																			_lowerCamelCase	:     List[Any]       =       torch.bfloataa
															_lowerCamelCase	:     Any       =       MixedPrecision(param_dtype=lowercase    , reduce_dtype=lowercase    , buffer_dtype=lowercase    )
															self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy    , lowercase    )
															if mp_dtype == FPaa:
																			self.assertTrue(isinstance(accelerator.scaler    , lowercase    )    )
															elif mp_dtype == BFaa:
																			self.assertIsNone(accelerator.scaler    )
															AcceleratorState._reset_state(lowercase    )






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

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






@require_fsdp
@require_multi_gpu
@slow
class lowerCAmelCase__      (	lowercase		):






			'''simple docstring'''



			def    A_       (    self    ):
							super().setUp()
							_lowerCamelCase	:     str       =       0.82
							_lowerCamelCase	:     Tuple       =       [
							    'fsdp_shard_grad_op_transformer_based_wrap',
							    'fsdp_full_shard_transformer_based_wrap',
							]
							_lowerCamelCase	:     Tuple       =       {
							    '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
							}
							_lowerCamelCase	:     Tuple       =       160
							_lowerCamelCase	:     int       =       160

							_lowerCamelCase	:     List[Any]       =       inspect.getfile(accelerate.test_utils    )
							_lowerCamelCase	:     Tuple       =       os.path.sep.join(mod_file.split(os.path.sep    )[:-1] + ['scripts', 'external_deps']    )
			def    A_       (    self    ):
							_lowerCamelCase	:     List[Any]       =       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	:     Tuple       =       cmd.copy()
											for i, strategy in enumerate(lowercase    ):
															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(lowercase    , env=os.environ.copy()    )
			def    A_       (    self    ):
							_lowerCamelCase	:     Optional[Any]       =       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(lowercase    ):
											_lowerCamelCase	:     int       =       cmd.copy()
											cmd_config.append(F'''--fsdp_sharding_strategy={i+1}'''    )
											if strategy != "FULL_SHARD":
															continue
											_lowerCamelCase	:     Optional[Any]       =       len(lowercase    )
											for state_dict_type in FSDP_STATE_DICT_TYPE:
															_lowerCamelCase	:     List[str]       =       cmd_config[:state_dict_config_index]
															cmd_config.append(F'''--fsdp_state_dict_type={state_dict_type}'''    )
															cmd_config.extend(
															    [
															        self.test_file_path,
															        F'''--output_dir={self.tmpdir}''',
															        '--partial_train_epoch=1',
															    ]    )
															with patch_environment(omp_num_threads=1    ):
																			execute_subprocess_async(lowercase    , env=os.environ.copy()    )

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






			def    A_       (    self    ):
							_lowerCamelCase	:     List[Any]       =       os.path.join(self.test_scripts_folder    , 'test_peak_memory_usage.py'    )
							_lowerCamelCase	:     Union[str, 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	:     Tuple       =       cmd.copy()
											if "fp16" in spec:
															cmd_config.extend(['--mixed_precision=fp16']    )
											else:
															cmd_config.extend(['--mixed_precision=no']    )

											if "multi_gpu" in spec:
															continue
											else:
															cmd_config.extend(['--use_fsdp']    )
															for i, strategy in enumerate(lowercase    ):
																			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(lowercase    , env=os.environ.copy()    )





"""simple docstring"""






import os
from typing import Dict, List, Union

import tensorflow as tf
from keras_nlp.tokenizers import BytePairTokenizer
from tensorflow_text import pad_model_inputs

from .tokenization_gpta import GPTaTokenizer






class lowerCAmelCase__      (	tf.keras.layers.Layer		):






			'''simple docstring'''



			def __init__(    self    , lowercase    , lowercase    , lowercase = None    , lowercase = None    ):
							super().__init__()
							_lowerCamelCase	:     Optional[Any]       =       pad_token_id
							_lowerCamelCase	:     Dict       =       max_length
							_lowerCamelCase	:     Tuple       =       vocab
							_lowerCamelCase	:     List[str]       =       merges
							_lowerCamelCase	:     int       =       BytePairTokenizer(lowercase    , lowercase    , sequence_length=lowercase    )
			@classmethod
			def    A_       (    cls    , lowercase    , *lowercase    , **lowercase    ):
							_lowerCamelCase	:     str       =       [' '.join(lowercase    ) for m in tokenizer.bpe_ranks.keys()]
							_lowerCamelCase	:     Union[str, Any]       =       tokenizer.get_vocab()
							return cls(lowercase    , lowercase    , *lowercase    , **lowercase    )
			@classmethod
			def    A_       (    cls    , lowercase    , *lowercase    , **lowercase    ):
							_lowerCamelCase	:     Optional[int]       =       GPTaTokenizer.from_pretrained(lowercase    , *lowercase    , **lowercase    )
							return cls.from_tokenizer(lowercase    , *lowercase    , **lowercase    )
			@classmethod
			def    A_       (    cls    , lowercase    ):
							return cls(**lowercase    )
			def    A_       (    self    ):
							return {
							    "vocab": self.vocab,
							    "merges": self.merges,
							    "max_length": self.max_length,
							    "pad_token_id": self.pad_token_id,
							}






			def    A_       (    self    , lowercase    , lowercase = None    ):
							_lowerCamelCase	:     int       =       self.tf_tokenizer(lowercase    )
							_lowerCamelCase	:     Optional[int]       =       tf.ones_like(lowercase    )

							if self.pad_token_id is not None:
											# pad the tokens up to max length
											_lowerCamelCase	:     str       =       max_length if max_length is not None else self.max_length

											if max_length is not None:
															_lowerCamelCase,					_lowerCamelCase	:     Union[str, Any]       =       pad_model_inputs(
															    lowercase    , max_seq_length=lowercase    , pad_value=self.pad_token_id    )

							return {"attention_mask": attention_mask, "input_ids": input_ids}