CarperAI / CarperAI/trlx

Question about trainer.save_pretrained

Open
#412 1 comment 0 reactions 0 assignees View on GitHub
feature request
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.