XLA (त्वरित रैखिक बीजगणित) रैखिक बीजगणित के लिए एक डोमेन-विशिष्ट संकलक है जो संभावित रूप से बिना किसी स्रोत कोड परिवर्तन के TensorFlow मॉडल को गति दे सकता है।
जैसे बर्ट में: परिणाम की गति और स्मृति के उपयोग में सुधार कर रहे हैं MLPerf 8 वोल्टा V100 GPUs 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(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 एपीआई चाहिए संकलन अर्थ विज्ञान है: या तो पूरे समारोह 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
ऑटो क्लस्टरिंग वर्तमान GPU वर्कलोड के लिए अनुकूलित है, लेकिन इसके साथ ही ध्वज का उपयोग करके यह भी CPU पर सक्रिय किया जा सकता --tf_xla_cpu_global_jit :
$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program
एक विस्तृत उपयोग उदाहरण के लिए देख ऑटो क्लस्टरिंग ट्यूटोरियल colab ।
साथ सीपीयू के लिए AOT (आगे-ऑफ-समय) संकलन tfcompile
तुम भी एक स्टैंडअलोन उपयोग कर सकते हैं tfcompile उपकरण है, जो (केवल x86-64 सीपीयू के लिए) निष्पादन योग्य कोड में 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 फ़ाइलें।
आप 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 कार्यक्रम के लिए एक बग को अलग करने की कोशिश replay_computation और iteratively उत्पन्न कार्यक्रमों पर यह चल रहा है।
अग्रिम पठन
- ज्ञात समस्याओं सूची XLA साथ ज्ञात समस्याओं की
- XLA वास्तुकला : XLA वास्तुकला का अवलोकन
- XLA - TensorFlow, संकलित : Google डेवलपर ब्लॉग पढ़ें
- की जाँच करें XLA स्रोत Github पर!
एक्सएलए फ्रंटेंड
TensorFlow के अलावा, XLA प्रोग्राम इसके द्वारा उत्पन्न किए जा सकते हैं:
- JAX : अजगर + NumPy कार्यक्रमों की composable परिवर्तनों
- जूलिया : वैज्ञानिक कंप्यूटिंग के लिए जूलिया भाषा
- PyTorch : PyTorch ढांचा
- NX : अमृत प्रोग्रामिंग भाषा के लिए संख्यात्मक कंप्यूटिंग पुस्तकालय