





'''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


__A	: Dict              =       get_logger(__name__)





def 			lowerCAmelCase_						(    a					:					Dict   , a					:					Union[str, Any]   , a					:					str   , a					:					Tuple   , a					:					Dict=0				):
							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_						(    a					:					int   , a					:					int   , a					:					str   , a					:					Optional[Any]   , a					:					Any=0				):
							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_						(    a					:					Dict   , a					:					str   , a					:					Optional[Any]   , a					:					Optional[int]   , a					:					int   , a					:					List[Any]=0				):
							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_						(    a					:					Optional[Any]   , a					:					Union[str, Any]   , a					:					List[Any]   , a					:					Any   , a					:					str   , a					:					Any=0				):
							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_				)







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


_lowercase											=    {
    """configuration_falcon""": ["""FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP""", """FalconConfig"""],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _lowercase											=    [
     """FALCON_PRETRAINED_MODEL_ARCHIVE_LIST""",
     """FalconForCausalLM""",
     """FalconModel""",
     """FalconPreTrainedModel""",
     """FalconForSequenceClassification""",
     """FalconForTokenClassification""",
     """FalconForQuestionAnswering""",
 ]


if TYPE_CHECKING:
 from .configuration_falcon import FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP, FalconConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_falcon import (
      FALCON_PRETRAINED_MODEL_ARCHIVE_LIST,
      FalconForCausalLM,
      FalconForQuestionAnswering,
      FalconForSequenceClassification,
      FalconForTokenClassification,
      FalconModel,
      FalconPreTrainedModel,
  )


else:
 import sys

 _lowercase											=    _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
