XProf ব্যবহার করে OOM ত্রুটিগুলি ডিবাগ করুন

অ্যাক্সিলারেটরের (GPU বা TPU) হাই ব্যান্ডউইথ মেমরি (HBM) ক্ষমতা শেষ হয়ে গেলে মেমরির বাইরে (OOM) ত্রুটি দেখা দেয়। OOM সমস্যা এবং ডিবাগিং কৌশলগুলির কিছু সাধারণ কারণ E1000 - কম্পাইল টাইম HBM OOM ডকুমেন্টেশন এবং GPU মেমরি বরাদ্দের উপর JAX ডকুমেন্টেশনে বিস্তারিতভাবে বর্ণনা করা হয়েছে।

এই পৃষ্ঠায় XProf এর মেমোরি ভিউয়ার টুল ব্যবহার করে আপনার JAX প্রোগ্রামের মেমোরি ব্যবহারের দৃশ্যমানতা, সর্বোচ্চ ব্যবহারের উদাহরণ সনাক্তকরণ এবং OOM ত্রুটিগুলি ডিবাগ করার পদ্ধতি বর্ণনা করা হয়েছে। এর জন্য নিম্নলিখিত পদক্ষেপগুলি অন্তর্ভুক্ত রয়েছে:

1. প্রোফাইলটি ক্যাপচার করার জন্য jax.profiler.trace দিয়ে আপনার প্রোগ্রামটি চালান।
2. ব্যাকগ্রাউন্ডে 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()

একটি TPU-মেশিনে, এই প্রোগ্রামটি নিম্নলিখিত ক্ষেত্রে ব্যর্থ হয়:

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 চালান

xprof ইনস্টল করুন ( pip install xprof ), এবং প্রোফাইলটি কোথায় সংরক্ষণ করা হবে তা নির্দিষ্ট করে একটি XProf ইনস্ট্যান্স শুরু করুন:

xprof --logdir=/tmp/xprof/ --port=6006

ইনস্ট্যান্সে যান (স্থানীয় মেশিনে, http://localhost:6006 এ)। Tools ড্রপডাউনে, Memory Viewer নির্বাচন করুন, এবং Memory Viewer টুল উইন্ডোতে, Memory Types ড্রপডাউনে (সাধারণত ডিফল্টরূপে নির্বাচিত) HBM নির্বাচন করুন।

XProf: মেমোরি ভিউয়ার টুল ডকুমেন্টেশন টুলের উপাদান এবং উপস্থাপিত তথ্য বর্ণনা করে।

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

Op সম্পর্কে আরও বিশদ জানতে আপনি বাফার চার্টগুলিতে হোভার করতে পারেন যেমন এর আকার, আকৃতি, বরাদ্দের ধরণ এবং আরও অনেক কিছু। এটি আপনাকে Ops সনাক্ত করতে এবং মূল্যায়ন করতে সাহায্য করতে পারে যেগুলিতে উচ্চ বা দীর্ঘস্থায়ী অস্থায়ী থাকতে পারে, যে কোনও বড় ইনপুট/ইন্টারমিডিয়েট/আউটপুট টেনসর যার অদক্ষ প্যাডিং আছে, ইত্যাদি, যা সর্বোচ্চ মেমোরিতে অবদান রাখছে এবং সামঞ্জস্য বা অপ্টিমাইজ করা প্রয়োজন।

E1000: ডিবাগিং- এ নির্দিষ্ট ডিবাগিং কৌশল শিখুন।