تحلیل فنی: توزیع حافظه در خوشه‌های پردازشی، گلوگاه مدل‌های عظیم را می‌زداید

اگر قصد آموزش یک مدل ۷۰ میلیاردی را دارید، متوجه خواهید شد که حتی قدرتمندترین تراشه‌های بازار هم به‌تنهایی برای این کار کوچک هستند. برای آموزش چنین مدلی به بیش از ۱ ترابایت حافظه نیاز است، در حالی که استاندارد طلایی فعلی یعنی GPUهای صنعتی تنها ۸۰ گیگابایت حافظه دارند.

این شکاف عظیم، «دیوار حافظه» نام دارد که آموزش مدل‌های بزرگ را بدون معماری‌های توزیع‌شده غیرممکن می‌کند. شریجیت ونکاترامانا (Shrijith Venkatramana) در تحلیل فنی مورخ ۲۴ اوت ۲۰۲۶ اعلام کرد که راه حل این مشکل، عبور از این تصور است که هر GPU باید نسخه کاملی از مدل را در اختیار داشته باشد.

همان‌طور که در تحلیل قبلی ما درباره‌ی نقش مهندسان آموزش و نقشه‌های راه آن‌ها اشاره کردیم، گلوگاه فعلی هوش مصنوعی تعداد تراشه‌ها نیست، بلکه نحوه مدیریت حافظه میان آن‌هاست. بسیاری از توسعه‌دهندگان ابتدا به سراغ موازی‌سازی داده‌های توزیع‌شده (Distributed Data Parallel یا DDP) می‌روند. در DDP، هر GPU یک کپی کامل از پارامترها، گرادیان‌ها و وضعیت‌های بهینه‌ساز را دریافت می‌کند. این روش ساده است اما به‌شدت اتلافی است، چون یک وضعیت عظیم را روی تک‌تک دستگاه‌های خوشه تکرار می‌کند.

تاریخچه رشد مدل‌ها

در اواخر دهه ۲۰۱۰، اندازه مدل‌های شبکه عصبی (Neural Network) — شبیه نقشه مترویی که سیگنال را از ورودی به جواب می‌رساند — به‌شدت افزایش یافت. مدل BERT-Large حدود ۳۴۰ میلیون پارامتر داشت و GPT-2 به ۱.۵ میلیارد رسید. تا سال ۲۰۲۰، GPT-3 با ۱۷۵ میلیارد پارامتر معرفی شد.

در این مقیاس، افزودن GPUهای بیشتر مشکل را حل نکرد. هرچند موازی‌سازی داده‌ها حافظه کل دسته را افزایش داد، اما خودِ مدل همچنان باید در هر GPU جای می‌گرفت. این تکرار بیهوده، هدف پژوهشگرانی چون سامیام راجبانداری و همکارانش شد. آن‌ها در مقاله سال ۲۰۱۹ خود با عنوان «ZeRO: بهینه‌سازی‌های حافظه برای آموزش مدل‌های تریلیونی»، نشان دادند که می‌توان مدل‌های بالای ۱۰۰ میلیارد پارامتر را روی ۴۰۰ پردازنده آموزش داد و این مسیر را تا مدل‌های تریلیونی گسترش داد.

ریاضیات دیوار حافظه

برای درک شکست DDP در مقیاس بالا، باید هزینه حافظه به ازای هر پارامتر را بررسی کرد. در آموزش با دقت ترکیبی (Mixed Precision) و استفاده از بهینه‌ساز (Optimizer) Adam، هزینه‌ها به این صورت است:

• پارامترها: ۲ بایت
• گرادیان‌ها: ۲ بایت
• وزن‌های اصلی FP32: ۴ بایت
• گشتاور اول Adam: ۴ بایت
• گشتاور دوم Adam: ۴ بایت

مجموع این‌ها حدود ۱۶ بایت به ازای هر پارامتر است. برای یک مدل ۷۰ میلیاردی، محاسبه ساده است: ۷۰ میلیارد ضرب‌در ۱۶ بایت، یعنی ۱.۱۲ ترابایت حافظه برای هر نسخه GPU.

باید توجه داشت که این عدد تنها یک خط پایه است و هزینه‌های حیاتی دیگر را شامل نمی‌شود، مواردی مثل:

• فعال‌سازها (Activations)
• بافرهای موقت
• فضای کاری CUDA
• بافرهای ارتباطی
• تکه‌تکه شدن حافظه (Memory Fragmentation)
• حافظه مربوط به مکانیزم توجه (Attention)
• هزینه‌های نقطه بازرسی (Checkpointing)

درک دقیق این هزینه‌ها برای محیط‌های محلی نیز حیاتی است، چرا که محاسبه دقیق VRAM برای مدل‌های محلی نشان می‌دهد که نیازهای حافظه بسیار فراتر از مجموع وزن‌های مدل است.

