5 ms·

Show HN: Autograd.c – A tiny ML framework built from scratch

85 points by sueszli 10mo ago

built a tiny pytorch clone in c after going through prof. vijay janapa reddi's mlsys book: mlsysbook.ai/tinytorch/ perfect for learning how ml frameworks work under the hood :)

sueszli 10mo ago
woah, this got way more attention than i expected. thanks a lot. if you are interested in the technical details, the design specs are here: https://github.com/sueszli/autograd.c/blob/main/docs/design.md https://github.com/sueszli/autograd.c/blob/main/docs/design.... if you are working on similar mlsys or compiler-style projects and think there could be overlap, please reach out: https://sueszli.github.io/ https://sueszli.github.io/
- deleted 10mo ago
  [deleted]
spwa4 10mo ago
Cool. But this makes me wonder. This negates most of the advantages of C. Is there a compiler-autograd "library"? Something that would compile into C specifically to execute as fast as possible on CPUs with no indirection at all.
- thechao 10mo ago
  At best you'd be restricted to the forward mode, which would still double stack pressure. If you needed reverse mode you'd need 2x stack, and the back sweep over the stack based tape would have the nearly perfectly unoptimal "grain". If you allows the higher order operators (both push out and pull back), you're going to end up with Jacobians & Hessians over nontrivial blocks. That's going to need the heap. It's still better than an unbounded loop tape, though. We had all these issues back in 2006 when my group was implementing autograd for C++ and, later, a computer algebra system called Axiom. We knew it'd be ideal for NN; I was trying to build this out for my brother who was porting AI models to GPUs. (This did not work in 2006 for both HW & math reasons.)
  - spwa4 10mo ago
    Why not recompile every iteration? Weights are only updated at the end of the batch size at the earliest, and for distributed training, n batch sizes at the fastest, and generally only at the end of an iteration. In either case the cost of recompiling would be negligeable, no?
    - thechao 10mo ago
      You'd pay the cost of the core computation O(n) times. Matrix products under the derivative fibration (jet; whatever your algebra calls it) are just more matrix products. A good sized NN is already in the heap. Also, the hard part is finding the ideal combination of fwd vs rev transforms (it's NP hard). This is similar to the complexity of finding the ideal subblock matrix multiply orchestration. So, the killer cost is at compile time, not runtime, which is fundamental to the underlying autograd operation. On the flip side, it's 2025, not 2006, so pro modern algorithms & heuristics can change this story quite a bit. All of this is spelled out in Griewank's work (the book).
      - spwa4 10mo ago
        This one? https://epubs.siam.org/doi/book/10.1137/1.9780898717761 https://epubs.siam.org/doi/book/10.1137/1.9780898717761
        thechao 10mo ago
        Yep. You can find used copies at some online places? Powell's in Portland (online store) sometimes has it for 25 or 30 $s.
- sueszli 10mo ago
  a heap-free implementation could be a really cool direction to explore. thanks! i think you might be interested in MLIR/IREE: https://github.com/openxla/iree https://github.com/openxla/iree
- attractivechaos 10mo ago
  > Is there a compiler-autograd "library"? Do you mean the method theano is using? Anyway, the performance bottleneck often lies in matrix multiplication or 2D-CNN (which can be reduced to matmul). Compiler autograd wouldn't save much time.
- marcthe12 10mo ago
  We would need to mirror jax architecture more. Since the jax is sort of jit arch wise. Basically you somehow need a good way to convert computational graph to machine code while at compile time also perform a set of operations on the graph.
- justinnk 10mo ago
  I believe Enzyme comes close to what you describe. It works on the LLVM IR level. https://enzyme.mit.edu https://enzyme.mit.edu
PartiallyTyped 10mo ago
Any reason for creating a new tensor when accumulating grads over updating the existing one? Edit: I asked this before I read the design decisions. Reasoning is, as far as I understand, that for simplificity no in-place operations hence accumulating it done on a new tensor.
- sueszli 10mo ago
  yeah, exactly. it's for explicit ownership transfer. you always own what you receive, sum it, release both inputs, done. no mutation tracking, no aliasing concerns. https://github.com/sueszli/autograd.c/blob/main/src/autograd.c#L170 https://github.com/sueszli/autograd.c/blob/main/src/autograd... i wonder whether there is a more clever way to do this without sacrificing simplicity.

