Rust 框架 DenseCore 称与 PyTorch 逐位一致
一名计算机专业学生在 r/MachineLearning 发帖介绍其用 Rust 编写的轻量深度学习框架 DenseCore，它提供类似 PyTorch 的 Python API，体积约 2 MB，而 PyTorch 约 580 MB。作者称，在不改动 Karpathy 的 nanoGPT 代码的情况下，DenseCore 在测试的 50 步中给出的初始权重、损失、梯度和更新后权重与 PyTorch 完全逐位相同，开启 dropout 时也是如此。作者同时表示，他是在复现 PyTorch CPU 注意力所用的 flash attention 内核、Adam 尾部循环的舍入方式以及 clip_grad_norm_ 用倒数乘法代替除法等细节后才做到的，并且该一致性目前仅限 Apple Silicon Mac。
作者称在 26 个测试算子中只有 1 个在 Mac 与 Linux x86 之间得到相同比特，说明跨机器复现 PyTorch 结果本身就存在差异，而逐位对齐需要复刻许多未文档化的内核行为。
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)