इस दस्तावेज़ में, XLA कॉम्पोनेंट बनाने का तरीका बताया गया है.
अगर आपने XLA रिपॉज़िटरी को क्लोन नहीं किया है या Bazel को इंस्टॉल नहीं किया है, तो XLA डेवलपर गाइड के शुरुआती सेक्शन देखें.
Linux
कॉन्फ़िगर करें
XLA बिल्ड, रिपॉज़िटरी के रूट डायरेक्ट्री में मौजूद .bazelrc फ़ाइल से कॉन्फ़िगर किए जाते हैं. ./configure.py स्क्रिप्ट का इस्तेमाल, सामान्य सेटिंग में बदलाव करने के लिए किया जा सकता है.
अगर आपको कॉन्फ़िगरेशन बदलना है, तो रिपॉज़िटरी की रूट डायरेक्ट्री से ./configure.py स्क्रिप्ट चलाएं. इस स्क्रिप्ट में, XLA डिपेंडेंसी की जगह और बिल्ड कॉन्फ़िगरेशन के अतिरिक्त विकल्पों (उदाहरण के लिए, कंपाइलर फ़्लैग) के लिए फ़्लैग होते हैं. ज़्यादा जानकारी के लिए, सैंपल सेशन सेक्शन देखें.
सीपीयू सपोर्ट
हमारा सुझाव है कि XLA को बनाने और उसकी जांच करने के लिए, सही Docker इमेज का इस्तेमाल करें. जैसे, ml-build. इसका इस्तेमाल GitHub पर XLA के सीआई वर्कफ़्लो में भी किया जाता है. ml-build इमेज में, Clang 18 पहले से इंस्टॉल होता है.
docker run -itd --rm \
--name xla \
-w /xla \
-v $PWD:/xla \
us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest \
bash
Docker कंटेनर का इस्तेमाल करके, सीपीयू सपोर्ट के साथ XLA बनाया जा सकता है. इसके लिए, ये निर्देश चलाएं:
docker exec xla ./configure.py --backend=CPU
docker exec xla bazel build \
--spawn_strategy=sandboxed \
--test_output=all \
//xla/...
अगर आपको सीपीयू सपोर्ट के साथ XLA टारगेट Docker का इस्तेमाल किए बिना बनाने हैं, तो आपको Clang इंस्टॉल करना होगा. XLA को फ़िलहाल CI में Clang 18 के साथ बनाया गया है. हालांकि, पुराने वर्शन भी काम करने चाहिए.
टारगेट कॉन्फ़िगर करने और बनाने के लिए, ये कमांड चलाएं:
./configure.py --backend=CPU
bazel build \
--spawn_strategy=sandboxed \
--test_output=all \
//xla/...
GPU के साथ काम करने की सुविधा
हमारा सुझाव है कि ऊपर बताए गए Docker कंटेनर का इस्तेमाल करके, GPU सपोर्ट के साथ XLA बनाएं.
सभी जीपीयू का ऐक्सेस रखने वाले Docker कंटेनर को शुरू करने के लिए, यह निर्देश चलाएं:
docker run -itd --rm \
--gpus all \
--name xla_gpu \
-w /xla \
-v $PWD:/xla \
us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest \
bash
जीपीयू की सुविधा के साथ XLA बनाने के लिए, ये कमांड चलाएं:
docker exec xla_gpu ./configure.py --backend=CUDA
docker exec xla_gpu bazel build \
--spawn_strategy=sandboxed \
--test_output=all \
//xla/...
- Docker कंटेनर शुरू करते समय,
--gpus allफ़्लैग का इस्तेमाल न करें. - CUDA की कंप्यूटिंग क्षमताओं को मैन्युअल तरीके से तय करें. उदाहरण के लिए:
docker exec xla_gpu ./configure.py --backend=CUDA \
--cuda_compute_capabilities="9.0"
TensorFlow के GPU डॉकर इमेज के बारे में ज़्यादा जानकारी के लिए, यह दस्तावेज़ पढ़ें.
GPU की सुविधा के साथ XLA टारगेट, Docker के बिना भी बनाए जा सकते हैं. इन कमांड का इस्तेमाल करके, टारगेट कॉन्फ़िगर और बिल्ड करें:
./configure.py --backend=CUDA
bazel build \
--spawn_strategy=sandboxed \
--test_output=all \
//xla/...
हर्मेटिक CUDA के बारे में ज़्यादा जानने के लिए, यह दस्तावेज़ पढ़ें.
JAX CI/रिलीज़ कंटेनर का इस्तेमाल करके, CUDA/cuDNN की सुविधा के साथ XLA बनाना
XLA, एक कंपाइलर है. इसका इस्तेमाल JAX में किया जाता है. JAX को PyPI व्हील के ज़रिए डिस्ट्रिब्यूट किया जाता है. JAX के कंटिन्यूअस इंटिग्रेशन के दस्तावेज़ में बताया गया है कि tensorflow/ml-build:latest Docker कंटेनर का इस्तेमाल करके, JAX व्हील कैसे बनाए जाते हैं.
हम इन निर्देशों का इस्तेमाल करके, JAX कंटेनर में XLA टारगेट भी बना सकते हैं. इससे यह पक्का होता है कि XLA टारगेट का बिल्ड कॉन्फ़िगरेशन, JAX/XLA के बिल्ड कॉन्फ़िगरेशन के साथ काम करता है. यह तब काम आ सकता है, जब हमें XLA टूल का इस्तेमाल करके वर्कलोड के नतीजों को फिर से जनरेट करना हो. ये टूल, मूल रूप से JAX में बनाए गए थे.
JAX CI कंटेनर में XLA टारगेट बनाना
- JAX के डेटा स्टोर करने की जगह को क्लोन करें और 'jax' डायरेक्ट्री पर जाएं
git clone <a href="https://github.com/jax-ml/jax.git">https://github.com/jax-ml/jax.git</a>
cd jax
- JAX CI/रिलीज़ डॉकर कंटेनर शुरू करने के लिए, यह कमांड चलाएं:
./ci/utilities/run_docker_container.sh
इससे 'jax' नाम का Docker कंटेनर शुरू हो जाएगा.
- कंटेनर में jax-cuda-plugin टारगेट बनाने के लिए, इसका इस्तेमाल करें:
docker exec jax ./ci/build_artifacts.sh jax-cuda-plugin
इससे .jax_configure.bazelrc फ़ाइल बन जाएगी. इसमें ज़रूरी बिल्ड कॉन्फ़िगरेशन शामिल होगा. जैसे, CUDA/cuDNN के साथ काम करने की सुविधा
- कंटेनर में इंटरैक्टिव शेल को ऐक्सेस करें:
docker exec -ti jax /bin/bash
अब आपको कंटेनर में मौजूद /jax डायरेक्ट्री में होना चाहिए
- इस कमांड का इस्तेमाल करके, XLA टारगेट बनाएं. उदाहरण के लिए:
/usr/local/bin/bazel build \
--config=cuda_libraries_from_stubs \
--verbose_failures=true \
@xla//xla/tools/multihost_hlo_runner:hlo_runner_main
अगर चाहें, तो HERMETIC envs को बदला जा सकता है. उदाहरण के लिए:
--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"
- ज़रूरत पड़ने पर, होस्ट ओएस से ऐक्सेस करने के लिए, नतीजे के तौर पर मिले आर्टफ़ैक्ट को
/jax/distमें कॉपी करें
cp bazel-bin/external/xla/xla/tools/multihost_hlo_runner/hlo_runner_main \
./dist/
- इंटरैक्टिव शेल से बाहर निकलें:
exit