ZeRO এর বিভিন্ন স্তর: Stage 1, Stage 2, এবং Stage 3

ZeRO (Zero Redundancy Optimizer) এর বিভিন্ন স্তর: Stage 1, Stage 2, এবং Stage 3

ZeRO (Zero Redundancy Optimizer) মডেল Training-এর মেমোরি ব্যবস্থাপনা উন্নত করার জন্য বিভিন্ন স্তর বা স্টেজে কাজ করে। প্রতিটি স্তর মডেলের মেমোরি ব্যবহারের একটি নির্দিষ্ট অংশ কমিয়ে Training প্রক্রিয়াকে আরও কার্যকরী করে। ZeRO-এর তিনটি স্তর হলো Stage 1, Stage 2, এবং Stage 3। প্রতিটি স্তর মডেলের বিভিন্ন অংশ যেমন Optimizer State, Gradient এবং Parameter-এর মেমোরি ব্যবস্থাপনা করে। নিচে প্রতিটি স্তর বিস্তারিতভাবে আলোচনা করা হলো।

Stage 1: Optimizer State Sharding

Stage 1 হলো ZeRO-এর প্রথম স্তর, যা মডেলের Optimizer State বিভিন্ন GPU-তে ভাগ করে সংরক্ষণ করে।

কাজের প্রক্রিয়া:

• মডেলের Optimizer State, যেমন ওজন (weights), বায়াস (biases), এবং মোমেন্টাম (momentum) একাধিক GPU তে ভাগ করে রাখা হয়।
• প্রতিটি GPU শুধুমাত্র তার দায়িত্বের Optimizer State সংরক্ষণ করে এবং অন্যান্য GPU-এর সাথে একত্রিত করে মডেল Training চালায়।

সুবিধা:

• Optimizer State শার্ডিংয়ের মাধ্যমে মডেল Training-এর জন্য মেমোরি প্রয়োজনীয়তা অনেক কমে যায়।
• মডেল Training-এর মেমোরি ব্যবস্থাপনা উন্নত হয়, যা GPU মেমোরি কম্পিউটেশনে সহায়ক।

উদাহরণ:

উদাহরণ হিসেবে, একটি বড় ভাষা মডেল Training করার সময় Optimizer State শার্ডিং ব্যবহার করলে GPU মেমোরি ব্যবহারের ৫০% এর বেশি সাশ্রয় করা যায়।

Stage 2: Gradient Sharding

Stage 2 এ Gradient Sharding করা হয়, যেখানে গ্রেডিয়েন্টগুলোকেও Optimizer State-এর মতো শার্ড করা হয়।

কাজের প্রক্রিয়া:

• Stage 2-এ প্রতিটি GPU শুধুমাত্র তার গ্রেডিয়েন্টের একটি নির্দিষ্ট অংশ সংরক্ষণ করে।
• Training প্রক্রিয়ায় বিভিন্ন GPU তাদের গ্রেডিয়েন্ট শেয়ার করে আপডেট করে।

সুবিধা:

• Gradient Sharding-এর ফলে GPU মেমোরির আরও সাশ্রয় হয়, কারণ প্রতিটি GPU-তে সম্পূর্ণ গ্রেডিয়েন্ট স্টোর না করে কেবলমাত্র তার অংশটি রাখা হয়।
• Stage 2 ব্যবহারে মডেল Training-এর জন্য প্রয়োজনীয় GPU মেমোরি আরও কমে আসে এবং Training-এর স্পিডও বৃদ্ধি পায়।

উদাহরণ:

Stage 2-এর Gradient Sharding ব্যবহার করে বড় মডেল যেমন BERT বা GPT Training-এর মেমোরি খরচ আরও কমানো সম্ভব হয়।

Stage 3: Parameter Sharding

Stage 3 হলো ZeRO-এর সবচেয়ে উন্নত স্তর, যেখানে মডেলের প্রতিটি Parameter অর্থাৎ মডেলের ওজন (weights) শার্ড করা হয় এবং GPU তে ভাগ করে রাখা হয়।

কাজের প্রক্রিয়া:

• মডেলের প্রতিটি Parameter (প্যারামিটার) বিভিন্ন GPU তে ভাগ করে রাখা হয়।
• Parameter Sharding করার ফলে GPU-তে পূর্ণ মডেলের Parameter না রেখে কেবলমাত্র তার একটি নির্দিষ্ট অংশ সংরক্ষণ করা হয়।
• Training প্রক্রিয়ায় GPU-গুলো একত্রে কাজ করে এবং Parameter শেয়ার করে Training সম্পন্ন করে।

সুবিধা:

• Stage 3 ব্যবহার করলে মডেলের সবকিছু, অর্থাৎ Optimizer State, Gradient এবং Parameter শার্ড হয়ে GPU তে ভাগ হয়, যা GPU মেমোরি প্রয়োজনীয়তা উল্লেখযোগ্যভাবে কমায়।
• বড় মডেল যেমন GPT-3 বা মেগা-মডেল Training এর জন্য Stage 3 বিশেষভাবে উপযোগী।

উদাহরণ:

Stage 3 ব্যবহার করে বিশাল আকারের মডেল যেমন GPT-3 এবং Turing-NLG মডেল Training করা সহজ হয়, কারণ এটি GPU মেমোরির উপর চাপ কমিয়ে দেয় এবং Training কার্যকর করে।

ZeRO-এর বিভিন্ন স্তরের তুলনা

ZeRO-এর কনফিগারেশন উদাহরণ

DeepSpeed কনফিগারেশন ফাইল ব্যবহার করে ZeRO-এর বিভিন্ন স্তর প্রয়োগ করা যায়। নিচে Stage 1, Stage 2 এবং Stage 3 এর জন্য কনফিগারেশন দেখানো হলো:

{
  "train_batch_size": 32,
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 3
  }
}

এখানে zero_optimization.stage ফিল্ডটি ৩ উল্লেখ করা হয়েছে, যা ZeRO Stage 3 সক্রিয় করে। এভাবে Stage 1 বা Stage 2 নির্ধারণ করতে zero_optimization.stage এর মান ১ বা ২ সেট করতে হবে।

সারসংক্ষেপ

• Stage 1: Optimizer State শার্ডিং করে মডেল Training এর মেমোরি ব্যবস্থাপনা উন্নত করে।
• Stage 2: Gradient Sharding যুক্ত করে, যা আরও বেশি মেমোরি সাশ্রয় করে।
• Stage 3: Parameter Sharding করে মডেলের সব অংশকে GPU-তে ভাগ করে দেয়, যা বড় মডেল Training এর জন্য বিশেষভাবে উপযোগী।

ZeRO-এর এই তিনটি স্তর বড় মডেল Training এর মেমোরি ব্যবস্থাপনাকে কার্যকরী করে এবং GPU মেমোরি সাশ্রয় করে, যা বড় মডেল যেমন GPT-3, BERT Training-এর জন্য অত্যন্ত গুরুত্বপূর্ণ।