跳到正文
r/MachineLearning· /u/Mental-Feature-2377·· 4 小时前AI 评分51

Rust 框架 DenseCore 让 nanoGPT 训练结果与 PyTorch 逐位一致

[P] my rust framework trains nanoGPT with the exact same bits as pytorch. worth continuing? [P]

AI 导读

一名计算机专业学生发布 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