آموزش DDP و FSDP در PyTorch از صفر

بسیاری از توسعه‌دهندگان هنگام مقیاس‌پذیری مدل‌های یادگیری عمیق به سراغ بسته‌های آماده‌ای مثل `DistributedDataParallel` یا FSDP در PyTorch می‌روند، اما وقتی فرایند آموزش به دلایل نامعلومی قفل می‌کند یا همگرایی مدل دچار مشکل می‌شود، عیب‌یابی آن به یک کابوس تبدیل خواهد شد. درک ریاضی و سازوکار پشت پرده‌ی عملیات اشتراکی (Collectives) کلید ساخت و اجرای بدون خطای سیستم‌های توزیع‌شده است.

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

در این درس یاد می‌گیرید چگونه یک لایه‌ی DDP و طرح اولیه‌ی FSDP را فقط با استفاده از ابزارهای بومی پایتورک بسازید، پارامترها را میان فرایندها همگام کنید و به‌صورت ریاضی برابری نتایج موازی را با آموزش تک‌پردازنده‌ای بسنجید.

آموزش چندرتبه‌ای (Multi-rank training) به دو عمل اشتراکی (Collective) و یک قانون خلاصه می‌شود: در ابتدای کار، پارامترها را همه‌پخشی (Broadcast) کنید؛ بعد از پس‌انتشار (Backward)، از گرادیان‌ها میانگین بگیرید؛ و هرگز نگذارید رتبه‌ها (Ranks) بر سر این‌که در کدام گام قرار دارند، با هم اختلاف پیدا کنند.

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

• راه‌اندازی یک گروه فرایندی (Process group) میان N رتبه با بک‌اند gloo، بدون نیاز به سخت‌افزار خاص.
• پیاده‌سازی یک لایه‌ی پوششی (Wrapper) ساده برای DDP که هنگام ساخت مدل، پارامترها را همه‌پخشی کند و پس از پس‌انتشار، روی گرادیان‌ها عمل all-reduce انجام دهد.
• اثبات این‌که all-reduce روی گرادیان‌های مجزای هر رتبه، خروجیِ کاملاً یکسانی با گرادیان محاسبه‌شده توسط یک فرایند تک روی داده‌های متصل‌شده تولید می‌کند.
• ترسیم طرح کلی شاردینگ پارامترها در FSDP: هر رتبه برشی از پارامترها را نگه می‌دارد، تنسور کامل فقط برای گذر رو به جلو (Forward pass) گردآوری و بلافاصله پس از آن رها می‌شود.

صورت مسئله

مدل درون یک دستگاه جا می‌شود، اما کل مجموعه‌داده نه. بودجه‌ی بهینه‌سازی ایجاب می‌کند که در هر ثانیه‌ی زمان واقعی، N برابر نمونه‌ی بیشتری ببینید. اهرم اول، موازی‌سازی داده (Data parallel) است: هر رتبه همان مدل یکسان را روی برش متفاوتی از دسته (Batch) اجرا می‌کند و پیش از گام بهینه‌ساز، گرادیان‌ها میانگین‌گیری می‌شوند. اهرم دوم، FSDP است: زمانی که مدل روی یک دستگاه جا نمی‌شود، هر رتبه تنها کسری از هر پارامتر را در حافظه نگه می‌دارد و تنسورهای کامل را در حین گذر رو به جلو، لایه به لایه بازسازی می‌کند.

بخش دشوار کار، حساب‌وکتاب و همگام‌سازی است. اگر پارامترها بین رتبه‌ها حتی کمی فاصله بگیرند، کل فرایند آموزش در سکوت خراب می‌شود. اگر از گرادیان‌ها میانگین بگیرید اما از تابع زیان (Loss) نه، ارقام روی داشبورد دروغ می‌گویند. اگر بک‌اندِ عملیات اشتراکی نتواند سر ساختار شبکه (Topology) به توافق برسد، اجرا برای همیشه قفل می‌کند. راه چاره این است: یک بار این اعمال اشتراکی را دستی بنویسید و دیگر به هیچ لایه‌ی پوششیِ آماده‌ای که توانایی پیاده‌سازی مجددش را ندارید، اعتماد نکنید.

این درس روی پردازنده‌ی اصلی اجرا می‌شود و فرضی روی وجود CUDA ندارد. بک‌اند gloo در همه‌ی نسخه‌های PyTorch گنجانده شده است و از ورکر‌های torch.multiprocessing پشتیبانی می‌کند؛ همین کد روی یک سیستم چند GPU، بدون تغییر ساختاری، با سوییچ به nccl کار خواهد کرد.

مبانی مفهومی

دو عمل اشتراکی که واقعاً اهمیت دارند

قرارداد پایه‌ایِ DDP، اجرای broadcast در زمان ساخت مدل و all_reduce بعد از پس‌انتشار است. در طرح اولیه‌ی FSDP، یک عمل all_gather نیز قبل از گذر رو به جلوی هر لایه اضافه می‌شود.

