آموزش پیاده‌سازی DDP از صفر در PyTorch برای موازی‌سازی داده

آموزش مدل‌های هوش مصنوعی بزرگ روی یک پردازنده‌ی تک معمولاً ناممکن یا بسیار زمان‌بر است. موازی‌سازی داده (Data Parallelism) اساسی‌ترین تکنیک برای توزیع محاسبات بین پردازنده‌ها به شمار می‌رود، اما بسیاری از توسعه‌دهندگان کتابخانه‌ی DistributedDataParallel (به‌اختصار DDP) را مثل یک جعبه‌ی سیاه در نظر می‌گیرند. درک دقیق سازوکار داخلی DDP مانع از افت کارایی و بروز خطاهای پیچیده در سیستم‌های توزیع‌شده می‌شود.

در این درس نشان می‌دهیم که DDP جادو نیست؛ بلکه لایه‌ای حدوداً 200 خطی روی عملیات جمعی allreduce است که با مقداردهی اولیه‌ی یکسان و میانگین‌گیری گرادیان‌ها در هر گام، وزن‌های مدل را در تمام فرآیندها همگام نگه می‌دارد. پیاده‌سازی کامل این ساختار روی CPU با بک‌اند gloo انجام می‌شود؛ بنابراین برای آزمودن و درک مفاهیم پایه‌ای آن هیچ نیازی به پردازنده‌های گرافیکی گران‌قیمت یا کلاسترهای پیچیده نخواهید داشت.

ماژول DistributedDataParallel صرفاً یک هوک (hook) روی عملیات allreduce است. کافی است مدل را درون آن بپیچید (wrap کنید)، پارامترهای اولیه را از رتبه (rank) 0 برودکست (broadcast) کنید تا همه‌ی رتبه‌ها با وزن‌های کاملاً یکسان کار را شروع کنند، روی تک‌تک پارامترها یک هوک پس‌انتشار (backward hook) بگذارید که دستور allreduce گرادیان را صادر کند؛ ادامه‌ی مسیر همان گرادیان کاهشی (gradient descent) استاندارد است. کل این الگو در 200 خط پیاده می‌شود.

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

• پیاده‌سازی یک لایه‌ی پوششی مشابه DistributedDataParallel که پارامترهای اولیه را برودکست کند و بعد از پس‌انتشار، روی گرادیان‌ها allreduce بزند.
• اجرای N رتبه‌ی پردازنده (CPU ranks) با تابع torch.multiprocessing.spawn روی بک‌اند gloo به کمک راندوو مبتنی بر فایل (file-based rendezvous).
• اثبات درستی همگام‌سازی گرادیان‌ها از طریق آموزش ترتیبی همان مدل روی همان داده‌ها و نمایش برابری گام‌به‌گام پارامترها.
• دفاع از دلایل استفاده از باکت‌بندی (تلفیق گرادیان‌ها یا gradient fusion) و همپوشانی (ارتباط هم‌زمان با پس‌انتشار) به‌عنوان دو تغییری که یک DDP ساده را به یک DDP آماده‌ی محیط عملیاتی (پروداکشن) تبدیل می‌کنند.

صورت مسئله

یک مدل 1 میلیارد پارامتری با 12 گیگابایت فعال‌سازی (activation) در حافظه‌ی یک پردازنده‌ی گرافیکی معمولی جا نمی‌شود. حتی اگر جا شود هم آموزش آن هفته‌ها طول می‌کشد. در موازی‌سازی داده (Data Parallelism)، دسته‌ی داده‌ها (batch) بین N رتبه تقسیم می‌شود، هر رتبه گذر رفت (forward) و برگشت (backward) را روی قطعه‌ی داده‌ی (shard) اختصاصی خود حساب می‌کند و در هر گام، گرادیان‌های همه‌ی رتبه‌ها با هم جمع می‌شوند تا هر N نسخه کاملاً یکسان بمانند. در نهایت بهینه‌ساز بر اساس همین گرادیانِ جمع‌شده گام برمی‌دارد.

اگر گرادیان‌ها همگام نشوند، این N نسخه تا همان گام 2 از هم فاصله می‌گیرند و واگرا می‌شوند. در این حالت دیگر با «یک مدل آموزش‌دیده روی داده‌های بیشتر» طرف نیستیم، بلکه N مدل مجزا داریم که برحسب تصادف وزن‌های اولیه‌شان یکی بوده است. از طرف دیگر، اگر همگام‌سازی گرادیان‌ها بد انجام شود (یک allreduce به ازای هر پارامتر، بدون همپوشانی و بدون باکت‌بندی)، شبکه گلوگاه سیستم می‌شود و پردازنده‌های گرافیکی بیکار می‌مانند تا داده از سیم‌های ارتباطی رد شود. هنر DDP این است که هزینه‌ی همگام‌سازی گرادیان‌ها را نسبت به محاسبات تقریباً به صفر برساند. نسخه‌ی استاندارد PyTorch DDP این کار را با باکت‌بندی گرادیان‌ها، همپوشانی allreduce با پس‌انتشار لایه‌ی بعدی و بهره‌گیری از NCCL روی بستر NVLink انجام می‌دهد. ما همه‌ی این سه مورد را روی CPU با استفاده از gloo پیاده می‌کنیم و دقیقاً همان مفاهیم را یاد می‌گیریم.