حتی خوشه‌ای با ۱۶ پردازنده ۸۰ گیگابایتی، مجموعاً ۱.۲۸ ترابایت حافظه دارد، اما DDP نمی‌تواند از این مقدار به‌عنوان یک استخر واحد استفاده کند. DDP از هر GPU می‌پرسد: «آیا تو به‌تنهایی ۱.۱۲ ترابایت جا داری؟» و پاسخ منفی است. حافظه مجموع با حافظه قابل‌استفاده متفاوت است.

ZeRO: حذف تکرارهای بیهوده

پژوهشگران مایکروسافت در سال ۲۰۱۹ با معرفی بهینه‌ساز حذف تکرار صفر (Zero Redundancy Optimizer یا ZeRO) این مشکل را حل کردند. ایده اصلی ساده بود: اگر GPUها در حال ارتباط هستند، به‌جای تکرار وضعیت آموزش، آن را تقسیم کنند. ZeRO این کار را در سه مرحله انجام می‌دهد:

• مرحله ۱: فقط وضعیت‌های بهینه‌ساز را تکه‌تکه (Shard) می‌کند.
• مرحله ۲: هم وضعیت‌های بهینه‌ساز و هم گرادیان‌ها را تقسیم می‌کند.
• مرحله ۳: همه چیز — پارامترها، گرادیان‌ها و وضعیت‌های بهینه‌ساز — را تکه‌تکه می‌کند.

در ZeRO-3، هزینه حافظه هر GPU از ۱۶P به ۱۶P/N کاهش می‌یابد (P پارامترها و N تعداد GPUهاست). برای مدل ۷۰ میلیاردی روی ۶۴ پردازنده، نیاز حافظه از ۱.۱۲ ترابایت به ۱۷.۵ گیگابایت می‌رسد. این کار فضای لازم برای فعال‌سازها و فضای کاری CUDA را باز می‌کند.

FSDP و ترفند «مادّی‌سازی»

قابلیت موازی‌سازی داده‌های کاملاً تکه‌تکه شده (Fully Sharded Data Parallel یا FSDP) در PyTorch، پیاده‌سازی بومی این ایده‌هاست و از ZeRO-3 الهام گرفته است. جادوی FSDP در این است که مدل کامل را در طول محاسبات در حافظه نگه نمی‌دارد، بلکه از فرآیند «مادّی‌سازی موقت» استفاده می‌کند.

تصور کنید یک ترنسفورمر (Transformer) با ۸۰ لایه دارید. وقتی در حال محاسبه لایه ۱۷ هستید، نیازی نیست پارامترهای تمام ۸۰ لایه روی GPU باشند. در لحظه محاسبه لایه ۱۷، GPUها عملیات all-gather را انجام می‌دهند. هر GPU تکه کوچک خود را می‌فرستد و برای لحظه‌ای کوتاه، پارامترهای کامل فقط برای همان لایه روی GPU ایجاد می‌شوند. پس از اتمام محاسبه لایه ۱۷، این کپی آزاد شده و مدل به حالت تکه‌تکه بازمی‌گردد.

به همین دلیل کلمه «کاملاً» در FSDP اهمیت دارد؛ چون نه فقط دسته‌ها، بلکه کل وضعیت آموزش مدل را تکه‌تکه می‌کند. FSDP مدل برنامه‌نویسی موازی‌سازی داده‌های معمولی را حفظ می‌کند و توسعه‌دهنده نیازی به تقسیم دستی لایه‌ها ندارد.

موازنه: حافظه در برابر ارتباطات

در آموزش توزیع‌شده هیچ غذای رایگانی وجود ندارد. FSDP و ZeRO با کاهش مصرف حافظه، حجم داده‌های در حال حرکت در شبکه را افزایش می‌دهند. چرخه آموزش به این صورت در می‌آید:

• گذر پیشرو (Forward Pass): پارامترهای تکه‌تکه $\rightarrow$ جمع‌آوری $\rightarrow$ پارامترهای کامل $\rightarrow$ محاسبه $\rightarrow$ آزاد کردن/تکه‌تکه کردن.
• گذر پس‌رو (Backward Pass): جمع‌آوری پارامترها $\rightarrow$ محاسبه $\rightarrow$ توزیع گرادیان‌ها $\rightarrow$ تکه‌تکه کردن.

این تغییر، اتصالات با پهنای‌باند بالا (مثل خوشه‌های A100 یا H100) را حیاتی می‌کند. اگر GPUها بیشتر از محاسبه، منتظر شبکه بمانند، سخت‌افزار گران‌قیمت بیکار می‌ماند. مدل ذهنی مهندسی این است: «کارایی آموزش $\approx$ کارایی محاسبات / (سربار ارتباطات + همگام‌سازی)».

