Rust 框架 DenseCore 让 nanoGPT 训练结果与 PyTorch 逐位一致
[P] my rust framework trains nanoGPT with the exact same bits as pytorch. worth continuing? [P]
一名计算机专业学生发布 Rust 深度学习框架 DenseCore,提供类 PyTorch 的 Python API,把 import torch 指向它并原样运行 Karpathy 的 nanoGPT,在测试的 50 步内初始权重、loss、梯度和权重都与 PyTorch 逐位相同,开启 dropout 时也一致。
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)
submitted by /u/Mental-Feature-2377
[link] [留言]
来源:r/MachineLearning · reddit.com