پیاده‌سازی عملیات جمعی در آموزش توزیع‌شده از صفر

آموزش توزیع‌شده (distributed training) مدل‌های یادگیری عمیق اغلب شبیه یک جعبه‌ی سیاه به نظر می‌رسد؛ جایی که فریم‌ورک‌هایی مانند PyTorch DDP یا DeepSpeed در پس‌زمینه وزن‌ها را همگام می‌کنند، اما هنگام بروز افت کارایی یا گیر کردن فرآیند آموزش، ریشه‌یابی مشکل دشوار می‌شود. هسته‌ی اصلی تمام این سیستم‌ها تنها چهار عملیات جمعی ارتباطی یعنی allreduce، broadcast، allgather و reduce_scatter است.

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

در این درس، این چهار عملیات را بدون تکیه بر زیرساخت‌های پیچیده سخت‌افزاری و صرفاً با استفاده از `multiprocessing.Queue` در پایتون روی CPU پیاده‌سازی می‌کنیم و درستی خروجی‌ها را مستقیماً در برابر بک‌اند gloo در کتابخانه `torch.distributed` می‌سنجیم؛ بنابراین برای درک و اجرای کامل این درس نیازی به دسترسی به کلاسترهای گران‌قیمت یا سرویس‌های ابری نخواهید داشت.

چهار عملیات جمعی (collective operations) که آموزش توزیع‌شده (distributed training) را سرپا نگه می‌دارند عبارتند از: allreduce، broadcast، allgather و reduce_scatter. هر جزء سازنده‌ی دیگری که فریم‌ورک‌های آموزشی ارائه می‌دهند، صرفاً پوسته‌ای دور همین چهار عملیات است. اگر این‌ها را یک بار روی مشی از multiprocessing.Queue بسازید و درستی‌شان را در برابر یک پیاده‌سازی مرجع بسنجید، باقی این مسیر صرفاً کارهای اتصال و لوله‌کشی خواهد بود.

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

• پیاده‌سازی ring allreduce در دو گام (ابتدا reduce-scatter و سپس allgather) و اثبات این‌که حجم ارتباطات به ازای هر رتبه (rank) برابر با 2(N-1)/N بایت به ازای هر عنصر است.
• ساخت broadcast، allgather و reduce_scatter بر پایه‌ی ارسال‌های نقطه‌به‌نقطه (point-to-point) روی multiprocessing.Queue.
• اعتبارسنجی تمام این اجزای سازنده در مقایسه با مرجع gloo در torch.distributed برای ورودی یکسان.
• دفاع از انتخاب میان ساختار حلقه‌ای (ring) و درختی (tree) بر اساس توپولوژی کلاستر، کف تاخیر (latency floor) و سقف پهنای باند (bandwidth ceiling).

صورت مسئله

یک allreduce ساده‌لوحانه روی N رتبه، تنسور را N بار به یک گره ریشه (root) می‌فرستد و N بار آن را به بقیه برمی‌گرداند (broadcast). در این حالت، پهنای باند به ازای هر رتبه با مقیاس $O(N)$ رشد می‌کند، گره ریشه به گلوگاه تبدیل می‌شود و حداقل زمان اجرای واقعی (wall-clock floor) برابر با کندترین پیوند ارتباطی ضرب‌در N خواهد بود. روش ring allreduce این فرایند را به $2(N-1)$ تکه با اندازه‌ی $T/N$ می‌شکند؛ در نتیجه بایت‌های ارسالی به ازای هر رتبه به $2T(N-1)/N$ کاهش می‌یابد که با رشد N به $2T$ نزدیک می‌شود و فرقی نمی‌کند کلاستر چقدر بزرگ باشد، هرگز از آن فراتر نمی‌رود. روش tree allreduce در Nهای کوچک و پیوندهای با تاخیر بالا برنده است، زیرا عمق ارتباطی آن به جای $2(N-1)$ گام لغزش (hop)، برابر با $\log_2(N)$ است. اگر توپولوژی اشتباهی را متناسب با ساختار کلاستر خود انتخاب کنید، کندترین GPU زمان هر گام (step time) را دیکته خواهد کرد.

تمام فریم‌ورک‌های آموزش توزیع‌شده‌ای که در این بخش بررسی می‌کنید به این چهار جزء سازنده وابسته‌اند. PyTorch DDP گرادیان‌ها را با یک allreduce به ازای هر باکت (bucket) پارامتر همگام می‌کند. ZeRO وضعیت بهینه‌ساز (optimiser state) را با reduce_scatter قطعه‌بندی (shard) می‌کند و پارامترهای به‌روزرسانی‌شده را با allgather انتشار می‌دهد. FSDP کل مرحله‌ی پیش‌رو (forward) را به یک allgather به همراه reduce_scatter تبدیل می‌کند. موازی‌سازی پایپ‌لاینی (pipeline parallel) برای جابه‌جایی اکتیویشن‌ها در میان گروه‌های مراحل پایپ‌لاین به broadcast نیاز دارد. اگر نتوانید این چهار عملیات جمعی را پیاده‌سازی کنید، نمی‌توانید سر در بیاورید چرا فرایند آموزش قفل می‌کند، چرا عدم تطابق گرادیان در رتبه‌ی 3 رخ می‌دهد یا چرا تعویض توپولوژی باعث دو برابر شدن حباب پایپ‌لاین (pipeline bubble) می‌شود.

