Skip to content

Update TRL documentation #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bebechien merged 4 commits into from
Oct 17, 2026
+279 −910
Merged

Update TRL documentation #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d5147ad
Drop outdated TRL workarounds from the text tuning guides
qgallouedec Oct 16, 2026
58cfb53
Let SFTTrainer and the pipelines load the model in the text tuning gu…
qgallouedec Oct 16, 2026
92fd907
Keep the default max_length in the QLoRA guide so no sample is truncated
qgallouedec Oct 16, 2026
4e3471f
Let SFTTrainer load the model in the FunctionGemma fine-tuning guide
qgallouedec Oct 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
135 changes: 61 additions & 74 deletions docs/models/functiongemma/fine-tuning.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,23 +197,13 @@ When using `shuffle=False` to your own custom datasets, ensure your source data

You are now ready to fine-tune your model. Hugging Face TRL [SFTTrainer](https://huggingface.co/docs/trl/sft_trainer) makes it straightforward to supervise fine-tune open LLMs. The `SFTTrainer` is a subclass of the `Trainer` from the `transformers` library and supports all the same features,

The following code loads the FunctionGemma model and tokenizer from Hugging Face.
The following code loads the FunctionGemma tokenizer from Hugging Face and shows how the chat template formats a training example. The `SFTTrainer` loads the model for you.

```python { .notebook-cell }
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
base_model,
dtype="auto",
device_map="auto",
attn_implementation="eager"
)
tokenizer = AutoTokenizer.from_pretrained(base_model)
from transformers import AutoTokenizer

print(f"Device: {model.device}")
print(f"DType: {model.dtype}")
# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model)

# Print formatted user prompt
print("--- dataset input ---")
Expand All @@ -224,8 +214,6 @@ print(debug_msg)
```

``` { .no-copy .output }
Device: cuda:0
DType: torch.bfloat16
--- dataset input ---
{
"messages": [
Expand Down Expand Up @@ -311,12 +299,10 @@ What is the reimbursement limit for travel meals?<end_of_turn>
<start_function_call>call:search_knowledge_base{query:<escape>travel meal reimbursement limit policy<escape>}<end_function_call><start_function_response>
```

## Before fine-tune

The output below shows that the out-of-the-box capabilities may not be good enough for this use case.
The following function measures how often a model calls the expected tool on the test set. You will use it before and after training.

```python { .notebook-cell }
def check_success_rate():
def check_success_rate(model):
success_count = 0
for idx, item in enumerate(dataset['test']):
messages = [
Expand Down Expand Up @@ -347,8 +333,61 @@ def check_success_rate():
print(f" -> ❌ wrong (hallucinated '{other_tool}')")

print(f"Success : {success_count} / {len(dataset['test'])}")
```

## Training

Before you can start your training, you need to define the hyperparameters you want to use in a `SFTConfig` instance.

```python { .notebook-cell }
from trl import SFTConfig

check_success_rate()
args = SFTConfig(
output_dir=checkpoint_dir, # directory to save and repository id
max_length=512, # max sequence length for model and packing of the dataset
packing=False, # Groups multiple samples in the dataset into a single sequence
num_train_epochs=8, # number of training epochs
per_device_train_batch_size=4, # batch size per device during training
gradient_checkpointing=False, # the model is small, so skip gradient checkpointing for speed
optim="adamw_torch_fused", # use fused adamw optimizer
logging_steps=1, # log every step
#save_strategy="epoch", # save checkpoint every epoch
eval_strategy="epoch", # evaluate checkpoint every epoch
learning_rate=learning_rate, # learning rate
lr_scheduler_type="constant", # use constant learning rate scheduler
push_to_hub=True, # push model to hub
report_to="tensorboard", # report metrics to tensorboard
)
```

You now have every building block you need to create your `SFTTrainer` to start the training of your model.

```python { .notebook-cell }
from trl import SFTTrainer

# Create Trainer object
trainer = SFTTrainer(
model=base_model,
args=args,
train_dataset=dataset['train'],
eval_dataset=dataset['test'],
)
```

``` { .no-copy .output }
Tokenizing train dataset: 0%| | 0/20 [00:00<?, ? examples/s]
Truncating train dataset: 0%| | 0/20 [00:00<?, ? examples/s]
Tokenizing eval dataset: 0%| | 0/20 [00:00<?, ? examples/s]
Truncating eval dataset: 0%| | 0/20 [00:00<?, ? examples/s]
The model is already on multiple devices. Skipping the move to device specified in `args`.
```

## Before fine-tune

The trainer has loaded the base model. The output below shows that its out-of-the-box capabilities may not be good enough for this use case.

```python { .notebook-cell }
check_success_rate(trainer.model)
```

``` { .no-copy .output }
Expand Down Expand Up @@ -415,58 +454,6 @@ check_success_rate()
Success : 2 / 20
```

## Training

Before you can start your training, you need to define the hyperparameters you want to use in a `SFTConfig` instance.

```python { .notebook-cell }
from trl import SFTConfig

torch_dtype = model.dtype

args = SFTConfig(
output_dir=checkpoint_dir, # directory to save and repository id
max_length=512, # max sequence length for model and packing of the dataset
packing=False, # Groups multiple samples in the dataset into a single sequence
num_train_epochs=8, # number of training epochs
per_device_train_batch_size=4, # batch size per device during training
gradient_checkpointing=False, # Caching is incompatible with gradient checkpointing
optim="adamw_torch_fused", # use fused adamw optimizer
logging_steps=1, # log every step
#save_strategy="epoch", # save checkpoint every epoch
eval_strategy="epoch", # evaluate checkpoint every epoch
learning_rate=learning_rate, # learning rate
fp16=True if torch_dtype == torch.float16 else False, # use float16 precision
bf16=True if torch_dtype == torch.bfloat16 else False, # use bfloat16 precision
lr_scheduler_type="constant", # use constant learning rate scheduler
push_to_hub=True, # push model to hub
report_to="tensorboard", # report metrics to tensorboard
)
```

You now have every building block you need to create your `SFTTrainer` to start the training of your model.

```python { .notebook-cell }
from trl import SFTTrainer

# Create Trainer object
trainer = SFTTrainer(
model=model,
args=args,
train_dataset=dataset['train'],
eval_dataset=dataset['test'],
processing_class=tokenizer,
)
```

``` { .no-copy .output }
Tokenizing train dataset: 0%| | 0/20 [00:00<?, ? examples/s]
Truncating train dataset: 0%| | 0/20 [00:00<?, ? examples/s]
Tokenizing eval dataset: 0%| | 0/20 [00:00<?, ? examples/s]
Truncating eval dataset: 0%| | 0/20 [00:00<?, ? examples/s]
The model is already on multiple devices. Skipping the move to device specified in `args`.
```

Start training by calling the `train()` method.

```python { .notebook-cell }
Expand Down Expand Up @@ -528,7 +515,7 @@ plt.show()
After the training is done, you'll want to evaluate and test your model. You can load different samples from the test dataset and evaluate the model on those samples.

```python { .notebook-cell }
check_success_rate()
check_success_rate(trainer.model)
```

``` { .no-copy .output }
Expand Down
83 changes: 8 additions & 75 deletions docs/notebooks/full-finetune.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -189,37 +189,6 @@
"\n",
"You are now ready to fine-tune your model. Hugging Face TRL [SFTTrainer](https://huggingface.co/docs/trl/sft_trainer) makes it straightforward to supervise fine-tune open LLMs. The `SFTTrainer` is a subclass of the `Trainer` from the `transformers` library and supports all the same features,\n",
"\n",
"The following code loads the Gemma model and tokenizer from Hugging Face."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"language": "python"
},
"outputs": [],
"source": [
"import torch\n",
"from transformers import AutoTokenizer, AutoModelForCausalLM\n",
"\n",
"# Load model and tokenizer\n",
"model = AutoModelForCausalLM.from_pretrained(\n",
" base_model,\n",
" torch_dtype=\"auto\",\n",
" device_map=\"auto\",\n",
" attn_implementation=\"eager\"\n",
")\n",
"tokenizer = AutoTokenizer.from_pretrained(base_model)\n",
"\n",
"print(f\"Device: {model.device}\")\n",
"print(f\"DType: {model.dtype}\")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Before fine-tune\n",
"\n",
"The output below shows that the out-of-the-box capabilities may not be good enough for this use case."
Expand All @@ -238,8 +207,8 @@
"from random import randint\n",
"import re\n",
"\n",
"# Load the model and tokenizer into the pipeline\n",
"pipe = pipeline(\"text-generation\", model=model, tokenizer=tokenizer)\n",
"# Load the base model into the pipeline\n",
"pipe = pipeline(\"text-generation\", model=base_model, dtype=\"auto\")\n",
"\n",
"# Load a random sample from the test dataset\n",
"rand_idx = randint(0, len(dataset[\"test\"])-1)\n",
Expand Down Expand Up @@ -335,29 +304,21 @@
"source": [
"from trl import SFTConfig\n",
"\n",
"torch_dtype = model.dtype\n",
"\n",
"args = SFTConfig(\n",
" output_dir=checkpoint_dir, # directory to save and repository id\n",
" max_length=512, # max sequence length for model and packing of the dataset\n",
" packing=False, # Groups multiple samples in the dataset into a single sequence\n",
" num_train_epochs=5, # number of training epochs\n",
" per_device_train_batch_size=4, # batch size per device during training\n",
" gradient_checkpointing=False, # Caching is incompatible with gradient checkpointing\n",
" gradient_checkpointing=False, # the model is small, so skip gradient checkpointing for speed\n",
" optim=\"adamw_torch_fused\", # use fused adamw optimizer\n",
" logging_steps=1, # log every step\n",
" save_strategy=\"epoch\", # save checkpoint every epoch\n",
" eval_strategy=\"epoch\", # evaluate checkpoint every epoch\n",
" learning_rate=learning_rate, # learning rate\n",
" fp16=True if torch_dtype == torch.float16 else False, # use float16 precision\n",
" bf16=True if torch_dtype == torch.bfloat16 else False, # use bfloat16 precision\n",
" lr_scheduler_type=\"constant\", # use constant learning rate scheduler\n",
" push_to_hub=True, # push model to hub\n",
" report_to=\"tensorboard\", # report metrics to tensorboard\n",
" dataset_kwargs={\n",
" \"add_special_tokens\": False, # Template with special tokens\n",
" \"append_concat_token\": True, # Add EOS token as separator token between examples\n",
" }\n",
")\n"
]
},
Expand All @@ -380,11 +341,10 @@
"\n",
"# Create Trainer object\n",
"trainer = SFTTrainer(\n",
" model=model,\n",
" model=base_model,\n",
" args=args,\n",
" train_dataset=dataset['train'],\n",
" eval_dataset=dataset['test'],\n",
" processing_class=tokenizer,\n",
")\n"
]
},
Expand Down Expand Up @@ -466,36 +426,9 @@
"\n",
"After the training is done, you'll want to evaluate and test your model. You can load different samples from the test dataset and evaluate the model on those samples.\n",
"\n",
"For this particular use case, the best model is a matter of preference. Interestingly, what we'd normally call 'overfitting' can be very useful for a game NPC. It forces the model to forget general information and instead lock onto the specific persona and characteristics it was trained on, ensuring it stays consistently in character."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"language": "python"
},
"outputs": [],
"source": [
"from transformers import AutoTokenizer, AutoModelForCausalLM\n",
"\n",
"model_id = checkpoint_dir\n",
"For this particular use case, the best model is a matter of preference. Interestingly, what we'd normally call 'overfitting' can be very useful for a game NPC. It forces the model to forget general information and instead lock onto the specific persona and characteristics it was trained on, ensuring it stays consistently in character.\n",
"\n",
"# Load Model\n",
"model = AutoModelForCausalLM.from_pretrained(\n",
" model_id,\n",
" torch_dtype=\"auto\",\n",
" device_map=\"auto\",\n",
" attn_implementation=\"eager\"\n",
")\n",
"tokenizer = AutoTokenizer.from_pretrained(model_id)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's load all questions from the test dataset and generate outputs."
"Let's load the fine-tuned model, then generate outputs for all questions from the test dataset."
]
},
{
Expand All @@ -508,8 +441,8 @@
"source": [
"from transformers import pipeline\n",
"\n",
"# Load the model and tokenizer into the pipeline\n",
"pipe = pipeline(\"text-generation\", model=model, tokenizer=tokenizer)\n",
"# Load the fine-tuned model into the pipeline\n",
"pipe = pipeline(\"text-generation\", model=checkpoint_dir, dtype=\"auto\")\n",
"\n",
"def test(test_sample):\n",
" # Convert as test example into a prompt with the Gemma template\n",
Expand Down
Loading
Loading