fix: fsdp sharded state dict wont work for save_only_model knob

Signed-off-by: Mehant Kammakomati <[email protected]>

Author: kmehant

src/transformers/trainer.py

@@ -5175,6 +5175,12 @@ def create_accelerator_and_postprocess(self):
             raise ValueError(
                 "`auto_find_batch_size` isn't supported yet with DeepSpeed Zero-3. Please consider using Zero-2, Zero-1, or FSDP"
             )
+        if (
+            self.args.save_only_model
+            and self.is_fsdp_enabled
+            and "SHARDED_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type)
+        ):
+            raise ValueError("save_only_model option is not compatible with FSDP state dict type 'SHARDED_STATE_DICT'")
 
     def propagate_args_to_deepspeed(self, auto_find_batch_size=False):
         """