Efficient MultiModal Data Pipeline

You've got everything ready - data, model, a beefy GPU setup. You hit "run" and... wait. And wait some more. Your GPUs are barely breaking a sweat while your wallet's getting lighter by the hour.

Sound familiar? We've been there. After some detective work on our nanoVLM project, we discovered the real culprit wasn't our model or hardware, it was our data pipeline being incredibly wasteful.

Here's what we found:

1. Idle GPUs: Our model was literally waiting around for data to show up
2. Padding hell: Every batch was stuffed with useless padding tokens that contributed nothing to training

In this post we build an efficient pipeline in five stages. In each stage we add or remove from the previous step and comment on what went right and what did not.

Table of Contents:

• Stage 0: Pre Requisites
• Stage 1: Visualising the Dataset
• Stage 2: Naive Padding
• Stage 3: Constrained Padding
• Stage 4: Packing Smarter with Knapsacks
• Stage 5: Knapsack for Multimodal Data
• Conclusion

[Stage 0] Preparation

To make it easier to follow the data preparation tasks, we created a separate repo laser-focused on the data pipeline only. We hope this will be much easier to understand that reading the code once integrated with the nanoVLM repository. In addition, this could be useful to bootstrap other data pipelines!

Repository: https://github.com/ariG23498/mmdp

To follow along, all you need to do is clone the repository. It contains the final data preparation tasks, but it's designed to showcase each step of the way.

$ git clone https://github.com/ariG23498/mmdp.git

[Stage 1] Visualising the Dataset

Before optimizing anything, we need to understand what we are working with. Our multimodal dataset has images, text prompts, and responses.

$ uv run 01_check_dataset.py

Getting familiar with your training data is crucial for success. The previous script shows a random sample each time you run it; you may want to copy the snippet to a notebook and run it multiple times to get a feeling about the data.