Hacker News Reader

Comment by liuliu | Hacker News Reader

And 10k ops is not a lot! SDXL has around 4k ops (from my memory) on the forward pass, from this test, that means you spend 40ms each iteration on autodiff!

View on HN

The author's workload is sort of different than the usual ML workload since the author's expression tree is large (10k nodes), while a modern neural net has a relatively smaller expression tree, maybe fewer than 100 for the larger neural nets?

Another commenter mentioned Dr.Jit which seems to be designed for this use case. This is a quote from their project page.

> Why did we create Dr.Jit, when dynamic derivative compilation is already possible using Python-based ML frameworks like JAX, Tensorflow, and PyTorch along with backends like XLA and TorchScript?

> The reason is related to the typical workloads: machine learning involves small-ish computation graphs that are, however, made of arithmetically intense operations like convolutions, matrix multiplications, etc. The application motivating Dr.Jit (differentiable rendering) creates giant and messy computation graphs consisting of 100K to millions of "trivial" nodes (elementary arithmetic operations). In our experience, ML compilation backends use internal representations and optimization passes that are too rich for this type of input, causing them to crash or time out during compilation. If you have encountered such issues, you may find Dr.Jit useful.

My number of ops == your number of expression nodes. Trying to say that for certain models, the number of expression nodes is not trivial (in the range of a few thousands). LLM have smaller number of expression nodes / parameter count ratio though. I think 7b models typically have around ~500 nodes give or take when break everything down.

Reply on news.ycombinator.com