API کامپایلر

پس زمینه

ما فرض می‌کنیم که خوانندگان حداقل با اصول اولیه نمایش تقسیم‌بندی آشنا هستند، که توضیح می‌دهد چگونه تقسیم یک تانسور می‌تواند در Shardy بیان شود. این سند نشان می‌دهد که چگونه می‌توان از نمایش‌های اشتراک‌گذاری در یک برنامه استفاده کرد، به عنوان مثال برای اتصال یک تقسیم‌بندی به یک تانسور خاص برنامه.

انتشار شاردینگ فرآیند تصمیم گیری در مورد تقسیم بندی برای هر تانسور در یک برنامه با توجه به محدودیت های تقسیم بندی برای زیر مجموعه ای از تانسورها است. API کامپایلر Shardy راه‌های مختلفی را برای تأثیر/کنترل انتشار اشتراک‌گذاری نشان می‌دهد. علاوه بر این، به کاربران اجازه می دهد تا محاسبات خرد شده دستی را در برنامه های خود قرار دهند.

هدف

این سند طراحی چنین مولفه‌های API را در Shardy توصیف می‌کند و رفتار و تغییرات ثابت آنها را توضیح می‌دهد. توجه داشته باشید که در حالی که این API برای کنترل انتشار به اشتراک گذاری استفاده می شود، این سند قرار نیست درباره رفتار انتشار و نحوه طراحی آن بحث کند.

نمای کلی

• اشتراک گذاری های ورودی/خروجی - یک شاردینگ را به ورودی یا خروجی تابع اصلی وصل کنید تا نشان دهد که تانسور ورودی/خروجی باید در هنگام داده شدن به/برگرداندن از تابع، به این ترتیب تقسیم شود.

• محدودیت Sharding - یک شاردینگ را به یک تانسور میانی (مثلاً نتیجه یک matmul) وصل کنید تا نشان دهد که این تانسور یا زیرمجموعه ای از کاربردهای آن باید به این ترتیب تقسیم شود.

• گروه شاردینگ - تانسورهای متعدد را با یک شناسه گروه بندی کنید تا نشان دهید که آنها باید به همان روش خرد شوند.

• محاسبات دستی - محاسبات فرعی را در بر می گیرد که به صورت دستی با استفاده از زیرمجموعه ای از محورهای مش پارتیشن بندی می شود، جایی که تقسیم بندی ها در امتداد آن محورهای دستی برای همه ورودی ها و خروجی ها مشخص شده است، و در داخل محاسبات فرعی، انواع تانسورها به صورت محلی در آن تقسیم بندی ها هستند.

طراحی دقیق

اشتراک گذاری ورودی/خروجی

به کاربران اجازه می‌دهد تا برای ورودی‌ها و خروجی‌های تابع اصلی یک اشتراک گذاری مشخص کنند.

در MLIR، ویژگی‌ها را می‌توان به آرگومان‌ها و نتایج تابع متصل کرد و بنابراین کاربران می‌توانند از این طریق ویژگی‌های اشتراک‌گذاری را به تابع متصل کنند.

به عنوان مثال:

@mesh_xy = <["x"=2, "y"=2]>

// The 1st input has a sharding specified, but the 2nd input doesn't.
// The output has a sharding specified.
func @main(%arg0: tensor<8x8xf32>
            {sdy.sharding = #sdy.sharding<@mesh_xy, [{"x"}, {}]>},
            %arg1: tensor<8x16xf32>)
    -> (tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_xy, [{}, {"y"}]>}) {
  ...
}

محدودیت شاردینگ

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

این یک عملیات MLIR است که تانسور را به عنوان ورودی می گیرد و یک ویژگی اشتراک گذاری به آن متصل است. این عملیات می تواند:

• بدون استفاده (آویزان) - به این معنی که شاردینگ متصل به این صورت است که خود تانسور باید خرد شود.
• دارای کاربردها باشد - به این معنی که تقسیم بندی متصل به این صورت است که چگونه استفاده از عملیات محدودیت تقسیم بندی باید تقسیم شود، در حالی که سایر کاربردهای تانسور ورودی ممکن است تقسیم بندی متفاوتی داشته باشند (اگر تانسور ورودی کاربرد دیگری نداشته باشد، رفتار مشابه با بدون مورد استفاده). انتشار، تقسیم بندی خود تانسور را تعیین می کند و در صورت لزوم آن را مجدداً تغییر می دهد.

می‌تواند دارای تقسیم‌بندی‌های بعد باز باشد، به این معنی که عملوند را می‌توان در امتداد محورهای موجود بیشتر خرد کرد.

@mesh_xy = <["x"=2, "y"=2]>

%0 = ... : tensor<8x8xf32>
%1 = sdy.sharding_constraint %0 <@mesh_xy, [{"x"}, {?}]> : tensor<8x8xf32>

گروه شاردینگ

در مواردی که هیچ وابستگی داده ای وجود ندارد یا وابستگی داده ای قوی بین دو یا چند تانسور وجود ندارد، در حالی که کاربران می دانند که آن تانسورها باید به روش های مشابه یا مشابه تقسیم شوند، API Shardy راهی برای مشخص کردن این رابطه ارائه می دهد. این به کاربران این آزادی را می دهد که به صراحت مشخص کنند که تانسورها باید به عنوان یکدیگر تقسیم شوند.

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

