یادگیری مجزای مفاهیم موازی‌سازی مانند همگام‌سازی گرادیانها یا شارد کردن وضعیت بهینهساز یک مسئله است، اما کنار هم قرار دادن درست آن‌ها بدون بروز واگرایی، بنبست (deadlock) یا نشت حافظه چالشی کاملاً متفاوت است. در این درس، تمام مفاهیم درس‌های گذشته را در قالب یک حلقه‌ی آموزش توزیع‌شده سرتاسری (end-to-end distributed training) یکپارچه می‌کنیم.
این پیاده‌سازی با طراحی یک مدل ترنسفورمر زبانی کوچک به نام MiniGPT، همگام‌سازی اولیه پارامترها با DDP، توزیع وضعیت بهینهساز با الگوریتم ZeRO-1 و ذخیره و بازیابی بایت‌به‌بایت چک‌پوینت‌های شاردشده را روی 4 رتبه (rank) شبیه‌سازی‌شده به نمایش می‌گذارد.
از آنجا که این پیاده‌سازی از بک‌اند gloo استفاده می‌کند، نیازی به کارت‌های گرافیک اختصاصی یا سرویس‌های ابری خارجی ندارد و به‌راحتی روی پردازنده‌ی معمولی (CPU) سیستم شخصی شما در عرض چند ثانیه اجرا و راستی‌آزمایی میشود.
در درس‌های 76 تا 80، هر بار یکی از قطعات پازل را ساختیم. این درس محل اتصال همه‌ی آن‌هاست: یک GPT کوچک که روی 4 رتبه‌ی شبیه‌سازی‌شده با DDP برای همگام‌سازی گرادیان، ZeRO-1 برای شارد کردن وضعیت بهینهساز و یک چک‌پوینت شاردشده در میانه‌ی مسیر آموزش داده میشود. این دمو 20 گام اجرا میشود، خودش به کارش پایان می‌دهد، نمودار تابع زیان (loss) و پروفایل حافظه را چاپ می‌کند و در نهایت یک چک‌پوینت با قابلیت ادامه‌ی آموزش ذخیره می‌نماید.
اهداف یادگیری
- ترکیب DDP (درس 77)، ZeRO-1 (درس 78) و چک‌پوینت‌های شاردشده (درس 80) در قالب یک حلقه‌ی آموزش واحد.
- آموزش یک مدل زبانی ترنسفورمر 2 لایه روی یک پیکره‌ی ساختگی کوچک به مدت 20 گام روی 4 رتبه‌ی شبیه‌سازی‌شده.
- چاپ جدول loss به تفکیک گام‌ها، پروفایل حافظه به تفکیک هر رتبه و یک مانیفست چک‌پوینت که با همان تعداد رتبه، بایت‌به‌بایت یکسان بازیابی شود.
- دفاع فنی از این ترکیب: هر بخش در درس‌های پیشین به شکل مستقل آزموده شده و این درس اثبات می‌کند که در کنار هم نیز به‌درستی کار می‌کنند.
صورت مسئله
پروژه‌ی پایانی اثباتی است بر اینکه قطعات سیستم به‌خوبی در کنار هم می‌نشینند. درس 76 توابع ارتباط جمعی (collectives) را پیاده کرد؛ درس 77 آن‌ها را در قالب DDP بسته‌بندی کرد؛ درس 78 وضعیت بهینهساز را با reduce_scatter شارد نمود؛ درس 79 پایپ‌لاین را تحلیل کرد؛ و درس 80 یک چک‌پوینت شاردشده ذخیره نمود. هر درس با تست اختصاصی خودش مستقل بود؛ اما در یک اجرای آموزشی واقعی، همه‌ی این اجزای سازنده هم‌زمان به کار گرفته می‌شوند. اگر ترکیب آن‌ها ایرادی داشته باشد، loss واگرا میشود، چک‌پوینت لود نمی‌شود یا حافظه‌ی هر رتبه به‌جای کاهش، افزایش می‌یابد.
این درس دموی سرتاسری را اجرا کرده و چهار ویژگی پایدار را راستی‌آزمایی می‌کند: (الف) loss در طول 20 گام (با در نظر گرفتن خطای محاسبات ممیز شناور) به‌طور یکنواخت کاهش یابد، (ب) همه‌ی رتبهها در هر گام نُرم پارامتر یکسانی داشته باشند، (ج) حافظه‌ی بهینهساز در هر رتبه با فرمول 12P/N بایت در ZeRO-1 برابر باشد، و (د) چک‌پوینت گام 10 پس از راه‌اندازی مجدد، بایت‌به‌بایت یکسان بازیابی شود. این دمو به‌طور خودکار خاتمه می‌یابد: 20 گام، یک دستور واحد و خروج با کد 0.
مفاهیم پایه
مدل GPT کوچک (mini GPT)
این مدل عمداً کوچک طراحی شده است: 2 بلوک ترنسفورمر، بُعد امبدینگ 32، 4 سرِ اتنشن، اندازه‌ی واژگان 64، طول توالی 16 و اندازه‌ی دسته (batch) برابر با 4؛ یعنی در مجموع چند هزار پارامتر. به اندازه‌ای بزرگ است که همه‌ی تصمیم‌های فنی اتصالات را درگیر کند (اتنشن چندسر از مسیر ماسک‌شده‌ی استاندارد استفاده می‌کند؛ LayerNorm وزن‌هایی دارد که باید همگام شوند؛ و سرِ LM یک تصویر کردن خطیِ مجزا به فضای واژگان است)؛ و در عین حال به اندازه‌ای کوچک است که اجرای 20 گام روی 4 رتبه‌ی CPU ظرف چند ثانیه تمام شود.
قواعد ترکیب اجزا
| بخش مربوط به درس | مسئولیت مستقیم | وظیفه‌ای که به حلقه‌ی آموزش واگذار میشود |
|---|---|---|
| برودکست DDP | همگام‌سازی اولیه‌ی پارامترها | یک‌بار فراخوانی در زمان ساخت |
| گام ZeRO-1 | همگام‌سازی گرادیان، به‌روزرسانی نسخه‌ی مرجع، برودکست پارامترها | یک‌بار فراخوانی در هر گام به‌جای optimiser.step |
| چک‌پوینت شاردشده | ذخیره‌سازی وضعیت هر رتبه، ساخت مانیفست با sha256 | فراخوانی در رتبه‌ی 0 با داده‌های جمع‌آوری‌شده از طریق allgather |
| حلقه‌ی آموزش | رفت، برگشت، ثبت لاگ loss | فراخوانی مرتب سه مورد بالا |
حلقه‌ی آموزش چیزی از reduce_scatter یا فایل‌های rendezvous نمی‌داند. ماژول‌های ZeRO و چک‌پوینت رابط‌های برنامه‌نویسی محدودی ارائه می‌دهند که حلقه آن‌ها را در کنار هم می‌نشاند.
چرا یک GPT کوچک و نه یک MLP ساده؟
مدل MLP درس 77 برای سنجش همگام‌سازی گرادیان کافی بود؛ اما GPT کوچک سه بخش تازه اضافه می‌کند: یک سرِ LM مجزا روی واژگان (در این درس برای شفافیت بیشتر، مستقل از امبدینگ توکن در نظر گرفته شده، گرچه معمولاً در GPT کامل این دو به هم گره می‌خورند)، ترکیب سافت‌مکس و آنتروپی متقاطع برای loss (که حالت‌های حدی عددی بیشتری نسبت به MSE دارد) و مسیر رفتِ نامتقارن (امبدینگ، سپس اتنشن و در نهایت MLP در هر لایه). اگر برای این پروژه‌ی پایانی به همان MLP بسنده می‌کردیم، مشخص نمی‌شد که آیا این ترکیب با LayerNorm یا ابعاد گرادیان لایهی امبدینگ درست کار می‌کند یا خیر.
منظور از خاتمه‌ی خودکار: خروج با کد 0
این حلقه دقیقاً 20 گام اجرا میشود و پایان می‌یابد؛ بدون while True، بدون دخالت انسان و بدون نیاز به بازیابی از وضعیتی بیرونی. پروژه‌ی پایانی‌ای که بتوانید بدون مراقبت رهایش کنید و پس از اتمام لاگ کاملش را تحویل بگیرید، ثابت می‌کند که اتصالات سیستم بدون نقص پیاده شده است. اگر بخشی دچار بنبست (deadlock) شود، برنامه هرگز خروجی نمی‌دهد و سامانه‌ی تست بلافاصله متوجه خطا میشود.
پیاده‌سازی از صفر
فایل code/main.py این بخش‌ها را پیاده‌سازی می‌کند:
MiniGPT: ترنسفورمر 2 لایه با خوداتنشن ماسک‌شده و سرِ LM مجزا.make_corpus(seed, total_tokens): تولید داده‌های قطعی برای پیش‌بینی توکن بعدی._train_worker: فرایندی که برای هر رتبه اجرا میشود؛ پارامترهای اولیه را برودکست می‌کند، حلقه را پیش می‌برد، گام ZeRO را صدا می‌زند و در گام 10 چک‌پوینت شاردشده را ذخیره می‌کند.verify_resume: پس از اجرای اصلی، چک‌پوینت گام 10 را درون فرایند بازخوانی کرده و بررسی می‌کند که شاردهای ذخیره‌شده‌ی نسخه‌ی مرجع با اسنپ‌شات حافظه بایت‌به‌بایت یکسان باشند.main: مدیریت کل دمو، چاپ جدول loss، پروفایل حافظه و نتیجه‌ی اعتبارسنجی.
اجرای برنامه:
python3 code/main.pyخروجی: جدولی با 20 ردیف از مقادیر loss، پروفایل حافظه‌ی 4 ردیفه به تفکیک رتبهها، مانیفست چک‌پوینت و در صورت موفقیت، چاپ خط "RESUME VERIFIED".
الگوهای محیط عملیاتی در دنیای واقعی
برای پیاده‌سازی این معماری در اجراهای واقعی محیط عملیاتی (پروداکشن)، سه الگو کار را کامل می‌کنند.
چک‌پوینت‌گیری بر اساس زمان (هر K دقیقه)، نه تعداد گام‌ها. مدت‌زمان هر گام با توجه به طول توالی و تعداد میکرودسته‌ها تغییر می‌کند. چک‌پوینت‌گیری با بازه‌ی زمانی 10 دقیقه‌ای، صرف‌نظر از اندازه‌ی مدل، حجم پردازش ثابتی را پوشش می‌دهد. در این درس برای سادگی از تعداد گام استفاده کردیم؛ اما در محیط عملیاتی زمان واقعی مبنا قرار میگیرد.
شناسایی زودهنگام واگرایی. در محیط عملیاتی، بعد از پاس برگشت محافظ مقادیر NaN و آشکارساز جهش ناگهانی loss اضافه میشود؛ اگر loss در یک گام بیش از 2 برابر جهش کند، به‌جای اینکه بگذاریم بهینهساز وارد وضعیتی تباهیده شود، به چک‌پوینت قبلی بازگشت به عقب (rollback) می‌زنیم. نمودار loss در این درس هموار است و نیازی به این محافظ نیست، اما قلاب (hook) آن در کد باقی مانده است.
تجمیع پروفایل حافظه در میان تمام رتبهها. در اجراهای واقعی، حافظه‌ی مصرفی در رتبه‌های مختلف متفاوت است (رتبه‌ای که بزرگ‌ترین مرحله‌ی پایپ‌لاین را دارد، اکتیویشن‌های بیشتری نگه می‌دارد). در محیط عملیاتی، حداکثر و میانگین حافظه در میان رتبهها ثبت میشود؛ در این درس برای نشان دادن تطابق دقیق با فرمول، مقدار حافظه‌ی هر رتبه جداگانه چاپ شده است.
استفاده از ابزارهای واقعی
الگوهای کاربردی در محیط‌های واقعی:
- DeepSpeed: همه‌ی موارد DDP + ZeRO + پایپ‌لاین + چک‌پوینت‌گذاری اکتیویشن را زیر چتر یک پیکربندی واحد ادغام می‌کند. ساختار پیاده‌شده در این درس، همان مدل DeepSpeed در مقیاس کوچک است.
- PyTorch FSDP: معادل بومی پایتورچ. حالت
FullyShardedDataParallelهمراه باShardingStrategy.SHARD_GRAD_OPدر واقع همان ZeRO-2 است. - NeMo و Megatron-LM: برای بزرگ‌ترین مدل‌ها، توازی تنسور (tensor parallel) را هم اضافه می‌کنند؛ در غیر این صورت ساختار کلی ترکیب دقیقاً به همین شکل است.
خروجی این درس
این بخش از دوره در اینجا به پایان می‌رسد. مجموع این 6 درس، همان زیرسیستم آموزش توزیع‌شده‌ای است که یک تیم فنی پیش از به‌کارگیری ابزاری مانند DeepSpeed باید طراحی کند؛ درستی این انتزاع روی gloo اثبات شده و حالت‌های شکست آن بررسی گردیده‌اند. در فاز 17 (زیرساخت و محیط عملیاتی)، این ساختار را به یک کلاستر واقعی منتقل خواهیم کرد.
تمرین‌ها
- تقسیم سرِ اتنشن با توازی تنسور را پیاده کنید و بررسی کنید که loss خروجی با خط مبنای تک‌رتبه‌ای یکسان باشد. دو رتبه: نیمی از سرها در هر رتبه و اجرای allreduce روی خروجی اتنشن.
- انباشت گرادیان (gradient accumulation) را در طول 4 میکرودسته پیاده کنید و اثبات کنید که گرادیان به‌دست‌آمده با گرادیان یک دسته‌ی بزرگ برابر است.
- مسیری برای ادامه‌ی آموزش از گام 10 اضافه کنید که آموزش را تا گام 20 ادامه دهد و نشان دهید که loss نهایی دقیقاً با خروجی اجرای پیوسته‌ی اولیه یکسان میشود.
- خروجی لاگ متریک‌ها (loss، نُرم گرادیان، زمان گام) را در قالب JSONL صادر کنید تا بعداً بتوان روند اجرا را مصورسازی کرد.
- یک محافظ NaN اضافه کنید تا در صورت جهش ناگهانی loss، به چک‌پوینت قبلی بازگردد؛ سپس با چندبرابر کردن نرخ یادگیری در یک گام، جهش مصنوعی ایجاد کرده و فرایند بازگشت به عقب را بیازمایید.
اصطلاحات کلیدی
| اصطلاح | در گفت‌وگوهای روزمره | معنای دقیق فنی |
|---|---|---|
| سرتاسری (End-to-end) | «سرهم کردن همه‌ی بخش‌ها» | اجرای یکپارچه‌ی همه‌ی اجزا در یک برنامه، نه اجرای تست واحد برای هر بخش |
| پروفایل حافظه (Memory profile) | «گیگابایت به ازای هر رتبه» | حجم بایت‌های اشغال‌شده در هر رتبه برای پارامترها، گرادیانها و وضعیت بهینهساز |
| توافق‌نامه‌ی بازیابی (Resume contract) | «ذخیره و بازیابی» | تطابق بایت‌به‌بایت وضعیت هر رتبه پس از چرخه‌ی ذخیره و لود چک‌پوینت |
| خاتمه‌ی خودکار (Self-terminating) | «اجرای محدود» | تعداد گام مشخص، خروج با کد 0 پس از اتمام و بدون نیاز به دخالت انسان |
برای مطالعه‌ی بیشتر
- راهنمای جامع آموزش سرتاسری در DeepSpeed
- آموزش پیشرفته‌ی PyTorch FSDP
- مرجع اسکریپت‌های آموزش Megatron-LM
- درس‌های 76 تا 80 از فاز 19؛ بررسی تک‌تک اجزایی که در این درس ترکیب شدند
- فاز 17؛ پیاده‌سازی این ترکیب روی یک کلاستر واقعی