مفاهیم پایه

الگوریتم Ring allreduce در دو گام

تنسور را به N تکه‌ی مساوی با اندیس‌های 0 تا N-1 تقسیم کنید. مالکیت تکه‌ای که اندیس آن با شماره‌ی رتبه برابر است به همان رتبه می‌رسد. گام اول یعنی reduce-scatter در N-1 مرحله اجرا می‌شود. در مرحله‌ی s، رتبه‌ی r تکه‌ی (r - s) mod N را به رتبه‌ی (r + 1) mod N می‌فرستد و تکه‌ی (r - s - 1) mod N را از رتبه‌ی (r - 1) mod N دریافت می‌کند و تکه‌ی دریافت‌شده را با نسخه‌ی محلی خود جمع می‌زند. پس از N-1 مرحله، رتبه‌ی r مجموع کامل تکه‌ی r را در اختیار دارد. گام دوم یعنی allgather در N-1 مرحله‌ی دیگر اجرا می‌شود و تکه‌های کامل‌شده را دور حلقه می‌چرخاند تا زمانی که تمام رتبه‌ها مجموع کامل همه‌ی تکه‌ها را داشته باشند.

شبکه‌ی صف به عنوان جایگزین NCCL

کتابخانه‌ی NCCL روی بسترهای PCIe و NVLink با کاهش‌های سخت‌افزاری اجرا می‌شود. روی CPU چنین امکانی ندارید. ایجاد یک multiprocessing.Queue به ازای هر یال در این حلقه، تحویل مرتب نقطه‌به‌نقطه را با یک تولیدکننده و یک مصرف‌کننده‌ی واحد برایتان فراهم می‌کند. عملیات کاهش (reduction) در فضای کاربری (user space) رخ می‌دهد؛ بنابراین سربار پایتون را می‌پردازید، اما الگوی ارتباطی دقیقاً همانند ring allreduce در NCCL است. وقتی منطق درستی عملکرد را روی نسخه‌ی مبتنی بر صف اثبات کنید، رفتار کلاستر واقعی هم دقیقاً از همین الگو پیروی می‌کند.

اعتبارسنجی در برابر gloo

برای هر جزء سازنده یک تست واحد (unit test) در نظر گرفته شده که خروجی آن را با torch.distributed راه‌اندازی‌شده با بک‌اند gloo روی همان تنسور و با تعداد فرایندهای (world size) یکسان مقایسه می‌کند. اگر ring allreduce پیاده‌سازی‌شده بیش از اپسیلون float32 با gloo اختلاف داشته باشد، تست رد می‌شود. اعتبارسنجی در برابر یک پیاده‌سازی مرجع غیرقابل مذاکره است؛ بدون آن، جزء سازنده درست به نظر می‌رسد تا زمانی که در مرحله‌ی 10000 از یک آموزش واقعی در محیط عملیاتی به خطا بخورید.

پیاده‌سازی

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

• کلاس Mesh که N نمونه از multiprocessing.Queue را در قالب یک حلقه به هم وصل می‌کند و متدهای send(dst, tensor) و recv(src) را در دسترس هر رتبه می‌گذارد.
• تابع ring_allreduce(mesh, rank, world_size, tensor) که الگوریتم دوگامه را اجرا می‌کند.
• تابع broadcast(mesh, rank, world_size, tensor, src) روی یک درخت لگاریتمی.
• تابع allgather(mesh, rank, world_size, tensor) با استفاده از N-1 چرخش.
• تابع reduce_scatter(mesh, rank, world_size, tensor) به عنوان نیمه‌ی اول allreduce.
• تابع _gloo_reference(op, world_size, tensor) که همان ورودی را برای مقایسه‌ی دقیق بایتی، با gloo در torch.distributed اجرا می‌کند.

اجرای کد:

python3 code/main.py

خروجی: جدول اعتبارسنجی هر جزء سازنده در مقایسه با خروجی‌های شبکه‌ی صف و gloo، به همراه شمارنده‌ی بایت به ازای هر رتبه که مقیاس‌پذیری 2T(N-1)/N را اثبات می‌کند.

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

سه الگو این اجزای سازنده را برای استقرار در محیط عملیاتی آماده و پایدار می‌کنند.

باکت‌بندی گرادیان‌ها پیش از allreduce. یک مدل 1 میلیارد پارامتری ده‌ها هزار تنسور گرادیان دارد. اجرای یک allreduce به ازای هر تنسور، کف تاخیر شبکه را N بار به سیستم تحمیل می‌کند. کتابخانه‌ی DDP گرادیان‌ها را در دسته‌های حدوداً 25 مگابایتی باکت‌بندی می‌کند و برای هر باکت یک allreduce می‌فرستد؛ بنابراین تنسورهای کوچک سوار بر تنسورهای بزرگ جابه‌جا می‌شوند. بدون باکت‌بندی، سربار تاخیر بر زمان گام محاسباتی چیره می‌شود.

