Visualize LLM training progress with Wandb Tables

In this example we will see how to instrument a custom callback for the huggingface Trainer to periodically visualize model predictions using Weight & Biases Tables

🤔 Why should I use W&B?

• Unified dashboard: Central repository for all your model metrics and predictions
• Lightweight: No code changes required to integrate with Hugging Face
• Accessible: Free for individuals and academic teams
• Secure: All projects are private by default
• Trusted: Used by machine learning teams at OpenAI, Toyota, Lyft and more

Think of W&B like GitHub for machine learning models— save machine learning experiments to your private, hosted dashboard. Experiment quickly with the confidence that all the versions of your models are saved for you, no matter where you're running your scripts.

W&B lightweight integrations works with any Python script, and all you need to do is sign up for a free W&B account to start tracking and visualizing your models.

🚀 Install, Import, and Log in

To get started with this example you will need to Install the Transformers, Weights & Biases, datasets libraries.

• Hugging Face Transformers
• Weights & Biases
• Huggingface Datasets

Uncomment the following cell install the libraries.

🖊️ Sign up for a free account →

🔑 Put in your API key

Once you've signed up, run the next cell. You'll be prompted to create a new API key at wandb.ai/settings if you haven't already. Store your API key securely. It can only be viewed once when created.

Train a language model

In this notebook, we'll see how to train a 🤗 Transformers model on the Causal language modeling task i.e. the model has to predict the next token in the sentence (so the labels are the same as the inputs shifted to the right). To make sure the model does not cheat, it gets an attention mask that will prevent it to access the tokens after token i when trying to predict the token i+1 in the sentence.

We will see how to load and preprocess the dataset for the task and train a model on it using the Trainer API. We will be building a custom trainer callback to visualize the model predictions using Weights & Biases Tables by priodically logging the predictions to the table.

Prepare the dataset

we will use the Wikitext 2 dataset as an example for this task.

For causal language modeling (CLM) task we are going to take all the texts in our dataset and concatenate them after they are tokenized. Then we split them in examples of a certain sequence length. This way the model will receive chunks of contiguous text that may look like:

part of text 1

or

end of text 1 [BOS_TOKEN] beginning of text 2

depending on whether they span over several of the original texts in the dataset or not. The labels will be the same as the inputs, shifted to the left.

We will use the gpt2 architecture for this example. You can pick any of the checkpoints listed here instead.

We can now call the tokenizer on all our texts using the map method from the Datasets library.

Next, we need to concatenate all our texts together then split the result in small chunks of a certain block_size. Here we will use a block_size of 128.

Model & TrainingArguments

Now that we have our dataset prepared we are readt to instantiate the model and the Training Arguments. For simpilicity we will train the model for 3 epochs and log the model predictions and metrics after each epoch.

WandbPredictionProgressCallback

To periodically visualize the results we will subclass the WandbCallback from the transformers library. This callback already instrumented to log model metrics, checkpoints and system metrics to Weights & Biases.

Here we will customize the callback to periodically log model predictions and labels to a wandb.Table so that we can visualize the model predictions as the training progresses. To do this, we will also need to pass the trainier and tokenizer to our callback in order to predict over the validation dataset.

Trainer

First we inistantiate the [Trainer] class with the model, training arguments, and the train, eval datasets. Since our callback needs to store predictions after each evaluation we will be passing the trainer to callback and then adding the callback to the trainer.

Note: Here we donot add the callback while inistatiating the Trainer but instead we will use the add_callback method to include the callback in the trainer after instantiation.

Training

And that's it, we are ready to train the model and visualize the predictions.

👀 Visualize results in dashboard

Click the link printed out above, or go to wandb.ai to see your results stream in live. The link to see your run in the browser will appear after all the dependencies are loaded — look for the following output:

Tracking run with wandb version <wandb-verison>
Run data is saved locally in <local-logs-path>
Syncing run <run-name> to Weights & Biases (docs)
View project at <project-url>
View run at <run-url>

Click on the to visualize the sample model predictions epoch. You should see a table similar to the one shown in the screenshot below.

📈 Track key information effortlessly by default

Weights & Biases saves a new run for each experiment. Here's the information that gets saved by default:

• Hyperparameters: Settings for your model are saved in Config
• Model Metrics: Time series data of metrics streaming in are saved in Log
• Terminal Logs: Command line outputs are saved and available in a tab
• System Metrics: GPU and CPU utilization, memory, temperature etc.

🤓 Learn more!

• Documentation: docs on the Weights & Biases and Hugging Face integration
• Videos: tutorials, interviews with practitioners, and more on our YouTube channel
• Contact: Message us at contact@wandb.com with questions