Google I/O-তে TensorFlow-এ যোগ দিন, 11-12 মে এখনই নিবন্ধন করুন

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

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

যেমন বার্ট মধ্যে: ফলাফল গতি এবং মেমরির ব্যবহার উন্নতি হয় MLPerf 8 পশ্চিমী V100 জিপিইউ XLA ব্যবহার করে একটি ~ 7 গুণ কর্মক্ষমতা বৃদ্ধি অর্জন এবং ~ 5x ব্যাচ আকার উন্নতি হয়েছে ব্যবহার জমা:

ভূমিকা

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

XLA চলমান মডেলগুলির একটি বিকল্প মোড সরবরাহ করে: এটি প্রদত্ত মডেলের জন্য বিশেষভাবে তৈরি করা গণনা কার্নেলের একটি ক্রম অনুসারে টেনসরফ্লো গ্রাফকে সংকলন করে। কারণ এই কার্নেলগুলি মডেলের জন্য অনন্য, তারা অপ্টিমাইজেশনের জন্য মডেল-নির্দিষ্ট তথ্য ব্যবহার করতে পারে। উদাহরণস্বরূপ, আসুন একটি অপ্টিমাইজেশান দেখুন 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 এপিআই আবশ্যক কম্পাইল শব্দার্থবিদ্যা রয়েছে: হয় সমগ্র ফাংশন XLA সঙ্গে কম্পাইল করা হয়, অথবা একটি errors.InvalidArgumentError ব্যতিক্রম ফেলে দেওয়া হয়। যে, যদি সমগ্র গণনার চলমান ছাড়া সব tensors এর মাত্রা অনুমান করা সম্ভব নয়: XLA বর্তমানে ফাংশন যেখানে মাত্রা inferrable নয় কম্পাইল করতে পারবে না। উদাহরণস্বরূপ, নিম্নলিখিত ফাংশন কম্পাইল হবে না:

@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]))

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

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

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

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

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

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

বিস্তারিত ব্যবহার উদাহরণস্বরূপ দ্রষ্টব্য স্বয়ংক্রিয় ক্লাস্টারিং টিউটোরিয়াল colab

সঙ্গে CPU- র জন্য AOT (এগিয়ে অফ সময়) সংকলন tfcompile

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

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

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 intrinsics।

  • 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 ডিরেক্টরি (উপরে উল্লিখিত)।

যদি সম্ভব হয়, ব্যবহার করে একটি একক XLA প্রোগ্রাম একটি বাগ বিছিন্ন করার চেষ্টা replay_computation এবং iteratively এই উত্পন্ন প্রোগ্রাম চলমান।

আরও পড়া

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

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

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

কথা বলে

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

XLA ওভারভিউ