LucilleAI
یک فریمورک یادگیری ماشین به زبان Rust که تنسورها را روی GPU نگه میدارد و هر عملیات را برای آموزش ضبط میکند
یک فریمورک یادگیری ماشین که از صفر به Rust نوشته شده، با بکاند CUDA، و یک مدل واقعی را آموزش میدهد.
چیست
میخواستم بفهمم بین یک عملیات روی تنسور و آپدیت گرادیان واقعاً چه اتفاقی میافتد، برای همین کل پشته را ساختم: کتابخانه تنسور، autodiff، مجموعه لایههای شبکه عصبی، بهینهساز، توکنایزر و حلقه آموزش. بعد روی آن یک مدل نوشتم و دیدم loss پایین میآید.
این بیان صادقانه ماجراست. ارزشش «رقابت با PyTorch» نیست — رقابت نمیکند و ادعایش را هم ندارم. ارزشش این است که هر لایه را میتوانم باز کنم و بخوانم، و به مشکلات مشخصی برخوردم که فریمورکها سرِشان را میپوشانند.
چه کار میکند
- تنسورها روی GPU زندگی میکنند و autodiff از هر عملیات یک ضبط نگه میدارد تا گرادیانها بتوانند به عقب پخش شوند.
- cuBLAS روی استریم CUDA خودش اجرا میشود، همزمان با کد من روی استریم دیگر. باارزشترین بخش این پروژه: ضربهای ماتریسی با کرنلهای اطرافشان همپوشانی پیدا میکنند، پس GPU هیچوقت منتظر کاری که هنوز به آن ندادهام نمیماند.
- فریمورک کرش نکرد — عدد غلط داد. یک ردیابی نشان داد «NaN شبحی» فقط روی سختافزار نسل تورینگ ظاهر میشود،
جایی که
tf32ورودیهای tensor-core را بیصدا گرد میکند. همان یک باگ از همه چیزهای دیگر این پروژه با هم بیشتر وقت گرفت، و راهحلش چیزی است که فقط کسی پیدا میکند که کد تولیدشده را بخواند. - توکنایزر یک باگ بازتولیدپذیری داشت. توکنها را در یک انباشتگر اعشار شناور میشمرد که نمیتواند ترتیب
جمعها را عوض کند، برای همین یک متن یکسان در اجراهای مختلف واکبوردی با اندازه متفاوت میداد.
u64حلش کرد. - یک مدل واقعی روی متن واقعی آموزش دید: ۳٫۷ میلیون پارامتر، کانتکست ۱۲۸ توکن، loss از حالت اولیه پایین آمد، و توکنایزری که متن فارسی را بدون خراب کردن رفتوبرگشت میکند.
اعداد
| ۸۱٬۲۰۰ | خط کد — ۷۶٬۷۰۰ Rust و ۴٬۵۰۰ CUDA |
| ۲۴ | کرنل CUDA دستنویس |
| ۳۲ | عملیات قابل مشتقگیری |
| ۳٫۷ میلیون | پارامتر مدل |
| ۳۹۶ | گام پیشآموزش، loss برابر ۰٫۰۴۸۶ |
| ۹۹٫۲٪ | دقت روی مجموعه نگهداشتهشده (۶۵۶ از ۶۶۱) |
چرا ارزش دیدن دارد
دو دلیل.
اولی پل استریمها. همپوشانی کار کتابخانه با کار خودت یک تکنیک واقعی است و گرفتنش سخت است — باید مالکیت حافظه را مدیریت کنی، چون تخصیصی که استریم دیگر هنوز استفاده میکند قابل استفاده مجدد نیست. دومی NaN است. با خواندن کد پیدایش نکردم؛ با ریختن trace هر کرنل و نگاه کردن به جایی که یک عدد متناهی به «نعدد» تبدیل میشد پیدایش کردم. فقط روی تورینگ رخ میداد، فقط در مسیرهای مشخص tensor-core، و هر بررسی مبتنی بر خطا در دنیا از آن عبور میکرد.
وضعیت فعلی
- AMP در هر دو اجرا خاموش بود. آموزش با دقت ترکیبی اصلاً آزمایش نشده.
- ارزیابی خودِ پروژه اعلام میکند overfit شده، و باید همینطور باشد. عدد ۹۹٪ از یک مجموعه آزمون که خودم نوشتهام یعنی آزمون زیادی ساده بوده، نه اینکه مدل عالی است.
- مدل کوچک است — ۳٫۷ میلیون پارامتر یک نمایش است، نه یک مدل زبانی.
هیچکدام از اینها جذابیتش را کم نمیکند. آن را همان چیزی میکند که هست: یک پشته ML کامل، خوانا و از صفر، که هر عدد این صفحه از کد و اجراهای خودم آمده.