تطابق میانگین‌گیری گرادیان با گرادیان تک‌فرایندی

مدلی که روی دسته‌ای با B نمونه در N رتبه آموزش می‌بیند، باید دقیقاً همان گرادیانی را تولید کند که یک فرایند تک با آموزش روی دسته‌ای به اندازه‌ی N*B تولید می‌کرد. ترفند ماجرا این‌جاست: جمع زدن گرادیان‌های هر رتبه و تقسیم آن بر N، گرادیانِ زیان میانگین را به دست می‌دهد؛ یعنی دقیقاً همان نتیجه‌ای که انتروپی متقاطع (Cross-entropy) با کاهش میانگین (Mean reduction) روی کل دسته ایجاد می‌کرد. کد این درس این موضوع را با بررسی شرط max-abs-diff < 1e-3 بین گرادیان all-reduce دستی و گرادیان مرجعِ تک‌فرایندی اثبات می‌کند.

طرح کلی FSDP

صرفه‌جویی در حافظه دقیق است: حافظه‌ی مورد نیاز برای نگهداری پارامترها در هر رتبه به اندازه‌ی 1/N کاهش می‌یابد. هزینه‌ی این کار، عمل گردآوری (Gather) است که در هر گذر رو به جلو پرداخت می‌شود. در محیط عملیاتی (پروداکشن)، FSDP عمل گردآوری را با محاسبات لایه‌ی قبلی همپوشانی می‌کند؛ بنابراین هزینه‌ی زمانی بسیار کمتر از محاسبات ساده‌انگارانه‌ی اولیه خواهد بود. در این درس، فرایند رفت و برگشت را روی تک‌تک پارامترها پیاده می‌کنیم و برابری بیتی بازسازی با نسخه‌ی اصلی را می‌سنجیم.

پردازنده‌ی اصلی و بک‌اند gloo

با اینکه محیط عملیاتی بر پایه‌ی CUDA است، همان مسیرهای کد روی پردازنده‌ی اصلی نیز وجود دارند. gloo بک‌اند عملیات اشتراکی برای CPU است. این بک‌اند چند مرتبه از nccl روی پردازنده‌های گرافیکی کندتر است، اما رابط برنامه‌نویسی (API) کاملاً یکسانی دارد. گروه فرایندی در این درس با backend="gloo" راه‌اندازی می‌شود و رتبه‌ها به جای torchrun با torch.multiprocessing اجرا می‌شوند؛ هر دو روش در نهایت به همان فراخوانی‌های torch.distributed ختم می‌شوند. در سیستمی با چند GPU، تنها تغییرات لازم عبارت‌اند از: استفاده از backend="nccl"، تنسورهای روی دستگاه (Device tensors) و راه‌اندازی با torchrun.

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

فایل code/main.py دست‌سازه‌ی اجرایی این درس است.

گام 1: بالا آوردن گروه فرایندی

os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = str(port)
dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)

دو متغیر MASTER_ADDR و MASTER_PORT نقش نقطه‌ی ملاقات (Rendezvous) را دارند: همه‌ی رتبه‌ها به یک پورت یکسان روی یک هاست متصل می‌شوند. در این درس با ترفند Bind کردن سوکت و بستن آن، پورت آزادی انتخاب می‌کنیم تا زمان اجرای هم‌زمان چند برنامه روی یک ماشین، تداخلی پیش نیاید.

گام 2: همه‌پخشی در زمان ساخت مدل

متد MinimalDDP.__init__ تمام پارامترها و بافرها را پیمایش کرده و dist.broadcast(tensor, src=0) را صدا می‌زند. مقادیر رتبه‌ی 0 به مقداردهی اولیه‌ی مرجع تبدیل می‌شوند. بدون این کار، هر رتبه با بذر تصادفی (Seed) خودش شروع به کار می‌کند و رتبه‌ها از همان گام اول واگرا می‌شوند.

گام 3: اجرای all-reduce روی گرادیان‌ها پس از پس‌انتشار

def all_reduce_grads_(module, world_size):
    for p in module.parameters():
        if p.grad is None:
            p.grad = torch.zeros_like(p.data)
        dist.all_reduce(p.grad.data, op=dist.ReduceOp.SUM)
        p.grad.data.div_(world_size)

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

گام 4: اثبات برابری

تابع manual_all_reduce_matches_single_process همان مدل را روی رتبه‌ی 0 می‌سازد و گرادیان پس از all-reduce را با گرادیانی که یک فرایندِ واحد روی ورودی‌های متصل‌شده محاسبه می‌کرد، مقایسه می‌کند. بیشینه‌ی قدر مطلق اختلاف (max-abs-diff) رقمی در حدود 1e-8 خواهد بود.

گام 5: چرخه‌ی کامل FSDP

