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

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

शुरुआती जानकारी

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

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

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

XLA के बिना चलने पर, ग्राफ़ तीन कर्नेल लॉन्च करता है: एक गुणा के लिए, एक जोड़ने के लिए और एक कम करने के लिए. हालांकि, XLA ग्राफ़ को ऑप्टिमाइज़ कर सकता है, ताकि यह नतीजे की गिनती एक ही कर्नेल लॉन्च में कर सके. ऐसा करने के लिए, यह फ़ंक्शन एक ही जीपीयू कर्नेल में जोड़, गुणा, और छोटा करके "फ़्यूज़" करता है. साथ ही, यह फ़्यूजन ऑपरेशन y*z और x+y*z से जनरेट हुई इंटरमीडिएट वैल्यू को मेमोरी में नहीं लिखता है. इसके बजाय, यह इन इंटरमीडिएट कंप्यूटेशन के नतीजों को सीधे उनके उपयोगकर्ताओं के साथ "स्ट्रीम" करता है, जबकि उन्हें पूरी तरह से जीपीयू रजिस्टर में रखा जाता है. फ़्यूज़न, 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 एपीआई में कंपाइल करना सिमैंटिक होता है: या तो पूरे फ़ंक्शन को 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 के इस्तेमाल के बारे में ट्यूटोरियल वीडियो देखें.

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

Keras मॉडल के लिए, 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 एनवायरमेंट वैरिएबल सेट करके चालू की जा सकती है:

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

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

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

इस्तेमाल के उदाहरण के लिए, ऑटो-क्लस्टरिंग से जुड़े ट्यूटोरियल कोलैब देखें.

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 की गड़बड़ी की रिपोर्ट सबमिट करते समय, इन्हें अटैच करना बहुत मददगार होता है!

  • module_XXXX.ir-*.ll NVPTX में मूल सुविधाओं के साथ, LLVM इंटरमीडिएट वर्शन में जनरेट की गई फ़ाइलें.

  • 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 डायरेक्ट्री का कॉन्टेंट अटैच करें (ऊपर बताया गया है).

अगर मुमकिन हो, तो किसी गड़बड़ी को एक XLA प्रोग्राम से अलग करने की कोशिश करें. इसके लिए, run_hlo_module का इस्तेमाल करें और उसे जनरेट किए गए प्रोग्राम पर बार-बार चलाएं.

इसके बारे में और पढ़ें

एक्सएलए फ़्रंटएंड

TensorFlow के अलावा, XLA प्रोग्राम इन तरीकों से जनरेट किए जा सकते हैं:

  • JAX: Python+NumPy प्रोग्राम के आसान ट्रांसफ़ॉर्मेशन
  • जूलिया: वैज्ञानिक कंप्यूटिंग के लिए जूलिया भाषा
  • PyTorch: PyTorch फ़्रेमवर्क
  • Nx: Elixir प्रोग्रामिंग भाषा के लिए न्यूमेरिकल कंप्यूटिंग लाइब्रेरी

भाषण

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

XLA की खास जानकारी