Fastest autograd in the West | Hacker News Reader

Fastest autograd in the West

arogozhnikov.github.io

90 pointstplrbv46 comments

Am I the only one surprised that the author of einops is looking for work? In an era of an AI arms race between many big labs? If you’re rolling your own networks, I’d definitely reach out to this guy!

Fun exercise… but for such use cases, people use Casadi and DrJit

https://github.com/casadi/casadi https://github.com/mitsuba-renderer/drjit

DrJit is made by same author of pybind11 and nanobind.

And for those of us just tuning in, what is an "autograd"? From skimming the article and some quick web searching, it looks like some sort of math function that gets used a lot in ML (and I assume other places) so there are implementations in a bunch of major libraries?

reverse-mode automatic differentiation to compute the gradient of a scalar with respect to some large vector of inputs that it's computed from, which is critically important for the most popular and most efficient mathematical optimization algorithms such as adam

you can't really implement it in a library in a normal programming language, because it has to introspect on how your program is computing everything, which is not something that a library can normally do; you have to implement it in a programming language, at least an embedded domain-specific language but in some cases something like a fortran compiler

of course you can implement it in a library. You just have to make sure the users are using your array objects with appropriately overloaded operators so you can keep track of the computational graph.

That's how pretty much how all python DL frameworks are built. And thats always why you can't call .numpy() on a torch tensor which requires a gradient. Doing so would break the computational graph.

that's why the comment you're replying to says 'at least an embedded domain-specific language', kepler boy

I still consider "you can't really implement it in a library in a normal programming language" as misleading.

I'll second this, as I consider Python a "normal programming language" and this is available as a library there.

i could be wrong, but as far as i know there isn't an automatic differentiation library in python which you can use to compute the gradient of an arbitrary python function, just the gradient of an arbitrary function expressed in an embedded domain-specific language (usually one implemented by the library itself). kepler boy said specifically that this is true of pytorch but as far as i know it's true of all of them

Maybe, I guess our difference is in what we mean by "DSL". I think we're disagreeing on whether "this library provides functions you can arrange your graph out of" is "DSL" or "normal language functionality".

there's a fuzzy boundary between 'embedded dsl' and 'just a library' (similar to the problem of 'what is a framework'), but at the point that a library provides functions you can use to build your computational graph out of, it's clearly not just an automatic differentiation library; it's some kind of computational-graph-composing library which perhaps sports automatic differentiation as one of its features. and what i'm saying is that, if you want to write an automatic differentiation library in most languages, it's basically going to have to be a compiler that statically transforms your algorithm, whether it's a compiler for fortran or for an embedded dsl

this is less true for forward-mode automatic differentiation, which you can do a pretty reasonable job of in python as just a dual-number class which you can pass to many, though not all, ordinary numerical computation functions. unfortunately forward-mode automatic differentiation is unusably slow for gradients with respect to large parameter vectors

It’s a DSL by definition because a function provided by an external library cannot be a language primitive in the host language, or else it wouldn’t need to be an external library! Unless we’re talking about a LISP but we’re not.

All libraries are DSLs, then? The initial claim was:

> you can't really implement it in a library in a normal programming [...] you have to implement it in a programming language, at least an embedded domain-specific language

So which is it?

No lol not all libraries are DSLs, we’re talking specifically about libraries which provide functions allowing you to construct a computational graph. These libraries are DSLs, not host language primitives.

What's the difference between a library function and a DSL library? What's a library function that's also a host language primitive?

A DSL lib is a subset of “all possible function libs” and there is no such thing as a function library which is a host language primitive by definition. You can’t be a primitive of a language if you’re implemented in a library for that language.

regexps, html, and css2 are also dsls despite not providing functions allowing you to construct a computational graph, but i agree that at the point where you have a graph of computations you have accidentally a programming language

That's what "embedded domain-specific language" means and what makes it different from other (unembedded) DSLs:

> A DomainSpecificLanguage that is defined as a library for a generic "host" programming language. The embedded DSL inherits the generic language constructs of its host language https://wiki.c2.com/?EmbeddedDomainSpecificLanguage

He’s saying that PyTorch in essence is a DSL embedded in python which introspects itself to perform automatic differentiation. Libraries can’t normally peek at what ops are performed on variables in a given language, only the compiler/runtime for that language can. In the end it’s semantics - is an overloaded operator a DSL or native language construct? I’d argue it’s a DSL because its definition is in the host language.

You can implement autograd as a library. Just take a look at this

https://github.com/sradc/SmallPebble

The first line of the description is:

  > SmallPebble is a minimal automatic differentiation and deep learning library written from scratch in Python, using NumPy/CuPy.

The author of this library wrote a blog post, which is in my opinion the best introduction to reverse mode automatic differentiation that exists:

https://sidsite.com/posts/autodiff/

