7 दिसंबर को महिला एमएल संगोष्ठी में भाग लें अभी पंजीकरण करें

XLA: मशीन लर्निंग के लिए ऑप्टिमाइज़िंग कंपाइलर

XLA (त्वरित रैखिक बीजगणित) रैखिक बीजगणित के लिए एक डोमेन-विशिष्ट संकलक है जो संभावित रूप से बिना किसी स्रोत कोड परिवर्तन के TensorFlow मॉडल को गति दे सकता है।

परिणाम गति और स्मृति उपयोग में सुधार हैं: उदाहरण के लिए BERT MLPerf सबमिशन में 8 वोल्टा V100 GPU का उपयोग करके XLA का उपयोग करके ~ 7x प्रदर्शन सुधार और ~ 5x बैच आकार में सुधार हासिल किया है:

परिचय

जब एक 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 के साथ स्पष्ट संकलन tf.function(jit_compile=True)

स्पष्ट संकलन एपीआई यह चुनने के लिए एक बढ़िया नियंत्रण प्रदान करता है कि कौन से कार्यों को संकलित किया जाना चाहिए। उदाहरण के लिए, निम्नलिखित 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 में सेमेन्टिक्स को संकलित करना चाहिए : या तो संपूर्ण फ़ंक्शन errors.InvalidArgumentError के साथ संकलित किया गया है, या कोई त्रुटि है। 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 . के साथ प्रयोग

केरस मॉडल के लिए, 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 का उपयोग करके संकलित और निष्पादित किया जा सकता है। GPU पर ऑटो-क्लस्टरिंग को TF_XLA_FLAGS परिवेश चर सेट करके सक्षम किया जा सकता है:

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

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

$ 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 में पा सकते हैं:

आप 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 डायरेक्टरी (ऊपर संदर्भित) की सामग्री संलग्न करें।

यदि संभव हो, तो replay_computation का उपयोग करके और इसे जेनरेट किए गए प्रोग्रामों पर पुनरावृत्त रूप से चलाकर बग को एकल XLA प्रोग्राम में अलग करने का प्रयास करें।

अग्रिम पठन

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

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

  • JAX : Python+NumPy प्रोग्राम का कंपोज़ेबल ट्रांसफ़ॉर्मेशन
  • जूलिया : वैज्ञानिक कंप्यूटिंग के लिए जूलिया भाषा
  • PyTorch : PyTorch फ्रेमवर्क
  • एनएक्स : अमृत प्रोग्रामिंग भाषा के लिए संख्यात्मक कंप्यूटिंग पुस्तकालय

बाते

jit_compile=True . का उपयोग करके TF से XLA का उपयोग करना

एक्सएलए अवलोकन