কম্পাইলার API

পটভূমি

আমরা অনুমান করি পাঠকরা অন্ততপক্ষে শার্ডিং উপস্থাপনার মূল বিষয়গুলির সাথে পরিচিত, যা বর্ণনা করে যে কীভাবে একটি টেনসরের শার্ডিং শার্ডিতে প্রকাশ করা যেতে পারে। এই নথিটি দেখায় যে কীভাবে একটি প্রোগ্রামে শার্ডিং উপস্থাপনাগুলি ব্যবহার করা যেতে পারে, যেমন প্রোগ্রামের একটি নির্দিষ্ট টেনসরের সাথে একটি শার্ডিং সংযুক্ত করতে।

টেনসরের একটি উপসেটের জন্য শার্ডিং সীমাবদ্ধতা প্রদত্ত একটি প্রোগ্রামে প্রতিটি টেনসরের জন্য একটি শার্ডিংয়ের সিদ্ধান্ত নেওয়ার প্রক্রিয়া হল শার্ডিং প্রচার। শার্ডি'স কম্পাইলার API শার্ডিং প্রচারকে প্রভাবিত/নিয়ন্ত্রিত করার বিভিন্ন উপায় প্রকাশ করে। উপরন্তু এটি ব্যবহারকারীদের তাদের প্রোগ্রামে ম্যানুয়ালি শার্ড করা গণনা সন্নিবেশ করার অনুমতি দেয়।

উদ্দেশ্য

এই ডক শার্ডিতে এই জাতীয় API উপাদানগুলির নকশা বর্ণনা করে এবং তাদের আচরণ এবং পরিবর্তনগুলি ব্যাখ্যা করে৷ মনে রাখবেন যে যখন এই APIটি শার্ডিং প্রচার নিয়ন্ত্রণ করতে ব্যবহৃত হয়, এই ডকটি প্রচারের আচরণ বা এটি কীভাবে ডিজাইন করা হয়েছে সে সম্পর্কে কিছু আলোচনা করতে যাচ্ছে না ।

ওভারভিউ

• ইনপুট/আউটপুট শার্ডিংস - প্রধান ফাংশনের একটি ইনপুট বা আউটপুটে একটি শার্ডিং সংযুক্ত করুন, যাতে নির্দেশ করা যায় যে ফাংশন থেকে দেওয়া-কে/রিটার্ন করার সময় ইনপুট/আউটপুট টেনসরকে এভাবে শার্ড করা উচিত।

• শার্ডিং সীমাবদ্ধতা - একটি মধ্যবর্তী টেনসরের সাথে একটি শার্ডিং সংযুক্ত করুন (যেমন একটি মাটমুলের ফলাফল) বোঝাতে যে এইভাবে সেই টেনসর, বা এর ব্যবহারগুলির একটি উপসেটকে শার্ড করা উচিত।

• শার্ডিং গ্রুপ - একটি আইডি দ্বারা একাধিক টেনসরকে গোষ্ঠীবদ্ধ করে নির্দেশ করে যে সেগুলিকে একইভাবে শার্ড করা উচিত।

• ম্যানুয়াল কম্পিউটেশন - একটি সাব-কম্পিউটেশন ঘেরা যা জাল অক্ষের একটি উপসেট ব্যবহার করে ম্যানুয়ালি পার্টিশন করা হয়, যেখানে সেই ম্যানুয়াল অক্ষ বরাবর শার্ডিংগুলি সমস্ত ইনপুট এবং আউটপুটগুলির জন্য নির্দিষ্ট করা হয় এবং সাব-কম্পিউটেশনের ভিতরে টেনসরের ধরনগুলি সেই শার্ডিংগুলির জন্য স্থানীয়।

বিস্তারিত ডিজাইন

ইনপুট/আউটপুট শার্ডিং

ব্যবহারকারীদের প্রধান ফাংশনের ইনপুট এবং আউটপুটগুলির জন্য একটি শার্ডিং নির্দিষ্ট করার অনুমতি দেয়৷

এমএলআইআর-এ, বৈশিষ্ট্যগুলি ফাংশন আর্গুমেন্ট এবং ফলাফলের সাথে সংযুক্ত করা যেতে পারে এবং তাই ব্যবহারকারীরা এইভাবে ফাংশনের সাথে শার্ডিং বৈশিষ্ট্য সংযুক্ত করতে পারে।

যেমন:

@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>

শেয়ারিং গ্রুপ

যে ক্ষেত্রে দুটি বা ততোধিক টেনসরের মধ্যে কোন ডেটা নির্ভরতা বা কোন শক্তিশালী ডেটা নির্ভরতা নেই, যখন ব্যবহারকারীরা জানেন যে সেই টেনসরগুলিকে একই বা একইভাবে বিভাজন করা উচিত, Shardy API এই সম্পর্কটি নির্দিষ্ট করার একটি উপায় অফার করে। এটি ব্যবহারকারীদের স্পষ্টভাবে উল্লেখ করার স্বাধীনতা দেয় যে টেনসরগুলি একে অপরের মতো বিভাজন করা উচিত।

