কাস্টম কল প্রোফাইলিং, কাস্টম কল প্রোফাইলিং

XLA কাস্টম কল আপনাকে এমন কাস্টম কার্নেল বা অপারেশন সম্পাদন করার সুযোগ দেয় যা XLA দ্বারা স্বাভাবিকভাবে সমর্থিত নয়। ট্রেস ভিউয়ার-এর মধ্যে এই কাস্টম কলগুলির পারফরম্যান্স সম্পর্কে বিস্তারিত জানতে, আপনি নির্দিষ্ট XLA ফ্ল্যাগ ব্যবহার করে বিশদ ট্রেসিং এবং LLO (লো-লেভেল অপটিমাইজার) ডিবাগ তথ্য সক্রিয় করতে পারেন।

⚠️ পরীক্ষামূলক ফিচার : লো লেভেল অপটিমাইজার (LLO) অ্যানালাইসিস এবং কাস্টম কল প্রোফাইলিং পরীক্ষামূলক পর্যায়ে রয়েছে। এই ফিচারগুলো এবং সমস্ত CLI অ্যানালাইসিস টুল ( get_kernel_stats , get_llo_analysis , get_llo_debug_string ) ব্যবহার করতে, xprof-nightly ইনস্টল করুন । xprof স্ট্যান্ডার্ড PyPI রিলিজ (2.23.1)-এ এই সাবকমান্ডগুলো নেই।

পূর্বশর্ত এবং টুলচেইন প্রয়োজনীয়তা

LLO ট্রেস ক্যাপচার করার আগে, যাচাই করুন যে আপনার পরিবেশ নিম্নলিখিত প্রয়োজনীয়তাগুলি পূরণ করে:

• পাইথন ৩.১১+ (পাইথন ৩.১২ প্রস্তাবিত) : ডিফল্ট ক্লাউড টিপিইউ ভিএম ইমেজ (উবুন্টু ২২.০৪) সিস্টেম পাইথন ৩.১০.১২ সহ আসে, যা নীরবে JAX-কে ০.৬.২ সংস্করণে সীমাবদ্ধ করে এবং libtpu ০.০.১৭ পুল করে। পুরোনো libtpu বিল্ডগুলিতে LLO ফ্ল্যাগ ডেফিনিশনের অভাব থাকে, যার ফলে ERROR: Unknown command line flag দেখা দেয় অথবা খালি LLO প্রোফাইল রিটার্ন করে। পাইথন ৩.১২ ভার্চুয়াল এনভায়রনমেন্ট ম্যানেজ করার জন্য uv ব্যবহার করার জন্য দৃঢ়ভাবে সুপারিশ করা হচ্ছে।

• প্যাকেজ ইনস্টলেশন ( xprof-nightly ) : jax[tpu] এর পাশাপাশি xprof-nightly ইনস্টল করুন :

  # 1. Setup Python 3.12 environment
  pip install uv
  uv python install 3.12
  uv venv --python 3.12 ~/venvs/v312
  
  # 2. Install xprof-nightly and JAX
  uv pip install --python ~/venvs/v312/bin/python \
      'jax[tpu]>=0.11.0' xprof-nightly numpy ml_dtypes absl-py fire
  

• JAX >= 0.11.0 : প্রস্তাবিত টুলচেইন সংস্করণ ( libtpu >= 0.0.44 )। উল্লেখ্য যে, কম্পাইল-টাইম LLO ডিবাগ তথ্য ( --xla_xprof_register_llo_debug_info=true ) jax >= 0.10.2 ( libtpu >= 0.0.42 ) থেকে সমর্থিত, অন্যদিকে অপ্ট-ইন রানটাইম কাস্টম কল ট্রেসিং ( --xla_xprof_enable_custom_call_tracing=true ) এর জন্য jax >= 0.11.0 ( libtpu >= 0.0.44 ) প্রয়োজন।