برای مقابله با این مشکل، FSDP از ترفندهایی برای هم‌زمان کردن ارتباطات و محاسبات استفاده می‌کند. مثلاً سیستم می‌تواند لایه ۱۰ را محاسبه کند و هم‌زمان پارامترهای لایه ۱۱ را دریافت کند. این کار از طریق روش‌هایی مثل پیش‌خوانی (Prefetching) و بافرهای ارتباطی نامتقارن انجام می‌شود.

اقتصاد زیرساختی

این تحول، آموزش هوش مصنوعی را از یک خرید سخت‌افزاری به یک مسئله اقتصاد زیرساختی تبدیل می‌کند. شما در واقع حافظه GPU را با پهنای‌باند شبکه معاوضه می‌کنید.

دو معماری را مقایسه کنید: معماری A از GPUهای عظیمی استفاده می‌کند که کل مدل را جای می‌دهند؛ برنامه‌نویسی ساده است و ارتباطات کم، اما هزینه بسیار بالاست. معماری B از تعداد زیادی GPU با حافظه کمتر استفاده می‌کند؛ نیاز حافظه هر تراشه کم می‌شود اما پیچیدگی سیستم توزیع‌شده و نیاز به شبکه افزایش می‌یابد.

هدف دیگر فقط «خرید سریع‌ترین GPU» نیست، بلکه یافتن ارزان‌ترین سیستمی است که واحدهای پردازشی را فعال نگه دارد. اگر ارتباطات باعث شود بهره‌وری GPU از ۵۰٪ به ۲۰٪ برسد، خرید GPUهای بیشتر عملاً سیستم را از نظر اقتصادی بدتر می‌کند.

همچنین FSDP یک راهکار قطعی برای همه مشکلات نیست. برای مدل‌های پیشرو، این روش با سایر انواع موازی‌سازی ترکیب می‌شود:

۱. موازی‌سازی تنسور (Tensor Parallelism): تقسیم عملیات ماتریسی برای لایه‌هایی که حتی برای یک تراشه هم زیاد هستند.
۲. موازی‌سازی خط لوله (Pipeline Parallelism): قرار دادن لایه‌های مختلف روی GPUهای مختلف برای توزیع عمق مدل.
۳. نقطه بازرسی فعال‌سازها (Activation Checkpointing): ذخیره حافظه با محاسبه مجدد فعال‌سازها در گذر پس‌رو به‌جای ذخیره همه آن‌ها.

برای کسانی که قصد دارند از سطح کتابخانه‌های آماده فراتر روند، پیاده‌سازی هسته‌های CUDA و ROCm امکان بهینه‌سازی‌های عمیق‌تری را در سطح سخت‌افزار برای مدیریت این جریان‌های داده فراهم می‌کند.

از ترفند به استاندارد

آنچه در سال ۲۰۱۹ به عنوان یک مقاله پژوهشی شروع شد، اکنون زیرساخت پیش‌فرض LLMهاست. مسیر تکامل از تکه‌تکه کردن ساده به نسخه‌های تهاجمی‌تر رسید:

• ZeRO-Offload: انتقال وضعیت بهینه‌ساز به حافظه CPU که ارزان‌تر و حجیم‌تر است. این کار آموزش مدل ۱۰ میلیاردی را روی یک V100 ممکن کرد.
• ZeRO-Infinity: گسترش ایده به سلسله‌مراتب حافظه ناهمگن شامل GPU، CPU و فضای ذخیره‌سازی.
• FSDP: ادغام این انتزاع‌ها در هسته PyTorch که مقیاس‌پذیری خطی در TFLOPS را گزارش کرده و از مدل‌های تا ۱ تریلیون پارامتر پشتیبانی می‌کند.

تغییر بنیادین این است که حافظه را به‌جای ویژگی یک تراشه، به‌عنوان یک منبع توزیع‌شده در کل خوشه می‌بینیم.

گام بعدی شما

• اگر از کتابخانه DeepSpeed استفاده می‌کنید، مراحل مختلف ZeRO را تست کنید تا نقطه بهینه بین مصرف VRAM و سرعت آموزش را بیابید.
• مستندات PyTorch FSDP را برای پیاده‌سازی مدل‌های بزرگ‌تر از ظرفیت تک‌تراشه مطالعه کنید.
• در طراحی خوشه، روی پهنای‌باند شبکه (Interconnect) به اندازه حافظه GPU سرمایه‌گذاری کنید.

اما داستان سخت‌افزاری این تحول حتی شگفت‌انگیزتر است — به تحلیل ما درباره‌ی تراشه‌های Blackwell مراجعه کنید.