XLA: মেশিন লার্নিং এর জন্য অপ্টিমাইজিং কম্পাইলার

OpenXLA হল রৈখিক বীজগণিতের জন্য একটি ডোমেন-নির্দিষ্ট কম্পাইলার যা সম্ভাব্য কোনো সোর্স কোড পরিবর্তন ছাড়াই TensorFlow মডেলকে ত্বরান্বিত করতে পারে।

ভূমিকা

যখন একটি টেনসরফ্লো প্রোগ্রাম চালানো হয়, তখন টেনসরফ্লো নির্বাহক দ্বারা সমস্ত ক্রিয়াকলাপ পৃথকভাবে সম্পাদিত হয়। প্রতিটি টেনসরফ্লো অপারেশনে একটি পূর্বে সংকলিত GPU কার্নেল বাস্তবায়ন থাকে যা নির্বাহক প্রেরণ করে।

এক্সএলএ চলমান মডেলগুলির একটি বিকল্প মোড সরবরাহ করে: এটি প্রদত্ত মডেলের জন্য বিশেষভাবে তৈরি করা গণনা কার্নেলের একটি ক্রমানুসারে টেনসরফ্লো গ্রাফকে সংকলন করে। কারণ এই কার্নেলগুলি মডেলের জন্য অনন্য, তারা অপ্টিমাইজেশনের জন্য মডেল-নির্দিষ্ট তথ্য ব্যবহার করতে পারে। উদাহরণস্বরূপ, আসুন একটি অপ্টিমাইজেশান দেখুন XLA একটি সাধারণ TensorFlow গণনার প্রসঙ্গে:

def model_fn(x, y, z):
  return tf.reduce_sum(x + y * z)

XLA ছাড়াই চালান, গ্রাফ তিনটি কার্নেল চালু করে: একটি গুণের জন্য, একটি যোগ করার জন্য এবং একটি হ্রাসের জন্য। যাইহোক, XLA গ্রাফটিকে অপ্টিমাইজ করতে পারে যাতে এটি একটি একক কার্নেল লঞ্চে ফলাফল গণনা করে। এটি একটি একক GPU কার্নেলে যোগ, গুণ এবং হ্রাসকে "ফিউজিং" করে করে। অধিকন্তু, এই ফিউজড অপারেশনটি মেমরিতে y*z এবং x+y*z দ্বারা উত্পাদিত মধ্যবর্তী মানগুলি লিখতে পারে না; পরিবর্তে এটি এই মধ্যবর্তী গণনার ফলাফলগুলিকে সরাসরি তাদের ব্যবহারকারীদের কাছে "স্ট্রিম" করে এবং তাদের সম্পূর্ণভাবে GPU রেজিস্টারে রাখে। ফিউশন হল XLA এর একক সবচেয়ে গুরুত্বপূর্ণ অপ্টিমাইজেশন। মেমরি ব্যান্ডউইথ সাধারণত হার্ডওয়্যার এক্সিলারেটরের সবচেয়ে কম সম্পদ, তাই মেমরি অপারেশনগুলি অপসারণ করা কর্মক্ষমতা উন্নত করার অন্যতম সেরা উপায়।

TensorFlow মডেলের জন্য XLA সক্ষম করুন

tf.function(jit_compile=True) সহ স্পষ্ট সংকলন

স্পষ্ট সংকলন API কোন ফাংশন কম্পাইল করা উচিত তা নির্বাচন করার জন্য একটি সূক্ষ্ম-দানাযুক্ত নিয়ন্ত্রণ অফার করে। উদাহরণস্বরূপ, নিম্নলিখিত TensorFlow ফাংশন যা MNIST প্রশিক্ষণ সম্পাদন করে XLA এর সাথে কম্পাইল করা হয়েছে:

@tf.function(jit_compile=True)
def train_mnist(images, labels):
    images, labels = cast(images, labels)

    with tf.GradientTape() as tape:
      predicted_labels = layer(images)
      loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=predicted_labels, labels=labels
      ))
    layer_variables = layer.trainable_variables
    grads = tape.gradient(loss, layer_variables)
    optimizer.apply_gradients(zip(grads, layer_variables))

jit_compile API-এর অবশ্যই-সংকলন শব্দার্থবিদ্যা আছে: হয় পুরো ফাংশনটি XLA দিয়ে কম্পাইল করা হয়েছে, অথবা একটি errors.InvalidArgumentError ব্যতিক্রম নিক্ষেপ করা হয়েছে৷ XLA বর্তমানে ফাংশনগুলি কম্পাইল করতে পারে না যেখানে মাত্রাগুলি অনুমানযোগ্য নয় : অর্থাৎ, যদি সম্পূর্ণ গণনা না চালিয়ে সমস্ত টেনসরের মাত্রা অনুমান করা সম্ভব না হয়৷ উদাহরণস্বরূপ, নিম্নলিখিত ফাংশন কম্পাইল হবে না:

@tf.function
def not_compilable(x):
  return tf.unique(x)

আকৃতি যদিও রান জুড়ে পরিবর্তিত হতে পারে:

@tf.function(jit_compile=True)
def recompiled_on_launch(a, b):
  return a + b

recompiled_on_launch(tf.ones([1, 10]), tf.ones([1, 10]))
recompiled_on_launch(tf.ones([1, 100]), tf.ones([1, 100]))

আরও বিস্তারিত ব্যবহারের উদাহরণের জন্য টিউটোরিয়াল কোল্যাব দেখুন, এবং jit_compile=True use এর একটি টিউটোরিয়াল ভিডিও দেখুন।

কেরাসের সাথে ব্যবহার

কেরাস মডেলের জন্য, jit_compile=True model.compile এ একটি আর্গুমেন্ট হিসাবে সেট করা যেতে পারে:

model.compile(optimizer="adam", jit_compile=True)

বিতরণ কৌশল সঙ্গে ব্যবহার

XLA:GPU টিএফ ডিস্ট্রিবিউটেড স্ট্র্যাটেজি ( MirroredStrategy বা MultiWorkerMirroredStrategy ) এর সাথে jit_compile=True এর সাথে স্টেপ ফাংশন টীকা দিয়ে ব্যবহার করা যেতে পারে।

@tf.function(jit_compile=True)
def step_fn():
  t = tf.ones(shape=[100], dtype=tf.float32)
  ctx = tf.distribute.get_replica_context()
  return ctx.all_reduce(tf.distribute.ReduceOp.SUM, t)

@tf.function
def run_fn():
  return strategy.run(step_fn)

স্বয়ংক্রিয় ক্লাস্টারিং

কোনো পরিবর্তন ছাড়াই TensorFlow মডেলে XLA ব্যবহার শুরু করার একটি সহজ উপায় হল স্বয়ংক্রিয়-ক্লাস্টারিং সক্ষম করা, যা স্বয়ংক্রিয়ভাবে TensorFlow ফাংশনগুলির মধ্যে ক্লাস্টার (সংযুক্ত সাবগ্রাফ) খুঁজে পায় যা XLA ব্যবহার করে সংকলিত এবং কার্যকর করা যেতে পারে। TF_XLA_FLAGS এনভায়রনমেন্ট ভেরিয়েবল সেট করে GPU-তে অটো-ক্লাস্টারিং সক্ষম করা যেতে পারে:

$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program

অটো-ক্লাস্টারিং বর্তমানে GPU ওয়ার্কলোডের জন্য অপ্টিমাইজ করা হয়েছে, তবে এটি CPU-তে অতিরিক্ত পতাকা ব্যবহার করে সক্রিয় করা যেতে পারে --tf_xla_cpu_global_jit :

$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program

বিস্তারিত ব্যবহারের উদাহরণের জন্য অটো-ক্লাস্টারিং টিউটোরিয়াল কোলাব দেখুন।