• কঠোর এনভায়রনমেন্ট অর্ডারিং : import jax ঠিক আগে LIBTPU_INIT_ARGS অবশ্যই শেল-এ এক্সপোর্ট করতে হবে অথবা os.environ এ কনফিগার করতে হবে। libtpu JAX-এর সর্বপ্রথম ইম্পোর্টের সময় ইনিশিয়ালাইজেশন ফ্ল্যাগগুলো পার্স করে; import jax পরে এগুলো সেট করলে কোনো এক্সেপশন তৈরি না করেই নীরবে কোনো প্রভাব ফেলে না।

হার্ডওয়্যার সামঞ্জস্য ম্যাট্রিক্স

ফ্ল্যাগ প্রাপ্যতা ডায়াগনস্টিক

ওয়ার্কলোড চালু করার আগে আপনার ইনস্টল করা libtpu বাইনারিতে প্রয়োজনীয় ফ্ল্যাগ সংজ্ঞাগুলো আছে কিনা তা যাচাই করতে, এই ডায়াগনস্টিক কোডটি চালান:

import glob
import os
import libtpu

so_paths = glob.glob(os.path.dirname(libtpu.__file__) + "/*libtpu*.so")
if so_paths:
  blob = open(so_paths[0], "rb").read()
  for flag in (
      b"xla_xprof_register_llo_debug_info",
      b"xla_xprof_enable_custom_call_tracing",
      b"tpu_enable_periodic_counter_sampling",
  ):
    print(flag.decode(), "PRESENT" if flag in blob else "ABSENT")

ট্রেসিং কীভাবে চালু করবেন

প্রস্তাবিত ডিফল্ট: শুধুমাত্র LLO ডিবাগ তথ্য

স্ট্যাটিক LLO বিশ্লেষণ এবং HLO/কার্নেল-স্তরের প্রোফাইলিং-এর জন্য, --xla_xprof_register_llo_debug_info=true দিয়ে LLO ডিবাগ তথ্য রেজিস্টার করুন। এটি সম্পূর্ণ HLO অপ স্ট্রিমকে অক্ষত রাখে, ফলে get_hlo_stats , get_roofline_model , get_top_hlo_ops , এবং get_kernel_stats সবগুলোই কোনো অতিরিক্ত রানটাইম ট্রেস ওভারহেড ছাড়াই কাজ করে এবং একই সাথে get_llo_analysis ও get_llo_debug_string দ্বারা ব্যবহৃত সম্পূর্ণ কম্পাইল-টাইম LLO সোর্স ম্যাপ তৈরি করে।

import os

# Flags MUST precede any jax / libtpu import
os.environ["LIBTPU_INIT_ARGS"] = "--xla_xprof_register_llo_debug_info=true"

import jax
# Workload definition and tracing...

• --xla_xprof_register_llo_debug_info=true : XProf ভিজ্যুয়ালাইজেশনের জন্য LLO ডিবাগ তথ্য, অপকোড এবং মেটাডেটা নিবন্ধন করে।

সূক্ষ্ম-স্তরের রানটাইম LLO বান্ডেল ট্রেসিং

• --xla_xprof_enable_custom_call_tracing=true : এটি একটি ক্যানোনিকাল ফ্ল্যাগ যা রানটাইম LLO এক্সিকিউশনের সূক্ষ্ম বিবরণ (ট্রেস ভিউয়ারে Pallas Primitives , LLO Ops এবং প্রতি-ইউনিট ইন্সট্রাকশন লেন) সক্ষম করে এবং স্বয়ংক্রিয়ভাবে ইন্সট্রাকশন বান্ডেল ইন্সট্রুমেন্টেশন সক্রিয় করে ( xla_tpu_bundle_instrumentation_options যার ডিফল্ট trace_best_effort_frequency=10 এবং trace_guaranteed_frequency=10 )।

--xla_xprof_enable_custom_call_tracing=true ব্যবহার করার সময়, যদি বর্ধিত ট্রেস সাইজ হার্ডওয়্যার ট্রেস বাফারকে ওভারফ্লো করে, তাহলে vtrace ফ্রিকোয়েন্সি টিউন করুন ( xla_tpu_bundle_instrumentation_options এ থাকা trace_best_effort_frequency এবং trace_guaranteed_frequency ; নিচে 'How to Tune' দেখুন)।

উদাহরণ ট্রেস ভিউয়ার

