واجهة برمجة التطبيقات Compiler API

الخلفية

نفترض أنّ القراء على دراية بأساسيات تمثيل التجزئة على الأقل، والتي تصف كيفية تعبير تجزئة مصفوفة تينسور في Shardy. يوضّح هذا المستند كيفية استخدام تمثيلات التحليل إلى أقسام في برنامج، مثلاً لربط عملية تقسيم إلى أقسام بأحد مصفوفات البرنامج.

إنّ انتشار التجزئة هو عملية تحديد عملية تجزئة لكلّ مصفوفة كثيفة في برنامج معيّن استنادًا إلى قيود التجزئة لمجموعة فرعية من المصفوفات الكثيفة. توفّر واجهة برمجة التطبيقات compiler API في Shardy عدة طرق للتأثير في عملية نشر التجزئة أو التحكّم فيها. بالإضافة إلى ذلك، يسمح للمستخدمين بإدراج عمليات حسابية مجزّأة يدويًا في برامجهم.

الهدف

يصف هذا المستند تصميم مكونات واجهة برمجة التطبيقات هذه في Shardy ويوضّح سلوكها وثوابت الأداء فيها. يُرجى العلم أنّه على الرغم من أنّ واجهة برمجة التطبيقات هذه تُستخدَم للتحكّم في نشر التجزئة، لن تتناول هذه المقالة أيّ شيء عن سلوك النشر أو طريقة تصميمه.

نظرة عامة

• عمليات تقسيم الإدخال/الإخراج: يمكنك إرفاق عملية تقسيم بأحد مدخلات الدالة الرئيسية أو مخرجاتها للإشارة إلى أنّه يجب تقسيم ملف كثيف معالجة الإدخال/الإخراج بهذه الطريقة عند إرساله إلى الدالة أو استرجاعه منها.

• قيد التجزئة: يمكنك إرفاق عملية تجزئة بأحد المتسلسلات الوسيطة (مثل نتيجة matmul) للإشارة إلى أنّه ينبغي تجزئة هذه المتسلسلة أو مجموعة فرعية من استخداماتها.

• مجموعة التجزئة: تجميع مصفوفات متعددة حسب معرّف لتحديد أنّه يجب تقسيمها بالطريقة نفسها

• الحساب اليدوي: يضمّ حسابًا فرعيًا يتم تقسيمه يدويًا باستخدام مجموعة فرعية من محاور الشبكة، حيث يتم تحديد عمليات التجزئة على طول هذه المحاور اليدوية لجميع الإدخالات والمخرجات، وداخل الحساب الفرعي تكون أنواع الن tensors محلية بالنسبة إلى عمليات التجزئة هذه.

التصميم التفصيلي

تقسيم الإدخال/الإخراج

يسمح للمستخدمين بتحديد تقسيم لمدخلات ومخرجات الدالة الرئيسية.

في 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 تأخذ مصفوفة كثافة Tensor كمدخل، وتتضمّن سمة تجزئة. يمكن أن تكون العملية:

• أن تكون غير مستخدَمة (غير مرتبطة) - ما يعني أنّ التقسيم المرفق هو الطريقة التي يجب بها تقسيم المتجه نفسه
• أن يكون لها استخدامات: ويعني ذلك أنّ التقسيم المرفق هو الطريقة التي يجب تقسيم استخدامات عملية قيد التقسيم بها، في حين أنّ استخدامات ناقلات المصفوفات المدخلة الأخرى قد يكون لها تقسيم مختلف (إذا لم تكن ناقلات المصفوفات المدخلة لها استخدامات أخرى، يكون السلوك هو نفسه في حالة عدم توفّر أي استخدامات). سيحدِّد التوسّع عملية تقسيم مصفوفة السلاسل المتعدّدة ذاتها ويعيد تقسيمها إذا لزم الأمر.

يمكن أن تتضمّن تقسيمات سمات مفتوحة، ما يعني أنّه يمكن تقسيم الم Operand بشكلٍ إضافي على طول المحاور المتاحة.

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

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

مجموعة التجزئة

في الحالات التي لا تتوفّر فيها أيّ تبعيات للبيانات أو تبعيات قوية للبيانات بين مصفوفة تينسور أو أكثر، مع معرفة المستخدمين بأنّه يجب تقسيم مصفوفات التنسور بالطريقة نفسها أو بطرق مشابهة، توفّر واجهة برمجة التطبيقات Shardy API طريقة لتحديد هذه العلاقة. يمنحك هذا حرية تحديد بوضوح أنّه يجب تقسيم مصفوفات Tensor على نحو متطابق.

لتحقيق ذلك، نقدّم مفهوم مجموعات الأجزاء، حيث تحتوي كل مجموعة على أي عدد من التعليمات المرتبطة بمعرّف مجموعة الأجزاء نفسه. تفرض مجموعات التجزئة أن تكون عمليات التجزئة ضمن المجموعة نفسها متماثلة.

على سبيل المثال، في برنامج افتراضي للمستخدِم مثل ما هو موضّح أدناه، نريد تقسيم مخرجات البرنامج بالطريقة نفسها المستخدَمة في إدخال البرنامج في حال عدم توفّر أيّ تداخلات بين البيانات.

في حال تنفيذ هذا البرنامج، لن تتمكّن عملية نشر التجزئة من الاستنتاج بشأن تجزئة مصفوفات %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>
  }
}

في هذا المثال البسيط أعلاه، كان بإمكاننا بدلاً من ذلك تحديد القسمة نفسها في الإخراج كما في الإدخال، ما سيؤدي إلى التأثير نفسه، لأنّنا سبق أن عرفنا الشريحة التي نريد تعيينها للإدخال قبل الوقت، ولكن في الحالات الأكثر واقعية، نستخدم الشريحة للحفاظ على توحيد sharding لعناصر متعددة بدون معرفة الشريحة لأي منها بالضرورة، بينما سيتولى Shardy الباقي ويجد أفضل تقسيم لتخصيصه.

الحساب اليدوي

قد يريد المستخدمون التحكّم بشكل صريح في كيفية تقسيم أجزاء من عمليات الحساب التي يجرونها، والمجموعات التي يتم استخدامها. على سبيل المثال، يريد بعض المستخدمين تطبيق matmul المجمّع يدويًا (من واجهة برمجة التطبيقات للواجهة الأمامية) بدلاً من تأخير التحويل إلى المُجمِّع. نوفّر واجهة برمجة تطبيقات Manual Computation 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. منطقة/جزء الحساب هو الحساب المحلي (مثلاً، بما في ذلك المجموعات التي يحدّدها المستخدِم). يجب أن يكون محليًا في ما يتعلّق بتقسيم البيانات إلى مجموعات في/خارج الإطار على طول المحاور اليدوية (راجِع الملاحظة أعلاه).

تداخل العمليات الحسابية اليدوية

يمكنك إدراج عمليات حسابية يدوية متعددة داخل بعضها ما دام كلّ منها يعمل على مجموعة فريدة من المحاور اليدوية.