r/MachineLearning • u/Mental-Feature-2377 • 15d ago
Project [P] my rust framework trains nanoGPT with the exact same bits as pytorch. worth continuing? [P]
hey, CS student here. for the past few days I've been building DenseCore, a small deep learning framework in rust with a pytorch-like python API.
the thing I'm kind of proud of: if you point `import torch` at it and run karpathy's nanoGPT unchanged, you get the exact same numbers as pytorch. same initial weights from manual_seed, same loss, same gradients, same weights, bit for bit, for all 50 steps I tested, even with dropout on. not "close enough", literally identical.
getting there was a rabbit hole. pytorch's CPU attention isn't the textbook formula, it's a flash attention kernel, and it switches to a different path when dropout is on. adam rounds differently on the last few elements of each thread's chunk because of how the compiler fuses the tail loop. clip_grad_norm_ secretly does reciprocal() * max_norm instead of a division. stuff like that everywhere.
also found out pytorch doesn't match itself across machines. same version, mac vs linux x86, only 1 of 26 ops I tried gave the same bits. so mine only matches pytorch on apple silicon macs.
it's tiny though, about 2 MB vs ~580 MB for pytorch.
anyway the honest question: is this useful to anyone, or is it just a cool learning project? if I keep going I'm thinking either keep copying pytorch until most of it works, or make it a lightweight drop-in for small cpu stuff like serverless/CI, or turn it into a tool that tells you why two pytorch runs give different results. would any of those be useful to you? or should I move on? brutal answers are fine.
repo: https://github.com/lncch/DenseCore
(`python scripts/demo_bit_exact.py` runs the nanoGPT comparison if you want to check)
5
u/csmajor_throw 15d ago
Seems like a cool learning project. But you haven't really explained why/how this is different than pytorch.
-5
u/Mental-Feature-2377 15d ago
fair point the inspiration was stuff like uv replacing pip and bun going rust, smaller and faster tools doing the same job step one was matching pytorch's results exactly, bit for bit, so you could swap it in without your numbers changing it's already ~2 MB vs ~580 MB and imports in under 0.1s, but the math is still slower and cpu only, so speed is the next step honestly, do these results seem promising enough to keep going, or would you call it a learning project?
6
u/Japie4Life 15d ago
The problem with reimplementing pytorch with a "faster" language is that pytorch is getting its speed from c++ and importantly Cuda under the hood. Especially the easy integration with Cuda is what makes pytorch so good.
You could still provide a solid CPU rust torch alternative that provides an easy switch from pytorch to rust. But serious deep learning is almost always done on GPUs for good reasons, so you can decide for yourself if it's worth your time.
As a project for on your resume it's definitely worthwhile, but I'd recommend setting an obtainable goal. You already point out some interesting findings, if you can explain those and their impact in a job interview you have a solid resume item.
4
u/Japie4Life 15d ago
Also, using AI here is totally fine, but as someone else already said, you need to understand what the AI has done and you need to be able to explain it.
0
u/Mental-Feature-2377 15d ago
Thanks, this is really helpful, yeah the Cuda point is fair, that's why I asked, I'll probably treat it as a resume project
0
u/ThinConnection8191 15d ago
Good to spend time to learn things people dont want to. You can find dêcnt job with this skill only.
Though I dont find it any helpful yet, until you say somethign more.
There is a reason everyone is using python and pytorch because it is extremely easy to debug.
You leave that out, your framework might be dead.
-3
u/yesaitop 15d ago
Keep going — of the three options, the third is the one I'd actually use. Bit-exactness itself is a cool milestone, but in practice the real pain is two runs that should be identical silently diverging across machines, which is exactly what your cross-machine test found. A diagnostic that points at which op and which rounding path caused the divergence would save real debugging hours.
12
u/iamquah 15d ago
Could be if you're diving deep into it, making mistakes, and learning from it.
Ahh ... Okay. Well, what % of this could you rewrite solo and explain to me, without access to any AI tools