Xprof ট্রেস ভিউয়ারে LLO ট্রেসগুলো দেখতে কেমন হয় তার একটি উদাহরণ এখানে দেওয়া হলো:

উন্নত প্যারামিটার (ইভেন্ট ড্রপ পরিচালনা)

আপনি যদি Xprof-এ ইভেন্ট ড্রপ বা বাফার ওভারফ্লো দেখতে পান, তার মানে হলো ট্রেস পয়েন্টগুলো খুব ঘন ঘন ট্রিগার হচ্ছে, যা হার্ডওয়্যার ট্রেস বাফারগুলোকে ওভারলোড করে ফেলছে। আপনি অ্যাডভান্সড প্যারামিটার ব্যবহার করে LLO ট্রেস ইনসারশনের ফ্রিকোয়েন্সি টিউন করতে পারেন।

এই প্যারামিটারগুলো xla_tpu_bundle_instrumentation_options এর মাধ্যমে কনফিগার করা হয়। আপনি নিয়ন্ত্রণ করতে পারেন কত ঘন ঘন ট্রেসগুলো ইন্সট্রাকশন বান্ডেলে প্যাক করা হবে।

মূল পরামিতি

• trace_best_effort_frequency (ডিফল্ট: ১০): বিদ্যমান বান্ডেলের মধ্যে সুযোগসন্ধানী ট্রেস সন্নিবেশ করার জন্য লক্ষ্যমাত্রা ব্যবধান (বান্ডেলের হিসাবে)। কম্পাইলার এই নির্দিষ্ট সময় অন্তর একটি ট্রেস সন্নিবেশ করার চেষ্টা করবে, কিন্তু এর জন্য নতুন বান্ডেল তৈরি করবে না ।
• trace_guaranteed_frequency (ডিফল্ট: ১০): দুটি ট্রেসের মধ্যে অনুমোদিত বান্ডেলের সর্বোচ্চ সংখ্যা। এটি একটি নিশ্চয়তা। যখন আমরা বিদ্যমান বান্ডেলগুলিতে ট্রেসগুলি প্যাক করে এটি পূরণ করতে পারি না, তখন আমরা একটি নতুন বান্ডেল তৈরি করব এবং সেখানে একটি ট্রেস (একাই) রাখব।

কীভাবে টিউন করবেন

• যদি আপনি 'ইভেন্ট ড্রপস' দেখতে পান : কম ঘন ঘন ট্রেস করার জন্য মানগুলি বাড়িয়ে দিন (যেমন, ৫০ বা ১০০-তে সেট করুন), যা তৈরি হওয়া ট্রেস ডেটার পরিমাণ কমিয়ে দেবে।
• যদি আপনার আরও সূক্ষ্ম বিশ্লেষণের প্রয়োজন হয় : আরও ঘন ঘন ট্রেস করার জন্য মানগুলি হ্রাস করুন (এর ফলে অতিরিক্ত ওভারহেড এবং সম্ভাব্য বাফার ওভারফ্লো হওয়ার ঝুঁকি থাকবে)।

নির্দেশনা চক্র গণনা কীভাবে হিসাব করা হয়

যেহেতু প্রতিটি নির্দেশনার সময় নয়, বরং সুযোগ বুঝে ট্রেস পয়েন্ট যুক্ত করা হয়, তাই আনুমানিক হার্ডওয়্যার সাইকেল খরচের উপর ভিত্তি করে অন্তর্বর্তী টাইমস্ট্যাম্পগুলো ইন্টারপোলেট করা হয়।

কম্পাইলার টার্গেট টিপিইউ জেনারেশন এবং যে এক্সিকিউশন ইউনিট এটিকে সমাধান করছে তার উপর ভিত্তি করে প্রতিটি এলএলও ইনস্ট্রাকশনের অন্তর্নিহিত হার্ডওয়্যার সাইকেল কস্ট গণনা করে। এই সাইকেল সংখ্যাগুলো এক্সিকিউশন থ্রুপুট এবং ল্যাটেন্সি ডিলে নির্দেশ করে।

উচ্চ-স্তরের প্রবাহ

