XLA (त्वरित रैखिक बीजगणित) रैखिक बीजगणित के लिए एक डोमेन-विशिष्ट संकलक है जो संभावित रूप से बिना किसी स्रोत कोड परिवर्तन के TensorFlow मॉडल को गति दे सकता है।
परिणाम गति और मेमोरी उपयोग में सुधार हैं: उदाहरण के लिए BERT MLPerf सबमिशन में XLA का उपयोग करके 8 वोल्टा V100 GPU का उपयोग करके ~ 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 ग्राफ़ को अनुकूलित कर सकता है ताकि यह एक कर्नेल लॉन्च में परिणाम की गणना कर सके। यह एक ही जीपीयू कर्नेल में जोड़, गुणा और कमी को "फ़्यूज़िंग" करके करता है। इसके अलावा, यह फ़्यूज्ड ऑपरेशन मेमोरी में 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 में आवश्यक-संकलित शब्दार्थ है: या तो संपूर्ण फ़ंक्शन 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 usage पर एक ट्यूटोरियल वीडियो देखें।
केरस के साथ प्रयोग
केरस मॉडल के लिए, jit_compile=True model.compile के तर्क के रूप में सेट किया जा सकता है:
model.compile(optimizer="adam", jit_compile=True)
वितरित रणनीति के साथ प्रयोग
XLA: jit_compile=True के साथ स्टेप फंक्शन को एनोटेट करके GPU को TF डिस्ट्रीब्यूटेड स्ट्रैटेजी ( MirroredStrategy या MultiWorkerMirroredStrategy ) के साथ इस्तेमाल किया जा सकता है:
@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 के साथ 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 निर्देशिका (ऊपर संदर्भित) की सामग्री संलग्न करें।
यदि संभव हो, तो run_hlo_module उपयोग करके और उत्पन्न प्रोग्राम पर इसे पुनरावृत्त रूप से चलाकर बग को एकल XLA प्रोग्राम में अलग करने का प्रयास करें।
अग्रिम पठन
- ज्ञात मुद्दे XLA के साथ ज्ञात समस्याओं की सूची
- XLA आर्किटेक्चर : XLA आर्किटेक्चर का अवलोकन
- XLA - TensorFlow, Compiled : Google Developers Blog पर पढ़ें
- Github पर XLA स्रोत देखें!
एक्सएलए फ्रंटेंड्स
TensorFlow के अलावा, XLA प्रोग्राम इनके द्वारा भी तैयार किए जा सकते हैं:
- JAX : Python+NumPy प्रोग्राम का कंपोज़ेबल ट्रांसफ़ॉर्मेशन
- जूलिया : वैज्ञानिक कंप्यूटिंग के लिए जूलिया भाषा
- PyTorch : PyTorch फ्रेमवर्क
- एनएक्स : एलिक्सिर प्रोग्रामिंग लैंग्वेज के लिए न्यूमेरिकल कंप्यूटिंग लाइब्रेरी