شاردی

Shardy یک سیستم پارتیشن بندی تانسور مبتنی بر MLIR برای همه گویش ها است. ساخته شده از همکاری هر دو تیم GSPMD و PartIR ، بهترین های هر دو سیستم و تجربه مشترک تیم ها و کاربران را در خود جای داده است.

مزایا

• کنترل و قابلیت پیش بینی بیشتر برای کاربران با ترکیب انتشار GSPMD با پارتیشن بندی افزایشی PartIR.
• ویژگی‌های جدید ناشی از تجربه مشترک، به‌عنوان مثال، پشتیبانی جدید از تغییر شکل‌هایی که به‌طور بدنامی ارتباط بیشتری ایجاد می‌کنند، مگر اینکه کاربران بدانند چگونه با آنها کار کنند.
• قابلیت استفاده و اشکال زدایی بهتر برای افزایش سرعت کاربر نهایی، به عنوان مثال با استفاده از نمایش اشتراک گذاری مبتنی بر محور.
• یک پایگاه کد منبع باز ساده با استفاده از MLIR، با مجموعه گسترده تری از مشارکت کنندگان فعال (داخلی، خارجی و در مناطق زمانی مختلف) برای پشتیبانی از کاربران.

اجزاء

• نمایش شاردینگ: یک نمایش تقسیم بندی مبتنی بر محور که به یک مش منطقی خاص (خارج از مش های بالقوه چندگانه) متصل است و از تقسیم بندی ها و محورهای محدود کننده ابعاد، محورهای تقسیم برای عملیات هایی مانند تغییر شکل، اولویت ها برای پارتیشن بندی افزایشی و موارد دیگر پشتیبانی می کند.
• APIهای کامپایلر: مجموعه‌ای از مؤلفه‌های کامپایلر که می‌توانند در کنار نمایش شاردینگ برای تأثیرگذاری بر انتشار شاردینگ استفاده شوند.
  • اشتراک گذاری های ورودی/خروجی - یک شاردینگ را به ورودی یا خروجی تابع اصلی وصل کنید تا نشان دهد که تانسور ورودی/خروجی باید در هنگام داده شدن به/برگرداندن از تابع، به این ترتیب تقسیم شود.
  • محدودیت Sharding - یک شاردینگ را به یک تانسور میانی (مثلاً نتیجه یک matmul) وصل کنید تا نشان دهد که این تانسور یا زیرمجموعه ای از کاربردهای آن باید به این ترتیب تقسیم شود.
  • Shard As/Like - تانسورهای متعدد را با یک شناسه گروه بندی کنید تا نشان دهید که آنها باید به همان روش خرد شوند.
  • محاسبات دستی - محاسبات فرعی را در بر می گیرد که به صورت دستی با استفاده از زیرمجموعه ای از محورهای مش پارتیشن بندی می شود، که در آن تقسیم بندی ها در امتداد آن محورهای دستی برای همه ورودی ها و خروجی ها مشخص شده است، و در داخل محاسبات فرعی، انواع تانسورهای محلی با آن تقسیم بندی ها هستند.
• Sharding Propagation: یک الگوریتم انتشار که اولویت های کاربر و محدودیت های اشتراک گذاری را با مدل های هزینه کامپایلر و اکتشافی ترکیب می کند:
  • اولویت های تعریف شده توسط کاربر، به عنوان مثال انجام موازی سازی دسته ای و سپس صفر
  • اولویت‌های مبتنی بر عملیات، به‌عنوان مثال، ابتدا عملیات‌های عنصری و سپس matmuls و غیره.
  • اکتشافی های ریزدانه بیشتر، به عنوان مثال، ابعاد دسته ای را ترجیح می دهند.
• پارتیشن‌کننده SPMD: مؤلفه‌ای است که با پارتیشن‌بندی برنامه به یک برنامه SPMD، افزودن حرکت/قالب‌بندی داده‌ها و عملیات جمعی در فرآیند، تصمیمات انتشار اشتراکی را کاهش می‌دهد.
  • کوتاه مدت، اجرای اولیه از پارتیشن‌کننده فعلی GSPMD SPMD استفاده می‌کند.
  • دراز مدت، ما قصد داریم یک پارتیشن‌کننده SPMD جدید مبتنی بر MLIR ایجاد کنیم.

مخزن کد

پروژه Shardy در حال توسعه فعال است و ما به دنبال بازخورد از جامعه منبع باز هستیم. کد Shardy در https://github.com/openxla/shardy موجود است