




import os

import torch

from ..logging import get_logger
from .constants import FSDP_PYTORCH_VERSION, MODEL_NAME, OPTIMIZER_NAME
from .versions import is_torch_version


if is_torch_version(">=", FSDP_PYTORCH_VERSION):
  import torch.distributed.checkpoint as dist_cp
  from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
  from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
  from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
  from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType


_lowerCAmelCase						:   List[str]		=   get_logger(__name__)



def 						lowerCAmelCase		(		_lowerCAmelCase							:   List[str]						,							_lowerCAmelCase							:   Dict						,							_lowerCAmelCase							:   Optional[int]						,							_lowerCAmelCase							:   Union[str, Any]						,							_lowerCAmelCase							:   List[str]=0	):




 """simple docstring"""







 os.makedirs(_lowerCAmelCase						,							exist_ok=_lowerCAmelCase	)
 with FSDP.state_dict_type(
     _lowerCAmelCase						,							fsdp_plugin.state_dict_type						,							fsdp_plugin.state_dict_config						,							fsdp_plugin.optim_state_dict_config	):
  UpperCAmelCase__							 =						model.state_dict()
  if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
   UpperCAmelCase__							 =						F'''{MODEL_NAME}.bin''' if model_index == 0 else F'''{MODEL_NAME}_{model_index}.bin'''
   UpperCAmelCase__							 =						os.path.join(_lowerCAmelCase						,							_lowerCAmelCase	)
   if accelerator.process_index == 0:
    logger.info(F'''Saving model to {output_model_file}'''	)
    torch.save(_lowerCAmelCase						,							_lowerCAmelCase	)
    logger.info(F'''Model saved to {output_model_file}'''	)
  elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
   UpperCAmelCase__							 =						(
       F'''{MODEL_NAME}_rank{accelerator.process_index}.bin'''
       if model_index == 0
       else F'''{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'''
   )
   UpperCAmelCase__							 =						os.path.join(_lowerCAmelCase						,							_lowerCAmelCase	)
   logger.info(F'''Saving model to {output_model_file}'''	)
   torch.save(_lowerCAmelCase						,							_lowerCAmelCase	)
   logger.info(F'''Model saved to {output_model_file}'''	)
  elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
   UpperCAmelCase__							 =						os.path.join(_lowerCAmelCase						,							F'''{MODEL_NAME}_{model_index}'''	)
   os.makedirs(_lowerCAmelCase						,							exist_ok=_lowerCAmelCase	)
   logger.info(F'''Saving model to {ckpt_dir}'''	)
   UpperCAmelCase__							 =						{"model": state_dict}

   dist_cp.save_state_dict(
       state_dict=_lowerCAmelCase						,							storage_writer=dist_cp.FileSystemWriter(_lowerCAmelCase	)						,							planner=DefaultSavePlanner()						,							)
   logger.info(F'''Model saved to {ckpt_dir}'''	)



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Union[str, Any]						,							_lowerCAmelCase							:   Optional[Any]						,							_lowerCAmelCase							:   str						,							_lowerCAmelCase							:   List[str]						,							_lowerCAmelCase							:   Optional[Any]=0	):




 """simple docstring"""







 accelerator.wait_for_everyone()
 with FSDP.state_dict_type(
     _lowerCAmelCase						,							fsdp_plugin.state_dict_type						,							fsdp_plugin.state_dict_config						,							fsdp_plugin.optim_state_dict_config	):
  if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
   if type(_lowerCAmelCase	) != FSDP and accelerator.process_index != 0:
    if not fsdp_plugin.sync_module_states:
     raise ValueError(
         "Set the `sync_module_states` flag to `True` so that model states are synced across processes when "
         "initializing FSDP object"	)
    return
   UpperCAmelCase__							 =						F'''{MODEL_NAME}.bin''' if model_index == 0 else F'''{MODEL_NAME}_{model_index}.bin'''
   UpperCAmelCase__							 =						os.path.join(_lowerCAmelCase						,							_lowerCAmelCase	)
   logger.info(F'''Loading model from {input_model_file}'''	)
   UpperCAmelCase__							 =						torch.load(_lowerCAmelCase	)
   logger.info(F'''Model loaded from {input_model_file}'''	)
  elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
   UpperCAmelCase__							 =						(
       F'''{MODEL_NAME}_rank{accelerator.process_index}.bin'''
       if model_index == 0
       else F'''{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'''
   )
   UpperCAmelCase__							 =						os.path.join(_lowerCAmelCase						,							_lowerCAmelCase	)
   logger.info(F'''Loading model from {input_model_file}'''	)
   UpperCAmelCase__							 =						torch.load(_lowerCAmelCase	)
   logger.info(F'''Model loaded from {input_model_file}'''	)
  elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
   UpperCAmelCase__							 =						(
       os.path.join(_lowerCAmelCase						,							F'''{MODEL_NAME}_{model_index}'''	)
       if F'''{MODEL_NAME}''' not in input_dir
       else input_dir
   )
   logger.info(F'''Loading model from {ckpt_dir}'''	)
   UpperCAmelCase__							 =						{"model": model.state_dict()}
   dist_cp.load_state_dict(
       state_dict=_lowerCAmelCase						,							storage_reader=dist_cp.FileSystemReader(_lowerCAmelCase	)						,							planner=DefaultLoadPlanner()						,							)
   UpperCAmelCase__							 =						state_dict["model"]
   logger.info(F'''Model loaded from {ckpt_dir}'''	)
  model.load_state_dict(_lowerCAmelCase	)



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Any						,							_lowerCAmelCase							:   str						,							_lowerCAmelCase							:   List[str]						,							_lowerCAmelCase							:   int						,							_lowerCAmelCase							:   Dict						,							_lowerCAmelCase							:   List[str]=0	):




 """simple docstring"""







 os.makedirs(_lowerCAmelCase						,							exist_ok=_lowerCAmelCase	)
 with FSDP.state_dict_type(
     _lowerCAmelCase						,							fsdp_plugin.state_dict_type						,							fsdp_plugin.state_dict_config						,							fsdp_plugin.optim_state_dict_config	):
  UpperCAmelCase__							 =						FSDP.optim_state_dict(_lowerCAmelCase						,							_lowerCAmelCase	)
  if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
   if accelerator.process_index == 0:
    UpperCAmelCase__							 =						(
        F'''{OPTIMIZER_NAME}.bin''' if optimizer_index == 0 else F'''{OPTIMIZER_NAME}_{optimizer_index}.bin'''
    )
    UpperCAmelCase__							 =						os.path.join(_lowerCAmelCase						,							_lowerCAmelCase	)
    logger.info(F'''Saving Optimizer state to {output_optimizer_file}'''	)
    torch.save(_lowerCAmelCase						,							_lowerCAmelCase	)
    logger.info(F'''Optimizer state saved in {output_optimizer_file}'''	)
  else:
   UpperCAmelCase__							 =						os.path.join(_lowerCAmelCase						,							F'''{OPTIMIZER_NAME}_{optimizer_index}'''	)
   os.makedirs(_lowerCAmelCase						,							exist_ok=_lowerCAmelCase	)
   logger.info(F'''Saving Optimizer state to {ckpt_dir}'''	)
   dist_cp.save_state_dict(
       state_dict={"optimizer": optim_state}						,							storage_writer=dist_cp.FileSystemWriter(_lowerCAmelCase	)						,							planner=DefaultSavePlanner()						,							)
   logger.info(F'''Optimizer state saved in {ckpt_dir}'''	)



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Optional[int]						,							_lowerCAmelCase							:   int						,							_lowerCAmelCase							:   Union[str, Any]						,							_lowerCAmelCase							:   Tuple						,							_lowerCAmelCase							:   Tuple						,							_lowerCAmelCase							:   Optional[int]=0	):




 """simple docstring"""







 accelerator.wait_for_everyone()
 with FSDP.state_dict_type(
     _lowerCAmelCase						,							fsdp_plugin.state_dict_type						,							fsdp_plugin.state_dict_config						,							fsdp_plugin.optim_state_dict_config	):
  if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
   UpperCAmelCase__							 =						None
   # below check should work but currently it isn't working (mostly opytorch issue),
   # in the meantime disabling it at the cost of excess memory usage
   # if accelerator.process_index == 0 or not fsdp_plugin.optim_state_dict_config.rank0_only:
   UpperCAmelCase__							 =						(
       F'''{OPTIMIZER_NAME}.bin''' if optimizer_index == 0 else F'''{OPTIMIZER_NAME}_{optimizer_index}.bin'''
   )
   UpperCAmelCase__							 =						os.path.join(_lowerCAmelCase						,							_lowerCAmelCase	)
   logger.info(F'''Loading Optimizer state from {input_optimizer_file}'''	)
   UpperCAmelCase__							 =						torch.load(_lowerCAmelCase	)
   logger.info(F'''Optimizer state loaded from {input_optimizer_file}'''	)
  else:
   UpperCAmelCase__							 =						(
       os.path.join(_lowerCAmelCase						,							F'''{OPTIMIZER_NAME}_{optimizer_index}'''	)
       if F'''{OPTIMIZER_NAME}''' not in input_dir
       else input_dir
   )
   logger.info(F'''Loading Optimizer from {ckpt_dir}'''	)
   UpperCAmelCase__							 =						load_sharded_optimizer_state_dict(
       model_state_dict=model.state_dict()						,							optimizer_key="optimizer"						,							storage_reader=dist_cp.FileSystemReader(_lowerCAmelCase	)						,							)
   UpperCAmelCase__							 =						optim_state["optimizer"]
   logger.info(F'''Optimizer loaded from {ckpt_dir}'''	)
  UpperCAmelCase__							 =						FSDP.optim_state_dict_to_load(_lowerCAmelCase						,							_lowerCAmelCase						,							_lowerCAmelCase	)
  optimizer.load_state_dict(_lowerCAmelCase	)