تابع fsdp_round_trip_sketch هر پارامتر را تخت (Flatten) می‌کند، طول آن را برای هماهنگی با مضربی از world_size با صفر پر (Pad) می‌کند، برش می‌زند، از طریق all-gather گردآوری می‌کند و داده‌های پد را برمی‌دارد. نسخه‌ی بازسازی‌شده در هر رتبه، با داده‌ی اولیه برابریِ کامل دارد. این، گام خروج از شارد (Unshard) است؛ عملیات معکوس (شاردینگ مجدد بعد از گذر رو به جلو) به سادگیِ جدا کردن یک برش از تنسورِ گردآوری‌شده خواهد بود.

کد را اجرا کنید:

python3 code/main.py

اندازه‌ی پیش‌فرض جهان (World size) برابر 2 است. دو فرایندِ CPU ایجاد می‌شوند، از طریق gloo با هم ارتباط برقرار می‌کنند و با کد وضعیت صفر به پایان می‌رسند. خروجی در فایل outputs/ddp-demo.json ثبت می‌شود که شامل این موارد است: مجموع پارامترها در هر رتبه، نُرم گرادیان بعد از all-reduce، نتیجه‌ی چرخه‌ی کامل FSDP و اختلاف گرادیان دستی در برابر گرادیان مرجع.

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

پشته‌های آموزش در محیط عملیاتی (پروداکشن) همین اجزای سازنده را به کار می‌گیرند. کلاس DistributedDataParallel در PyTorch چند قابلیت مهم اضافه می‌کند: قلاب‌های گرادیان پس از پس‌انتشار که all-reduce را با ادامه‌ی خود پس‌انتشار همپوشانی می‌کنند، all-reduce دسته‌بندی‌شده (Bucketed) که چند گرادیان کوچک را درون یک عملیات اشتراکی ادغام می‌کند و کانتکست no_sync که در درس 46 با آن کار کردید.

ماژول FSDP در PyTorch امکانات دیگری می‌افزاید: یک نمای پارامتر تخت برای هر لایه به‌طوری که هر رتبه یک بافر پیوسته در حافظه داشته باشد، همپوشانی خروج از شارد در لایه‌ی بعدی با محاسباتِ لایه‌ی جاری، و امکان تخلیه روی پردازنده‌ی اصلی برای شاردها.

الگوی کلی همواره ثابت است: همه‌پخشی در راه‌اندازی، کاهش (Reduce) بعد از پس‌انتشار، و شارد کردن پارامترها به محض اینکه دیگر درون حافظه جا نگرفتند.

خروجی این درس

فایل outputs/skill-distributed-fsdp-ddp.md شامل دستورالعمل آماده برای نوشتن اسکریپت‌های آموزشی جدید است: راه‌اندازی گروه فرایندی با gloo برای پردازنده‌ی اصلی و nccl برای GPU، بسته‌بندی مدل درون یک پوسته‌ی DDP که پارامترها را در زمان ساخت همه‌پخشی کرده و بعد از پس‌انتشار کاهش می‌دهد، و در صورت نیاز، شارد کردن پارامترها با استفاده از الگوی all_gather طبق طرح FSDP.

تمرین‌ها

1. کد را با گزینه‌ی --world-size 4 اجرا کنید و مطمئن شوید اختلاف پارامترها در طول اجرا زیر 1e-3 باقی می‌ماند.
2. میانگین‌گیری دستی را با dist.all_reduce(op=dist.ReduceOp.AVG) جایگزین کرده و تفاوت زمانی را اندازه بگیرید.
3. یک قلاب پس از پس‌انتشار به لایه‌ی پوششی DDP اضافه کنید تا all-reduce با ادامه‌ی پس‌انتشار همپوشانی داشته باشد؛ میزان بهبود زمان واقعی را ثبت کنید.
4. گام شاردینگ مجدد FSDP را پیاده کنید: بعد از گذر رو به جلو، تنسور کامل را دوباره با همان شارد محلی جایگزین کنید و تأیید کنید که حافظه‌ی هر رتبه کاهش می‌یابد.
5. بک‌اند را روی سیستمی مجهز به CUDA به nccl تغییر دهید. بررسی کنید کدام متغیرهای محیطی دست‌نخورده باقی می‌مانند و کدام تغییر می‌کنند.

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

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

• مستندات torch.distributed در PyTorch برای آشنایی با مفاهیم پایه‌ایِ عملیات اشتراکی که در این درس به کار رفت.
• فهرست عملیات اشتراکی کتابخانه‌ی gloo که ساختاری کاملاً مشابه با اجزای سازنده‌ی nccl روی CUDA دارد.
• درس 46 از فاز 19 برای یادگیری الگوی انباشت گرادیان که all-reduce در DDP را با no_sync مدیریت می‌کند.
• درس 47 از فاز 19 برای ساختار چک‌پوینت‌هایی که در طول اجراهای DDP و FSDP سالم باقی می‌مانند.
• مستندات PyTorch FSDP جهت مشاهده‌ی پیاده‌سازی عملیاتی شاردینگ پارامترها که در این درس بررسی شد.