Set pretrained model name correctly, if provided, in HF Checkpointer (#…

…1407)

llmfoundry/callbacks/hf_checkpointer.py

@@ -217,6 +217,14 @@ def __init__(
             )
 
         self.mlflow_logging_config = mlflow_logging_config
+        if 'metadata' in self.mlflow_logging_config:
+            self.pretrained_model_name = self.mlflow_logging_config[
+                'metadata'].get(
+                    'pretrained_model_name',
+                    None,
+                )
+        else:
+            self.pretrained_model_name = None
 
         self.huggingface_folder_name_fstr = os.path.join(
             'huggingface',
@@ -529,6 +537,16 @@ def tensor_hook(
                 original_tokenizer,
             )
 
+            # Ensure that the pretrained model name is correctly set on the saved HF checkpoint.
+            if self.pretrained_model_name is not None:
+                new_model_instance.name_or_path = self.pretrained_model_name
+                if self.using_peft:
+                    new_model_instance.base_model.name_or_path = self.pretrained_model_name
+                    for k in new_model_instance.peft_config.keys():
+                        new_model_instance.peft_config[
+                            k
+                        ].base_model_name_or_path = self.pretrained_model_name
+
             log.debug('Saving Hugging Face checkpoint to disk')
             # This context manager casts the TE extra state in io.BytesIO format to tensor format
             # Needed for proper hf ckpt saving.
@@ -624,10 +642,7 @@ def tensor_hook(
                     # Upload the license file generated by mlflow during the model saving.
                     license_filename = _maybe_get_license_filename(
                         local_save_path,
-                        self.mlflow_logging_config['metadata'].get(
-                            'pretrained_model_name',
-                            None,
-                        ),
+                        self.pretrained_model_name,
                     )
                     if license_filename is not None:
                         mlflow_logger._mlflow_client.log_artifact(