Question about trainer.save_pretrained
- Dominant language
- Python
- Stars
- 4.8k
- Forks
- 487
- PR merge metrics
- No merged PRs in 30d
Description
### 🚀 The feature, motivation, and pitch
Here is the reply for #365 :
> Assume that you having checkpoint output is `best_checkpoint/pytorch_model/mp_rank_00_model_states.pt`. I guess you can try something like this:
>
> ```python
> import torch
> # Import the model architecture used during training and load the weights
> from trlx.models.modeling_ppo import AutoModelForCausalLMWithHydraValueHead
> model = AutoModelForCausalLMWithValueHead.from_pretrained("...")
> model.load_state_dict(torch.load("best_checkpoint/pytorch_model/mp_rank_00_model_states.pt")["module"])
> ```
>
> Another way that you can save your model directly to huggingface format by this refer this https://github.com/CarperAI/trlx#save-the-resulting-model-to-a-hugging-face-pretrained-language-model-ready-to-upload-to-the-hub.
And I have another question about this issue, when executing:
```python
trainer = trlx.train(config=config, reward_fn=lambda samples, **kwargs: [float(int(sample)) for sample in samples])
trainer.save_pretrained('/path/to/output/folder/')
```
Is the trainer storing the last checkpoint or the best checkpoint?
I suspect that it is the last checkpoint? If so, how can I save the best checkpoint so that I can load it using:
```python
AutoModelForCausalLM.from_pretrained(path)
```
### Alternatives
_No response_
### Additional context
_No response_
Contributor guide
Research direction
Start at the trainer.save_pretrained entry point and trace how checkpoints are selected and written. Reproduce the shown call with a small training run, then verify which checkpoint can be loaded through AutoModelForCausalLM.from_pretrained; done means the last-versus-best behavior and supported best-checkpoint workflow are clear.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- huggingface, python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 30/100