Exploring TensorFlow Model Prediction Issues | by Adam Brownell

Steps to debug BERT’s (and other LLMs’) slow prediction times on a personal computer

This all started when I was playing around with BERT models, and I got the ominous message all Data Scientists hope to avoid:

The dreaded “Kernel Died” message 💀

This happened to me while I was running my TensorFlow BERT model on my Jupyter Notebook. Training large language models (LLMs) notoriously takes a large amount of data and compute, so it could make sense for my comparably puny laptop to crash here…

… except this crash occurred during prediction, rather than training, which was strange given my assumption that more memory was used during training than prediction.

The “Kernel Died” error provided is unfortunately not very descriptive, and debugging line-by-line through the TensorFlow sounded like a daunting exercise.

A few quick searches around Stack Overflow did not completely answer my outstanding questions either. But I still needed a path forward.

This is my exploration of the Kernel dying problem and how I found a solution. 🚀

Given the only thing I knew about my issue was that the kernel died, I had to gather more context. From a few other threads, it seemed clear that the reason for the kernel dying was the my model prediction required more RAM than my CPU…