Method save

distributed/FSDP2/checkpoint.py:199–209

(self, model: FSDPModule, optim: torch.optim.Optimizer)

            return {}

    def save(self, model: FSDPModule, optim: torch.optim.Optimizer):
        model_state_dict = self._get_full_model_state_dict(model)
        optim_state_dict = self._get_full_optimizer_state_dict(model, optim)
        if torch.distributed.get_rank() == 0:
            new_training_time = int(time.time() * 1000)
            new_checkpoint_folder = f"{self.folder}/{'dcp_api' if self.dcp_api else 'dtensor_api'}/{new_training_time}"
            new_model_checkpoint = f"{new_checkpoint_folder}/{MODEL_CHECKPOINT}"
            new_optim_checkpoint = f"{new_checkpoint_folder}/{OPTIM_CHECKPOINT}"
            os.makedirs(new_checkpoint_folder, exist_ok=True)
            torch.save(model_state_dict, new_model_checkpoint)
            torch.save(optim_state_dict, new_optim_checkpoint)