এমএল কমিউনিটি দিবস 9 নভেম্বর! TensorFlow, Jax থেকে আপডেটের জন্য আমাদের সাথে যোগ দিন, এবং আরও আরও জানুন

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

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

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

ভূমিকা

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

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

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

এক্সএলএ ছাড়াই চালান, গ্রাফ তিনটি কার্নেল চালু করে: একটি গুণের জন্য, একটি সংযোজনের জন্য এবং আরেকটি হ্রাসের জন্য। যাইহোক, এক্সএলএ গ্রাফটি অপ্টিমাইজ করতে পারে যাতে এটি একটি একক কার্নেল লঞ্চে ফলাফল গণনা করে। এটি একটি একক GPU কার্নেলে সংযোজন, গুণ এবং হ্রাসকে "ফিউজিং" করে এটি করে। তাছাড়া, এই নিলীন অপারেশন দ্বারা উত্পাদিত অন্তর্বর্তী মান লিখতে না y*z এবং x+y*z মেমরি; পরিবর্তে এটি এই মধ্যবর্তী গণনার ফলাফলগুলি সরাসরি তাদের ব্যবহারকারীদের কাছে "স্ট্রিম" করে এবং তাদের সম্পূর্ণরূপে জিপিইউ রেজিস্টারে রাখে। ফিউশন হল 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_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 ফাইল।

আপনি টেন্সরফ্লো গ্রাফের অভ্যন্তরে এক্সএলএ ক্লাস্টারগুলির এম্বেডিং কল্পনা করে গ্রাফটি ডাম্প করতে পারেন:

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

প্রজননযোগ্য বাগ রিপোর্ট

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

$ 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 এই উত্পন্ন প্রোগ্রাম চলমান।

আরও পড়া

এক্সএলএ ফ্রন্টএন্ডস

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

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

কথা বলে

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

এক্সএলএ ওভারভিউ