বিভাগ: রানটাইম: প্রোগ্রাম বরাদ্দ ব্যর্থতা
এই ত্রুটিটি নির্দেশ করে যে, একটি TPU ডিভাইসের XLA রানটাইম একটি কম্পাইল করা XLA প্রোগ্রাম এক্সিকিউটেবলকে TPU-এর HBM-এ লোড করতে ব্যর্থ হয়েছে।
নমুনা ত্রুটি বার্তা:
XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).
XLA ব্যাকএন্ড: TPU
সংক্ষিপ্ত বিবরণ
এই ত্রুটিটি সাধারণত নিম্নলিখিত কারণগুলির মধ্যে কোনো একটির কারণে ঘটে থাকে:
- প্রোগ্রামের আকার উপলব্ধ HBM অতিক্রম করেছে: কম্পাইল করা XLA প্রোগ্রামটি, তার নির্দেশাবলী, স্ট্যাটিক ডেটা এবং যেকোনো এমবেডেড কনস্ট্যান্ট সহ, সেই নির্দিষ্ট TPU কোর(গুলি)তে বর্তমানে উপলব্ধ মোট ফ্রি HBM-এর পরিমাণের চেয়ে বড়, যেখানে প্রোগ্রামটি লোড করা হচ্ছে।
- এইচবিএম খণ্ডীকরণ: যদিও ডিভাইসে থাকা মোট খালি এইচবিএম সমষ্টিগতভাবে পর্যাপ্ত হতে পারে, তবে এটি সম্পূর্ণ প্রোগ্রামটি ধারণ করার মতো যথেষ্ট বড় একটি একক, অবিচ্ছিন্ন ব্লক হিসেবে উপলব্ধ নয়।
টিপিইউ রানটাইম কীভাবে মেমরিকে অগ্রাধিকার দেয়, তা বোঝা গুরুত্বপূর্ণ। লোড হওয়া প্রোগ্রামের চেয়ে বাফার অ্যালোকেশনকে বেশি অগ্রাধিকার দেওয়া হয়। যদি কোনো বাফার অ্যালোকেশন ব্যর্থ হয়, তাহলে জায়গা খালি করার জন্য রানটাইম এইচবিএম (HBM) থেকে ইতিমধ্যে লোড হওয়া প্রোগ্রামগুলোকে সরিয়ে দেয়। এর ফলে এমন পরিস্থিতি তৈরি হতে পারে যেখানে আগে সফলভাবে লোড হওয়া একটি প্রোগ্রাম এখন ওওএম (OOM) ত্রুটির কারণে ব্যর্থ হয়, কারণ এইচবিএম এখন আরও বেশি ডেটা বাফার দ্বারা পূর্ণ হয়ে গেছে।
ডিবাগিং
- বাফার মেমরি ফুটপ্রিন্ট হ্রাস করুন: ডেটা বাফার দ্বারা ব্যবহৃত মেমরি খালি করলে প্রোগ্রামটির জন্য আরও জায়গা পাওয়া যাবে।
- ব্যাচ সাইজ কমান: অ্যাক্টিভেশনের জন্য ব্যবহৃত মেমরির পরিমাণ কমানোর এটি অন্যতম কার্যকর উপায়।
- প্যারামিটার শার্ডিং: খুব বড় মডেলের ক্ষেত্রে, মডেলের প্যারামিটার এবং গণনা একাধিক TPU কোর বা হোস্টের মধ্যে বন্টন করার জন্য মডেল প্যারালেলিজম বা শার্ডিং কৌশল (যেমন FSDP বা মেগাস্কেল) ব্যবহার করুন। মনে রাখবেন যে, একাধিক চিপের মধ্যে টেনসর ভাগ করার কারণে এটি নেটওয়ার্ক কমিউনিকেশন ওভারহেড বাড়িয়ে দিতে পারে।
- সিকোয়েন্স/কন্টেক্সটের দৈর্ঘ্য কমানো: যেসব মডেল সিকোয়েন্সিয়াল ডেটা প্রসেস করে (যেমন, এনএলপি মডেল), তাদের ক্ষেত্রে সিকোয়েন্সের দৈর্ঘ্য কমালে মেমোরি ব্যবহার উল্লেখযোগ্যভাবে হ্রাস পেতে পারে।
- বাফার ডোনেশন: ফ্রেমওয়ার্কের ফিচার (যেমন,
jax.jit(..., donate_argnums=...)) ব্যবহার করে XLA-কে আউটপুট সংরক্ষণের জন্য ইনপুট বাফারের মেমরি পুনরায় ব্যবহার করার সুযোগ দিন, যা সর্বোচ্চ মেমরি ব্যবহার কমিয়ে আনে।
- প্রোগ্রামের টেম্পোরারি ভেরিয়েবলের জন্য প্রয়োজনীয় মেমরি হ্রাস করুন।
-
tpu_shared_memory_percentফ্ল্যাগ ব্যবহার করে টেম্পোরারি ভেরিয়েবলের জন্য প্রোগ্রামের মেমরি ব্যবহার কমান। মনে রাখবেন, এটি পারফরম্যান্সে নেতিবাচক প্রভাব ফেলতে পারে।
-
- কার্য সম্পাদনের কৌশল উন্নত করুন/পরিষেবার চাপ কমান
- প্রোগ্রাম লোডিং পরিচালনা করুন: আপনি যদি একাধিক ফাংশন JIT-কম্পাইল করেন, তবে মনে রাখবেন যে প্রতিটি ফাংশনের ফলে একটি প্রোগ্রাম লোড হতে পারে। একই সাথে লোড হওয়া প্রোগ্রামের সংখ্যা সর্বনিম্ন রাখতে আপনার ওয়ার্কলোড সাজানোর চেষ্টা করুন।
- মেমরি লিক যাতে না হয় তা নিশ্চিত করুন।
-
jax.Arrayঅবজেক্টের রেফারেন্সগুলো যেন উদ্দেশ্যের চেয়ে বেশি সময় ধরে রাখা না হয়, তা নিশ্চিত করুন।jax.Arrayঅবজেক্ট ধরে রাখলে, প্রোগ্রাম কম্পাইলেশন সম্পন্ন হওয়ার পরেও স্বয়ংক্রিয় ডি-অ্যালোকেশন বাধাগ্রস্ত হতে পারে।
-
প্রতিটি প্রোগ্রামের মেমরি ব্যবহার কমানোর জন্য আপনি আরও যেসব কৌশল অবলম্বন করতে পারেন, তার জন্য এরর কোড: E1000 দেখুন।
টুলিং
-
tpu_log_allocations_on_oomফ্ল্যাগটি সক্রিয় করুন, যার ফলে OOM ঘটলে অ্যালোকেটরটি সমস্ত চলমান অ্যালোকেশনের একটি বিস্তারিত রিপোর্ট ডাম্প করবে, যা ডিবাগিংয়ের জন্য অত্যন্ত মূল্যবান হতে পারে। - আপনার প্রোগ্রামের প্রোফাইল তৈরি করুন: সময়ের সাথে সাথে আপনার প্রোগ্রামের মেমরি ব্যবহারের একটি বিস্তারিত চিত্র পেতে JAX মেমরি প্রোফাইলার অথবা TensorFlow প্রোফাইলার ব্যবহার করুন। এটি মেমরি ব্যবহারের অপ্রত্যাশিত বৃদ্ধি শনাক্ত করতে সাহায্য করতে পারে।