যখন অ্যাক্সিলারেটরের (GPU বা TPU) হাই ব্যান্ডউইথ মেমোরি (HBM)-এর ধারণক্ষমতা শেষ হয়ে যায়, তখন আউট অফ মেমোরি (OOM) ত্রুটি দেখা দেয়। OOM সমস্যার কিছু সাধারণ কারণ এবং ডিবাগিং কৌশল E1000 - কম্পাইল টাইম HBM OOM ডকুমেন্টেশন এবং GPU মেমোরি অ্যালোকেশন সম্পর্কিত JAX ডকুমেন্টেশনে বিস্তারিতভাবে বর্ণনা করা হয়েছে।
এই পৃষ্ঠায় বর্ণনা করা হয়েছে কীভাবে আপনার JAX প্রোগ্রামের মেমরি ব্যবহার পর্যবেক্ষণ করতে, সর্বোচ্চ ব্যবহারের মুহূর্তগুলো শনাক্ত করতে এবং OOM ত্রুটি ডিবাগ করতে XProf-এর মেমরি ভিউয়ার টুলটি ব্যবহার করতে হয়। এর জন্য নিম্নলিখিত ধাপগুলো অনুসরণ করতে হয়:
- প্রোফাইলটি ক্যাপচার করতে আপনার প্রোগ্রামটি
jax.profiler.traceসহ চালান। - ব্যাকগ্রাউন্ডে এক্সপ্রফ (XProf) চালু করুন এবং মেমরি ব্যবহারের বিস্তারিত তথ্য দেখতে মেমরি ভিউয়ার টুলটি ব্যবহার করুন।
উদাহরণ প্রোগ্রাম
নিম্নলিখিত JAX প্রোগ্রামটি একটি OOM ত্রুটির কারণ হয়:
import jax
from jax import random
import jax.numpy as jnp
@jax.profiler.trace("/tmp/xprof")
@jax.jit
def oom():
a = random.normal(random.PRNGKey(1), (327680, 327680), dtype=jnp.bfloat16)
return a @ a
if __name__ == "__main__":
oom()
একটি টিপিইউ-মেশিনে, এই প্রোগ্রামটি নিম্নলিখিত কারণে ব্যর্থ হয়:
XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=107374182400) would exceed memory (size=17179869184) :: #allocation7 [shape = 'u8[327680,327680]{1,0:T(8,128)(4,1)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of xor_convert_fusion@{}'] :: <no-hlo-instruction>
(একটি GPU-মেশিনে, ত্রুটিটি দেখতে এইরকম হয়: XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 214748364800 bytes. )
রান এক্সপ্রফ
xprof ইনস্টল করুন ( pip install xprof ), এবং প্রোফাইলটি যেখানে সংরক্ষিত আছে সেই ডিরেক্টরিটি উল্লেখ করে একটি XProf ইনস্ট্যান্স চালু করুন:
xprof --logdir=/tmp/xprof/ --port=6006
ইনস্ট্যান্সটিতে যান (লোকাল মেশিনে, http://localhost:6006 ঠিকানায়)। Tools ড্রপডাউন থেকে Memory Viewer নির্বাচন করুন, এবং Memory Viewer টুল উইন্ডোতে, Memory Types ড্রপডাউন থেকে HBM নির্বাচন করুন (যা সাধারণত ডিফল্টরূপে নির্বাচিত থাকে)।
XProf: Memory Viewer টুলের ডকুমেন্টেশনে টুলটির উপাদানসমূহ এবং উপস্থাপিত তথ্যের বর্ণনা দেওয়া হয়েছে।
'HLO Ops at Peak Memory Allocation' অংশটির উপর মনোযোগ দিন, যেখানে সর্বোচ্চ মেমরি ব্যবহারের মুহূর্তে তিনটি বাফার চার্ট দেখানো হয়েছে। বাফারটিতে অন্তর্ভুক্ত রয়েছে:
- প্রোগ্রামের ইনপুট এবং আউটপুট: ট্রেনিং ব্যাচ, অপটিমাইজারের অবস্থা, ইত্যাদি।
- TensorCore এবং SparseCore টেম্পোরারি: অন্তর্বর্তীকালীন গণনার (যেমন অ্যাক্টিভেশন, গ্রেডিয়েন্ট, ইত্যাদি) জন্য প্রয়োজনীয় ডাইনামিক মেমরি।
আপনি বাফার চার্টের উপর মাউস হোভার করে Op-টির আকার, আকৃতি, অ্যালোকেশন টাইপ এবং আরও অনেক বিস্তারিত তথ্য পেতে পারেন। এটি আপনাকে এমন Op-গুলো শনাক্ত ও মূল্যায়ন করতে সাহায্য করতে পারে, যেগুলোতে উচ্চ বা দীর্ঘস্থায়ী টেম্পোরারি থাকতে পারে, অথবা অদক্ষ প্যাডিংযুক্ত কোনো বড় ইনপুট/ইন্টারমিডিয়েট/আউটপুট টেনসর থাকতে পারে, ইত্যাদি, যা সর্বোচ্চ মেমরি ব্যবহারের জন্য দায়ী এবং যেগুলোকে সামঞ্জস্য বা অপ্টিমাইজ করা প্রয়োজন।
E1000: ডিবাগিং কোর্সে নির্দিষ্ট ডিবাগিং কৌশল শিখুন।