this is a nice blog post, but, two things:

- it's very much taking the edsl approach

- it's only implementing forward-mode

It is reverse mode, not forward mode.

As for the edsl, you can claim it is, but why? There is no parser, no syntax, nothing that would make any user think they are learning a new language. Would you call numpy an edsl ?

the text says it's reverse mode, but it's wrong. what the code in the post implements is forward mode

embedded dsls don't have parsers; that's what makes them embedded. numpy is solidly in the center of the embedded dsl concept

> the text says it's reverse mode, but it's wrong. what the code in the post implements is forward mode

I see. It's actually worse than both forward and reverse mode.

Let's say you have the simple program x1+x2+x3+x4+x5. Internally the program calculates x6=x1+x2, x7 = x6+x3, x8=x7+x4 and result = x9 = x8+x5. In forward mode you get all the partial derivatives for all intermediate variables w.r.t the inputs x1 through x5, so you get things like dx8/dx2, but you never calculate (or store in memory) dx8/dx6. In reverse mode you get the derivatives of the result w.r.t. all intermediate variables, so things like dx9/dx8, dx9/dx7, etc, but again you never need to calculate dx8/dx6. In this library you calculate recursively all dxj/dxi and store them, as long as xj has a direct or indirect dependency on xi. This is a strict superset of the derivatives calculated in both forward and reverse mode.

Good catch.

hmm, really? i didn't realize that

> you can't really implement it in a library in a normal programming language, because it has to introspect on how your program is computing everything

AFAIK autograd generically refers to a computational graph's gradient being automatically derived from those of it's constituent computational nodes/operators. It's not limited to cases like PyTorch that presents itself as "differentiable programming" where the graph is implicitly built by code execution.

The alternative to a differentiable language is something like the original Lua-based Torch, or TensorFlow's original non-eager mode, where the computational graph is explicitly built by library calls out of predefined nodes/operators, each of which knows it's own gradient.

agreed

Jax is unrolling for-loops during compilation. I don't know about the hypothetical performance of canonical Jax code for this example, but the comparison as it is seems unfair.

https://jax.readthedocs.io/en/latest/faq.html#jit-decorated-...

> 1. we test many computation graphs (graph is changing constantly) 2. many-many scalar operations with roughly 10k—100k nodes in each graph 3. every graph should be compiled and ran around 10k times both forward and backward

Now I'm curious to know where that is useful. Some kind of meta-learning approach?

Maybe instruction-level autograd. When you use a differentiable programming language.

Based on the "we TEST" wording, maybe evaluating graphs for fitness - some sort of genetic/evolutionary approach perhaps?

Automatic differentiation isn't just for machine learning, it's useful in several other applications where you have an optimization problem built of many variables interacting in differentiable ways. One possible guess could be structure-from- motion or SLAM, where the graph is often dynamic and several thousands of nodes isn't uncommon (the nodes would be things such as camera poses and 3D landmarks for example). However, in that case there are other frameworks that are built with that scenario in mind (Ceres, GTSAM), and would probably be better baselines.

Optimizing something with AD is machine learning.

i want to point out that in many cases you can compute a differentiable relaxation of a nondifferentiable problem to get something you can optimize more easily (with, for example, automatic differentiation) and then use the relaxation to solve the original nondifferentiable problem, for example with branch-and-bound

Robotic planning of liquid handling. There are many potential alternative ways to achieve the same result, each resulting in its own computation graph.

I imagine any kind of trajectory planning for many agents faces similar challenges (e.g. robots in a factory).

But these computation graphs are similar every time. You would/should write these as a small graph with conditionals on large tensors, not as a fast changing graph of scalars. The latter will always be slower when you have an accelerator like a GPU.

> But these computation graphs are similar every time.

For now I'm testing one-after-another because failure of previous (overflow, too-small-transfer, etc) gives a hint about which part can be changed in the graph.

I've been thinking about vectorizing over multiple guesses at once, but don't see a fast reliable way to merge several graphs (also that induces significant burden on other parts of the project - autograd is just one component).

Related: my post about this: https://bernsteinbear.com/blog/compiling-ml-models/

I recommend trying out TCC, which compiles C very fast.

Instead of doing a switch having separate loops for each operation would be much better / faster (branch free).

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!

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.

I think tinygrad is faster.

https://github.com/tinygrad/tinygrad

Maybe he thinks San Diego (George Hotz) is not west of Moscow (or maybe Alex moved now to London).

According to LinkedIn, the author lives in SF Bay Area since 2019: https://www.linkedin.com/in/alexrogozhnikov/

I recently had a similar need for my “real-time” Python program to do some autograd operations – but not on big tensors – and also ended up leveraging Rust as part of my solution. But I didn’t know about rustimport (or cppimport). That’s really slick, and I’m going to file that nugget away for later.

Reply on news.ycombinator.com