موازی‌سازی پایپ‌لاین و تحلیل حباب (Pipeline Parallelism)

وقتی مدل‌های زبانی به ده‌ها میلیارد پارامتر می‌رسند، حجم وزن‌ها و حالت‌های بهینه‌ساز از گنجایش حافظه‌ی یک کارت گرافیک فراتر می‌رود. در چنین سناریوهایی، موازی‌سازی داده به‌تنهایی کارساز نیست و باید سراغ خرد کردن مدل در طول لایه‌ها بروید؛ رویکردی که موازی‌سازی پایپ‌لاین (Pipeline Parallelism) نامیده می‌شود.

چالش اصلی در موازی‌سازی پایپ‌لاین، زمان‌های بیکاری یا همان «حباب» (Bubble) است که به دلیل وابستگی‌های ترتیبی میان لایه‌ها به وجود می‌آید. تسلط بر نحوه زمان‌بندی ریزدسته‌ها (Microbatches) و تفاوت‌های معماری روش‌هایی مثل GPipe و 1F1B تعیین می‌کند که پردازنده‌ها چه میزان از چرخه‌ها را واقعاً صرف پردازش مفید می‌کنند.

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

موازی‌سازی تنسور (Tensor Parallelism) ضرب ماتریسی را بین رتبه‌ها تقسیم می‌کند. موازی‌سازی پایپ‌لاین (Pipeline Parallelism) کل مدل را بین رتبه‌ها می‌شکند؛ یعنی هر رتبه یک مرحله (stage) را بر عهده می‌گیرد. ریزدسته‌ها (Microbatches) درون پایپ‌لاین جریان می‌یابند. زمان‌های خالی در آغاز و پایان، همان حباب (bubble) هستند؛ تمام هنر این روش، به حداقل رساندن همین زمان‌های تلف‌شده است.

اهداف یادگیری

• تقسیم یک مدل ترتیبی به N مرحله و شبیه‌سازی پایپ‌لاین فاز رفت (forward) روی N رتبه.
• زمان‌بندی M ریزدسته در طول پایپ‌لاین با استفاده از زمان‌بندی GPipe (پر کردن پایپ‌لاین فقط با فاز رفت، سپس آغاز فاز برگشت) و محاسبه‌ی سهم حباب (Bubble fraction).
• مقایسه‌ی سهم حباب این روش با زمان‌بندی متناوب 1F1B که در Megatron-LM و PipeDream به کار می‌رود.
• دفاع از شیوه‌ی تخصیص مراحل: برابری حجم محاسبات در هر مرحله بسیار مهم‌تر از برابری تعداد پارامترها در هر مرحله است.

صورت مسئله

یک مدل 70 میلیارد پارامتری با دقت fp16 فقط برای نگهداری وزن‌هایش به 140 گیگابایت حافظه نیاز دارد؛ هیچ پردازنده‌ی گرافیکی تجاری چنین فضایی ندارد. روش ZeRO-3 پارامترها را بین رتبه‌ها خرد می‌کند، اما همچنان برای هر گام فاز رفت نیاز دارد که همه‌ی رتبه‌ها کل لایه را با عملیات allgather جمع‌آوری کنند، که به ازای هر لایه هزینه‌ی $\log(N)$ گام لغزش (Hop) را تحمیل می‌کند. موازی‌سازی پایپ‌لاین مسیر متفاوتی را در پیش می‌گیرد: مدل را به N مرحله می‌شکند و هر مرحله را روی یک رتبه قرار می‌دهد. فاز رفتِ لایه‌ی 1 روی رتبه‌ی 0 تمام می‌شود و تنسور اکتیویشن را به رتبه‌ی 1 تحویل می‌دهد؛ رتبه‌ی 1 لایه‌ی 2 را اجرا می‌کند و خروجی را به رتبه‌ی 2 می‌سپارد، و همین روند ادامه پیدا می‌کند. فاز برگشت (backward) نیز در جهت معکوس جریان می‌یابد. حافظه‌ی مصرفی به‌صورت خطی کاهش می‌یابد چون هر رتبه فقط یک مرحله را نگه می‌دارد؛ اما محاسبات حالت ترتیبی پیدا می‌کنند و مسئله‌ی حباب دقیقاً از همین‌جا شکل می‌گیرد.

