Save checkpoints locally to /mnt/local_storage before reporting to tr…

…ain (ray-project#42157) 70B finetuning ran into OOD issues because it saved the checkpoints to `/tmp`. We should save it to `/mnt/local_storage` instead. Signed-off-by: Huaiwei Sun <[email protected]>
Mark2000 · Jan 3, 2024 · 3759098 · 3759098
1 parent 4df79bb
commit 3759098
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/doc/source/templates/04_finetuning_llms_with_deepspeed/finetune_hf_llm.py b/doc/source/templates/04_finetuning_llms_with_deepspeed/finetune_hf_llm.py
@@ -487,7 +487,7 @@ def training_function(kwargs: dict):
             "learning_rate": lr_scheduler.get_lr()[0],
         }
 
-        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+        with tempfile.TemporaryDirectory(dir=args.output_dir) as temp_checkpoint_dir:
             accelerator.print(f"Saving the model locally at {temp_checkpoint_dir}")
             accelerator.wait_for_everyone()