تکثیر

نمای کلی

انتشار شاردینگ از خرده‌کاری‌های مشخص شده توسط کاربر برای استنباط تقسیم‌بندی نامشخص تانسورها (یا بعد خاص تانسورها) استفاده می‌کند. جریان داده (زنجیره های استفاده-دف) نمودار محاسباتی را در هر دو جهت طی می کند تا زمانی که به یک نقطه ثابت برسد، به عنوان مثال، تقسیم بندی دیگر نمی تواند بدون لغو تصمیمات اشتراک گذاری قبلی تغییر کند.

انتشار را می توان به مراحل تجزیه کرد. هر مرحله شامل بررسی یک عملیات خاص و انتشار بین تانسورها (عملگرها و نتایج)، بر اساس ویژگی های آن عملیات است. با در نظر گرفتن یک matmul به عنوان مثال، بین بعد غیر انقباضی lhs یا rhs به بعد متناظر نتیجه یا بین بعد انقباضی lhs و rhs انتشار می دهیم.

ویژگی‌های یک عملیات ارتباط بین ابعاد متناظر در ورودی و خروجی آن را تعیین می‌کند و می‌تواند به عنوان یک قانون تقسیم عملیاتی انتزاع شود.

بدون حل تعارض، یک مرحله انتشار به سادگی تا آنجا که می تواند منتشر می شود و در عین حال محورهای متضاد را نادیده می گیرد. ما از آن به عنوان (طولانی ترین) محورهای اصلی شاردینگ سازگار یاد می کنیم.

طراحی دقیق

سلسله مراتب حل تعارض

ما چندین استراتژی حل تعارض را در یک سلسله مراتب ترکیب می کنیم:

1. اولویت های تعریف شده توسط کاربر در Sharding Representation ، توضیح دادیم که چگونه می‌توان اولویت‌ها را به تقسیم‌بندی ابعاد متصل کرد تا امکان پارتیشن‌بندی تدریجی برنامه فراهم شود، به عنوان مثال، انجام موازی‌سازی دسته‌ای -> مگاترون -> تقسیم‌بندی صفر. این با اعمال انتشار در تکرارها به دست می‌آید - در تکرار i همه تقسیم‌بندی‌های ابعادی را که اولویت <=i دارند منتشر می‌کنیم و بقیه را نادیده می‌گیریم. ما همچنین مطمئن می‌شویم که انتشار، اشتراک‌گذاری‌های تعریف‌شده توسط کاربر با اولویت پایین‌تر ( >i ) را لغو نمی‌کند، حتی اگر در تکرارهای قبلی نادیده گرفته شوند.
2. اولویت های عملیاتی ما بر اساس نوع عملیات، شاردینگ ها را منتشر می کنیم. عملیات "عبور" (به عنوان مثال، عملیات عنصر عاقلانه و تغییر شکل) بالاترین اولویت را دارند، در حالی که عملیات با تبدیل شکل (مانند نقطه و کاهش) اولویت کمتری دارند.
3. انتشار تهاجمی شاردینگ ها را با یک استراتژی تهاجمی منتشر کنید. استراتژی اصلی فقط خرده‌کاری‌ها را بدون درگیری منتشر می‌کند، در حالی که استراتژی تهاجمی تضادها را حل می‌کند. تهاجمی بالاتر می تواند ردپای حافظه را به قیمت ارتباطات بالقوه کاهش دهد.
4. تکثیر اساسی این پایین ترین استراتژی انتشار در سلسله مراتب است، که هیچ حل تعارضی را انجام نمی دهد و در عوض محورهایی را منتشر می کند که بین همه عملوندها و نتایج سازگار هستند.

این سلسله مراتب را می توان به عنوان حلقه های تو در تو تعبیر کرد. به عنوان مثال، برای هر اولویت کاربر، یک انتشار با اولویت کامل اعمال می شود.

قانون تقسیم عملیات

قانون تقسیم بندی انتزاعی از هر عملیاتی را معرفی می کند که الگوریتم انتشار واقعی را با اطلاعاتی که برای انتشار تقسیم بندی ها از عملوندها به نتایج یا در بین عملوندها و غیره نیاز دارد، بدون نیاز به استدلال در مورد انواع عملیات خاص و ویژگی های آنها ارائه می دهد. این اساساً منطق op-specific را فاکتور می‌کند و یک نمایش مشترک (ساختار داده) برای همه عملیات‌ها فقط به منظور انتشار ارائه می‌کند. در ساده ترین شکل خود، فقط این تابع را ارائه می دهد:

GetOpShardingRule(Operation *) -> OpShardingRuleAttr

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

بیایید به مثال ماتمول برگردیم.

رمزگذاری که اطلاعات مورد نیاز در حین انتشار، یعنی روابط بین ابعاد را در بر می گیرد، می تواند به شکل نماد einsum نوشته شود:

(i, k), (k, j) -> (i, j)

در این رمزگذاری، هر بعد به یک عامل منفرد نگاشت می شود.

