Correct saving with FSDP

llama3.1-70b-bnb/llama3.1-70b_train.py

@@ -120,6 +120,8 @@ if local_process_index == 0:
 
 # Save model in first process only:
 if process_index == 0:
+    if trainer.is_fsdp_enabled:
+        trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
     trainer.save_model()
 
 # trainer.push_to_hub()