


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							(			UpperCamelCase__			):

							def       lowerCamelCase_ (							self:	str							)			->   int:





									"""simple docstring"""







									super().setUp()

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

							def       lowerCamelCase_ (							self:	Optional[Any]							)			->   List[str]:





									"""simple docstring"""







									from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

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

							def       lowerCamelCase_ (							self:	str							)			->   str:





									"""simple docstring"""







									from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

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





									"""simple docstring"""







									from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

									for i, state_dict_type in enumerate(UpperCamelCase_							):
											lowercase__            =      self.dist_env.copy()
											lowercase__            =      state_dict_type
											with mockenv_context(**UpperCamelCase_							):
													lowercase__            =      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:	Tuple							)			->   int:





									"""simple docstring"""







									lowercase__            =      AutoModel.from_pretrained(UpperCamelCase_							)
									for policy in FSDP_AUTO_WRAP_POLICY:
											lowercase__            =      self.dist_env.copy()
											lowercase__            =      policy
											if policy == "TRANSFORMER_BASED_WRAP":
													lowercase__            =      '''BertLayer'''
											elif policy == "SIZE_BASED_WRAP":
													lowercase__            =      '''2000'''
											with mockenv_context(**UpperCamelCase_							):
													lowercase__            =      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							)

									lowercase__            =      self.dist_env.copy()
									lowercase__            =      '''TRANSFORMER_BASED_WRAP'''
									lowercase__            =      '''T5Layer'''
									with mockenv_context(**UpperCamelCase_							):
											lowercase__            =      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							)							)

									lowercase__            =      self.dist_env.copy()
									lowercase__            =      '''SIZE_BASED_WRAP'''
									lowercase__            =      '''0'''
									with mockenv_context(**UpperCamelCase_							):
											lowercase__            =      FullyShardedDataParallelPlugin()
											fsdp_plugin.set_auto_wrap_policy(UpperCamelCase_							)
											self.assertIsNone(fsdp_plugin.auto_wrap_policy							)

							def       lowerCamelCase_ (							self:	Union[str, Any]							)			->   int:





									"""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:
											lowercase__            =      self.dist_env.copy()
											lowercase__            =      mp_dtype
											with mockenv_context(**UpperCamelCase_							):
													lowercase__            =      Accelerator()
													if mp_dtype == "fp16":
															lowercase__            =      torch.floataa
													elif mp_dtype == "bf16":
															lowercase__            =      torch.bfloataa
													lowercase__            =      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       lowerCamelCase_ (							self:	Any							)			->   str:





									"""simple docstring"""







									from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

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


@require_fsdp
@require_multi_gpu
@slow
class        _a							(			UpperCamelCase__			):

							def       lowerCamelCase_ (							self:	Any							)			->   str:





									"""simple docstring"""







									super().setUp()
									lowercase__            =      0.82
									lowercase__            =      [
									    '''fsdp_shard_grad_op_transformer_based_wrap''',
									    '''fsdp_full_shard_transformer_based_wrap''',
									]
									lowercase__            =      {
									    '''multi_gpu_fp16''': 3_200,
									    '''fsdp_shard_grad_op_transformer_based_wrap_fp16''': 2_000,
									    '''fsdp_full_shard_transformer_based_wrap_fp16''': 1_900,
									    # 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
									}
									lowercase__            =      160
									lowercase__            =      160

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

							def       lowerCamelCase_ (							self:	Optional[int]							)			->   Any:





									"""simple docstring"""







									lowercase__            =      os.path.join(self.test_scripts_folder    ,						'''test_performance.py'''							)
									lowercase__            =      ['''accelerate''', '''launch''', '''--num_processes=2''', '''--num_machines=1''', '''--machine_rank=0''', '''--use_fsdp''']
									for config in self.performance_configs:
											lowercase__            =      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       lowerCamelCase_ (							self:	Any							)			->   Tuple:





									"""simple docstring"""







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

													lowercase__            =      cmd_config[:-1]
													lowercase__            =      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       lowerCamelCase_ (							self:	Any							)			->   List[str]:





									"""simple docstring"""







									lowercase__            =      os.path.join(self.test_scripts_folder    ,						'''test_peak_memory_usage.py'''							)
									lowercase__            =      [
									    '''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():
											lowercase__            =      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()							)






from pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


lowerCAmelCase       =							logging.get_logger(__name__)
def 			_a (      SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE     ):


		"""simple docstring"""




		try:
				with open(SCREAMING_SNAKE_CASE						,				'''rb'''     ) as flax_state_f:
						lowercase__            =      from_bytes(SCREAMING_SNAKE_CASE						,				flax_state_f.read()     )
		except UnpicklingError as e:
				try:
						with open(SCREAMING_SNAKE_CASE     ) as f:
								if f.read().startswith('''version'''     ):
										raise OSError(
										    '''You seem to have cloned a repository without having git-lfs installed. Please'''
										    ''' install git-lfs and run `git lfs install` followed by `git lfs pull` in the'''
										    ''' folder you cloned.'''     )
								else:
										raise ValueError from e
				except (UnicodeDecodeError, ValueError):
						raise EnvironmentError(f'Unable to convert {model_file} to Flax deserializable object. '     )

		return load_flax_weights_in_pytorch_model(SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE     )
def 			_a (      SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE     ):


		"""simple docstring"""




		try:
				import torch  # noqa: F401
		except ImportError:
				logger.error(
				    '''Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see'''
				    ''' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'''
				    ''' instructions.'''     )
				raise

		# check if we have bf16 weights
		lowercase__            =      flatten_dict(jax.tree_util.tree_map(lambda SCREAMING_SNAKE_CASE     : x.dtype == jnp.bfloataa						,				SCREAMING_SNAKE_CASE     )     ).values()
		if any(SCREAMING_SNAKE_CASE     ):
				# convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

				# and bf16 is not fully supported in PT yet.
				logger.warning(
				    '''Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` '''
				    '''before loading those in PyTorch model.'''     )
				lowercase__            =      jax.tree_util.tree_map(
				    lambda SCREAMING_SNAKE_CASE     : params.astype(np.floataa     ) if params.dtype == jnp.bfloataa else params						,				SCREAMING_SNAKE_CASE     )

		lowercase__            =      ''''''

		lowercase__            =      flatten_dict(SCREAMING_SNAKE_CASE						,				sep='''.'''     )
		lowercase__            =      pt_model.state_dict()

		# keep track of unexpected & missing keys
		lowercase__            =      []
		lowercase__            =      set(pt_model_dict.keys()     )

		for flax_key_tuple, flax_tensor in flax_state_dict.items():
				lowercase__            =      flax_key_tuple.split('''.'''     )

				if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
						lowercase__            =      flax_key_tuple_array[:-1] + ['''weight''']
						lowercase__            =      jnp.transpose(SCREAMING_SNAKE_CASE						,				(3, 2, 0, 1)     )
				elif flax_key_tuple_array[-1] == "kernel":
						lowercase__            =      flax_key_tuple_array[:-1] + ['''weight''']
						lowercase__            =      flax_tensor.T
				elif flax_key_tuple_array[-1] == "scale":
						lowercase__            =      flax_key_tuple_array[:-1] + ['''weight''']

				if "time_embedding" not in flax_key_tuple_array:
						for i, flax_key_tuple_string in enumerate(SCREAMING_SNAKE_CASE     ):
								lowercase__            =      (
								    flax_key_tuple_string.replace('''_0'''						,				'''.0'''     )
								    .replace('''_1'''						,				'''.1'''     )
								    .replace('''_2'''						,				'''.2'''     )
								    .replace('''_3'''						,				'''.3'''     )
								    .replace('''_4'''						,				'''.4'''     )
								    .replace('''_5'''						,				'''.5'''     )
								    .replace('''_6'''						,				'''.6'''     )
								    .replace('''_7'''						,				'''.7'''     )
								    .replace('''_8'''						,				'''.8'''     )
								    .replace('''_9'''						,				'''.9'''     )
								)

				lowercase__            =      '''.'''.join(SCREAMING_SNAKE_CASE     )

				if flax_key in pt_model_dict:
						if flax_tensor.shape != pt_model_dict[flax_key].shape:
								raise ValueError(
								    f'Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '
								    f'to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'     )
						else:
								# add weight to pytorch dict
								lowercase__            =      np.asarray(SCREAMING_SNAKE_CASE     ) if not isinstance(SCREAMING_SNAKE_CASE						,				np.ndarray     ) else flax_tensor
								lowercase__            =      torch.from_numpy(SCREAMING_SNAKE_CASE     )
								# remove from missing keys
								missing_keys.remove(SCREAMING_SNAKE_CASE     )
				else:
						# weight is not expected by PyTorch model
						unexpected_keys.append(SCREAMING_SNAKE_CASE     )

		pt_model.load_state_dict(SCREAMING_SNAKE_CASE     )

		# re-transform missing_keys to list
		lowercase__            =      list(SCREAMING_SNAKE_CASE     )

		if len(SCREAMING_SNAKE_CASE     ) > 0:
				logger.warning(
				    '''Some weights of the Flax model were not used when initializing the PyTorch model'''
				    f' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'
				    f' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'
				    ''' (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This'''
				    f' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'
				    ''' to be exactly identical (e.g. initializing a BertForSequenceClassification model from a'''
				    ''' FlaxBertForSequenceClassification model).'''     )
		if len(SCREAMING_SNAKE_CASE     ) > 0:
				logger.warning(
				    f'Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'
				    f' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'
				    ''' use it for predictions and inference.'''     )

		return pt_model


