def save(
        self,
        booster: Booster,
        save_dir: str,
        model: nn.Module = None,
        ema: nn.Module = None,
        optimizer: Optimizer = None,
        lr_scheduler: _LRScheduler = None,
        sampler=None,
        epoch: int = None,
        step: int = None,
        global_step: int = None,
        batch_size: int = None,
        lora: bool = False,
        actual_update_step: int = None,
        ema_shape_dict: dict = None,
        async_io: bool = True,
        include_master_weights: bool = False,
    ) -> str:
        """
        Save a checkpoint.

        Args:
            booster (Booster): The Booster object.
            save_dir (str): The directory to save the checkpoint to.
            model (nn.Module): The model to save the checkpoint from.
            ema (nn.Module): The EMA model to save the checkpoint from.
            optimizer (Optimizer): The optimizer to save the checkpoint from.
            lr_scheduler (_LRScheduler): The learning rate scheduler to save the checkpoint from.
            sampler: The sampler to save the checkpoint from.
            epoch (int): The epoch of the checkpoint.
            step (int): The step of the checkpoint.
            global_step (int): The global step of the checkpoint.
            batch_size (int): The batch size of the checkpoint.
            lora (bool): Whether the model is trained with LoRA.

        Returns:
            str: The path to the saved checkpoint
        """
        self._sync_io()
        save_dir = os.path.join(save_dir, f"epoch{epoch}-global_step{actual_update_step}")
        os.environ["TENSORNVME_DEBUG_LOG"] = os.path.join(save_dir, "async_file_io.log")
        if model is not None:
            if not lora:
                os.makedirs(os.path.join(save_dir, "model"), exist_ok=True)
                booster.save_model(
                    model,
                    os.path.join(save_dir, "model"),
                    shard=True,
                    use_safetensors=True,
                    size_per_shard=4096,
                    use_async=async_io,
                )
            else:
                os.makedirs(os.path.join(save_dir, "lora"), exist_ok=True)
                booster.save_lora_as_pretrained(model, os.path.join(save_dir, "lora"))
        if optimizer is not None:
            booster.save_optimizer(
                optimizer, os.path.join(save_dir, "optimizer"), shard=True, size_per_shard=4096, use_async=async_io
            )
            if include_master_weights:
                self._prepare_master_pinned_state_dict(model, optimizer)
                master_weights_gathering(model, optimizer, self.master_pinned_state_dict)
        if lr_scheduler is not None:
            booster.save_lr_scheduler(lr_scheduler, os.path.join(save_dir, "lr_scheduler"))
        if ema is not None:
            self._prepare_pinned_state_dict(ema, ema_shape_dict)
            model_gathering(ema, ema_shape_dict, self.pinned_state_dict)
        if dist.get_rank() == 0:
            running_states = {
                "epoch": epoch,
                "step": step,
                "global_step": global_step,
                "batch_size": batch_size,
                "actual_update_step": actual_update_step,
            }
            save_json(running_states, os.path.join(save_dir, "running_states.json"))

            if ema is not None:
                if async_io:
                    self.writer = async_save(os.path.join(save_dir, "ema.safetensors"), self.pinned_state_dict)
                else:
                    torch.save(ema.state_dict(), os.path.join(save_dir, "ema.pt"))

            if sampler is not None:
                # only for VariableVideoBatchSampler
                torch.save(sampler.state_dict(step), os.path.join(save_dir, "sampler"))

            if optimizer is not None and include_master_weights:
                self.master_writer = async_save(
                    os.path.join(save_dir, "master.safetensors"), self.master_pinned_state_dict
                )

        dist.barrier()
        return save_dir

        Save a checkpoint.

        Args:
            booster (Booster): The Booster object.
            save_dir (str): The directory to save the checkpoint to.
            model (nn.Module): The model to save the checkpoint from.
            ema (nn.Module): The EMA model to save the checkpoint from.
            optimizer (Optimizer): The optimizer to save the checkpoint from.
            lr_scheduler (_LRScheduler): The learning rate scheduler to save the checkpoint from.
            sampler: The sampler to save the checkpoint from.
            epoch (int): The epoch of the checkpoint.
            step (int): The step of the checkpoint.
            global_step (int): The global step of the checkpoint.
            batch_size (int): The batch size of the checkpoint.
            lora (bool): Whether the model is trained with LoRA.

        Returns:
            str: The path to the saved checkpoint
        
opensora/utils/ckpt.py