--- title: I Optimized Google Colab’s Free TPU to Train 44× Faster... description: A free Colab TPU running optimized JAX code trains a language model about 44 times faster than a free T4 GPU. One billion tokens takes roughly 3 hours on the... canonical: https://daily.dev/posts/i-optimized-google-colab-s-free-tpu-to-train-44-faster-than-t4-hbhm48uao twitter:card: summary_large_image twitter:site: @dailydotdev og:type: website og:site_name: daily.dev og:title: I Optimized Google Colab’s Free TPU to Train 44× Faster Than T4 | daily.dev og:description: A free Colab TPU running optimized JAX code trains a language model about 44 times faster than a free T4 GPU. One billion tokens takes roughly 3 hours on the... og:url: https://daily.dev/posts/i-optimized-google-colab-s-free-tpu-to-train-44-faster-than-t4-hbhm48uao og:image: https://api.daily.dev/og/posts/hbhM48Uao.png og:image:alt: I Optimized Google Colab’s Free TPU to Train 44× Faster Than T4 og:image:width: 1200 og:image:height: 630 og:locale: en --- > ## Documentation Index > Fetch the complete documentation index at: https://daily.dev/llms.txt > Use this file to discover all available pages before exploring further. # I Optimized Google Colab’s Free TPU to Train 44× Faster Than T4 **[Medium](https://daily.dev/sources/medium_js)** · 22 min read · 0 upvotes · 0 comments ## Summary A free Colab TPU running optimized JAX code trains a language model about 44 times faster than a free T4 GPU. One billion tokens takes roughly 3 hours on the TPU versus about 6 days on the T4. Most of the speedup comes from the TPU hardware itself and switching to JAX, not from custom kernels, which added less than 10%. Splash attention was the single biggest win at long sequences but did almost nothing at short ones. Fine-grained mixture-of-experts layers turned out to be expensive on the TPU, with tiny experts costing about 6 times a dense layer. The optimized TPU's performance is roughly in the same class as an RTX 5090 GPU, by estimate from published specs. ## Full article daily.dev links to this article rather than hosting it. Read it at the original source: ## Questions this post answers ### How much faster is a Colab TPU v5e than a T4 GPU for training with optimized JAX code? An optimized TPU v5e running JAX trains about 44 times faster than a free Colab T4 GPU running PyTorch. One billion tokens takes about 3 hours on the optimized TPU versus about 6 days on the T4. Most of the gain, roughly 7.7x, comes from switching hardware, another 4.5x from moving to JAX, and splash attention, Megablox and data-movement fixes contribute only about 1.08x more. _Anyone weighing free TPU versus GPU training time can track benchmarks like this on daily.dev._ ### Does splash attention actually speed up TPU training compared to flash attention? Yes, splash attention with block size 1024 and a fused backward pass ran about 8.8 times faster than default flash attention at sequence length 8192, taking 8.05 ms versus 70.6 ms. However, the gain is bottleneck-dependent: at a shorter sequence length of 512 used in a separate smaller model, splash attention contributed almost no speedup since attention was not the dominant cost there. _Developers tuning attention kernels can follow benchmark breakdowns like this on daily.dev._ ### Why are mixture-of-experts models with many small experts slow on TPUs? Tiny experts are expensive on TPU hardware because streaming expert weights from memory dominates the cost when each expert sees few tokens. Timing stayed flat at about 24 microseconds from 8 to 64 tokens per expert, and roughly 400 or more tokens per expert were needed to reach about 90% of peak efficiency. Even the best granularity setting still cost about 2.2 times a dense layer with equivalent active parameters, and finer experts cost far more, up to 6x at a width of 256. _Engineers designing MoE architectures for TPU can reference benchmarks like these on daily.dev._ ## Similar posts on daily.dev - [Writing an LLM from scratch, part 28 -- training a base model from scratch on an RTX 3090](https://daily.dev/posts/writing-an-llm-from-scratch-part-28----training-a-base-model-from-scratch-on-an-rtx-3090-qv0nonoeu) · Hacker News · 2 upvotes · 0 comments - [Google New TPU Generation is Specifically Designed for Agents and SOTA Model Training](https://daily.dev/posts/google-new-tpu-generation-is-specifically-designed-for-agents-and-sota-model-training-aznf1vfqp) · InfoQ · 0 upvotes · 0 comments - [Training large models on Ironwood TPUs](https://daily.dev/posts/training-large-models-on-ironwood-tpus-zrglcgqla) · Google Cloud · 1 upvotes · 0 comments --- Tags: [#data-science](https://daily.dev/tags/data-science), [#mixture-of-experts](https://daily.dev/tags/mixture-of-experts) [View this post on daily.dev](https://daily.dev/posts/i-optimized-google-colab-s-free-tpu-to-train-44-faster-than-t4-hbhm48uao) ```json {"@context":"https://schema.org","@graph":[{"@type":"Organization","@id":"https://daily.dev/#organization","name":"daily.dev","url":"https://daily.dev","logo":{"@type":"ImageObject","url":"https://daily.dev/apple-touch-icon.png","width":180,"height":180},"sameAs":["https://twitter.com/dailydotdev","https://github.com/dailydotdev","https://www.linkedin.com/company/daily-dev-ltd"]},{"@type":"WebSite","@id":"https://daily.dev/#website","url":"https://daily.dev","name":"daily.dev","publisher":{"@id":"https://daily.dev/#organization"},"potentialAction":{"@type":"SearchAction","target":{"@type":"EntryPoint","urlTemplate":"https://daily.dev/search?q={search_term_string}"},"query-input":"required name=search_term_string"}}]} {"@context":"https://schema.org","@type":"TechArticle","headline":"I Optimized Google Colab’s Free TPU to Train 44× Faster Than T4","url":"https://daily.dev/posts/i-optimized-google-colab-s-free-tpu-to-train-44-faster-than-t4-hbhm48uao","mainEntityOfPage":{"@type":"WebPage","@id":"https://daily.dev/posts/i-optimized-google-colab-s-free-tpu-to-train-44-faster-than-t4-hbhm48uao"},"datePublished":"2026-10-31T11:07:58.975Z","dateModified":"2026-10-31T11:09:09.768Z","description":"A free Colab TPU running optimized JAX code trains a language model about 44 times faster than a free T4 GPU. One billion tokens takes roughly 3 hours on the...","image":"https://media.daily.dev/image/upload/f_auto,q_auto/v1/posts/6eec1bdd9eb2c38510972c7a72463a63?_a=AQAEuop","thumbnailUrl":"https://media.daily.dev/image/upload/f_auto,q_auto/v1/posts/6eec1bdd9eb2c38510972c7a72463a63?_a=AQAEuop","isAccessibleForFree":true,"articleSection":"Medium","inLanguage":"en","publisher":{"@type":"Organization","name":"daily.dev","url":"https://daily.dev","logo":{"@type":"ImageObject","url":"https://daily.dev/apple-touch-icon.png","width":180,"height":180}},"author":{"@type":"Organization","name":"Medium","logo":"https://media.daily.dev/image/upload/t_logo,f_auto/v1/logos/medium","url":"https://daily.dev/sources/medium_js"},"commentCount":0,"discussionUrl":"https://daily.dev/posts/i-optimized-google-colab-s-free-tpu-to-train-44-faster-than-t4-hbhm48uao","interactionStatistic":[{"@type":"InteractionCounter","interactionType":{"@type":"LikeAction"},"userInteractionCount":0},{"@type":"InteractionCounter","interactionType":{"@type":"CommentAction"},"userInteractionCount":0}],"keywords":"data-science,mixture-of-experts","timeRequired":"PT22M"} {"@context":"https://schema.org","@type":"BreadcrumbList","itemListElement":[{"@type":"ListItem","position":1,"name":"Home","item":"https://daily.dev"},{"@type":"ListItem","position":2,"name":"Medium","item":"https://daily.dev/sources/medium_js"},{"@type":"ListItem","position":3,"name":"I Optimized Google Colab’s Free TPU to Train 44× Faster Than T4"}]} {"@context":"https://schema.org","@type":"FAQPage","@id":"https://daily.dev/posts/i-optimized-google-colab-s-free-tpu-to-train-44-faster-than-t4-hbhm48uao#faq","mainEntity":[{"@type":"Question","name":"How much faster is a Colab TPU v5e than a T4 GPU for training with optimized JAX code?","acceptedAnswer":{"@type":"Answer","text":"An optimized TPU v5e running JAX trains about 44 times faster than a free Colab T4 GPU running PyTorch. One billion tokens takes about 3 hours on the optimized TPU versus about 6 days on the T4. Most of the gain, roughly 7.7x, comes from switching hardware, another 4.5x from moving to JAX, and splash attention, Megablox and data-movement fixes contribute only about 1.08x more. Anyone weighing free TPU versus GPU training time can track benchmarks like this on daily.dev."}},{"@type":"Question","name":"Does splash attention actually speed up TPU training compared to flash attention?","acceptedAnswer":{"@type":"Answer","text":"Yes, splash attention with block size 1024 and a fused backward pass ran about 8.8 times faster than default flash attention at sequence length 8192, taking 8.05 ms versus 70.6 ms. However, the gain is bottleneck-dependent: at a shorter sequence length of 512 used in a separate smaller model, splash attention contributed almost no speedup since attention was not the dominant cost there. Developers tuning attention kernels can follow benchmark breakdowns like this on daily.dev."}},{"@type":"Question","name":"Why are mixture-of-experts models with many small experts slow on TPUs?","acceptedAnswer":{"@type":"Answer","text":"Tiny experts are expensive on TPU hardware because streaming expert weights from memory dominates the cost when each expert sees few tokens. Timing stayed flat at about 24 microseconds from 8 to 64 tokens per expert, and roughly 400 or more tokens per expert were needed to reach about 90% of peak efficiency. Even the best granularity setting still cost about 2.2 times a dense layer with equivalent active parameters, and finer experts cost far more, up to 6x at a width of 256. Engineers designing MoE architectures for TPU can reference benchmarks like these on daily.dev."}}]} ```