حباب همان زمان بیکاری پردازنده‌ها در ابتدای پایپ‌لاین (انتظار برای رسیدن اولین ریزدسته به آخرین مرحله) و در انتهای آن (انتظار برای تخلیه شدن آخرین ریزدسته در مسیر برگشت) است. با داشتن M ریزدسته و N مرحله، سهم حباب به ازای هر مرحله برابر با (N-1)/(M+N-1) خواهد بود. اگر M=8 و N=4 باشد، این مقدار 27% است. اما در M=64 و N=4 به 4.5% کاهش می‌یابد. حباب زمانی کوچک می‌شود که در هر گام، تعداد ریزدسته‌ها زیاد باشد؛ این یعنی اندازه‌ی دسته‌ی هر ریزدسته باید کوچک بماند، و این دقیقاً همان قیدی است که طراحی ریزدسته‌ها را هدایت می‌کند.

مفهوم

زمان‌بندی GPipe

پایپ‌لاین در فاز رفت ابتدا با تمام M ریزدسته پر می‌شود، پیش از آن‌که هیچ فاز برگشتی شروع شود؛ سپس فاز برگشت در جهت معکوس تخلیه می‌شود. اکتیویشن‌های هر ریزدسته باید تا زمان فاز برگشتِ همان دسته در حافظه بمانند، بنابراین مصرف حافظه به‌صورت خطی با M افزایش می‌یابد. فاز رفت M+N-1 چرخه طول می‌کشد و فاز برگشت نیز M+N-1 چرخه‌ی دیگر زمان می‌برد. کار مفید هر مرحله 2M چرخه و حباب هر مرحله 2(N-1) چرخه است. اگر هر فاز رفت و برگشت یک واحد زمانی طول بکشد، سهم حباب برابر با (N-1)/(M+N-1) می‌شود. انتخاب M بسیار بزرگ‌تر از N، بخش اعظم این حباب را پنهان می‌کند.

زمان‌بندی 1F1B

رویکرد متناوب (Interleave): به محض اینکه فاز رفت یک ریزدسته به آخرین مرحله رسید، فاز برگشت آن بلافاصله آغاز می‌شود و به سمت عقب جریان می‌یابد. این زمان‌بندی در هر مرحله، یک فاز رفت و یک فاز برگشت را یکی‌درمیان اجرا می‌کند. سهم حباب همچنان N-1 است، اما حافظه‌ی مورد نیاز برای اکتیویشن‌ها به عمق پایپ‌لاین محدود می‌شود، نه تعداد کل ریزدسته‌ها. پایپ‌لاین‌های محیط عملیاتی (پروداکشن) مانند Megatron و PipeDream از الگوی 1F1B استفاده می‌کنند. در این درس ابتدا GPipe را پیاده‌سازی می‌کنیم چون ساده‌تر است، و پیاده‌سازی 1F1B را به‌عنوان تمرین پیش می‌برید.

چرا برابری محاسبات در هر مرحله مهم است؟