به عنوان مثال، در یک برنامه کاربری فرضی مانند شکل زیر، ما می خواهیم خروجی برنامه را دقیقاً مشابه ورودی برنامه تقسیم کنیم در حالی که هیچ وابستگی داده ای بین این دو وجود ندارد.

اگر این برنامه را اجرا کنیم، انتشار شاردینگ قادر به استنتاج به اشتراک گذاری تانسورهای %1 و %2 نخواهد بود و در نهایت تکرار می شوند. با این حال، با پیوست کردن یک ویژگی shard_group که می‌گوید ورودی %0 و خروجی %2 در یک shard_group هستند، به اشتراک گذاری @mesh_xy, [{"x"},{"y"}]> اجازه می‌دهیم تا از ورودی %0 به خروجی %2 و به نوبه خود به بقیه نمودار که ثابت %1 در اینجا پخش می شود. ما می توانیم با عملیات sdy.sharding_group یک مقدار به یک گروه اختصاص دهیم.

@mesh_xy = <["x"=2, "y"=2]>

module @"jit_zeros_like" {
  func.func @main(%arg0: tensor<8x2xi64> {sdy.sharding = #sdy.sharding<@mesh_xy, [{"x"},{"y"}]>} }) -> (tensor<8x2xi64>) {
    %0 = sdy.sharding_group %arg0, id=0 : tensor<8x2xi64>
    %1 = stablehlo.constant dense<0> : tensor<8x2xi64>
    %2 = sdy.sharding_group %1, id=0 : tensor<8x2xi64>
    return %2 : tensor<8x2xi64>
  }
}

در این مثال ساده در بالا، به طور متناوب، می‌توانیم به‌صراحت همان تقسیم‌بندی را در خروجی به‌عنوان ورودی مشخص کنیم، که همان اثر را به دست می‌آورد، زیرا ما قبلاً می‌دانیم چه قطعه‌ای را می‌خواهیم به ورودی اختصاص دهیم، اما در در موارد واقع بینانه تر، ما از shard برای همگام نگه داشتن شاردینگ چندین تانسور بدون دانستن شاردینگ برای هیچ یک از آنها استفاده می کنیم، در حالی که Shardy از بقیه مراقبت می کند و بهترین تقسیم بندی را برای اختصاص به آنها پیدا می کند.

محاسبات دستی

کاربران ممکن است خواهان کنترل صریح در مورد نحوه تقسیم بندی بخش هایی از محاسبات آنها و استفاده از چه مجموعه هایی باشند. به عنوان مثال، برخی از کاربران می خواهند به جای تعویق به کامپایلر، matmul جمعی را به صورت دستی (از API frontend) اعمال کنند. ما یک API محاسبات دستی ارائه می دهیم که به آنها اجازه می دهد این کار را انجام دهند.

این عملیات MLIR با یک منطقه واحد برای محاسبه فرعی دستی است. کاربران با استفاده از یک زیرمجموعه (شامل احتمالاً همه) محورهای مش، تقسیم‌بندی‌های ورودی/خروجی را برای این محاسبه فرعی مشخص می‌کنند. محاسبات فرعی محلی/دستی با محورهای مش مشخص شده (معروف به محورهای دستی) و جهانی/تقسیم نشده محورهای نامشخص (معروف به محورهای آزاد) خواهد بود. محاسبات فرعی را می توان بیشتر در امتداد محورهای آزاد در طول انتشار به همان روشی که محاسبات خارج از این عملیات می تواند تقسیم کرد.

به عنوان مثال:

@mesh_name = <["data"=2, "model"=2]>

%0 = ... : tensor<16x32xf32>
%1 = sdy.manual_computation(%0)
    in_shardings=[<@mesh_name, [{"data"}, {"model",?}]>]
    out_shardings=[<@mesh_name, [{"data"}, {?}]>]
    manual_axes={"data"}
    (%arg1: tensor<8x32xf32>) {
  // body
  return %42 : tensor<8x32xf32>
} : (tensor<16x32xf32>) -> tensor<16x32xf32>

متغیرها

1. همه in_shardings ، out_shardings و manual_axes باید به یک مش اشاره داشته باشند. manual_axes بر اساس مش مرتب شده است.

2. manual_axes باید به طور صریح در همه تقسیم‌بندی‌های درون/خارج استفاده شوند، به‌عنوان مثال، برای هر تقسیم‌بندی، تمام محورهای دستی باید یا یک بعد را تقسیم کنند یا به صراحت تکرار شوند.

3. اگر یک محور آزاد (هر محور مش که در manual_axes نباشد) در یکی از تقسیم‌بندی‌های درون/خروجی وجود داشته باشد، باید نسبت به هر محور دستی در همان ابعاد کوچک‌تر باشد (در مثال بالا، یک بعد تقسیم‌بندی {"model", "data"} ​​نامعتبر خواهد بود).

4. ناحیه/ بدنه محاسبات، محاسبات محلی است (به عنوان مثال، از جمله مجموعه های مشخص شده توسط کاربر). باید به صورت محلی باشد و در امتداد محورهای دستی به صورت محلی باشد (به یادداشت بالا مراجعه کنید).

تودرتو محاسبات دستی

شما می توانید چندین محاسبات دستی را درون یکدیگر قرار دهید تا زمانی که هر یک بر روی مجموعه منحصر به فرد خود از محورهای دستی عمل کنند.