مفاهیم پایه

سه عملیات کلیدی مورد نیاز DDP

چرا میانگین و نه مجموع؟

حاصل Allreduce-SUM تقسیم بر world_size، میانگین گرادیان را به ما می‌دهد. میانگین نسبت به world_size پایدار است؛ یعنی نرخ یادگیری‌ای (learning rate) که برای یک رتبه تنظیم کرده‌اید روی 4 رتبه هم کار می‌کند، چون اندازه‌ی گرادیان در هر گام تغییری نکرده است. اگر Allreduce-SUM را بدون تقسیم استفاده کنید، با هر بار تغییر اندازه‌ی کلاستر مجبورید نرخ یادگیری را دوباره از اول تنظیم کنید. ماژول DDP عملیات SUM را درون خود قرار داده و حاصل را تقسیم می‌کند؛ در این درس ما هم دقیقاً همین کار را انجام می‌دهیم.

چرا باکت‌بندی گرادیان‌ها؟

یک ترنسفورمر هزاران تنسور پارامتر دارد. اجرای یک allreduce به ازای هر تنسور، یعنی پرداخت هزاران‌باره‌ی تاخیر پایه‌ی gloo. ساختار DDP گرادیان‌ها را در دسته‌های حدوداً 25 مگابایتی (باکت) گروه‌بندی می‌کند و برای هر باکت یک allreduce صادر می‌کند. در این روش حجم کل داده‌ی ارسالی روی شبکه تغییری نمی‌کند، اما تاخیر ارتباطی سرشکن می‌شود. برای مدل کوچک این درس، همه‌ی پارامترها را در یک باکت قرار می‌دهیم؛ آنچه اهمیت دارد و منتقل می‌شود درک این ساختار است.

اهمیت تثبیت سید (Seed)

هر رتبه برای بُر زدن داده‌ها (shuffling) باید torch.manual_seed(seed + rank) را صدا بزند، اما برای مقداردهی اولیه‌ی پارامترها باید از torch.manual_seed(seed) استفاده کند. اگر برای بُر زدن از یک سید مشترک استفاده کنید، همه‌ی رتبه‌ها دسته‌ی داده‌ها را با یک ترتیب مشابه می‌بینند (که نقض غرض موازی‌سازی داده است). برای پارامترها، برودکست اولیه‌ی رتبه‌ی 0 در زمان ساخت ماژول همان چیزی است که نسخه‌ها را کاملاً همسان می‌کند و مقداردهی اولیه‌ی هر رتبه را رونویسی (overwrite) خواهد کرد. یکسان بودن سید مقداردهی اولیه صرفاً یک پشتیبان است: اگر برودکست اجرا نشود یا مشکلی پیدا کند، مقداردهی‌های اختصاصی هر رتبه باعث می‌شوند وزن‌های اولیه تفاوت داشته باشند و همگام‌سازی گرادیان به‌تنهایی هرگز نمی‌تواند آن‌ها را دوباره به هم برساند. این الگوی سید را درست پیاده کنید، وگرنه آزمون برابری پارامترها در همان گام 1 شکست می‌خورد.

پیاده‌سازی از صفر

فایل code/main.py این بخش‌ها را پیاده کرده است:

• MiniMLP: یک شبکه‌ی پرسپترون چندلایه (MLP) سه‌لایه، آن‌قدر کوچک که در چند ثانیه همگرا شود و آن‌قدر کامل که ساختار کار را شفاف نشان دهد.
• DistributedDataParallel(model, world_size): در زمان ساخت، پارامترها را برودکست می‌کند و لایه‌ی پوششی‌ای برمی‌گرداند که متد sync_grads آن گرادیان‌های جمع‌شده از طریق allreduce را بر world_size تقسیم می‌کند.
• worker(rank, world_size, ...): حلقه‌ی کامل آموزش شامل مقداردهی اولیه‌ی torch.distributed روی gloo، گذر رفت، برگشت، همگام‌سازی و اعمال گام بهینه‌ساز.
• _reference_single_process_loop(...): همین مدل را به‌صورت ترتیبی روی همین داده‌ها و در یک فرایند آموزش می‌دهد؛ این متد در تست‌ها استفاده می‌شود تا برابری بایتی پارامترها بعد از هر گام تایید شود.

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

python3 code/main.py

خروجی: جدولی از روند آموزش در هر گام که تابع زیان (loss) و چک‌سام پارامترهای اجرای تک‌فرایندی را با اجرای DDP روی 4 رتبه مقایسه می‌کند. هر دو مسیر منحنی‌های زیان کاملاً یکسانی (تا دقت اعشار float epsilon) تولید می‌کنند که نشان می‌دهد همگام‌سازی گرادیان کاملاً درست است.

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

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

