বিভাগ: কম্পাইল সময়: HBM OOM
এই ত্রুটিটি ইঙ্গিত দেয় যে প্রোগ্রামটির জন্য TPU ডিভাইসে শারীরিকভাবে উপলব্ধ হাই ব্যান্ডউইথ মেমরি (HBM) এর চেয়ে বেশি প্রয়োজন।
নমুনা ত্রুটি বার্তা:
RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.
RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G
XLA ব্যাকএন্ড: TPU
সংক্ষিপ্ত বিবরণ
XLA পরীক্ষা করে নিশ্চিত করে যে সমস্ত প্রয়োজনীয় স্ট্যাটিক বরাদ্দের সামগ্রিক আকার ডিভাইসের HBM-এ ফিট করে।
কম্পাইলারটি বিভিন্ন ধরণের বরাদ্দের জন্য TPU-এর স্থির HBM ক্ষমতা পরিচালনা করে:
- প্রোগ্রামের ইনপুট এবং আউটপুট: প্রশিক্ষণ ব্যাচ, অপ্টিমাইজার অবস্থা ইত্যাদি।
- TensorCore + SparseCore Temporaries: মধ্যবর্তী গণনার জন্য প্রয়োজনীয় গতিশীল মেমরি (যেমন অ্যাক্টিভেশন, গ্রেডিয়েন্ট ইত্যাদি)।
- কম্পাইলড বাইনারি: টেনসরকোর (TC) এবং স্পার্সকোর (SC) উভয়ের জন্য মেশিন কোড।
- সিস্টেম ওভারহেড: XLA রানটাইমের জন্য সংরক্ষিত স্থান (যেমন পুরোনো TPU প্রজন্মের ইনফিড বাফার)।
- ধ্রুবক: HLO IR-এ এমবেড করা ধ্রুবক মানগুলি HBM-এ বরাদ্দ করা হয়।
- কম্পাইলার ইন্টার্নাল: প্রোগ্রাম লেভেল এবং প্রতি-এইচএলও বরাদ্দ (যেমন মেশে নোডের জন্য রাউটিং তথ্য)
এই ত্রুটিটি তখন ঘটে যখন XLA কম্পাইলার উপরের সমস্ত বরাদ্দ HBM ডিভাইসে ফিট করতে পারে না।
ডিবাগিং
নিচের কোন বিভাগটি আপনার ত্রুটিকে সবচেয়ে ভালোভাবে বর্ণনা করে তা নির্ধারণ করতে ত্রুটি বার্তা এবং লগগুলি সাবধানতার সাথে বিশ্লেষণ করুন:
- TensorCore (TC) + SparseCore (SC) HBM ব্যবহারের সীমা অতিক্রম করে : যদি ত্রুটিটি স্পষ্টভাবে ব্যবহারকে ভেঙে দেয়, যেমন, "TC Hbm ব্যবহার: X, SC Hbm ব্যবহার Y" । → বিভাগ 1 এ যান। ব্যালেন্স TC এবং SC HBM ব্যবহার ।
- অপ্রত্যাশিতভাবে বড় বরাদ্দ : যদি ত্রুটিটি "মেমরি স্পেসে মেমরি শেষ হয়ে গেছে HBM" লেখা থাকে, তাহলে HBM-এ সবচেয়ে বড় বরাদ্দের তালিকার জন্য লগগুলি পরীক্ষা করুন। যদি, এক বা একাধিক অপ্রত্যাশিতভাবে বড় টেনসর (যেমন HBM সীমার 50%) উপস্থিত থাকে → বিভাগ 2-এ যান। অপ্রত্যাশিতভাবে বড় বরাদ্দ ।
- সমষ্টিগত বরাদ্দ HBM সীমা অতিক্রম করেছে : যদি ত্রুটিটি "মেমরি স্পেসে মেমরি শেষ হয়ে গেছে HBM" লেখা থাকে কিন্তু লগে অপ্রত্যাশিতভাবে কোনও বড় টেনসর উপস্থিত না থাকে → বিভাগ 3 এ যান। সমষ্টিগত বরাদ্দ HBM সীমা অতিক্রম করেছে ।
বিভাগ ১. ব্যালেন্স টিসি এবং এসসি এইচবিএম ব্যবহার
যদি ত্রুটিটি স্পষ্টভাবে ব্যবহারকে ভেঙে দেয়, যেমন, "TC Hbm ব্যবহার: X, SC Hbm ব্যবহার Y" তাহলে বাধা সনাক্ত করতে দুটি মানের তুলনা করুন।
- উচ্চ স্পার্সকোর ব্যবহার:
- HBM স্ট্যাক ব্যবহার অপ্টিমাইজ করুন:
feature_width,max_unique_nz_per_rowএবংlogical_replica_countব্যবহার করে HBM স্ট্যাক মেমরি ব্যবহারের স্কেল তৈরি করুন।--xla_sc_num_serialized_tables_to_optimize_hbmফ্ল্যাগ টিউন করে আপনি সর্বোচ্চ স্ট্যাক ব্যবহার কমাতে পারেন যা টেবিলের প্রক্রিয়াকরণকে সিরিয়ালাইজ করে। এর জন্য সমান্তরালতা হ্রাসের খরচ দিতে হয়। - প্যাডিং ওভারহেড পরীক্ষা করুন: স্পার্সকোর এম্বেডিং টেবিলগুলিকে 32B (8 floats) এ সারিবদ্ধ করে। ছোট বৈশিষ্ট্য প্রস্থের (যেমন, < 8 floats) টেবিলগুলিতে উল্লেখযোগ্য প্যাডিং ওভারহেড খরচ হয়, যার ফলে HBM নষ্ট হয়।
- হিপ ব্যবহার হ্রাস করুন:
maximum_parallel_iterationsএর জন্য উচ্চ মান HBM হিপে প্রিফেচ করা ইনপুট ডেটার পরিমাণ বৃদ্ধি করে। এই মান কমিয়ে দিলে উল্লেখযোগ্য মেমরি খালি হতে পারে। - শেয়ারিং যাচাই করুন: নিশ্চিত করুন যে এম্বেডিং টেবিলগুলি সমস্ত চিপগুলিতে সঠিকভাবে মোড-শার্ড করা আছে। দেখুন কিভাবে limits টেবিলে অনুবাদ করে ।
- আরও ধারণার জন্য SC: পারফরম্যান্স এবং মেমরির বাধা দেখুন।
- HBM স্ট্যাক ব্যবহার অপ্টিমাইজ করুন:
- উচ্চ টেনসরকোর ব্যবহার:
- বিভাগ ২- এ যান।
- সুষম
- যদি কোনটিই পৃথকভাবে অতিরিক্ত না হয় কিন্তু যোগফল খুব বেশি হয়, তাহলে আপনি চিপের ধারণক্ষমতায় আছেন। আপনাকে অবশ্যই উভয় উপাদানের ব্যবহার কমানোর চেষ্টা করতে হবে। তিনটি বিভাগেই সুপারিশ অনুসরণ করুন।
বিভাগ ২। অপ্রত্যাশিতভাবে বড় বরাদ্দ
যদি লগগুলিতে এক বা একাধিক অপ্রত্যাশিতভাবে বড় বরাদ্দ (> HBM সীমার ৫০%) থাকে, তাহলে এটি প্রায় কখনই হার্ডওয়্যার ক্ষমতার সমস্যা নয়। এটি সাধারণত একটি কনফিগারেশন ত্রুটি। JAX সোর্স কোডের ইঙ্গিতের জন্য বৃহৎ বরাদ্দের XLA লেবেল (যদি থাকে) পরীক্ষা করুন।
- ডিবাগিং আর্টিফ্যাক্টগুলি সরান:
- বৃহৎ পরিসরে jax.debug.print() ব্যবহার করলে কম্পাইলার HBM-এর সম্পূর্ণ টেনসরকে CPU-তে স্থানান্তর করতে বাধ্য হতে পারে, যার ফলে ফিউশন ভেঙে যায় এবং সর্বোচ্চ মেমোরি ব্যবহার বৃদ্ধি পায়। বাকি থাকা
jax.debug.print()গুলি সরিয়ে ফেলুন।
- বৃহৎ পরিসরে jax.debug.print() ব্যবহার করলে কম্পাইলার HBM-এর সম্পূর্ণ টেনসরকে CPU-তে স্থানান্তর করতে বাধ্য হতে পারে, যার ফলে ফিউশন ভেঙে যায় এবং সর্বোচ্চ মেমোরি ব্যবহার বৃদ্ধি পায়। বাকি থাকা
- অদক্ষ জালের আকার বা ভাগাভাগি ঠিক করুন:
- ভুল মেশ শেপ বা অনুপস্থিত শার্ডিং অ্যানোটেশনের কারণে কম্পাইলারটি ডিফল্টভাবে রেপ্লিকেশনে চলে যেতে পারে - যার ফলে কম্পাইলারকে একটি একক চিপে সত্যিই বড় টেনসর লাগানোর চেষ্টা করতে বাধ্য করা হয়।
- বৃহৎ বরাদ্দের আকারগুলি পরীক্ষা করুন এবং যাচাই করুন যে শার্ডিং সঠিকভাবে নির্দিষ্ট করা হয়েছে এবং XLA দ্বারা প্রচারিত হয়েছে।
বিভাগ ৩। মোট বরাদ্দ HBM সীমা অতিক্রম করে
যদি HBM সীমা অতিক্রম করে মোট বরাদ্দের পরিমাণের কারণে প্রোগ্রামটির ক্ষমতা শেষ হয়ে যায়, তাহলে পিক ব্যবহারে অবদান রাখা নির্দিষ্ট বাফারগুলি সনাক্ত করার জন্য মেমরি প্রোফাইলটি কল্পনা করা প্রায়শই সহায়ক। পিক মেমরি অবদানকারীদের সনাক্ত করার জন্য ধাপে ধাপে নির্দেশিকাটির জন্য XProf দিয়ে OOM ত্রুটিগুলি ডিবাগ করুন দেখুন।
একবার আপনি কিছু শীর্ষ অবদানকারীকে চিহ্নিত করার পরে, মেমরি ফুটপ্রিন্ট অপ্টিমাইজ করতে নিম্নলিখিত পদক্ষেপগুলি ব্যবহার করুন।
A. টেনসর প্যাডিং এবং অ্যালাইনমেন্ট পরীক্ষা করুন
অদক্ষ টেনসর আকারগুলি TPU-তে OOM-এর একটি সাধারণ, নীরব কারণ। TPU-তে সর্বোচ্চ কর্মক্ষমতা পেতে, XLA টেনসরের মাত্রা প্যাড করে—সাধারণত ক্ষুদ্রতম মাত্রার জন্য 128 এবং দ্বিতীয় ক্ষুদ্রতম মাত্রার জন্য 8 এর গুণিতকে। এই প্যাডিং ইনপুট অ্যারে এবং মধ্যবর্তী টেনসর (HLO টেম্পোরারি) উভয়কেই প্রভাবিত করে, যা সম্ভাব্যভাবে মেমরির ব্যবহারকে উল্লেখযোগ্যভাবে স্ফীত করে, বিশেষ করে ছোট মাত্রার আকারের সাথে। অ্যারে লেআউট দেখুন।
- বৃহৎ বাফারের অডিট আকার: (ডিফল্ট লেআউট সহ TPU v5-এ)
- Xprof Memory Viewer- এ একটি বাফারের উপর ঘোরালে বাফার ডিটেইলস কার্ডটি দেখা যাবে যাতে প্যাডিং তথ্য সহ বাফার ডিটেইলস থাকে।
- উদাহরণ :
(129, 1024)এর একটি আকৃতি(256, 1024)তে প্যাড করা হতে পারে, যার ফলে প্রায় ৫০% মেমোরি নষ্ট হয়। - সংশোধন:
(128, 1024)আকৃতির জন্য কোনও প্যাডিং প্রয়োজন হয় না এবং ০% মেমরি নষ্ট হয়।
- মাত্রা সারিবদ্ধ করুন: নিশ্চিত করুন যে সমস্ত বৃহৎ টেনসরের মাত্রা (ব্যাচের আকার, এম্বেডিং মাত্রা, লুকানো আকার) 128 এর গুণিতক।
খ. কনফিগারেশন সামঞ্জস্য করুন
আপনি প্রায়শই এই কনফিগারেশন সমন্বয়গুলির মাধ্যমে OOM গুলি সমাধান করতে পারেন:
- ব্যাচের আকার হ্রাস করুন: মধ্যবর্তী অ্যাক্টিভেশন এবং গ্রেডিয়েন্টের জন্য প্রয়োজনীয় মেমোরি ব্যাচের আকারের সাথে সরাসরি সমানুপাতিক। ব্যাচের আকার হ্রাস করলে প্রায়শই মেমোরির ব্যবহার হ্রাস করতে সাহায্য করতে পারে।
- ইনপুট বাফার দান করুন:
jax.jitব্যবহার করার সময়, আপনার মডেল প্যারামিটারের জন্য donate_argnums নির্দিষ্ট করুন। এটি XLA কে আউটপুটের সাথে ইনপুট মেমোরি ওভাররাইট করতে দেয়। - মিশ্র স্পষ্টতা সক্ষম করুন (bfloat16): মডেল আর্কিটেকচার এবং মানের প্রয়োজনীয়তা অনুমতি দিলে প্রোগ্রামের বৃহত্তম টেনসরগুলির জন্য bfloat16 বা কোয়ান্টাইজেশন (int8 ইত্যাদি) ব্যবহার করুন।
গ. আর্কিটেকচার এবং শারডিং অপ্টিমাইজ করুন
যদি কনফিগারেশন পরিবর্তনগুলি পর্যাপ্ত না হয়, তাহলে বর্তমান হার্ডওয়্যার সেটআপের জন্য মডেল টপোলজিটি খুব বড় হতে পারে।
- নতুন প্রজন্মের TPU ব্যবহার করুন: নতুন প্রজন্মের TPU সাধারণত প্রতি চিপে বেশি HBM অফার করে; যদি পাওয়া যায় তবে নতুন প্রজন্মের TPU ব্যবহার করুন।
- একটি বৃহত্তর চিপ টপোলজি ব্যবহার করুন: যদি মডেলের ওজন বিদ্যমান টপোলজির জন্য খুব বেশি হয়, তাহলে আপনি আরও চিপগুলিতে ভাগ করে নেওয়ার চেষ্টা করতে পারেন।
- উন্নত শার্ডিং কৌশল বাস্তবায়ন করুন:
- আরও উন্নত ডেটা, টেনসর, অথবা পাইপলাইন সমান্তরাল পদ্ধতিগুলি অন্বেষণ করুন।
- মধ্যবর্তী মান এবং আউটপুটগুলির জন্য শারডিং ইঙ্গিতগুলি নির্দিষ্ট করুন।
- JAX হোস্ট অফলোডিং ব্যবহার করুন: হোস্ট CPU মেমোরিতে বড় টেনসর অফলোড করুন। যেমন অ্যাক্টিভেশন অফলোডিং এবং অপ্টিমাইজার স্টেট অফলোডিং ।
D. XLA ফ্ল্যাগগুলিকে প্রভাবিত করে এমন কী মেমোরি টিউন করুন:
কম মেমোরি ব্যবহারের জন্য কী মেমোরি ফ্ল্যাগগুলিকে ট্রেড-অফ পারফরম্যান্সের সাথে সামঞ্জস্য করা যেতে পারে। তবে এগুলিকে শেষ অবলম্বন হিসাবে ব্যবহার করা উচিত কারণ এটি কর্মক্ষমতাকে বিরূপভাবে প্রভাবিত করতে পারে।
ই. টিউন এক্সএলএ রিমেটেরিয়ালাইজেশন পাস / ম্যানুয়াল চেকপয়েন্টিং
যদি মডেলটি মেমোরিতে ফিট করার কাছাকাছি থাকে, তাহলে আপনি XLA::Rematerialization পাসকে মেমোরি সাশ্রয়কে অগ্রাধিকার দিতে বাধ্য করতে পারেন, সম্ভবত ধীর সংকলনের খরচে:
| পতাকা | বিবরণ | প্রভাব / বিনিময় বন্ধ |
|---|---|---|
--xla_tpu_max_hbm_size_mib | রিমেটেরিয়ালাইজেশন পাস দ্বারা ব্যবহৃত HBM আকারের সীমা ম্যানুয়ালি সেট করে। | কম্পাইলারকে প্রোগ্রামটিকে প্রকৃত ভৌত HBM এর চেয়ে ছোট সীমার মধ্যে ফিট করার জন্য আরও কঠোর পরিশ্রম করতে বাধ্য করে। |
--xla_tpu_rematerialization_algo=PEAK_PRIORITY | সর্বোচ্চ মেমরি ব্যবহারের বিন্দুতে প্রচেষ্টাকে কেন্দ্রীভূত করে। | ডিফল্ট অ্যালগরিদমের চেয়ে আক্রমণাত্মক মেমরি হ্রাসের জন্য আরও দক্ষ হতে পারে। |
--xla_tpu_rematerialization_max_block_size_limit=32 | একটি ব্লকে একবারে পুনঃউপাদান করা যেতে পারে এমন সর্বোচ্চ সংখ্যক নির্দেশাবলী নিয়ন্ত্রণ করে। | এটি বৃদ্ধি করলে মেমোরি সাশ্রয় হয়, বরং কম্পাইলের সময় উল্লেখযোগ্যভাবে বৃদ্ধি পায় । |
--xla_tpu_rematerialization_block_effort_factor=10.0 | ব্লকগুলিকে পুনঃম্যাটেরিয়ালাইজ করার জন্য অনুসন্ধান করতে ব্যয় করা প্রচেষ্টার পরিমাণ (কম্পাইল সময়) নির্ধারণ করে। | উচ্চতর মান বর্ধিত কম্পাইল সময়ের খরচে মেমরি সঞ্চয়ের জন্য আরও বিস্তৃত অনুসন্ধানের অনুমতি দেয়। |
--xla_tpu_pre_fusion_remat=true | ফিউশন পাসের আগে একটি অতিরিক্ত রিমেটেরিয়ালাইজেশন পাস সক্ষম করে। | আরও মেমরি সঞ্চয় খুঁজে পেতে পারে, তবে কম্পাইলের সময় বাড়ায় এবং সম্ভাব্যভাবে সংখ্যাসূচক স্থিতিশীলতার উপর প্রভাব ফেলতে পারে। |
বিকল্পভাবে, jax.grad সহ jax.checkpoint ডেকোরেটর ব্যবহার করে ম্যানুয়ালি নিয়ন্ত্রণ করুন কোন ইন্টারমিডিয়েটগুলি ফরোয়ার্ড পাসে সংরক্ষিত হবে বনাম ব্যাকওয়ার্ড পাসে পুনঃগণিত হবে, HBM-এর জন্য ট্রেডিং কম্পিউট চক্র।
চ. উন্নত প্রোফাইলিং টুল ব্যবহার করুন
XProf ব্যবহার করে OOM ত্রুটিগুলি ডিবাগ করুন। XProf মেমরি ভিউয়ার ব্যবহার করার জন্য একটি টিউটোরিয়াল প্রদান করে যা HBM ব্যবহারের কম্পাইলারের ভিউ কল্পনা করে।
এই টুলটি আপনাকে সর্বোচ্চ মেমোরি বরাদ্দ এবং বাফার লাইফটাইম দেখতে দেয়, যা সর্বোচ্চ ব্যবহারের সময় HBM ঠিক কী ব্যবহার করে তা বোঝার জন্য অত্যন্ত গুরুত্বপূর্ণ। সাধারণ প্রোফাইলিং সেটআপের জন্য, Xprof এবং TensorBoard প্রোফাইলিং দিয়ে শুরু করা দেখুন।