from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tensorflow_text_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_lowerCAmelCase						:   Union[str, Any]		=   {
    "configuration_bert": ["BERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "BertConfig", "BertOnnxConfig"],
    "tokenization_bert": ["BasicTokenizer", "BertTokenizer", "WordpieceTokenizer"],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   str		=   ["BertTokenizerFast"]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   int		=   [
      "BERT_PRETRAINED_MODEL_ARCHIVE_LIST",
      "BertForMaskedLM",
      "BertForMultipleChoice",
      "BertForNextSentencePrediction",
      "BertForPreTraining",
      "BertForQuestionAnswering",
      "BertForSequenceClassification",
      "BertForTokenClassification",
      "BertLayer",
      "BertLMHeadModel",
      "BertModel",
      "BertPreTrainedModel",
      "load_tf_weights_in_bert",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   Dict		=   [
      "TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST",
      "TFBertEmbeddings",
      "TFBertForMaskedLM",
      "TFBertForMultipleChoice",
      "TFBertForNextSentencePrediction",
      "TFBertForPreTraining",
      "TFBertForQuestionAnswering",
      "TFBertForSequenceClassification",
      "TFBertForTokenClassification",
      "TFBertLMHeadModel",
      "TFBertMainLayer",
      "TFBertModel",
      "TFBertPreTrainedModel",
  ]
try:
  if not is_tensorflow_text_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   str		=   ["TFBertTokenizer"]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase						:   Optional[Any]		=   [
      "FlaxBertForCausalLM",
      "FlaxBertForMaskedLM",
      "FlaxBertForMultipleChoice",
      "FlaxBertForNextSentencePrediction",
      "FlaxBertForPreTraining",
      "FlaxBertForQuestionAnswering",
      "FlaxBertForSequenceClassification",
      "FlaxBertForTokenClassification",
      "FlaxBertModel",
      "FlaxBertPreTrainedModel",
  ]

if TYPE_CHECKING:
  from .configuration_bert import BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, BertConfig, BertOnnxConfig
  from .tokenization_bert import BasicTokenizer, BertTokenizer, WordpieceTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_bert_fast import BertTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_bert import (
        BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        BertForMaskedLM,
        BertForMultipleChoice,
        BertForNextSentencePrediction,
        BertForPreTraining,
        BertForQuestionAnswering,
        BertForSequenceClassification,
        BertForTokenClassification,
        BertLayer,
        BertLMHeadModel,
        BertModel,
        BertPreTrainedModel,
        load_tf_weights_in_bert,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_bert import (
        TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFBertEmbeddings,
        TFBertForMaskedLM,
        TFBertForMultipleChoice,
        TFBertForNextSentencePrediction,
        TFBertForPreTraining,
        TFBertForQuestionAnswering,
        TFBertForSequenceClassification,
        TFBertForTokenClassification,
        TFBertLMHeadModel,
        TFBertMainLayer,
        TFBertModel,
        TFBertPreTrainedModel,
    )

  try:
    if not is_tensorflow_text_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_bert_tf import TFBertTokenizer

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_bert import (
        FlaxBertForCausalLM,
        FlaxBertForMaskedLM,
        FlaxBertForMultipleChoice,
        FlaxBertForNextSentencePrediction,
        FlaxBertForPreTraining,
        FlaxBertForQuestionAnswering,
        FlaxBertForSequenceClassification,
        FlaxBertForTokenClassification,
        FlaxBertModel,
        FlaxBertPreTrainedModel,
    )

else:
  import sys

  _lowerCAmelCase						:   Union[str, Any]		=   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