شناسایی پارامترهای بلااستفاده (Find unused parameters): در برخی شبکه‌ها، گذر رفت بسته به شرایط از روی برخی پارامترها عبور می‌کند (مانند خروج زودهنگام یا مسیریاب در مدل‌های مخلوطی از خبرگان / MoE). پارامترهای ردشده فاقد گرادیان هستند، اما هوک آماده‌سازی باکت در DDP همچنان منتظر گرادیان آن‌ها می‌ماند و کل allreduce دچار بن‌بست (deadlock) می‌شود. فعال کردن find_unused_parameters=True به DDP می‌گوید پیش از اجرای reduce، بررسی کند که کدام پارامترها واقعاً گرادیان گرفته‌اند. هزینه‌ی این کار پیمایش گراف در هر گام است؛ بنابراین تا زمانی که گذر رفت شاخه‌بندی شرطی (branch) ندارد، آن را خاموش بگذارید.

بهینه‌سازی گراف ایستا (Static graph optimisation): اگر ساختار گذر رفت در تمام گام‌ها ثابت بماند، با فعال کردن static_graph=True به DDP اجازه می‌دهید زمان‌بندی باکت‌ها را از پیش محاسبه کند. این بهینه‌سازی در مقیاس‌های بزرگ اهمیت زیادی دارد؛ ذخیره‌ی چند میلی‌ثانیه‌ای که با پیش‌محاسبه در هر گام به دست می‌آید، در طول 10000 گام تاثیر چشمگیری خواهد داشت.

دقت در انباشت گرادیان (Gradient Accumulation): انباشت گرادیان روی K ریزدسته (microbatch) بدون همگام‌سازی لحظه‌ای در هر ریزدسته، نرخ پردازش را تا 10 برابر بیشتر می‌کند. ماژول DDP یک مدیریت‌کننده‌ی زمینه (context manager) به نام no_sync() در اختیارتان می‌گذارد که همگام‌سازی allreduce بعد از پس‌انتشار را متوقف می‌کند. اگر استفاده از این زمینه را فراموش کنید، K بار allreduce بیهوده اجرا خواهید کرد و نرخ پردازش به کف می‌رسد.

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

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

• PyTorch DDP: پیاده‌سازی مرجع و استاندارد. ماژول torch.nn.parallel.DistributedDataParallel(model) امکانات باکت‌بندی، همپوشانی و زمینه‌ی no_sync را به‌صورت آماده پیاده‌سازی کرده است.
• HuggingFace Accelerate: یک لانچر (launcher) اضافه می‌کند که متغیرهای محیطی torchrun و بسته‌بندی مدل را مدیریت می‌کند؛ اما در لایه‌های زیرین از همان DDP استفاده می‌کند.
• موازی‌سازی داده در Megatron-LM: برای مدل‌های بزرگ، DDP را با موازی‌سازی تنسور (tensor parallel) ترکیب می‌کند؛ البته بخش موازی‌سازی داده‌ی آن از همان الگوی allreduce-after-backward پیروی می‌کند.

خروجی این درس

در درس 78 (خرد کردن به سبک ZeRO)، عملیات allreduce به ازای هر پارامتر را با reduce_scatter جایگزین می‌کنیم تا هر رتبه تنها قطعه‌ی اختصاصی خود از وضعیت بهینه‌ساز را ذخیره کند. در درس 81 نیز DDP را با ZeRO ترکیب می‌کنیم تا دمو را به‌صورت سرتاسری بسازیم.

تمرین‌ها

1. قابلیت ساخت باکت‌های گرادیان با اندازه‌ی قابل‌تنظیم را اضافه کرده و میزان بهبود سرعت را در مقایسه با روش «یک allreduce به ازای هر پارامتر» روی مدلی عمیق‌تر بسنجید.
2. سازوکار no_sync() را به‌صورت یک context manager پیاده کنید و مطمئن شوید حاصل انباشت گرادیان روی K ریزدسته، با خط مبنای تک‌فرایندی مطابقت دارد.
3. حالتی برای find_unused_parameters تعریف کنید که در آن گذر رفت گاهی از یکی از لایه‌های MLP عبور کند؛ بررسی کنید که بدون این فلگ، سیستم حتماً دچار بن‌بست شود.
4. gloo را بردارید و همگام‌سازی را فقط با torch.distributed.barrier() انجام دهید تا تفاوت کارایی همگام‌سازی مبتنی بر allreduce با روش مبتنی بر barrier را احساس کنید.
5. سهم تاخیر همگام‌سازی گرادیان را از کل زمان یک گام برای اندازه‌های دسته‌ی 1، 16 و 256 اندازه بگیرید و چگونگی مقیاس‌پذیری آن را توضیح دهید.

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

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

• مستندات DistributedDataParallel در PyTorch
• آموزش ساختار داخلی PyTorch DDP
• مقاله‌ی Li و همکاران درباره‌ی تجربیات سرعت‌بخشی به آموزش موازی داده در PyTorch
• فاز 19، درس 76: عملیات جمعی که DDP بر پایه‌ی آن‌ها بنا شده است
• فاز 19، درس 78: خرد کردن داده به سبک ZeRO و جایگزینی allreduce پارامترها با reduce_scatter