tfcompile সহ CPU-এর জন্য AOT (সময়ের আগে) সংকলন

আপনি একটি স্বতন্ত্র tfcompile টুলও ব্যবহার করতে পারেন, যা TensorFlow গ্রাফকে এক্সিকিউটেবল কোডে রূপান্তর করে (শুধুমাত্র x86-64 CPU-এর জন্য)।

সংকলিত প্রোগ্রাম পরিদর্শন করুন

XLA আত্মদর্শন সুবিধা প্রদান করে যা আপনাকে জেনারেট করা প্রোগ্রামগুলি পরিদর্শন করতে দেয়। জেনারেট করা প্রোগ্রামগুলি ডাম্প করতে, পরিবেশ পরিবর্তনশীল XLA_FLAGS ব্যবহার করুন:

$ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program

ডাম্পিং সঞ্চালিত হওয়ার পরে, আপনি নিম্নলিখিত ফাইলগুলি /tmp/generated এ খুঁজে পেতে পারেন:

  • module_XXXX.*_optimizations.txt জেনারেটেড XLA প্রোগ্রাম , প্রতি কম্পাইল করা ক্লাস্টার প্রতি একটি। XLA বাগ রিপোর্ট জমা দেওয়ার সময় তাদের সংযুক্ত করা অত্যন্ত সহায়ক!

  • module_XXXX.ir-*.ll LLVM ইন্টারমিডিয়েট রিপ্রেজেন্টেশানে জেনারেট করা ফাইল, NVPTX ইন্ট্রিনসিক্স সহ।

  • module_XXXX.ptx জেনারেটেড PTX ফাইল।

এছাড়াও আপনি টেনসরফ্লো গ্রাফের ভিতরে XLA ক্লাস্টারের এম্বেডিংকে ভিজ্যুয়ালাইজ করে গ্রাফটি ডাম্প করতে পারেন:

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"

পুনরুত্পাদনযোগ্য বাগ রিপোর্ট

একটি বাগ রিপোর্ট পুনরুত্পাদন করা অনেক সহজ যদি এতে জেনারেট করা XLA প্রোগ্রামগুলির জন্য ডাম্প এবং ব্যবহৃত স্বয়ংক্রিয়-ক্লাস্টারিং এম্বেডিং অন্তর্ভুক্ত থাকে। অটো-ক্লাস্টারিংয়ের সাথে চলমান একটি টেনসরফ্লো প্রোগ্রামের জন্য সেগুলি তৈরি করতে, লঞ্চ করুন:

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated \
  TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" \
  XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/tmp/generated" \
    my/tensorflow/program"

বাগ ফাইল করার সময়, /tmp/generated ডিরেক্টরির বিষয়বস্তু সংযুক্ত করুন (উপরে উল্লেখ করা হয়েছে)।

যদি সম্ভব হয়, run_hlo_module ব্যবহার করে একটি একক XLA প্রোগ্রামে একটি বাগ বিচ্ছিন্ন করার চেষ্টা করুন এবং এটিকে জেনারেট করা প্রোগ্রামগুলিতে পুনরাবৃত্তি করুন।

আরও পড়া

XLA ফ্রন্টেন্ডস

TensorFlow ছাড়াও, XLA প্রোগ্রামগুলি তৈরি করা যেতে পারে:

  • JAX : Python+NumPy প্রোগ্রামের সংমিশ্রণযোগ্য রূপান্তর
  • জুলিয়া : বৈজ্ঞানিক কম্পিউটিংয়ের জন্য জুলিয়া ভাষা
  • PyTorch : PyTorch ফ্রেমওয়ার্ক
  • Nx : এলিক্সির প্রোগ্রামিং ভাষার জন্য সংখ্যাসূচক কম্পিউটিং লাইব্রেরি

কথা বলে

jit_compile=True ব্যবহার করে TF থেকে XLA ব্যবহার করা

XLA ওভারভিউ