اگر اجرای مرحله‌ی 0 حدود 50 میلی‌ثانیه و مرحله‌ی 1 حدود 100 میلی‌ثانیه زمان ببرد، سرعت کل چرخه به مرحله‌ی 1 گره می‌خورد. سایر مراحل در هر چرخه 50 میلی‌ثانیه بیکار می‌مانند تا مرحله‌ی 1 آزاد شود. تقسیم بر اساس تعداد برابر پارامترها معیار اشتباهی است: بخش عمده‌ی محاسبات یک ترنسفورمر (Transformer) به اتنشن (Attention) به‌علاوه‌ی لایه‌های MLP در هر لایه وابسته است، در حالی که لایه‌های امبدینگ (Embedding) پارامترهای بسیار زیاد اما بار محاسباتی ناچیزی دارند. بنابراین، تخصیص مراحل باید بر اساس یکسان‌سازی تعداد عملیات ممیز شناور (FLOPs) در هر مرحله انجام شود، نه وزن‌ها در هر مرحله.

ریزدسته در برابر دسته

یک پایپ‌لاین، M ریزدسته با اندازه‌ی B را اجرا می‌کند. اندازه‌ی دسته‌ی مؤثر (Effective batch size) برابر با MB خواهد بود. گرادیان نهایی در پایان یک گام پایپ‌لاین، معادل گرادیان حاصل از مجموع MB نمونه است. سهم حباب به M وابسته است، در حالی که بهینه‌ساز (Optimizer) کل داده‌های M*B را می‌بیند. تنظیم M یعنی سبک‌سنگین کردن میان سهم حباب (کمتر شدن با افزایش M) و حافظه‌ی مصرفی هر ریزدسته (افزایش حافظه‌ی اکتیویشن‌ها در الگوی GPipe با بالا رفتن M).

پیاده‌سازی

فایل code/main.py موارد زیر را پیاده‌سازی می‌کند:

• کلاس PipelineStage: یک nn.Module کوچک که پارامترهای یک مرحله را نگه می‌دارد و متد forward(activation) را فراهم می‌کند.
• کلاس Pipeline(stages, num_microbatches): زمان‌بندی GPipe را روی مراحل شبیه‌سازی‌شده و بر مبنای زمان واقعی فرضی اجرا می‌کند.
• تابع bubble_fraction(num_stages, num_microbatches): محاسبه‌ی فرم بسته (Closed-form) با رابطه‌ی (N-1)/(M+N-1).
• یک نمونه‌ی نمایشی 4-مرحله‌ای که ردگیری محاسبات (Trace) هر ریزدسته و سهم حباب اندازه‌گیری‌شده را چاپ می‌کند.

اجرای برنامه:

python3 code/main.py

خروجی: یک نمودار گانت (Gantt chart) بر اساس مراحل و ریزدسته‌ها، به همراه درصد حباب اندازه‌گیری‌شده در مقایسه با تخمین فرم بسته.

الگوهای کاربردی در محیط عملیاتی

سه الگو وجود دارد که موازی‌سازی پایپ‌لاین را برای تحویل به محیط عملیاتی به بلوغ کافی می‌رساند:

ترکیب چک‌پوینت‌گذاری اکتیویشن با پایپ‌لاین. در الگوی GPipe با وجود M ریزدسته‌ی در حال پردازش، حافظه‌ی اکتیویشن M برابر یک ریزدسته می‌شود. چک‌پوینت‌گذاری اکتیویشن (Activation Checkpointing) فاز رفت را هنگام فاز برگشت دوباره محاسبه می‌کند؛ یعنی با فدا کردن محاسبات، حافظه آزاد می‌شود. همین ترفند است که استفاده از پایپ‌لاین را برای دنباله‌های طولانی ممکن می‌سازد.

توازن مراحل اندازه‌گیری می‌شود، نه حدس زده. تیم‌های فنی یک مرحله‌ی پروفایلینگ اجرا می‌کنند تا زمان واقعی و محاسبات هر لایه (FLOPs و زمان سپری‌شده) را روی سخت‌افزار نهایی بسنجند، سپس بر اساس همان داده‌های تجربی مدل را تقسیم می‌کنند. فلگ --num-layers-per-stage در Megatron-LM یک لیست می‌گیرد تا در صورت متفاوت بودن هزینه‌ی لایه‌ها، تعداد لایه‌های نابرابر به هر مرحله اختصاص یابد.