هم‌پوشانی ارتباطات با محاسبات. مرحله‌ی برگشت (backward)، گرادیان‌ها را لایه به لایه و به ترتیب معکوس محاسبه می‌کند. به محض آماده شدن گرادیان آخرین لایه، عملیات allreduce آن را شروع کنید در حالی که لایه‌ی بعدی هم‌چنان در حال محاسبه است. کتابخانه‌ی PyTorch DDP این فرایند را با قلاب‌های bucket-ready پیاده می‌کند. وقتی شبکه پهنای باند خالی داشته باشد، این هم‌پوشانی زمان مشهود ارتباطات را به نصف می‌رساند.

انتخاب حلقه یا درخت بر اساس اندازه‌ی پیام، نه تعصب فنی. کتابخانه‌ی NCCL شامل یک سیستم تشخیص توپولوژی است که برای پیام‌های بالاتر از حدود 1 مگابایت ساختار حلقه‌ای و برای پایین‌تر از آن ساختار درختی را انتخاب می‌کند. این نقطه‌ی تلاقی بر اساس رقابت پهنای باند در برابر تاخیر تعیین می‌شود: بالای 1 مگابایت، عبارت پهنای باند 2T(N-1)/N غالب است و ساختار حلقه برنده می‌شود؛ زیر 1 مگابایت، تعداد گام‌های log2(N) برتری دارد. هاردکد کردن یکی از این دو توپولوژی، نرخ پردازش (throughput) را در پیام‌هایی با اندازه‌ی نامناسب به باد می‌دهد.

کاربردها

الگوهای محیط عملیاتی (پروداکشن):

• PyTorch DDP. بعد از مرحله‌ی برگشت، متد dist.all_reduce را روی گرادیان‌های باکت‌بندی‌شده فراخوانی می‌کند. اندازه‌ی باکت قابل تنظیم است؛ مقدار پیش‌فرض 25 مگابایت برای شبکه‌های 100 گیگابیت اترنت مقداری منطقی است.
• DeepSpeed ZeRO. عملیات reduce_scatter را برای قطعه‌بندی گرادیان‌ها و allgather را برای بازسازی پارامترهای کامل پیش از مرحله‌ی پیش‌رو صادر می‌کند. اجزای سازنده‌ی این درس دقیقاً همان فراخوانی‌هایی هستند که ZeRO به کار می‌برد.
• FSDP. مرحله‌ی پیش‌رو با allgather برای خروج لایه از حالت قطعه‌بندی شروع می‌شود، محاسبات انجام می‌گیرد، سپس با reduce_scatter تجمیع شده و پارامترهای غیرقطعه‌بندی‌شده دور ریخته می‌شوند. همان اجزای سازنده، اما با زمان‌بندی متفاوت.

نحوه‌ی استقرار

از اجزای سازنده‌ی مش صف در درس‌های 77 تا 81 استفاده کنید. درس 77 متد allreduce را وارد DDP می‌کند. درس 78 متد reduce_scatter را به ZeRO متصل می‌کند. درس 79 اکتیویشن‌های پایپ‌لاین را با ارسال و دریافت نقطه‌به‌نقطه جابه‌جا می‌کند، نه با یک عملیات جمعی. درس 81 هر چهار جزء سازنده را برای ساخت دموی سرتاسری (end-to-end) کنار هم قرار می‌دهد.

تمرین‌ها

1. نسخه‌ای از tree allreduce اضافه کنید و بر اساس اندازه‌ی پیام، بین حلقه و درخت سوییچ کنید. نقطه‌ی تلاقی را اندازه بگیرید.
2. یک پارامتر recv_timeout_ms اضافه کنید تا رتبه‌ای که متوقف شده، به جای معلق ماندن همیشگی، خطای مهلت زمانی (deadline error) بالا بیاورد.
3. برای این چهار جزء سازنده، multiprocessing.Queue را با سوکت‌های TCP جایگزین کنید؛ همان تست‌ها روی کابل و شبکه‌ی واقعی.
4. یک قلاب ابزاردقیق (instrumentation hook) برای پهنای باند اضافه کنید تا شمارنده‌ی بایت هر رتبه را در قالب JSONL لاگ کند.
5. زمان اجرای واقعی (wall-clock time) حلقه را در برابر درخت روی 4 رتبه برای تنسورهایی با اندازه‌ی 1 کیلوبایت، 1 مگابایت و 16 مگابایت مقایسه کنید. نقطه‌ی تلاقی را با داده‌های تجربی اثبات کنید.

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

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

• PyTorch Distributed: عملیات‌های جمعی NCCL
• مقاله Horovod درباره ring allreduce
• انتخاب توپولوژی و الگوریتم در NCCL
• مقاله Patarasuk و Yuan درباره الگوریتم‌های بهینه از نظر پهنای باند برای allreduce
• فاز 10 درس 05 - بررسی اجمالی آموزش توزیع‌شده
• فاز 19 درس 77 - پیاده‌سازی DDP بر پایه‌ی همین اجزای سازنده