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 में पा सकते हैं:
module_XXXX.*_optimizations.txtजेनरेट किए गए XLA प्रोग्राम , प्रत्येक संकलित क्लस्टर के लिए एक। XLA बग रिपोर्ट सबमिट करते समय उन्हें संलग्न करना अत्यंत सहायक होता है!module_XXXX.ir-*.llएनवीपीटीएक्स इंट्रिनिक्स के साथ एलएलवीएम इंटरमीडिएट प्रतिनिधित्व में जेनरेट की गई फाइलें।module_XXXX.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 डायरेक्टरी (ऊपर संदर्भित) की सामग्री संलग्न करें।
यदि संभव हो, तो replay_computation का उपयोग करके और इसे जेनरेट किए गए प्रोग्रामों पर पुनरावृत्त रूप से चलाकर बग को एकल XLA प्रोग्राम में अलग करने का प्रयास करें।
अग्रिम पठन
- ज्ञात समस्याएँ XLA के साथ ज्ञात समस्याओं की सूची
- XLA आर्किटेक्चर : XLA आर्किटेक्चर का अवलोकन
- XLA - TensorFlow, संकलित : Google Developers Blog पर पढ़ें
- जीथब पर एक्सएलए स्रोत देखें!
एक्सएलए फ्रंटेंड
TensorFlow के अलावा, XLA प्रोग्राम इसके द्वारा उत्पन्न किए जा सकते हैं:
- JAX : Python+NumPy प्रोग्राम का कंपोज़ेबल ट्रांसफ़ॉर्मेशन
- जूलिया : वैज्ञानिक कंप्यूटिंग के लिए जूलिया भाषा
- PyTorch : PyTorch फ्रेमवर्क
- एनएक्स : अमृत प्रोग्रामिंग भाषा के लिए संख्यात्मक कंप्यूटिंग पुस्तकालय