XLA (Accelerated Linear Algebra) เป็นคอมไพเลอร์เฉพาะโดเมนสำหรับพีชคณิตเชิงเส้นที่สามารถเร่งโมเดล TensorFlow โดยอาจไม่มีการเปลี่ยนแปลงซอร์สโค้ด
ผลลัพธ์คือการปรับปรุงความเร็วและการใช้หน่วยความจำ: เช่น การส่ง BERT MLPerf โดยใช้ 8 Volta V100 GPUs โดยใช้ XLA ได้รับการปรับปรุงประสิทธิภาพ ~7x และการปรับปรุงขนาดชุด ~5x:

การแนะนำ
เมื่อรันโปรแกรม TensorFlow การดำเนินการทั้งหมดจะถูกดำเนินการแยกกันโดยตัวดำเนินการ TensorFlow การดำเนินการ TensorFlow แต่ละครั้งมีการติดตั้งเคอร์เนล GPU ที่คอมไพล์แล้วซึ่งตัวดำเนินการส่งไปให้
XLA ให้โหมดทางเลือกในการรันโมเดล: คอมไพล์กราฟ TensorFlow เป็นลำดับของเคอร์เนลการคำนวณที่สร้างขึ้นโดยเฉพาะสำหรับโมเดลที่กำหนด เนื่องจากเคอร์เนลเหล่านี้มีลักษณะเฉพาะสำหรับโมเดล จึงสามารถใช้ประโยชน์จากข้อมูลเฉพาะของโมเดลเพื่อเพิ่มประสิทธิภาพได้ ตัวอย่างเช่น ลองดูที่การเพิ่มประสิทธิภาพ XLA ในบริบทของการคำนวณ TensorFlow อย่างง่าย:
def model_fn(x, y, z):
return tf.reduce_sum(x + y * z)
เรียกใช้โดยไม่ใช้ XLA กราฟจะเปิดตัวเคอร์เนลสามเมล็ด: เมล็ดหนึ่งสำหรับการคูณ หนึ่งสำหรับการบวก และอีกอันสำหรับการลดลง อย่างไรก็ตาม XLA สามารถเพิ่มประสิทธิภาพกราฟเพื่อให้คำนวณผลลัพธ์ในการเรียกใช้เคอร์เนลเดียว มันทำได้โดยการ "หลอมรวม" การบวก การคูณ และการลดลงในเคอร์เนล GPU เดียว นอกจากนี้ การดำเนินการที่หลอมรวมนี้ไม่ได้เขียนค่ากลางที่สร้างโดย y*z
และ x+y*z
ไปยังหน่วยความจำ แทนที่จะ "สตรีม" ผลลัพธ์ของการคำนวณระดับกลางเหล่านี้โดยตรงไปยังผู้ใช้ในขณะที่เก็บไว้ในการลงทะเบียน GPU ทั้งหมด Fusion คือการเพิ่มประสิทธิภาพที่สำคัญที่สุดเพียงครั้งเดียวของ XLA แบนด์วิธของหน่วยความจำโดยทั่วไปเป็นทรัพยากรที่หายากที่สุดในตัวเร่งความเร็วของฮาร์ดแวร์ ดังนั้นการลบการทำงานของหน่วยความจำจึงเป็นวิธีที่ดีที่สุดวิธีหนึ่งในการปรับปรุงประสิทธิภาพ
เปิดใช้งาน XLA สำหรับโมเดล TensorFlow
การรวบรวมอย่างชัดเจนด้วย tf.function(jit_compile=True)
API การคอมไพล์อย่างชัดเจนมีการควบคุมแบบละเอียดสำหรับการเลือกฟังก์ชันที่ควรคอมไพล์ ตัวอย่างเช่น ฟังก์ชัน 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))
API jit_compile
มีความหมาย ที่ต้องคอมไพล์ : ฟังก์ชันทั้งหมดถูกคอมไพล์ด้วย 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]))
ดู colab บทช่วยสอน สำหรับตัวอย่างการใช้งานโดยละเอียดเพิ่มเติม และ วิดีโอบทช่วยสอน เกี่ยวกับ jit_compile=True
applications
การใช้งานกับ Keras
สำหรับโมเดล 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)
การจัดกลุ่มอัตโนมัติ
วิธีง่ายๆ ในการเริ่มใช้ XLA ในโมเดล TensorFlow โดยไม่มีการเปลี่ยนแปลงใดๆ คือการเปิดใช้ งานการทำคลัสเตอร์อัตโนมัติ ซึ่งจะค้นหา คลัสเตอร์ (กราฟย่อยที่เชื่อมต่อ) โดยอัตโนมัติภายในฟังก์ชัน TensorFlow ซึ่งสามารถรวบรวมและดำเนินการโดยใช้ XLA สามารถเปิดใช้งานการทำคลัสเตอร์อัตโนมัติบน 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 (ล่วงหน้า) สำหรับ CPU ด้วย tfcompile
คุณยังสามารถใช้เครื่องมือ 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
สร้างไฟล์ในการเป็นตัวแทนระดับกลาง ของ LLVM ด้วย NVPTX ที่แท้จริงmodule_XXXX.ptx
สร้างไฟล์ PTX
คุณยังสามารถดัมพ์กราฟแสดงภาพการฝังกลุ่ม XLA ภายในกราฟ TensorFlow ด้วย:
$ 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 เดียวโดยใช้ run_hlo_module
และรันโปรแกรมที่สร้างขึ้นซ้ำๆ
อ่านเพิ่มเติม
- ปัญหาที่ทราบ รายการปัญหาที่ทราบเกี่ยวกับ XLA
- สถาปัตยกรรม XLA : ภาพรวมของสถาปัตยกรรม XLA
- XLA - TensorFlow รวบรวม : อ่านในบล็อก Google Developers
- ตรวจสอบ แหล่ง XLA บน Github!
XLA ส่วนหน้า
นอกเหนือจาก TensorFlow แล้ว โปรแกรม XLA สามารถสร้างโดย:
- JAX : การแปลงที่ประกอบได้ของโปรแกรม Python+NumPy
- Julia : ภาษาจูเลียสำหรับการคำนวณทางวิทยาศาสตร์
- PyTorch : กรอบ PyTorch
- Nx : ไลบรารีการคำนวณเชิงตัวเลขสำหรับภาษาโปรแกรม Elixir