एक्सएलए (एक्सएलए): मशीन लर्निंग के लिए, कंपाइलर को ऑप्टिमाइज़ करना

OpenXLA, लीनियर ऐल्जेब्रा के लिए डोमेन-स्पेसिफ़िक कंपाइलर है. यह TensorFlow मॉडल को तेज़ कर सकता है. इसके लिए, सोर्स कोड में बदलाव करने की ज़रूरत नहीं होती.

परिचय

जब कोई TensorFlow प्रोग्राम चलाया जाता है, तो TensorFlow एक्सेक्यूटर सभी ऑपरेशन को अलग-अलग तरीके से लागू करता है. TensorFlow के हर ऑपरेशन में, पहले से संकलित GPU कोर लागू होता है. इसे एक्ज़ीक्यूटर डिस्पैच करता है.

XLA, मॉडल चलाने का एक अन्य मोड उपलब्ध कराता है: यह TensorFlow ग्राफ़ को, खास तौर पर दिए गए मॉडल के लिए जनरेट किए गए कंप्यूटेशन कर्नेल के क्रम में कंपाइल करता है. ये कर्नेल, मॉडल के हिसाब से अलग-अलग होते हैं. इसलिए, ये ऑप्टिमाइज़ेशन के लिए, मॉडल के हिसाब से जानकारी का फ़ायदा उठा सकते हैं. उदाहरण के लिए, TensorFlow के किसी सामान्य कैलकुलेशन के संदर्भ में, XLA के ऑप्टिमाइज़ेशन की सुविधा को देखें:

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) के साथ साफ़ तौर पर कॉन्टेंट कंपाइल करना

एक्सप्लिशिट कंपाइलेशन एपीआई, यह चुनने के लिए बेहतर कंट्रोल देता है कि किन फ़ंक्शन को कंपाइल किया जाना चाहिए. उदाहरण के लिए, MNIST ट्रेनिंग करने वाले इस TensorFlow फ़ंक्शन को 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 में must-compile सेमेटिक्स है: या तो पूरे फ़ंक्शन को 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]))

इस्तेमाल के उदाहरण के बारे में ज़्यादा जानने के लिए, ट्यूटोरियल वाला colab देखें. साथ ही, jit_compile=True के इस्तेमाल के बारे में ट्यूटोरियल वीडियो देखें.

Keras के साथ इस्तेमाल करना

Keras मॉडल के लिए, jit_compile=True को model.compile के लिए आर्ग्युमेंट के तौर पर सेट किया जा सकता है:

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

डिस्ट्रिब्यूट की गई रणनीति के साथ इस्तेमाल करना

XLA:GPU का इस्तेमाल, TF की डिस्ट्रिब्यूटेड रणनीति (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 एनवायरमेंट वैरिएबल सेट करके, जीपीयू पर अपने-आप क्लस्टर बनने की सुविधा चालू की जा सकती है:

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

फ़िलहाल, अपने-आप क्लस्टर बनने की सुविधा को GPU वर्कलोड के लिए ऑप्टिमाइज़ किया गया है. हालांकि, इसे सीपीयू पर भी चालू किया जा सकता है. इसके लिए, --tf_xla_cpu_global_jit फ़्लैग का इस्तेमाल करें:

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

इस्तेमाल के उदाहरण के बारे में ज़्यादा जानने के लिए, अपने-आप क्लस्टर बनाने के ट्यूटोरियल वाला Colab देखें.

tfcompile वाले सीपीयू के लिए, पहले से कंपाइल करने की सुविधा

इसके अलावा, स्टैंडअलोन tfcompile टूल का भी इस्तेमाल किया जा सकता है. यह टूल, TensorFlow ग्राफ़ को एक्ज़ीक्यूटेबल कोड में बदलता है. हालांकि, यह सिर्फ़ x86-64 सीपीयू के लिए उपलब्ध है.

कंपाइल किए गए प्रोग्राम की जांच करना

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 फ़ाइलें.

TensorFlow ग्राफ़ में XLA क्लस्टर को एम्बेड करने के बारे में जानकारी देने वाले ग्राफ़ को भी डंप किया जा सकता है. इसके लिए, इनका इस्तेमाल करें:

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

गड़बड़ी की ऐसी रिपोर्ट जिन्हें फिर से बनाया जा सकता है

गड़बड़ी की रिपोर्ट को फिर से जनरेट करना बहुत आसान होता है, अगर उसमें जनरेट किए गए XLA प्रोग्राम और इस्तेमाल किए गए ऑटो-क्लस्टरिंग एम्बेडिंग के लिए डंप शामिल होते हैं. अपने-आप क्लस्टर बनने की सुविधा के साथ चल रहे TensorFlow प्रोग्राम के लिए, उन्हें जनरेट करने के लिए:

$ 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 प्रोग्राम के कॉम्पोज़ेबल ट्रांसफ़ॉर्मेशन
  • Julia: वैज्ञानिक कंप्यूटिंग के लिए Julia भाषा
  • PyTorch: PyTorch फ़्रेमवर्क
  • Nx: Elixir प्रोग्रामिंग लैंग्वेज के लिए, संख्यात्मक कंप्यूटिंग लाइब्रेरी

भाषण

jit_compile=True का इस्तेमाल करके, TF से XLA का इस्तेमाल करना

XLA के बारे में खास जानकारी