1. LLO নির্দেশনা পার্স করুন : অপকোড এবং মেটাডেটা শনাক্ত করুন।
2. বেস হার্ডওয়্যার সাইকেল জানুন : টিপিইউ জেনারেশন (v5e/v5p, v6e/v7x, ইত্যাদি) অনুযায়ী সাইকেল নির্ধারণ করুন।
3. GTC টিক্সে রূপান্তর করুন : নিম্নলিখিত সূত্র ব্যবহার করে সাইকেলকে গ্লোবাল টাইমার কাউন্টার (GTC) টিক্সে রূপান্তর করুন: Cycles * (GTC_Freq * 16) / TC_Freq ।
4. টাইমলাইন স্প্যান তৈরি করুন : পরিচিত ট্রেস সীমানাগুলির মধ্যে মধ্যবর্তী ইভেন্টগুলিকে সমানভাবে ইন্টারপোলেট করুন।

ইউনিট এবং প্রজন্ম অনুসারে চক্রের অনুমান

বিভিন্ন এক্সিকিউশন ইউনিটের জন্য বেস হার্ডওয়্যার সাইকেলগুলো কীভাবে মডেল করা হয়, তার উদাহরণ নিচে দেওয়া হলো:

ম্যাট্রিক্স গুণন একক (MXU)

MXU সাইকেল গণনা ডেটা টাইপের ঘনত্বের উপর ভিত্তি করে থ্রুপুট প্রতিফলিত করে।

ট্রান্সপোজ ইউনিট (XLU)

সাইকেল কাউন্ট ট্রান্সপোজ মেমরি লেআউট এবং ক্রসবার ডিলে নির্দেশ করে।

এক্সিকিউশন ইউনিট পুল (EUP)

EUP নির্দেশাবলী ভেক্টর গাণিতিক ফাংশন (যেমন, tanh , log , exp ) উপস্থাপন করে।

পতাকা অভিবাসন ও পুনর্মিলন

XLA এবং TPU ডকুমেন্টেশনের পূর্ববর্তী সংস্করণগুলিতে --xla_enable_custom_call_region_trace=true লিগ্যাসি ফ্ল্যাগটির উল্লেখ ছিল।

• ক্যানোনিকাল ফ্ল্যাগ : --xla_xprof_enable_custom_call_tracing (রানটাইম ইন্ট্রা-কার্নেল বান্ডেল টাইমলাইন স্প্যান প্রয়োজন হলে এটি ক্যানোনিকাল নাম; ডিফল্টরূপে শুধু --xla_xprof_register_llo_debug_info=true ব্যবহার করুন)। এটি সক্রিয় করা হলে, এটি কাস্টম কল ট্রেসিং চালু করে এবং স্বয়ংক্রিয়ভাবে প্রয়োজনীয় ইন্সট্রাকশন বান্ডেল ইন্সট্রুমেন্টেশন ও ট্রেস ফ্রিকোয়েন্সি ( xla_tpu_bundle_instrumentation_options ) কনফিগার করে।
• লিগ্যাসি ফ্ল্যাগ : --xla_enable_custom_call_region_trace=true (অপ্রচলিত বিকল্প নাম)। যদিও এটি এখনও পুরোনো কম্পাইলার ব্যাকএন্ড দ্বারা সমর্থিত, যেসব ব্যবহারকারীর রানটাইম বান্ডেল টাইমলাইন স্প্যান প্রয়োজন, তাদের --xla_xprof_enable_custom_call_tracing এ স্থানান্তরিত হওয়া উচিত।

ডিফল্ট ক্যাপচার উদাহরণ (যা HLO এবং কার্নেল পরিসংখ্যান অক্ষত রেখে LLO ডিবাগ তথ্য নিবন্ধন করে):

export LIBTPU_INIT_ARGS="--xla_xprof_register_llo_debug_info=true"
python your_jax_workload.py

যখন LLO ডিবাগ তথ্যের পাশাপাশি কাস্টম কল ট্রেসিং সক্রিয় করা হয়, তখন কাস্টম কলটি সম্পাদনকারী প্রতিটি TPU কোর বা ডিভাইসের জন্য ট্রেস ভিউয়ারে একটি নতুন LLO ইউটিলাইজেশন লাইন প্রদর্শিত হবে।

