Shameless request for help: if anybody has experience with seq2seq on TPU, and you want to do a cool project to deploy a world class Pytorch image parsing model to TPU (and do this quickly), please contact me immediately for a well paid and interesting job opportunity at nico [at] mathpix.com.
How to scale your model: A systems view of LLMs on TPUs
11–20 of 31 posts
Re: How to scale your model: A systems view of LLMs on TPUs
#12How do they make these fancy animations?
Nothing fancy. I made these with some pretty simple hand written scripts in javascript rendering to canvas: lots of fiddly little boxes moving around are simpler to script than to hand animate. (If I were to do much more of this I might rewrite these in blender since it has much nicer authoring tooling and export control.)
Re: How to scale your model: A systems view of LLMs on TPUs
#13I am really looking forward for JAX to take over pytorch/cuda over the next years. The whole PTX kerfuffle with Deepseek team shows the value of investing in more low levels approaches to squeeze out the most out of your hardware.
I like JAX but I'm not sure how an ML framework debate like "JAX vs PyTorch" is relevant to DeepSeek/PTX. The JAX API is at a similar level of abstraction to PyTorch [0]. Both are Python libraries and sit a few layers of abstraction above PTX/CUDA and their TPU equivalents.
[0] Although PyTorch arguably encompasses 2 levels, with both a pure functional library like the JAX API, as well as a "neural network" framework on top of it. Whereas JAX doesn't have the latter and leaves that to separate libraries like Flax.
Re: How to scale your model: A systems view of LLMs on TPUs
#14Not strictly related, but does anyone know why JAX uses tracing and not AST via reflection?
Re: How to scale your model: A systems view of LLMs on TPUs
#15I am really looking forward for JAX to take over pytorch/cuda over the next years. The whole PTX kerfuffle with Deepseek team shows the value of investing in more low levels approaches to squeeze out the most out of your hardware.
You do understand that PTX is part of CUDA right?
Re: How to scale your model: A systems view of LLMs on TPUs
#16Any way to convert this Jekyll site to a PDF?
Re: How to scale your model: A systems view of LLMs on TPUs
#17Any way to convert this Jekyll site to a PDF?
There are plans to release a PDF version; need to fix some formatting issues + convert the animated diagrams into static images.
Re: How to scale your model: A systems view of LLMs on TPUs
#18Not strictly related, but does anyone know why JAX uses tracing and not AST via reflection?
The short answer is that tracing is way, way easier to implement in a predictable and reliably performant way. This especially matters for distributed computation and automatic differentiation, two areas where JAX shines.
AST parsing via reflection means your ML compiler needs to re-implement all of Python, which is not a small language. This is a lot of work and hard to do well with abstractions that are not designed for those use-cases. (I believe Julia's whole language auto-diff systems struggle for essential the same reason.)
Re: How to scale your model: A systems view of LLMs on TPUs
#19I am really looking forward for JAX to take over pytorch/cuda over the next years. The whole PTX kerfuffle with Deepseek team shows the value of investing in more low levels approaches to squeeze out the most out of your hardware.
The interesting thing about this comment is that JAX is actually higher-level even than pytorch generally. Since everything is compiled you just express a logcial program and let the compiler (XLA) worry about the rest.
Are you suggesting that XLA would be where this "lower level" approach would reside since it can do more automatic optimization?
Re: How to scale your model: A systems view of LLMs on TPUs
#20Shameless request for help: if anybody has experience with seq2seq on TPU, and you want to do a cool project to deploy a world class Pytorch image parsing model to TPU (and do this quickly), please contact me immediately for a well paid and interesting job opportunity at nico [at] mathpix.com.
if you're using tpu why are you using pytorch