Browse Source

only save training params on rank 0

jpgard 7 months ago
parent
commit
6954b16b3b
1 changed files with 1 additions and 1 deletions
  1. 1 1
      src/llama_recipes/utils/train_utils.py

+ 1 - 1
src/llama_recipes/utils/train_utils.py

@@ -255,7 +255,7 @@ def train(model, train_dataloader,eval_dataloader, tokenizer, optimizer, lr_sche
         results["metrics_filename"] = metrics_filename
         results["metrics_filename"] = metrics_filename
 
 
     #saving the training params including fsdp setting for reference.
     #saving the training params including fsdp setting for reference.
-    if train_config.enable_fsdp and not train_config.use_peft:
+    if train_config.enable_fsdp and not train_config.use_peft and rank==0:
         save_train_params(train_config, fsdp_config, rank)
         save_train_params(train_config, fsdp_config, rank)
 
 
     return results
     return results