`Trainer._save()` May Incorrectly Save Empty Model State (safetensors)
Author: ChenDaiwei-99
Pull request: no
Labels: name: bug

There appears to be a potential issue in the `save_model()` method of the `Trainer` class in the `Transformers` library.

When the model inherits from `PreTrainedModel`, the _save() function follows [this logic](https://github.com/huggingface/transformers/blob/10627c1a0f6877ce6715b9537afe7fafb2a89edd/src/transformers/trainer.py#L4001-L4020). However, it is possible that the `state_dict` variable is None, which can result in saving an empty state dict.

This can happen in the following case:

> `self.args.should_save` is True (as in lines [3913–3914](https://github.com/huggingface/transformers/blob/10627c1a0f6877ce6715b9537afe7fafb2a89edd/src/transformers/trainer.py#L3913-L3914))
> 
> All preceding conditional branches (lines [3879–3911](https://github.com/huggingface/transformers/blob/10627c1a0f6877ce6715b9537afe7fafb2a89edd/src/transformers/trainer.py#L3879-L3911)) are false
> 
> As a result, `self._save(output_dir)` is called without a state_dict, and state_dict remains None.

I suggest moving the following lines (currently in line [4002–4003](https://github.com/huggingface/transformers/blob/10627c1a0f6877ce6715b9537afe7fafb2a89edd/src/transformers/trainer.py#L4002C1-L4003C53)):

```
if state_dict is None:
    state_dict = self.model.state_dict()
```
outside the if-else block to ensure state_dict is always properly loaded.

**Please let me know if this analysis is correct, and whether I can submit a PR to fix it. Thank you!**


### Reproduction

1. Define and initialize a custom model that inherits from PreTrainedModel.

2. Instantiate a Trainer with the following TrainingArguments configuration:

- load_best_model_at_end=True
- fsdp disabled (is_fsdp_enabled=False)
- deepspeed disabled (is_deepspeed_enabled=False)

3. At the end of training, when the trainer tries to load the best model, it will return error msg that weights are missing.


Comments:
cc @SunMarc 

> I suggest moving the following lines (currently in line [4002–4003](https://github.com/huggingface/transformers/blob/10627c1a0f6877ce6715b9537afe7fafb2a89edd/src/transformers/trainer.py#L4002C1-L4003C53)):
> 
> if state_dict is None:
>     state_dict = self.model.state_dict()
> outside the if-else block to ensure state_dict is always properly loaded.

Thanks for the report @ChenDaiwei-99 ! 

In your case it should go through the following path no ?          
```python
else:
            self.model.save_pretrained(
                output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors
            )
```

In that case, it's not an issue if the state_dict passed is empty as we will get it in `save_pretrained` function. If you can share a minimal reproducer, I can have a quick look at why we are not saving the weights as expected. 

Thanks for your reply!

You're absolutely right — the following code does get triggered, as shown in lines [4017-4019](https://github.com/huggingface/transformers/blob/10627c1a0f6877ce6715b9537afe7fafb2a89edd/src/transformers/trainer.py#L4017):
```
else:
            self.model.save_pretrained(
                output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors
            )
```
However, I noticed a potential issue in an earlier part of the logic: when `self.args.should_save` is True (as in lines [3913–3914](https://github.com/huggingface/transformers/blob/10627c1a0f6877ce6715b9537afe7fafb2a89edd/src/transformers/trainer.py#L3913-L3914)):
```
        elif self.args.should_save:
            self._save(output_dir)

```
The `_save()` method is called without an explicit `state_dict` argument. As a result, `state_dict` defaults to `None`, and this can **_lead to an empty pytorch_model.bin (or .safetensors) being saved._**

Would you mind double-checking if this behavior is intentional, or if a fallback to `self.model.state_dict()` might be needed here?

Could you check the logic?

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.