removes missing safetensors.index if not attempting to save_in_hf

Author: wesleytruong

torchtitan/protocols/state_dict_adapter.py

@@ -67,7 +67,7 @@ def __init__(self, model_args: BaseModelArgs, hf_assets_path: str | None):
             except FileNotFoundError:
                 logger.warning(
                     f"model.safetensors.index.json not found at hf_assets_path: {mapping_path}. \
-                    Defaulting to saving a single safetensors file if checkpoint is saved in HF format."
+                    Defaulting to saving final checkpoint to a single safetensors file."
                 )
                 hf_safetensors_indx = None
 

torchtitan/train.py

@@ -309,7 +309,11 @@ def __init__(self, job_config: JobConfig):
             checkpoint_config=job_config.checkpoint,
             sd_adapter=(
                 self.train_spec.state_dict_adapter(
-                    model_args, job_config.model.hf_assets_path
+                    model_args,
+                    job_config.model.hf_assets_path
+                    if job_config.checkpoint.enable_checkpoint
+                    and job_config.checkpoint.last_save_in_hf
+                    else None,
                 )
                 if self.train_spec.state_dict_adapter
                 else None