Hacker News Board Hacker News, browsed like a classic message board
Live data from Hacker News

How to scale your model: A systems view of LLMs on TPUs

jax-ml.github.io

Post reply on HN
Print view View original

21–30 of 31 posts

Re: How to scale your model: A systems view of LLMs on TPUs

#21

Post by Scene_Cast2 »

jdeaton wrote: post #19
3abiton wrote: post #4

I 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?

I'm curious, what does paradigmatic JAX look like? Is there an equivalent of picoGPT [1] for JAX?

[1] https://github.com/jaymody/picoGPT/blob/main/gpt2.py

Re: How to scale your model: A systems view of LLMs on TPUs

#23

Post by jdeaton »

Scene_Cast2 wrote: post #21
jdeaton wrote: post #19

Earlier quoted context omitted.

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?

I'm curious, what does paradigmatic JAX look like? Is there an equivalent of picoGPT [1] for JAX? [1] https://github.com/jaymody/picoGPT/blob/main/gpt2.py

yeah it looks exactly like that file but replace "import numpy as np" with "import jax.numpy as np" :)

Re: How to scale your model: A systems view of LLMs on TPUs

#24

Post by swyx »

awongh wrote: post #10
mattjjatgoogle wrote: post #6

An author's tweet thread: https://x.com/jacobaustin132/status/1886844716446007300

Here in the thread he says: https://x.com/jacobaustin132/status/1886844724339675340 : `5 years ago, there were many ML architectures, but today, there is (mostly) only one [transformers].` To what degree is this actually true, and what else is on the horizon that might become as popular as transformers?

it's quite true. the convergence of all archs to transformers is well documented by karpathy. SSMs were once touted as transformer killers, but increasingly look like just optional supplements.

Re: How to scale your model: A systems view of LLMs on TPUs

#25

Post by almostgotcaught »

brap wrote: post #14

Not strictly related, but does anyone know why JAX uses tracing and not AST via reflection?

> AST via reflection

I literally am a paid ML compiler engineer and I have no idea what this means. You understand that reflection, ala looking in a mirror is about being about to identify a type's type at runtime. It has nothing to do with the AST.

Re: How to scale your model: A systems view of LLMs on TPUs

#26

Post by achierius »

3abiton wrote: post #4

I 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.

What PTX kerfuffle are you referring to?

Re: How to scale your model: A systems view of LLMs on TPUs

#27

Post by bronxbomber92 »

almostgotcaught wrote: post #25
brap wrote: post #14

Not strictly related, but does anyone know why JAX uses tracing and not AST via reflection?

> AST via reflection I literally am a paid ML compiler engineer and I have no idea what this means. You understand that reflection, ala looking in a mirror is about being about to identify a type's type at runtime . It has nothing to do with the AST.

An example of “AST via reflection”

https://docs.scala-lang.org/scala3/reference/metaprogramming...

Re: How to scale your model: A systems view of LLMs on TPUs

#28

Post by hustwindmaple1 »

jdeaton wrote: post #20
nicodjimenez wrote: post #11

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.

if you're using tpu why are you using pytorch

there is limited TPU support in pytorch via torch_xla

Re: How to scale your model: A systems view of LLMs on TPUs

#29

Post by brap »

almostgotcaught wrote: post #25
brap wrote: post #14

Not strictly related, but does anyone know why JAX uses tracing and not AST via reflection?

> AST via reflection I literally am a paid ML compiler engineer and I have no idea what this means. You understand that reflection, ala looking in a mirror is about being about to identify a type's type at runtime . It has nothing to do with the AST.

Congratulations, but that's not what reflection means.

Wikipedia: "reflection is the ability of a process to examine, introspect, and modify its own structure and behavior."

Would you say inspect.getsource(func) fits the definition of reflection?

Would you say ast.parse(inspect.getsource(func)) has something to do with the AST?

Post reply on HN