‏HRM-Text: پیش‌آموزش مدل زبانی با هزار دلار

HRM-Text یه مدل زبانی ۱ میلیارد پارامتری‌ه که روی معماری HRM (سلسله‌مراتب بازگشتی) بنا شده. مشکلی که حل می‌کنه خیلی مهمه: پیش‌آموزش مدل‌های پایه (foundation models) تا حالا به منابع محاسباتی سنگینی نیاز داشت که فقط شرکت‌های بزرگ بهش دسترسی داشتن. HRM-Text این مانع رو با کاهش ۱۳۰ تا ۶۰۰ برابری نیاز محاسباتی و ۱۵۰ تا ۹۰۰ برابری نیاز به داده از سر راه برمی‌داره.

دو سایز برای آموزش وجود داره: نسخه L با ۰.۶ میلیارد پارامتر که روی ۸ عدد H100 در حدود ۵۰ ساعت (~۸۰۰ دلار) آموزش می‌بینه، و نسخه XL با ۱ میلیارد پارامتر که به ۱۶ عدد H100 روی دو نود نیاز داره و حدود ۴۶ ساعت (~۱۴۷۲ دلار) طول می‌کشه. قیمت‌ها بر اساس نرخ ۲ دلار در ساعت برای هر H100 محاسبه شده.

نتایج بنچمارک‌ها هم قابل توجهن؛ مدل XL روی تست‌هایی مثل GSM8k به ۸۴.۷٪ و MATH به ۵۶.۵٪ و DROP به ۸۲.۳٪ می‌رسه. این اعداد برای مدلی با این حجم از محاسبات، خیلی چشمگیرن.

از نظر فنی، HRM-Text از چند تکنولوژی کلیدی استفاده می‌کنه: PrefixLM sequence packing برای بهره‌وری بیشتر از توکن‌ها، FlashAttention 3 برای سرعت بخشیدن به محاسبات Attention (که به GPU‌های نسل Hopper مثل H100 نیاز داره)، و PyTorch FSDP2 برای توزیع آموزش روی چند GPU. معماری HRM از دو ماژول H و L تشکیل شده که لایه‌های مدل رو بین خودشون تقسیم می‌کنن.

فرآیند آموزش پنج مرحله داره: آماده‌سازی داده با pipeline مجزای data_io، راه‌اندازی محیط (ترجیحاً با Docker)، اجرای آموزش با torchrun، ارزیابی روی بنچمارک‌ها، و در نهایت تبدیل چک‌پوینت به فرمت HuggingFace Transformers. برای اجرای multi-node، یه shared storage مشترک بین نودها لازمه تا چک‌پوینت‌ها درست ذخیره بشن.

torchrun --nproc_per_node=8 pretrain.py arch/size@arch=L lr=2.5e-4 global_batch_size=172032

دستور بالا آموزش نسخه L رو روی یه نود ۸ GPU‌ای شروع می‌کنه. پیکربندی مدل از طریق Hydra قابل تغییره و می‌شه معماری‌های مختلفی از جمله Transformer استاندارد، Universal Transformer، و چند baseline بازگشتی دیگه رو هم امتحان کرد.

نکات کلیدی:

• پیش‌آموزش یه مدل ۰.۶ میلیارد پارامتری با ~۸۰۰ دلار امکان‌پذیره
• به GPU‌های نسل Hopper (مثل H100) نیاز داره چون FlashAttention 3 وابسته به این معماریه
• کد کاملاً متن‌بازه و تحت لایسنس Apache 2.0 منتشر شده
• پشتیبانی native از vLLM هنوز در دست توسعه‌ست
• می‌شه چک‌پوینت‌ها رو به فرمت HuggingFace Transformers تبدیل کرد