এটি অর্জনের জন্য, আমরা শার্ড গ্রুপের একটি ধারণা প্রবর্তন করি, যেখানে প্রতিটি গ্রুপে একই শার্ড গ্রুপ আইডির সাথে সম্পর্কিত যেকোন সংখ্যক নির্দেশাবলী রয়েছে। শেয়ারিং গোষ্ঠী একই গ্রুপের মধ্যে শার্ডিংগুলিকে একই রকম হওয়ার জন্য প্রয়োগ করে।

উদাহরণস্বরূপ, একটি অনুমানমূলক ব্যবহারকারী প্রোগ্রামে যেমন নীচে দেখানো হয়েছে, আমরা প্রোগ্রামের আউটপুটকে প্রোগ্রামের ইনপুটের মতো ঠিক একই রকম করতে চাই যখন দুটির মধ্যে কোনও ডেটা নির্ভরতা নেই।

যদি আমরা এই প্রোগ্রামটি চালাই, শার্ডিং প্রচার টেনসর %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>
  }
}

উপরের এই সাধারণ উদাহরণে, বিকল্পভাবে আমরা ইনপুট হিসাবে আউটপুটে একই শার্ডিং স্পষ্টভাবে নির্দিষ্ট করতে পারতাম, যা একই প্রভাব অর্জন করবে, যেহেতু আমরা ইতিমধ্যেই জেনেছি যে আমরা আগে থেকেই ইনপুটটিতে কোন শার্ড বরাদ্দ করতে চাই কিন্তু আরও বাস্তবসম্মত ক্ষেত্রে, আমরা একাধিক টেনসরের শার্ডিংকে সিঙ্কে রাখার জন্য শার্ড ব্যবহার করি, অগত্যা সেগুলির কোনওটির জন্য শার্ডিং না জেনেই, শার্ডি বাকিগুলির যত্ন নেবে এবং খুঁজে পাবে তাদের বরাদ্দ করার জন্য সেরা শার্ডিং।

ম্যানুয়াল কম্পিউটেশন

ব্যবহারকারীরা তাদের গণনার অংশগুলিকে কীভাবে বিভাজন করা হয় এবং কোন সমষ্টি ব্যবহার করা হয় সে সম্পর্কে স্পষ্ট নিয়ন্ত্রণ চাইতে পারে। উদাহরণ স্বরূপ, কিছু ব্যবহারকারী কম্পাইলারকে স্থগিত করার পরিবর্তে ম্যানুয়ালি (ফ্রন্টএন্ড API থেকে) সমষ্টিগত ম্যাটমুল প্রয়োগ করতে চান। আমরা একটি ম্যানুয়াল কম্পিউটেশন API প্রদান করি যা তাদের এটি করতে দেয়।

ম্যানুয়াল সাব-কম্পিউটেশনের জন্য এটি একটি একক অঞ্চলের সাথে MLIR অপারেশন। ব্যবহারকারীরা জাল অক্ষের একটি উপসেট (সম্ভবত সমস্ত সহ) ব্যবহার করে এই সাব-কম্পিউটেশনে ইনপুট/আউটপুট শার্ডিংগুলি নির্দিষ্ট করবে। সাব-কম্পিউটেশন হবে স্থানীয়/ম্যানুয়াল হবে নির্দিষ্ট জাল অক্ষ (ওরফে ম্যানুয়াল অক্ষ), এবং গ্লোবাল/বিভাজনবিহীন wrt অনির্দিষ্ট (ওরফে ফ্রি অক্ষ)। এই অপারেশনের বাইরের গণনা যেভাবে হতে পারে সেভাবে প্রচারের সময় মুক্ত অক্ষ বরাবর সাব-কম্পিউটেশনকে আরও শার্ড করা যেতে পারে।

যেমন:

@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. গণনার অঞ্চল/বডি হল স্থানীয় গণনা (যেমন, ব্যবহারকারীর নির্দিষ্ট সমষ্টি সহ)। ম্যানুয়াল অক্ষের সাথে ইন/আউট শার্ডিংয়ের সাথে এটি অবশ্যই স্থানীয় হতে হবে (উপরের নোট দেখুন)।

নেস্টিং ম্যানুয়াল গণনা

আপনি একাধিক ম্যানুয়াল কম্পিউটেশনকে একে অপরের মধ্যে নেস্ট করতে পারেন যতক্ষণ না প্রত্যেকটি তাদের নিজস্ব অনন্য ম্যানুয়াল অক্ষের সেটে কাজ করে।