এলএলও ইউটিলাইজেশন লাইন

এলএলও ইউটিলাইজেশন লাইনটি একটি কাস্টম কল সম্পাদনের সময় হার্ডওয়্যার রিসোর্স কীভাবে ব্যবহৃত হয় তার একটি ভিজ্যুয়ালাইজেশন প্রদান করে। কাস্টম কার্নেলের (যেমন, প্যালাস বা মোজাইকে লেখা) অভ্যন্তরীণ বাধা বা বটলনেক শনাক্ত করার জন্য এটি বিশেষভাবে উপযোগী।

সর্বোত্তম অনুশীলন এবং মাঠের অপ্রত্যাশিত সমস্যা

• xprof-nightly ব্যবহার করুন : স্ট্যান্ডার্ড xprof 2.23.1-এ get_kernel_stats এবং LLO CLI সাবকমান্ডগুলো নেই (পাশাপাশি সংখ্যাগত প্যারিটি যাচাইয়ের জন্য স্বতন্ত্র xparity কনসোল স্ক্রিপ্টটিও নেই)। Google3-বহির্ভূত পরিবেশে, সর্বদা xprof-nightly ইনস্টল করুন।
• প্যালাস কার্নেলের জন্য মেট্রিক্সের ব্যাখ্যা (রুফলাইন ব্লাইন্ড স্পট) : XLA-এর tpu_custom_call এর জন্য কোনো কস্ট মডেল নেই। তাই, get_roofline_model এবং get_overview 0.0 GFLOP/s , "bound_by": "Unknown" এবং 0.0% MXU ইউটিলাইজেশন রিপোর্ট করবে, এমনকি যখন LLO ইনস্ট্রাকশনগুলো হার্ডওয়্যারে সম্পূর্ণরূপে ক্যাপচার এবং এক্সিকিউট হচ্ছে।
  • কার্নেলের সময়কাল ও লেটেন্সি জানতে, xprof get_kernel_stats <logdir> ব্যবহার করুন।
  • নিম্ন-স্তরের নির্দেশাবলী সম্পাদনের বিভাজন এবং চক্রের আনুমানিক হিসাবের জন্য, xprof get_llo_analysis <logdir> ব্যবহার করুন।
• লাইট প্রোটো ইনার লুপ বডি : get_llo_debug_string এ, ইনার লুপ বডিগুলোকে সংক্ষেপে এভাবে বর্ণনা করা হয়: // Loop body not available in lite proto । এটি পারিপার্শ্বিক মডিউল কাঠামো, রেজিস্টার বরাদ্দ এবং বাইরের নির্দেশনার ক্রম সরবরাহ করে।
• ট্রেস যাচাইকরণ হিউরিস্টিক : LLO ডেটা আছে কিনা তা নির্ধারণ করতে SALU / VALU / EUP / XLU / VLD / VST / MXU Instructions মতো ট্রেস লাইনের নাম পরীক্ষা করবেন না । বৈধ LLO ট্রেসগুলিতে এই লাইনের নামগুলি ব্যবহৃত হয় না। xprof get_llo_analysis <logdir> এক্সিকিউট করে এবং "success": true যাচাই করে LLO ক্যাপচারটি ভ্যালিডেট করুন।
• ক্যাপচার কার্নেলের আকার নির্ধারণ : টেস্ট/ক্যাপচার কার্নেলের আকার খুব বড় করলে কম্পাইলার ত্রুটি দেখা দিতে পারে, যেমন CompileTimeScopedVmemOom: Scoped allocation with size 32.81M and limit 32.00M exceeded scoped vmem limit । ক্যাপচার ম্যাট্রিক্সের আকার পরিমিত রাখুন (যেমন (512, 512, 1024) f32)।
• পৃথক ভার্চুয়াল এনভায়রনমেন্ট : বিভিন্ন JAX ভার্সনের মধ্যে কার্নেল পোর্ট করার সময় (যেমন, JAX 0.11+ এর বাতিল হওয়া pltpu.repeat এর মতো বিষয়গুলির ক্ষেত্রে), নিজস্ব ভার্চুয়াল এনভায়রনমেন্ট বজায় রাখুন।