Introduction to Weights & Biases for Fine-Tuning
Fine-tuning large language models and other machine learning models is an inherently experimental process. You tweak learning rates, adjust batch sizes, swap datasets, modify architectures, and rerun training dozens of times. Without a robust tracking system, you quickly lose visibility into which configuration produced the best results. This is where Weights & Biases (W&B) comes in — a popular experiment tracking platform that lets you log, visualize, compare, and share your machine learning experiments with minimal code overhead.
In this tutorial, you will learn what Weights & Biases is, why it is especially valuable for fine-tuning workflows, how to integrate it into your training scripts, and what best practices will keep your experiment tracking clean and useful.
What Is Weights & Biases?
Weights & Biases is a developer-first MLOps platform designed to help machine learning practitioners track experiments, version datasets, evaluate models, and collaborate across teams. At its core, it provides a lightweight Python SDK that you drop into your training code. As your model trains, W&B streams metrics, system stats, configuration parameters, and artifacts to a centralized dashboard where you can compare runs side by side.
The platform offers several key components:
- Runs: A single execution of your training script. Each run captures configs, metrics, logs, and output files.
- Projects: A collection of related runs, typically representing a single model or task.
- Artifacts: Versioned datasets and model checkpoints, useful for reproducibility.
- Reports: Shareable documents combining charts, markdown, and run comparisons.
- Sweeps: Automated hyperparameter search with grid, random, or Bayesian optimization strategies.
Why Tracking Matters for Fine-Tuning
Fine-tuning is fundamentally a search problem. You are searching for the combination of hyperparameters, data preprocessing, and training strategy that yields the best performance on your target task. This search is complicated by several factors:
- High cost: Fine-tuning large models consumes significant compute. You cannot afford to lose track of which configuration worked.
- Subtle differences: A change in learning rate from 2e-5 to 3e-5 can dramatically affect convergence. Without logging, these differences are invisible.
- Reproducibility: When you find a winning configuration, you need to know exactly what produced it — including random seeds, library versions, and data versions.
- Collaboration: Teams often fine-tune in parallel. A shared tracking system prevents duplicated effort and enables knowledge sharing.
By integrating W&B into your fine-tuning pipeline, every run becomes a queryable, comparable, and reproducible record. You can answer questions like "Which learning rate gave the lowest validation loss?" or "Did adding more training data improve F1 score?" in seconds rather than digging through scattered log files.
Setting Up Weights & Biases
Installation and Authentication
Start by installing the W&B Python package. It works with any Python 3.7+ environment and integrates natively with PyTorch, TensorFlow, JAX, Hugging Face Transformers, and most other ML frameworks.
pip install wandb
After installation, authenticate with your W&B account. You will need an API key, which you can generate from your account settings at wandb.ai/authorize.
wandb login
This command prompts you to paste your API key and stores it locally so future runs authenticate automatically. In CI/CD or cloud training environments, you can set the key as an environment variable instead:
export WANDB_API_KEY="your-api-key-here"
Creating Your First Tracked Run
The minimal integration requires just three lines of code: initialize a run, log metrics during training, and finish the run when done. Here is a bare-bones example:
import wandb
# Initialize a new run
wandb.init(
project="my-finetuning-project",
config={
"learning_rate": 2e-5,
"batch_size": 16,
"epochs": 3,
"model_name": "bert-base-uncased",
}
)
# Simulate a training loop
for epoch in range(wandb.config.epochs):
train_loss = 0.4 / (epoch + 1)
val_loss = 0.5 / (epoch + 1)
val_accuracy = 0.7 + 0.08 * epoch
wandb.log({
"epoch": epoch,
"train_loss": train_loss,
"val_loss": val_loss,
"val_accuracy": val_accuracy,
})
# Finish the run
wandb.finish()
When you run this script, W&B creates a new run in your project, streams the logged metrics to your dashboard in real time, and saves the configuration dictionary. You can view live loss curves, system resource usage, and the full config from your browser.
Integrating W&B with Hugging Face Transformers
Automatic Logging with the Trainer API
One of the most common fine-tuning scenarios involves Hugging Face Transformers. The Trainer class has built-in W&B support, so integration requires almost no extra code. When W&B is installed and authenticated, the Trainer automatically logs training and evaluation metrics, the model configuration, and training arguments.
from transformers import (
AutoModelForSequenceClassification,
AutoTokenizer,
TrainingArguments,
Trainer,
)
from datasets import load_dataset
import wandb
# Initialize a W&B run with a descriptive name and config
wandb.init(
project="sentiment-finetuning",
name="bert-base-lr2e5-bs16",
config={
"model_name": "bert-base-uncased",
"task": "sentiment-classification",
"dataset": "imdb",
"learning_rate": 2e-5,
"batch_size": 16,
"epochs": 3,
"max_seq_length": 256,
},
)
# Load tokenizer and model
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(
model_name, num_labels=2
)
# Load and tokenize dataset
dataset = load_dataset("imdb")
def tokenize_function(examples):
return tokenizer(
examples["text"],
padding="max_length",
truncation=True,
max_length=256,
)
tokenized_datasets = dataset.map(tokenize_function, batched=True)
train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(2000))
eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(500))
# Configure training arguments
training_args = TrainingArguments(
output_dir="./results",
eval_strategy="epoch",
save_strategy="epoch",
learning_rate=2e-5,
per_device_train_batch_size=16,
per_device_eval_batch_size=16,
num_train_epochs=3,
weight_decay=0.01,
logging_dir="./logs",
report_to="wandb",
load_best_model_at_end=True,
metric_for_best_model="eval_loss",
)
# Define a compute_metrics function
import numpy as np
from sklearn.metrics import accuracy_score, f1_score
def compute_metrics(eval_pred):
logits, labels = eval_pred
predictions = np.argmax(logits, axis=-1)
return {
"accuracy": accuracy_score(labels, predictions),
"f1": f1_score(labels, predictions),
}
# Create and run the trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
compute_metrics=compute_metrics,
)
trainer.train()
# Log a summary table of final metrics
eval_results = trainer.evaluate()
wandb.log({
"final_eval_loss": eval_results["eval_loss"],
"final_accuracy": eval_results["eval_accuracy"],
"final_f1": eval_results["eval_f1"],
})
wandb.finish()
The key line is report_to="wandb" in TrainingArguments. This tells the Trainer to send all its standard logging — loss, learning rate, gradient norms, evaluation metrics — to W&B. Combined with the explicit wandb.init() call, you get both automatic metric logging and a custom run name and configuration.
Logging Custom Metrics and Visualizations
Beyond automatic logging, you often want to track custom metrics that the Trainer does not capture by default. For example, you might want to log per-class precision and recall, confusion matrices, or sample predictions. W&B supports logging rich objects like tables and images.
import wandb
from sklearn.metrics import confusion_matrix
import numpy as np
# After running evaluation
predictions = trainer.predict(eval_dataset)
pred_labels = np.argmax(predictions.predictions, axis=-1)
true_labels = predictions.label_ids
# Log a confusion matrix as a W&B Table
cm = confusion_matrix(true_labels, pred_labels)
wandb.log({
"confusion_matrix": wandb.Table(
data=cm.tolist(),
columns=["predicted_negative", "predicted_positive"],
)
})
# Log sample predictions for qualitative inspection
sample_indices = np.random.choice(len(true_labels), size=10, replace=False)
sample_data = []
for idx in sample_indices:
sample_data.append([
eval_dataset[int(idx)]["text"][:200],
int(true_labels[idx]),
int(pred_labels[idx]),
])
wandb.log({
"sample_predictions": wandb.Table(
data=sample_data,
columns=["text", "true_label", "predicted_label"],
)
})
This kind of qualitative logging is invaluable during fine-tuning. Seeing where your model makes mistakes helps you decide whether to adjust the dataset, the learning rate, or the model architecture.
Tracking Model Checkpoints as Artifacts
Metrics tell you how well a model performed, but artifacts let you version and retrieve the actual model weights and datasets that produced those metrics. W&B Artifacts provide a versioning system that integrates directly with your training runs.
import wandb
wandb.init(project="sentiment-finetuning", name="bert-with-artifacts")
# Log the dataset as an artifact
dataset_artifact = wandb.Artifact(
name="imdb-subset",
type="dataset",
description="2000 train / 500 eval subset of IMDB",
metadata={
"train_size": len(train_dataset),
"eval_size": len(eval_dataset),
"max_seq_length": 256,
},
)
# If you saved the dataset to disk, add the files
# dataset_artifact.add_dir("./data")
wandb.log_artifact(dataset_artifact)
# After training, log the model checkpoint as an artifact
model_artifact = wandb.Artifact(
name="bert-sentiment-model",
type="model",
description="Fine-tuned BERT for sentiment classification",
metadata={
"base_model": "bert-base-uncased",
"epochs": 3,
"final_eval_loss": eval_results["eval_loss"],
},
)
model_artifact.add_dir("./results/best_model")
wandb.log_artifact(model_artifact)
wandb.finish()
Each time you log an artifact with the same name, W&B creates a new version (v0, v1, v2, and so on). You can then reference a specific version in downstream runs, ensuring full reproducibility. For example, you can load a specific dataset version before training:
import wandb
wandb.init(project="sentiment-finetuning", name="reuse-dataset-v2")
# Download and use a specific dataset artifact version
artifact = wandb.use_artifact("imdb-subset:v2", type="dataset")
artifact_dir = artifact.download()
# Load data from artifact_dir...
wandb.finish()
Using W&B Sweeps for Hyperparameter Search
Manually trying different hyperparameter combinations is tedious and inefficient. W&B Sweeps automate this process. You define a search space and a strategy, and W&B launches multiple runs with different configurations, visualizing all of them together so you can identify the best performing setup.
import wandb
import numpy as np
from transformers import (
AutoModelForSequenceClassification,
AutoTokenizer,
TrainingArguments,
Trainer,
)
sweep_config = {
"method": "bayes",
"metric": {
"name": "eval_loss",
"goal": "minimize",
},
"parameters": {
"learning_rate": {
"min": 1e-5,
"max": 5e-5,
},
"batch_size": {
"values": [8, 16, 32],
},
"num_train_epochs": {
"values": [2, 3, 4],
},
"weight_decay": {
"values": [0.0, 0.01, 0.1],
},
},
}
sweep_id = wandb.sweep(sweep_config, project="sentiment-finetuning-sweeps")
def train_with_config():
wandb.init()
config = wandb.config
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained(
"bert-base-uncased", num_labels=2
)
training_args = TrainingArguments(
output_dir="./sweep_results",
eval_strategy="epoch",
learning_rate=config.learning_rate,
per_device_train_batch_size=config.batch_size,
num_train_epochs=config.num_train_epochs,
weight_decay=config.weight_decay,
report_to="wandb",
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
)
trainer.train()
wandb.finish()
# Launch the sweep agent (runs 10 trials)
wandb.agent(sweep_id, function=train_with_config, count=10)
The Bayesian search strategy intelligently explores the hyperparameter space by using results from previous runs to guide subsequent trials. This is typically more efficient than random or grid search, especially when each run is expensive.
Best Practices for Tracking Fine-Tuning Experiments
Name Your Runs Descriptively
Default run names like "happy-otter-42" are cute but unhelpful when you have fifty runs in a project. Use a naming convention that encodes the key variables you are experimenting with, such as bert-lr2e5-bs16-ep3-wd01. This makes it easy to scan the dashboard and identify runs at a glance.
Log All Hyperparameters in the Config
Every variable that could affect the outcome should be in the config dictionary passed to wandb.init(). This includes learning rate, batch size, number of epochs, weight decay, warmup steps, random seed, model name, dataset version, and any preprocessing choices. The config becomes the ground truth for what produced a given set of metrics.
Use Tags and Notes for Organization
W&B supports tagging runs and adding notes. Use tags to mark runs by status (e.g., "baseline", "experiment", "production") or by theme (e.g., "lr-tuning", "data-augmentation"). Notes let you record observations that do not fit into structured fields, such as "this run diverged at step 500, possibly due to too high a learning rate."
wandb.init(
project="sentiment-finetuning",
name="bert-lr2e5-bs16-ep3",
tags=["baseline", "lr-tuning", "bert"],
notes="Initial baseline run with standard hyperparameters.",
config={...},
)
Log at Consistent Intervals
For meaningful comparison across runs, log metrics at consistent intervals. If one run logs every 10 steps and another logs every 100 steps, the resulting curves will be hard to compare. Use a fixed logging step interval in your training loop or rely on the Trainer's built-in logging_steps parameter.
Set Random Seeds Everywhere
Reproducibility requires controlling randomness. Set seeds for Python, NumPy, PyTorch, and any other libraries that use random number generators. Log the seed value in your W&B config so you can recreate a run exactly.
import random
import numpy as np
import torch
def set_seed(seed=42):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
set_seed(42)
wandb.init(config={"seed": 42, ...})
Avoid Logging Sensitive Data
Be cautious about logging raw text samples, especially when fine-tuning on proprietary or sensitive datasets. W&B stores everything you log, and it may be visible to team members or, in the case of public projects, anyone on the internet. Consider logging only metadata or anonymized samples.
Use Grouping for Related Runs
When running cross-validation or repeated experiments with the same configuration, use the group parameter to cluster related runs together in the dashboard. This keeps your project view clean and lets you aggregate metrics across the group.
wandb.init(
project="sentiment-finetuning",
group="5fold-cv",
name="fold-0",
config={...},
)
Clean Up Failed Runs
Not every run will complete successfully. Runs that crash early can clutter your dashboard and skew aggregate statistics. You can delete failed runs from the W&B UI, or mark them with a "failed" tag so they can be filtered out during analysis.
Conclusion
Weights & Biases transforms fine-tuning from a chaotic process of scattered log files and guesswork into a structured, queryable, and reproducible workflow. By investing a small amount of time upfront to integrate W&B into your training scripts, you gain the ability to compare runs visually, search through historical experiments, version your datasets and model checkpoints, and automate hyperparameter search. The practices outlined in this tutorial — descriptive run names, comprehensive config logging, artifact versioning, consistent logging intervals, and disciplined seed management — will help you get the most out of the platform. As your fine-tuning projects grow in complexity and your team grows in size, having a robust experiment tracking system becomes not just a convenience but a necessity for making informed, data-driven decisions about your models.