

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					:			Union[str, Any]  				=			get_logger(__name__)






def        UpperCamelCase_(							_snake_case   :  Optional[int]       ,     _snake_case   :  Optional[Any]       ,     _snake_case   :  Optional[int]       ,     _snake_case   :  List[Any]       ,     _snake_case   :  int=0	):




			"""simple docstring"""






			os.makedirs(_snake_case       ,     exist_ok=_snake_case	)
			with FSDP.state_dict_type(
			    _snake_case       ,     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(_snake_case       ,     _snake_case	)
									if accelerator.process_index == 0:
												logger.info(F'Saving model to {output_model_file}'	)
												torch.save(_snake_case       ,     _snake_case	)
												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(_snake_case       ,     _snake_case	)
									logger.info(F'Saving model to {output_model_file}'	)
									torch.save(_snake_case       ,     _snake_case	)
									logger.info(F'Model saved to {output_model_file}'	)
						elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
									__a        =os.path.join(_snake_case       ,     F'{MODEL_NAME}_{model_index}'	)
									os.makedirs(_snake_case       ,     exist_ok=_snake_case	)
									logger.info(F'Saving model to {ckpt_dir}'	)
									__a        ={'model': state_dict}

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






def        UpperCamelCase_(							_snake_case   :  Any       ,     _snake_case   :  Tuple       ,     _snake_case   :  Dict       ,     _snake_case   :  Optional[Any]       ,     _snake_case   :  Any=0	):




			"""simple docstring"""






			accelerator.wait_for_everyone()
			with FSDP.state_dict_type(
			    _snake_case       ,     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(_snake_case	) != 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(_snake_case       ,     _snake_case	)
									logger.info(F'Loading model from {input_model_file}'	)
									__a        =torch.load(_snake_case	)
									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(_snake_case       ,     _snake_case	)
									logger.info(F'Loading model from {input_model_file}'	)
									__a        =torch.load(_snake_case	)
									logger.info(F'Model loaded from {input_model_file}'	)
						elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
									__a        =(
									    os.path.join(_snake_case       ,     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=_snake_case       ,     storage_reader=dist_cp.FileSystemReader(_snake_case	)       ,     planner=DefaultLoadPlanner()       ,     )
									__a        =state_dict['model']
									logger.info(F'Model loaded from {ckpt_dir}'	)
						model.load_state_dict(_snake_case	)






def        UpperCamelCase_(							_snake_case   :  str       ,     _snake_case   :  Union[str, Any]       ,     _snake_case   :  Optional[Any]       ,     _snake_case   :  Optional[Any]       ,     _snake_case   :  Any       ,     _snake_case   :  str=0	):




			"""simple docstring"""






			os.makedirs(_snake_case       ,     exist_ok=_snake_case	)
			with FSDP.state_dict_type(
			    _snake_case       ,     fsdp_plugin.state_dict_type       ,     fsdp_plugin.state_dict_config       ,     fsdp_plugin.optim_state_dict_config	):
						__a        =FSDP.optim_state_dict(_snake_case       ,     _snake_case	)
						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(_snake_case       ,     _snake_case	)
												logger.info(F'Saving Optimizer state to {output_optimizer_file}'	)
												torch.save(_snake_case       ,     _snake_case	)
												logger.info(F'Optimizer state saved in {output_optimizer_file}'	)
						else:
									__a        =os.path.join(_snake_case       ,     F'{OPTIMIZER_NAME}_{optimizer_index}'	)
									os.makedirs(_snake_case       ,     exist_ok=_snake_case	)
									logger.info(F'Saving Optimizer state to {ckpt_dir}'	)
									dist_cp.save_state_dict(
									    state_dict={'optimizer': optim_state}       ,     storage_writer=dist_cp.FileSystemWriter(_snake_case	)       ,     planner=DefaultSavePlanner()       ,     )
									logger.info(F'Optimizer state saved in {ckpt_dir}'	)






def        UpperCamelCase_(							_snake_case   :  List[str]       ,     _snake_case   :  List[Any]       ,     _snake_case   :  Dict       ,     _snake_case   :  List[Any]       ,     _snake_case   :  Tuple       ,     _snake_case   :  str=0	):




			"""simple docstring"""






			accelerator.wait_for_everyone()
			with FSDP.state_dict_type(
			    _snake_case       ,     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(_snake_case       ,     _snake_case	)
									logger.info(F'Loading Optimizer state from {input_optimizer_file}'	)
									__a        =torch.load(_snake_case	)
									logger.info(F'Optimizer state loaded from {input_optimizer_file}'	)
						else:
									__a        =(
									    os.path.join(_snake_case       ,     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(_snake_case	)       ,     )
									__a        =optim_state['optimizer']
									logger.info(F'Optimizer loaded from {ckpt_dir}'	)
						__a        =FSDP.optim_state_dict_to_load(_snake_case       ,     _snake_case       ,     _snake_case	)
						optimizer.load_state_dict(_snake_case	)









from manim import *






class 	__magic_name__    (			lowerCAmelCase_      ):





			def 		__magic_name__    (		self					)   -> Any:




						'''simple docstring'''
						__a        =Rectangle(height=0.5					,  width=0.5					)
						__a        =Rectangle(height=0.25					,  width=0.25					)
						__a        =Rectangle(height=0.46					,  width=0.46					).set_stroke(width=0					)

						__a        =[mem.copy() for i in range(6					)]
						__a        =[mem.copy() for i in range(6					)]
						__a        =VGroup(*__snake_case					).arrange(__snake_case					,  buff=0					)
						__a        =VGroup(*__snake_case					).arrange(__snake_case					,  buff=0					)
						__a        =VGroup(__snake_case					,  __snake_case					).arrange(__snake_case					,  buff=0					)
						__a        =Text('CPU'					,  font_size=24					)
						__a        =Group(__snake_case					,  __snake_case					).arrange(__snake_case					,  buff=0.5					,  aligned_edge=__snake_case					)
						cpu.move_to([-2.5, -0.5, 0]					)
						self.add(__snake_case					)

						__a        =[mem.copy() for i in range(4					)]
						__a        =VGroup(*__snake_case					).arrange(__snake_case					,  buff=0					)
						__a        =Text('GPU'					,  font_size=24					)
						__a        =Group(__snake_case					,  __snake_case					).arrange(__snake_case					,  buff=0.5					,  aligned_edge=__snake_case					)
						gpu.move_to([-1, -1, 0]					)
						self.add(__snake_case					)

						__a        =[mem.copy() for i in range(6					)]
						__a        =VGroup(*__snake_case					).arrange(__snake_case					,  buff=0					)

						__a        =Text('Model'					,  font_size=24					)
						__a        =Group(__snake_case					,  __snake_case					).arrange(__snake_case					,  buff=0.5					,  aligned_edge=__snake_case					)
						model.move_to([3, -1.0, 0]					)
						self.add(__snake_case					)

						__a        =[]
						__a        =[]
						__a        =[]

						for i, rect in enumerate(__snake_case					):
									rect.set_stroke(__snake_case					)

									__a        =Rectangle(height=0.46 / 4					,  width=0.46 / 3					).set_stroke(width=0.0					).set_fill(__snake_case					,  opacity=0.7					)

									if i == 0:
												cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT					)					,  buff=0.02					,  direction=__snake_case					)
												cpu_target.set_x(cpu_target.get_x() + 0.1					)
									elif i == 3:
												cpu_target.next_to(model_cpu_arr[0]					,  direction=__snake_case					,  buff=0.0					)
									else:
												cpu_target.next_to(model_cpu_arr[i - 1]					,  direction=__snake_case					,  buff=0.0					)
									self.add(__snake_case					)
									model_cpu_arr.append(__snake_case					)

						self.add(*__snake_case					,  *__snake_case					,  *__snake_case					)

						__a        =[mem.copy() for i in range(6					)]
						__a        =VGroup(*__snake_case					).arrange(__snake_case					,  buff=0					)

						__a        =Text('Loaded Checkpoint'					,  font_size=24					)
						__a        =Group(__snake_case					,  __snake_case					).arrange(__snake_case					,  buff=0.5					,  aligned_edge=__snake_case					)
						checkpoint.move_to([3, 0.5, 0]					)

						self.add(__snake_case					)

						__a        =[]
						__a        =[]

						for i, rect in enumerate(__snake_case					):
									__a        =fill.copy().set_fill(__snake_case					,  opacity=0.7					)
									target.move_to(__snake_case					)
									ckpt_arr.append(__snake_case					)

									__a        =target.copy()
									if i < 5:
												cpu_target.move_to(cpu_left_col_base[i + 1]					)
									else:
												cpu_target.move_to(cpu_right_col_base[i - 5]					)
									ckpt_cpu_arr.append(__snake_case					)
						self.add(*__snake_case					,  *__snake_case					)

						__a        =Square(side_length=2.2					)
						key.move_to([-5, 2, 0]					)

						__a        =MarkupText(
						    f'<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'					,  font_size=18					,  )

						key_text.move_to([-5, 2.4, 0]					)

						self.add(__snake_case					,  __snake_case					)

						__a        =MarkupText(
						    f'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'					,  font_size=18					,  )

						blue_text.next_to(__snake_case					,  DOWN * 2.4					,  aligned_edge=key_text.get_left()					)
						self.add(__snake_case					)

						__a        =MarkupText(
						    f'Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device.'					,  font_size=24					,  )
						step_a.move_to([2, 2, 0]					)

						__a        =[meta_mem.copy() for i in range(6					)]
						__a        =[meta_mem.copy() for i in range(6					)]
						__a        =VGroup(*__snake_case					).arrange(__snake_case					,  buff=0					)
						__a        =VGroup(*__snake_case					).arrange(__snake_case					,  buff=0					)
						__a        =VGroup(__snake_case					,  __snake_case					).arrange(__snake_case					,  buff=0					)
						__a        =Text('Disk'					,  font_size=24					)
						__a        =Group(__snake_case					,  __snake_case					).arrange(__snake_case					,  buff=0.5					,  aligned_edge=__snake_case					)
						disk.move_to([-4.0, -1.25, 0]					)
						self.play(Write(__snake_case					,  run_time=3					)					,  Write(__snake_case					,  run_time=1					)					,  Create(__snake_case					,  run_time=1					)					)

						__a        =[]
						for i, rect in enumerate(__snake_case					):
									__a        =rect.copy()
									target.generate_target()
									target.target.move_to(disk_left_col_base[i]					).scale(0.5					)
									animations.append(MoveToTarget(__snake_case					,  run_time=1.5					)					)
						self.play(*__snake_case					)

						self.play(FadeOut(__snake_case					)					)

						__a        =MarkupText(f'Then, the checkpoint is removed from memory\nthrough garbage collection.'					,  font_size=24					)
						step_a.move_to([2, 2, 0]					)

						self.play(Write(__snake_case					,  run_time=3					)					)

						self.play(
						    FadeOut(__snake_case					,  __snake_case					,  *__snake_case					,  *__snake_case					)					,  )

						self.wait()






