Hugging Face Transformers
Track Hugging Face Trainer fine-tuning runs in InstantML with a Transformers callback that logs metrics and checkpoints.
Add im.TransformersCallback to the Hugging Face Trainer to log scalar metrics as training progresses and record each saved checkpoint as an artifact. The callback is rank-zero safe, so only the main process logs in distributed training.
Install
python -m pip install "instantml[frameworks]"Authenticate with `instantml login` or INSTANTML_API_KEY if you haven't already.
Add the callback to the Trainer
This complete example fine-tunes a tiny BERT on a synthetic classification dataset:
import instantml as im
import torch
from transformers import (
AutoModelForSequenceClassification,
AutoTokenizer,
Trainer,
TrainingArguments,
)
model_name = "prajjwal1/bert-tiny"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
texts = ["great run", "bad run"] * 32
labels = [1, 0] * 32
encodings = tokenizer(texts, truncation=True, padding=True)
class TinyDataset(torch.utils.data.Dataset):
def __init__(self, encodings, labels):
self.encodings = encodings
self.labels = labels
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
item = {k: torch.tensor(v[idx]) for k, v in self.encodings.items()}
item["labels"] = torch.tensor(self.labels[idx])
return item
run = im.init(project="sft", name="bert-tiny-demo", config={"lr": 5e-5})
trainer = Trainer(
model=model,
args=TrainingArguments(
output_dir="out",
num_train_epochs=2,
logging_steps=5,
report_to="none",
),
train_dataset=TinyDataset(encodings, labels),
callbacks=[im.TransformersCallback(run=run)],
)
trainer.train()
run.finish()Open the project at instantml.ai to watch the training loss chart by global_step. You can also attach the callback to an existing Trainer with trainer.add_callback(im.TransformersCallback(run=run)).
Let the callback create the run
If you do not pass run=, the callback creates one on first log. Extra keyword arguments are forwarded to im.init(...):
trainer.add_callback(
im.TransformersCallback(project="sft", tags=["lora"])
)When no project is given it falls back to a project attribute on the TrainingArguments, then to "transformers".
What gets logged
- Metrics — scalar values from the Trainer's
on_logevents (loss, learning rate, eval metrics) at the currentglobal_step. - Checkpoints — on each
on_save, theoutput_diris recorded as acheckpointartifact tied to the current step.