زمان‌بندی ارسال و دریافت باید مانع بن‌بست شود. پایپ‌لاینی که در آن همه‌ی مراحل قبل از دریافت، اقدام به ارسال داده کنند، روی بستر ارتباطی دچار بن‌بست (Deadlock) می‌شود. راهکار استاندارد اجرای متناوب است: مراحل با رتبه‌ی زوج اول ارسال و بعد دریافت می‌کنند؛ مراحل با رتبه‌ی فرد اول دریافت و بعد ارسال می‌کنند. در این درس، زمان‌بندی رتبه‌ها به‌شکل صریح پیاده می‌شود تا این الگو به‌وضوح دیده شود.

استفاده از ابزارهای واقعی

الگوهای محیط عملیاتی:

• Megatron-LM. مرجع اصلی موازی‌سازی پایپ‌لاین در مقیاس‌های بزرگ. از الگوی 1F1B استفاده می‌کند و از ترکیب هم‌زمان موازی‌سازی تنسور، پایپ‌لاین و داده پشتیبانی می‌کند.
• DeepSpeed Pipeline. با روش‌های ZeRO یکپارچه است؛ ترکیب ZeRO-1 به همراه پایپ‌لاین، انتخاب رایج بزرگ‌ترین مدل‌های متن‌باز به شمار می‌آید.
• PyTorch Pipe. رپر بومی پایپ‌لاین در PyTorch که بر پایه‌ی torch.distributed.pipeline.sync.Pipe بنا شده است.

تحویل به محیط عملیاتی

در درس 80، بخش‌های خردشده‌ی پارامترهای هر مرحله را درون یک چک‌پوینت چندبخشی ذخیره می‌کنیم. در درس 81، ترکیب DDP و ZeRO و پایپ‌لاین را در یک نمونه‌ی نمایشی سرتاسری (End-to-end) کنار هم قرار خواهیم داد (به‌لحاظ ساختاری؛ پایپ‌لاین برای محیط اجرا شبیه‌سازی‌شده باقی می‌ماند).

تمرین‌ها

1. زمان‌بندی 1F1B را پیاده‌سازی کنید و بررسی کنید که سهم حباب آن دقیقاً با GPipe هم‌خوانی داشته باشد اما حافظه‌ی اکتیویشن کراندار بماند.
2. زمان واقعی هر مرحله را روی مدلی عمیق‌تر پروفایل کنید و مراحل را بر مبنای زمان واقعی اندازه‌گیری‌شده بازتنظیم کنید.
3. قابلیت انباشت گرادیان (Gradient Accumulation) را در طول ریزدسته‌های پایپ‌لاین پیاده کنید و مطمئن شوید گرادیان نهایی با گرادیان حاصل از فاز رفت دسته‌ی کامل برابری می‌کند.
4. پایپ‌لاین را با چک‌پوینت‌گذاری اکتیویشن همراه کنید و میزان کاهش حافظه را در مقایسه با هزینه‌ی محاسبات بسنجید.
5. پایپ‌لاین را با DDP ترکیب کنید (هر رتبه‌ی پایپ‌لاین روی یک گروه موازی‌سازی داده تکرار شود) و نحوه‌ی زمان‌بندی دوبعدی آن را تحلیل کنید.

اصطلاحات کلیدی

برای مطالعه‌ی بیشتر

• هوانگ و همکاران، GPipe: Efficient Training of Giant Neural Networks
• نارایانان و همکاران، PipeDream: Generalized Pipeline Parallelism for DNN Training
• مستندات موازی‌سازی پایپ‌لاین در Megatron-LM
• بخش 19، درس 76: اجزای سازنده‌ی send/recv که در این زمان‌بندی به کار می‌روند
• بخش 19، درس 78: تکنیک ZeRO مستقل از پایپ‌لاین است و اغلب با آن ترکیب می‌شود