Leveraging TPUs in Colab — featuring Antigravity

I previously wrote about Google’s new Colab extension for VS Code. Shortly after that piece was posted, Google announced their new IDE, Antigravity. Oh, and they released Gemini 3 Pro too. So let’s talk about all those things!

Introduction

In my last post, while exploring the Colab for VS Code extension, I missed an opportunity to showcase the true power of Tensor Processing Units (TPUs) for deep learning. We saw some impressive numbers with NVIDIA GPUs, but the TPU performance was underwhelming. The bottleneck wasn’t the hardware; it was my sub-optimal input pipeline. In this post I will offer some simple corrections that can help improve your training pipeline.

Improvements on my Gemma fine-tuning notebook

I made a few changes to my notebook. Most notably, I’m using the TensorFlow dataset API with my in-memory data:

This is significantly better than training a model than iterating over a standard Python dictionary (which I was using before). Here are some reasons why:

• .from_tensor_slices() processes the holds the data as TensorFlow tensors, which eliminates the need to repeatedly convert Python dictionary data types to TF tensors during the training step.
• .shuffle(), .batch(), and .cache(), are all performance boosters to better streamline the processes that would otherwise incur added overhead using Python dictionary and NumPy operations.
• .prefetch(tf.data.AUTOTUNE) enables parallelism by allowing the CPU work on data prep work for the next training step while the GPU/TPU is performing model training in the current step. This overlap is seamlessly handled by the tf.data API is crucial for faster training performance.
• tf.data.Dataset also integrates perfectly with model.fit() because you just have to pass the dataset object and Keras handles the rest.

Other minor improvements include setting the Keras backend to JAX, enabling mixed_bfloat16 precision, and setting the steps_per_execution in the compile step. steps_per_execution is a powerful throughput optimization feature when training on TPUs to run multiple batches on a single compile, reducing Python overhead by keeping the TPU cores busier for longer periods without interruption.

Training results

TPU v6e: 10.4k dataset samples, rank 16, JAX backend

Notice I’m using the full dataset (10.4k for training, 400 for validation) with 2 epochs. On the NVIDIA A100, it took 1m 35s to go through 2000 samples for training one time.

Get Glen Yu’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

You can find a copy of my updated notebook here.

import antigravity

Source: https://xkcd.com/353/

Google Antigravity was recently announced and is currently in public preview. It is an agent-first IDE that is designed to enable developers operate at a higher level in an environment that is also very familiar. Visually, Antigravity is a fork of VS Code, with very simple and streamlined onboarding setup. You can just import all your extensions and settings over from VS Code.

Antigravity comes with an Agent Manager where you can manage multiple asynchronous, local agents that can work in parallel on all of your workspaces. It supports the Model Context Protocol (MCP) that allows your IDE to securely connect to a myriad of supported servers and tools.

Antigravity UI: looks just like VS Code!

At the moment there’s a very generous free tier access of Gemini 3 Pro within Antigravity, so take advantage of this now!

Taking Antigravity for a test drive

I asked Antigravity (powered by Gemini 3 Pro) to make some recommendations to optimize my LoRA tuning notebook, and as you can see, I implemented many of its suggestions:

Antigravity TPU-optimization implementation plan

I wish that, instead of modifying the code in place, Antigravity would create a separate branch for changes, similar to how Jules operates.

However, Jules and Antigravity are fundamentally different tools designed for different use cases. Jules was built to work asynchronously on existing GitHub repositories, allowing it to leverage established branching workflows. Antigravity, being an IDE, does not guarantee that the code being worked on originates from an existing repository.

Antigravity has certainly made a strong first impression. While I’m still just scratching the surface of its full capabilities, I’m excited to see how it integrates into my daily workflow over the upcoming Christmas holidays. 🎅

If you’re looking for a new IDE to try, I’d certainly recommend giving Antigravity a go.

Reference: Deep Learning with Python, 3rd edition