def register_hf_model(
        self,
        state: State,
        logger: Logger,
        temp_save_dir: str,
        original_tokenizer: PreTrainedTokenizerBase,
        use_temp_dir: bool,
        new_model_instance: Union[PreTrainedModel, 'PeftModel'],
    ):
        """Register the model to MLflow.

        Args:
            state (State): The training state.
            logger (Logger): The logger.
            temp_save_dir (str): The temporary save directory.
            original_tokenizer (PreTrainedTokenizerBase): The original tokenizer.
            use_temp_dir (bool): Whether to use a temporary directory.
            new_model_instance (Union[PreTrainedModel, PeftModel]): The new model instance.
        """
        assert new_model_instance is not None
        new_model_instance = self.transform_model_pre_registration(
            new_model_instance,
        )
        register_save_dir = os.path.join(
            temp_save_dir,
            'register_save',
        )
        new_model_instance.save_pretrained(
            register_save_dir,
            max_shard_size='1GB',
        )
        if original_tokenizer:
            original_tokenizer.save_pretrained(register_save_dir)

        self.pre_register_edit(register_save_dir)

        self.save_additional_contents(
            state,
            logger,
            register_save_dir,
        )

        for mlflow_logger in self.mlflow_loggers:
            if self.mlflow_registered_model_name:
                log.debug(
                    f'Registering model to UC at {mlflow_logger.model_registry_prefix}.{self.mlflow_registered_model_name}',
                )

            # Save the monitor process to be restored after registering the model.
            with _monitor_process_saver(mlflow_logger):
                process = SpawnProcess(
                    target=_log_model_with_multi_process,
                    kwargs={
                        'mlflow_logger':
                            mlflow_logger,
                        'python_logging_level':
                            logging.getLogger('llmfoundry').level,
                        'transformers_model': {
                            'model': new_model_instance,
                            'tokenizer': original_tokenizer,
                        } if self.using_peft else register_save_dir,
                        'artifact_path':
                            'final_model_checkpoint',
                        'pretrained_model_name':
                            self.pretrained_model_name,
                        'registered_model_name':
                            self.mlflow_registered_model_name,
                        'await_registration_for':
                            3600,
                        'mlflow_logging_config':
                            self.mlflow_logging_config,
                    },
                )

                process.start()
                self.register_processes.append(process)

        # Save the temporary directory to be cleaned up later.
        if use_temp_dir:
            self.temp_save_dir = temp_save_dir
Register the model to MLflow.

        Args:
            state (State): The training state.
            logger (Logger): The logger.
            temp_save_dir (str): The temporary save directory.
            original_tokenizer (PreTrainedTokenizerBase): The original tokenizer.
            use_temp_dir (bool): Whether to use a temporary directory.
            new_model_instance (Union[PreTrainedModel, PeftModel]): The new model instance.
        
llmfoundry/callbacks/hf_checkpointer.py