






"""simple docstring"""


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


lowercase_				     =     get_logger(__name__)



def 					lowerCAmelCase				(    __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase=0 ):




  """simple docstring"""







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

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



def 					lowerCAmelCase				(    __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase=0 ):




  """simple docstring"""







  accelerator.wait_for_everyone()
  with FSDP.state_dict_type(
      __UpperCamelCase		,       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(__UpperCamelCase ) != 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
      __A									= f'{MODEL_NAME}.bin' if model_index == 0 else f'{MODEL_NAME}_{model_index}.bin'
      __A									= os.path.join(__UpperCamelCase		,       __UpperCamelCase )
      logger.info(f'Loading model from {input_model_file}' )
      __A									= torch.load(__UpperCamelCase )
      logger.info(f'Model loaded from {input_model_file}' )
    elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
      __A									= (
          f'{MODEL_NAME}_rank{accelerator.process_index}.bin'
          if model_index == 0
          else f'{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'
      )
      __A									= os.path.join(__UpperCamelCase		,       __UpperCamelCase )
      logger.info(f'Loading model from {input_model_file}' )
      __A									= torch.load(__UpperCamelCase )
      logger.info(f'Model loaded from {input_model_file}' )
    elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
      __A									= (
          os.path.join(__UpperCamelCase		,       f'{MODEL_NAME}_{model_index}' )
          if f'{MODEL_NAME}' not in input_dir
          else input_dir
      )
      logger.info(f'Loading model from {ckpt_dir}' )
      __A									= {'''model''': model.state_dict()}
      dist_cp.load_state_dict(
          state_dict=__UpperCamelCase		,       storage_reader=dist_cp.FileSystemReader(__UpperCamelCase )		,       planner=DefaultLoadPlanner()		,       )
      __A									= state_dict['''model''']
      logger.info(f'Model loaded from {ckpt_dir}' )
    model.load_state_dict(__UpperCamelCase )



def 					lowerCAmelCase				(    __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase=0 ):




  """simple docstring"""







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





def 					lowerCAmelCase				(    __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase=0 ):




  """simple docstring"""







  accelerator.wait_for_everyone()
  with FSDP.state_dict_type(
      __UpperCamelCase		,       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:
      __A									= 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:
      __A									= (
          f'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else f'{OPTIMIZER_NAME}_{optimizer_index}.bin'
      )
      __A									= os.path.join(__UpperCamelCase		,       __UpperCamelCase )
      logger.info(f'Loading Optimizer state from {input_optimizer_file}' )
      __A									= torch.load(__UpperCamelCase )
      logger.info(f'Optimizer state loaded from {input_optimizer_file}' )
    else:
      __A									= (
          os.path.join(__UpperCamelCase		,       f'{OPTIMIZER_NAME}_{optimizer_index}' )
          if f'{OPTIMIZER_NAME}' not in input_dir
          else input_dir
      )
      logger.info(f'Loading Optimizer from {ckpt_dir}' )
      __A									= load_sharded_optimizer_state_dict(
          model_state_dict=model.state_dict()		,       optimizer_key='''optimizer'''		,       storage_reader=dist_cp.FileSystemReader(__UpperCamelCase )		,       )
      __A									= optim_state['''optimizer''']
      logger.info(f'Optimizer loaded from {ckpt_dir}' )
    __A									= FSDP.optim_state_dict_to_load(__UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase )
    optimizer.load_state_dict(__UpperCamelCase )












"""simple docstring"""


import os
import tempfile
import unittest
from pathlib import Path

from transformers import AutoConfig, is_torch_available
from transformers.testing_utils import require_torch, torch_device


if is_torch_available():
       from transformers import PyTorchBenchmark, PyTorchBenchmarkArguments


@require_torch
class 					snake_case		(							unittest.TestCase						):


   '''simple docstring'''

   def 				_SCREAMING_SNAKE_CASE				(	self :     Any,       _lowerCamelCase :     Optional[int]  ):

     '''simple docstring'''
     for model_result in results.values():
       for batch_size, sequence_length in zip(model_result['''bs'''],       model_result['''ss''']  ):
         __A									= model_result['''result'''][batch_size][sequence_length]
         self.assertIsNotNone(_lowerCamelCase  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[Any]  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[Any]  ):

     '''simple docstring'''
     __A									= '''sgugger/tiny-distilbert-classification'''
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       only_pretrain_model=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Any  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       torchscript=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )


   @unittest.skipIf(torch_device == '''cpu''',       '''Cant do half precision'''  )
   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[Any]  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       fpaa=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Dict  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= AutoConfig.from_pretrained(_lowerCamelCase  )
     # set architectures equal to `None`
     __A									= None
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase,       configs=[config]  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     int  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result  )
     self.check_results_dict_not_empty(results.memory_train_result  )


   @unittest.skipIf(torch_device == '''cpu''',       '''Can\'t do half precision'''  )
   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[int]  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       fpaa=_lowerCamelCase,       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result  )
     self.check_results_dict_not_empty(results.memory_train_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     str  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= AutoConfig.from_pretrained(_lowerCamelCase  )
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase,       configs=[config]  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     List[Any]  ):

     '''simple docstring'''
     __A									= '''sshleifer/tinier_bart'''
     __A									= AutoConfig.from_pretrained(_lowerCamelCase  )
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase,       configs=[config]  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[Any]  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= AutoConfig.from_pretrained(_lowerCamelCase  )
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase,       configs=[config]  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result  )
     self.check_results_dict_not_empty(results.memory_train_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Dict  ):

     '''simple docstring'''
     __A									= '''sshleifer/tinier_bart'''
     __A									= AutoConfig.from_pretrained(_lowerCamelCase  )
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase,       configs=[config]  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result  )
     self.check_results_dict_not_empty(results.memory_train_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[Any]  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     with tempfile.TemporaryDirectory() as tmp_dir:
       __A									= PyTorchBenchmarkArguments(
           models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       save_to_csv=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       inference_time_csv_file=os.path.join(_lowerCamelCase,       '''inf_time.csv'''  ),       train_memory_csv_file=os.path.join(_lowerCamelCase,       '''train_mem.csv'''  ),       inference_memory_csv_file=os.path.join(_lowerCamelCase,       '''inf_mem.csv'''  ),       train_time_csv_file=os.path.join(_lowerCamelCase,       '''train_time.csv'''  ),       env_info_csv_file=os.path.join(_lowerCamelCase,       '''env.csv'''  ),       multi_process=_lowerCamelCase,       )
       __A									= PyTorchBenchmark(_lowerCamelCase  )
       benchmark.run()
       self.assertTrue(Path(os.path.join(_lowerCamelCase,       '''inf_time.csv'''  )  ).exists()  )
       self.assertTrue(Path(os.path.join(_lowerCamelCase,       '''train_time.csv'''  )  ).exists()  )
       self.assertTrue(Path(os.path.join(_lowerCamelCase,       '''inf_mem.csv'''  )  ).exists()  )
       self.assertTrue(Path(os.path.join(_lowerCamelCase,       '''train_mem.csv'''  )  ).exists()  )
       self.assertTrue(Path(os.path.join(_lowerCamelCase,       '''env.csv'''  )  ).exists()  )






   def 				_SCREAMING_SNAKE_CASE				(	self :     Union[str, Any]  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''

     def _check_summary_is_not_empty(_lowerCamelCase :     List[Any]  ):
       self.assertTrue(hasattr(_lowerCamelCase,       '''sequential'''  )  )
       self.assertTrue(hasattr(_lowerCamelCase,       '''cumulative'''  )  )
       self.assertTrue(hasattr(_lowerCamelCase,       '''current'''  )  )
       self.assertTrue(hasattr(_lowerCamelCase,       '''total'''  )  )

     with tempfile.TemporaryDirectory() as tmp_dir:
       __A									= PyTorchBenchmarkArguments(
           models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       log_filename=os.path.join(_lowerCamelCase,       '''log.txt'''  ),       log_print=_lowerCamelCase,       trace_memory_line_by_line=_lowerCamelCase,       multi_process=_lowerCamelCase,       )
       __A									= PyTorchBenchmark(_lowerCamelCase  )
       __A									= benchmark.run()
       _check_summary_is_not_empty(result.inference_summary  )
       _check_summary_is_not_empty(result.train_summary  )
       self.assertTrue(Path(os.path.join(_lowerCamelCase,       '''log.txt'''  )  ).exists()  )




