← Back to DevBytes

How to Use Weights & Biases to Track Fine-Tuning Experiments

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:

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:

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.

— Ad —

Google AdSense will appear here after approval

← Back to all articles