85 points by sueszli 10mo ago

built a tiny pytorch clone in c after going through prof. vijay janapa reddi's mlsys book: mlsysbook.ai/tinytorch/ perfect for learning how ml frameworks work under the hood :)

sueszli 10mo ago

woah, this got way more attention than i expected. thanks a lot. if you are interested in the technical details, the design specs are here: https://github.com/sueszli/autograd.c/blob/main/docs/design.md https://github.com/sueszli/autograd.c/blob/main/docs/design.... if you are working on similar mlsys or compiler-style projects and think there could be overlap, please reach out: https://sueszli.github.io/ https://sueszli.github.io/

deleted 10mo ago
[deleted]

spwa4 10mo ago

Cool. But this makes me wonder. This negates most of the advantages of C. Is there a compiler-autograd "library"? Something that would compile into C specifically to execute as fast as possible on CPUs with no indirection at all.

thechao 10mo ago
At best you'd be restricted to the forward mode, which would still double stack pressure. If you needed reverse mode you'd need 2x stack, and the back sweep over the stack based tape would have the nearly perfectly unoptimal "grain". If you allows the higher order operators (both push out and pull back), you're going to end up with Jacobians & Hessians over nontrivial blocks. That's going to need the heap. It's still better than an unbounded loop tape, though. We had all these issues back in 2006 when my group was implementing autograd for C++ and, later, a computer algebra system called Axiom. We knew it'd be ideal for NN; I was trying to build this out for my brother who was porting AI models to GPUs. (This did not work in 2006 for both HW & math reasons.)
- spwa4 10mo ago
  Why not recompile every iteration? Weights are only updated at the end of the batch size at the earliest, and for distributed training, n batch sizes at the fastest, and generally only at the end of an iteration. In either case the cost of recompiling would be negligeable, no?
  - thechao 10mo ago
    You'd pay the cost of the core computation O(n) times. Matrix products under the derivative fibration (jet; whatever your algebra calls it) are just more matrix products. A good sized NN is already in the heap. Also, the hard part is finding the ideal combination of fwd vs rev transforms (it's NP hard). This is similar to the complexity of finding the ideal subblock matrix multiply orchestration. So, the killer cost is at compile time, not runtime, which is fundamental to the underlying autograd operation. On the flip side, it's 2025, not 2006, so pro modern algorithms & heuristics can change this story quite a bit. All of this is spelled out in Griewank's work (the book).
    - spwa4 10mo ago
      This one? https://epubs.siam.org/doi/book/10.1137/1.9780898717761 https://epubs.siam.org/doi/book/10.1137/1.9780898717761
      - thechao 10mo ago
        Yep. You can find used copies at some online places? Powell's in Portland (online store) sometimes has it for 25 or 30 $s.
sueszli 10mo ago
a heap-free implementation could be a really cool direction to explore. thanks! i think you might be interested in MLIR/IREE: https://github.com/openxla/iree https://github.com/openxla/iree
attractivechaos 10mo ago
> Is there a compiler-autograd "library"? Do you mean the method theano is using? Anyway, the performance bottleneck often lies in matrix multiplication or 2D-CNN (which can be reduced to matmul). Compiler autograd wouldn't save much time.
marcthe12 10mo ago
We would need to mirror jax architecture more. Since the jax is sort of jit arch wise. Basically you somehow need a good way to convert computational graph to machine code while at compile time also perform a set of operations on the graph.
justinnk 10mo ago
I believe Enzyme comes close to what you describe. It works on the LLVM IR level. https://enzyme.mit.edu https://enzyme.mit.edu

PartiallyTyped 10mo ago

Any reason for creating a new tensor when accumulating grads over updating the existing one? Edit: I asked this before I read the design decisions. Reasoning is, as far as I understand, that for simplificity no in-place operations hence accumulating it done on a new tensor.

sueszli 10mo ago
yeah, exactly. it's for explicit ownership transfer. you always own what you receive, sum it, release both inputs, done. no mutation tracking, no aliasing concerns. https://github.com/sueszli/autograd.c/blob/main/src/autograd.c#L170 https://github.com/sueszli/autograd.c/blob/main/src/autograd... i wonder whether there is a more clever way to do this without sacrificing simplicity.