def _save_checkpoint(
        self,
        state: State,
        logger: Logger,
        upload_to_save_folder: bool,
        register_to_mlflow: bool,
    ):
        """Save a HuggingFace formatted checkpoint.

        Args:
            state (State): The training state.
            logger (Logger): The logger.
            upload_to_save_folder (bool): Whether to upload the HF checkpoint to the save folder.
            register_to_mlflow (bool): Whether to register the model to MLFlow
        """
        save_dir = format_name_with_dist_and_time(
            str(
                Path(self.save_dir_format_str) /
                self.huggingface_folder_name_fstr,
            ),
            state.run_name,
            state.timestamp,
        )

        # Use a temporary directory if save_dir is remote.
        use_temp_dir = self.remote_ud is not None
        temp_save_dir = tempfile.mkdtemp() if use_temp_dir else save_dir

        new_model_instance, original_tokenizer = self._get_hf_model(state)

        dist.barrier()

        if dist.get_global_rank() == 0:
            assert new_model_instance is not None
            if upload_to_save_folder:
                # This context manager casts the TE extra state in io.BytesIO format to tensor format
                # Needed for proper hf ckpt saving.
                context_manager = te.onnx_export(
                    True,
                ) if is_te_imported and state.precision == Precision.AMP_FP8 else contextlib.nullcontext(
                )

                _LLM_FOUNDRY_SAVE_FOLDER_HF_MAX_SHARD_SIZE = os.environ.get(
                    'LLM_FOUNDRY_SAVE_FOLDER_HF_MAX_SHARD_SIZE',
                    '1GB',
                )
                with context_manager:
                    new_model_instance.save_pretrained(
                        temp_save_dir,
                        max_shard_size=
                        _LLM_FOUNDRY_SAVE_FOLDER_HF_MAX_SHARD_SIZE,
                    )
                if original_tokenizer is not None:  # type: ignore
                    assert isinstance(
                        original_tokenizer,
                        PreTrainedTokenizerBase,
                    )
                    original_tokenizer.save_pretrained( # type: ignore
                        temp_save_dir,
                    )

                # Only need to edit files for MPT because it has custom code
                if new_model_instance.config.model_type == 'mpt':  # type: ignore
                    log.debug('Editing MPT files for HuggingFace compatibility')
                    edit_files_for_hf_compatibility(
                        temp_save_dir,
                        self.flatten_imports,
                    )

                self.save_additional_contents(
                    state,
                    logger,
                    temp_save_dir,
                )

                if self.remote_ud is not None:
                    for filename in os.listdir(temp_save_dir):
                        remote_file_name = os.path.join(save_dir, filename)
                        remote_file_uri = self.remote_ud.remote_backend.get_uri(
                            remote_file_name,
                        )
                        log.info(
                            f'Uploading HuggingFace formatted checkpoint to {remote_file_uri}',
                        )
                        self.remote_ud.upload_file(
                            state=state,
                            remote_file_name=remote_file_name,
                            file_path=Path(
                                os.path.join(temp_save_dir, filename),
                            ),
                            overwrite=self.overwrite,
                        )

        dist.barrier()

        if dist.get_global_rank() == 0:
            assert new_model_instance is not None
            if self.using_peft:
                model_name = self.mlflow_logging_config.get('metadata', {}).get(
                    'pretrained_model_name',
                    None,
                )
                if model_name is not None:
                    from peft import PeftModel
                    assert isinstance(new_model_instance, PeftModel)
                    assert isinstance(
                        new_model_instance.model,
                        PreTrainedModel,
                    )
                    new_model_instance.name_or_path = model_name
                    new_model_instance.model.name_or_path = model_name
                    new_model_instance.base_model.name_or_path = model_name
            if register_to_mlflow:
                self.register_hf_model(
                    state,
                    logger,
                    temp_save_dir,
                    original_tokenizer,
                    use_temp_dir,
                    new_model_instance,
                )
            else:
                # Clean up the temporary directory if we don't need to register to mlflow.
                if use_temp_dir:
                    shutil.rmtree(temp_save_dir)
        dist.barrier()
Save a HuggingFace formatted checkpoint.

        Args:
            state (State): The training state.
            logger (Logger): The logger.
            upload_to_save_folder (bool): Whether to upload the HF checkpoint to the save folder.
            register_to_mlflow (bool): Whether to register the model to MLFlow
        
llmfoundry/callbacks/hf_checkpointer.py