How to Debug a ML Model: A Step-by-Step Case Study in NLP
By: Yada Pruksachatkun
Tags: AI, NLP, Data Science, Technology, Machine Learning

How to Debug a ML Model: A Step-by-Step Case Study in NLP

While there are so many articles out there on how to get started on NLP or teaching you a tutorial, one of the hardest lessons to learn is how to debug a model or task implementation.

Mental state before model/task implementation, StockSnap via Pixabay (CC0)

Mental state after model/task implementation, Free-Photo via Pixabay (CC0)

Not to worry! This article will go through the debugging process of a pretty subtle (and not so subtle) series of bugs, and how we fixed them, with a case study to walk you through the lessons. If you would like to just see a bulleted list of tips, scroll down to the end!

In order to do that, let me take you back a few months, to when we (my research collaborator Phu and I) were first implementing masked language modeling into jiant, which is an opensource NLP framework, with the goal of doing multi task training on a RoBERTa model. If this sounds like an alien language to you, I would first suggest you look into this article on transfer learning and multitask learning, and this article about the RoBERTa model.

Setting up the Scene

Masked language modeling is one of the pretraining objectives in BERT, RoBERTa, and many BERT-style variants. It consists of an input-noising objective, where given a text, the model has to predict 15% of the tokens given the context. To make things harder, these predicted tokens are 80% of the time replaced by “[MASK]”, 10% by another random token, and 10% is the correct, unreplaced token.

For example, the model will be shown the below

Example of a text changed for MLM training. Here, the model will learn to predict “tail” for the token currently occupied with “[MASK]”

Designing the Initial Implementation

We first looked into if other people had implemented MLM before, and found the original implementation by Google, and a Pytorch implementation by AllenNLP. We used mostly all of the Huggingface implementation (which has been moved since, since it seems like the file that used to be there no longer exists) for the forward function. Following the RoBERTa paper, we dynamically masked the batch at each time step. Furthermore, Huggingface exposes the pretrained MLM head here, which we utilized as below.

Thus, the MLM flow in our code became the below:

Load MLM data -> Preprocess and index data -> Load model -> In each step of the model training, we: 1. Dynamically mask batch 2. Compute NLL loss for each masked token

The jiant framework uses primarily AllenNLP for vocabulary creation and indexing, as well as instance and dataset management.

We first tested with a toy dataset of 100 dataset examples to make sure the loading was correct with AllenNLP. After we went through some pretty explicit bugs, such as some label type mismatch with AllenNLP, we came upon a bigger bug.

The First Signs of Trouble

After making sure our preprocessing code worked with AllenNLP, we found a strange bug.

TypeError: ~ (operator.invert) is only implemented on byte tensors. Traceback (most recent call last):

This was because the code we copy-pasted from Huggingface was written with an older version of Python, and in the Pytorch you needed to use .byte() instead of bool() .

Thus, we simply changed one line, from

indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices

to

bernoulli_mask = torch.bernoulli(torch.full(labels.shape, 0.8)).to( device=inputs.device, dtype=torch.uint8 )

Trouble Strikes

Now, finally, we were able to run a forward function without erroring out! After a few minutes of celebration, we got to work verifying more subtle bugs. We first tested the correctness of our implementation by calling model.eval() and running the model through the MLM forward functions. Since the model, in this case RoBERTa-large, has been pretrained with MLM, we would expect it to do very well on MLM. That was not the case, and we were getting very high losses.

It became clear why: the predictions were always 2 off from the gold labels. For example, if the token “tail” was assigned index 25, the label for “dog wagged its [MASK] when it saw the treat” and “[MASK]” would be 25, but the prediction would be 27.

We only discovered this after hitting this error.

`/pytorch/aten/src/THCUNN/ClassNLLCriterion.cu:105: void cunn_ClassNLLCriterion_updateOutput_kernel(Dtype *, Dtype *, Dtype *, long *, Dtype *, int, int, int, int, long) [with Dtype = float, Acctype = float]: block: [0,0,0], thread: [19,0,0] Assertion `t >= 0 && t < n_classes` failed.`

This error meant that the prediction space was larger than the number of classes.

After a lot of pdb tracing, we realized that we weren’t using AllenNLP tag/label namespace. In AllenNLP, you can keep track of all the vocabularies you need in an AllenNLP Vocabulary object using namespaces such as the label namespace and input namespace. We found that AllenNLP vocabulary object automatically inserts @@PADDING@@ and @@UNKOWN@@ tokens to index 0 and 1 for all namespaces except for label namespaces (which are all those ending in “_tag” or “_labels.” Since we did not use the label namespace, our indices were being shifted forward by two, and the prediction space (defined by the label vocabulary size) was larger by 2! After finding this out, we renamed the label index and this particular threat was curbed.

A Last Hidden Bug, and a pivot

By this point, we had thought we had caught all or most of the bugs, and that MLM was working correctly. However, while the model was getting low perplexity now, a week later, while looking through the code with a third person, we found another hidden bug.

if self._unk_id is not None:

ids = (ids — 2) * valid_mask + self._pad_id * pad_mask + self._unk_id * unk_mask

else:

ids = (ids — 2) * valid_mask + self._pad_id * pad_mask

Somewhere buried in a separate part of the code, which we had written a few months back, we had shifted the inputs of any Transformer-based model back by 2, since AllenNLP shifted it forward by 2. Thus, essentially the model was seeing gibberish, since it was seeing whatever vocabulary tokens were two indices away from the correct inputs’ indices.

How did we fix this?

We ended up reversing a previous fix for a previous bug and not using the label namespace for inputs, since everything was shifted back by 2 anyhow. and simply making sure that the dynamically generated mask_idx is shifted forward by 2 before being fed into the model. In order to fix the previous error with mismatch between prediction and label space sizes, we made the number of labels the size of the pretrained model’s tokenizer, since that includes all of the vocabulary that that model was pretrained on.

After many countless hours debugging and running preliminary experiments, we were finally out of bugs. Phew!

So just as a recap of things we did to make sure the code was bug free, and to also recap the types of bugs we saw, here’s a nifty list.

Key Takeaways for debugging a model

Start testing with a toy dataset. For this case study, preprocessing the entire dataset would’ve taken ~4 hours. Use already created infrastructure if possible. Beware, if you are using other people’s code, make sure you know exactly how it fits into your code and any incompatibilities that may arise, both subtle and not, from integrating outside code into your own. If you’re working with pretrained models, and if it makes sense, try to load a trained model and make sure that it does well in a task. Beware of differences in Pytorch versions (and other versions of dependencies) between code. Be very careful with indexing. Sometimes sketching ut the flow of the indexing is very messy, and can cause a lot of headaches of why your model isn’t performing well. Get other people to look at your code too. You may need to get deeper into the weeds and look at the source codes of packages you are using for preprocessing (such as AllenNLP) in order to understand the source of bugs. Create unit tests to keep track of subtle bugs and to ensure you don’t fall back on those bugs with a code change.

And there you have it, a debugging case study and some lessons. We’re all on a journey to get better at debugging together, and I know that I’m far from an expert at model debugging, but hopefully this post was helpful for you! Special thanks to Phu Mon Htut for editing this post.

If you’d like to see the final implementation, check it out here!