چگونه انتشار از این نگاشت استفاده می کند: اگر یک بعد از یک عملوند/نتیجه در امتداد یک محور تقسیم شود، انتشار عامل آن بعد را در این نگاشت جستجو می کند و سایر عملوندها/نتایج را در امتداد بعد مربوطه خود با همان عامل - و (موضوع) خرد می کند. به بحث قبلی در مورد همانندسازی) به طور بالقوه عملوندها/نتایج دیگری را نیز تکرار کنید که آن فاکتور را در امتداد آن محور ندارند.

عوامل مرکب: گسترش قاعده برای تغییر شکل

در بسیاری از عملیات ها، به عنوان مثال، matmul، ما فقط نیاز داریم که هر بعد را به یک عامل منفرد نگاشت کنیم. با این حال، برای تغییر شکل کافی نیست.

شکل زیر دو بعد را در یک بعد ادغام می کند:

%out = mhlo.reshape(%in) : (tensor<2x4x32xf32>) -> tensor<8x32xf32>

در اینجا هر دو بعد 0 و 1 ورودی با بعد 0 خروجی مطابقت دارند. فرض کنید با دادن فاکتورهایی به ورودی شروع می کنیم:

(i,j,k) : i=2, j=4, k=32

می بینید که اگر بخواهیم از عوامل یکسانی برای خروجی استفاده کنیم، برای ارجاع چند عامل به یک بعد واحد نیاز داریم:

(i,j,k) -> ((ij), k) : i=2, j=4, k=32

همین کار را می توان انجام داد اگر تغییر شکل یک بعد را تقسیم کند:

%out = mhlo.reshape(%in) : (tensor<8x32xf32>) -> tensor<2x4x32xf32> ((ij), k) -> (i,j,k) : i=2, j=4, k=32

بعد اندازه 8 در اینجا اساساً از عوامل 2 و 4 تشکیل شده است، به همین دلیل است که ما فاکتورها را فاکتور (i,j,k) می نامیم.

این عوامل همچنین می توانند با مواردی که هیچ بعد کاملی که با یکی از عوامل مطابقت دارد وجود ندارد کار کند:

%out = mhlo.reshape(%in) : (tensor<8x4xf32>) -> tensor<2x16xf32> ((ij), k) -> (i,(jk)) : i=2, j=4, k=4

این مثال همچنین تاکید می کند که چرا باید اندازه فاکتورها را ذخیره کنیم - زیرا نمی توانیم به راحتی آنها را از ابعاد مربوطه استنتاج کنیم.

الگوریتم انتشار هسته

تکه تکه ها را در طول فاکتورها منتشر کنید

در شاردی سلسله مراتب تانسورها، ابعاد و فاکتورها را داریم. آنها داده ها را در سطوح مختلف نشان می دهند. یک عامل یک بعد فرعی است. این یک سلسله مراتب داخلی است که در انتشار به اشتراک گذاری استفاده می شود. هر بعد ممکن است با یک یا چند عامل مطابقت داشته باشد. نگاشت بین بعد و عامل توسط OpShardingRule تعریف می شود.

Shardy محورهای شاردینگ را در امتداد فاکتورها به جای ابعاد منتشر می کند . برای انجام این کار، مانند شکل زیر سه مرحله داریم

1. پروژه DimSharding به FactorSharding
2. محورهای شاردینگ را در فضای FactorSharding منتشر کنید
3. برای دریافت DimSharding به روز شده FactorSharding را طراحی کنید

تجسم انتشار شاردینگ در امتداد عوامل

ما از جدول زیر برای تجسم مسئله و الگوریتم انتشار شاردینگ استفاده خواهیم کرد.

• هر ستون نشان دهنده یک عامل است. F0 به معنای عامل با شاخص 0 است. ما تکه تکه ها را در طول فاکتورها (ستون ها) منتشر می کنیم.
• هر ردیف نشان دهنده یک تانسور است. T0 به تانسور با شاخص 0 اشاره دارد. تانسورها همه عملوندها و نتایج مربوط به یک عملیات خاص هستند. محورها در یک ردیف نمی توانند همپوشانی داشته باشند. یک محور (یا محور فرعی) نمی تواند برای تقسیم یک تانسور چندین بار استفاده شود. اگر یک محور صریحاً تکرار شود، نمی‌توانیم از آن برای پارتیشن بندی تانسور استفاده کنیم.

بنابراین، هر سلول نشان دهنده یک اشتراک فاکتور است. یک عامل ممکن است در تانسورهای جزئی وجود نداشته باشد. جدول C = dot(A, B) در زیر آمده است. سلول های حاوی N نشان می دهد که فاکتور در تانسور نیست. به عنوان مثال، F2 در T1 و T2 است، اما در T0 نیست.

محورهای شاردینگ را جمع آوری و منتشر کنید

ما از یک مثال ساده در زیر برای تجسم انتشار استفاده می کنیم.

مرحله 1. محورهایی را برای انتشار در امتداد هر عامل پیدا کنید (با نام مستعار (طولانی ترین) محورهای اصلی شاردینگ سازگار). برای این مثال، ["a", "b"] را در امتداد F0، ["c"] را در امتداد F1، و هیچ چیز را در امتداد F2 منتشر نمی کنیم.

مرحله 2. تقسیم بندی فاکتورها را برای به دست آوردن نتیجه زیر گسترش دهید.