XLA: การเพิ่มประสิทธิภาพคอมไพเลอร์สำหรับแมชชีนเลิร์นนิง

OpenXLA เป็นคอมไพเลอร์เฉพาะของโดเมนสำหรับพีชคณิตเชิงเส้นที่ช่วยเร่งโมเดล TensorFlow ได้แม้ไม่มีการเปลี่ยนแปลงซอร์สโค้ด

เกริ่นนำ

เมื่อเรียกใช้โปรแกรม TensorFlow ตัวดำเนินการ TensorFlow จะดำเนินการการดำเนินการทั้งหมดทีละรายการ การดำเนินการของ TensorFlow แต่ละรายการมีการใช้งานเคอร์เนล GPU ที่คอมไพล์ไว้ล่วงหน้าซึ่งผู้ดำเนินการจัดส่งไปถึง

XLA ให้โหมดทางเลือกของโมเดลการเรียกใช้ โดยคอมไพล์กราฟ TensorFlow เป็นลำดับเคอร์เนลการคำนวณที่สร้างขึ้นสำหรับโมเดลที่ระบุโดยเฉพาะ เนื่องจากเคอร์เนลเหล่านี้มีความเฉพาะตัวของโมเดล จึงอาจใช้ประโยชน์จากข้อมูลเฉพาะรุ่นเพื่อเพิ่มประสิทธิภาพได้ ลองดูตัวอย่างการเพิ่มประสิทธิภาพ XLA ในบริบทของการคำนวณ TensorFlow แบบง่าย

def model_fn(x, y, z):
  return tf.reduce_sum(x + y * z)

เมื่อเรียกใช้โดยไม่มี XLA กราฟจะเปิดตัวเคอร์เนล 3 รายการ คือ เคอร์เนล 1 รายการสำหรับการคูณ 1 ครั้งสำหรับการเพิ่ม 1 คลิปสำหรับการลดทอน อย่างไรก็ตาม 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))

jit_compile API มีความหมายของต้องคอมไพล์ กล่าวคือ ฟังก์ชันทั้งหมดจะคอมไพล์ด้วย XLA หรือมีการยกเว้น errors.InvalidArgumentError ก็ได้ ปัจจุบัน XLA ไม่สามารถคอมไพล์ฟังก์ชันที่ขนาดไม่อนุมานได้ กล่าวคือ หากไม่สามารถอนุมานมิติข้อมูลของ Tensor ทั้งหมดโดยไม่ต้องเรียกใช้การคำนวณทั้งหมด เช่น ฟังก์ชันต่อไปนี้จะไม่คอมไพล์

@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

การใช้งานกับ 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 ได้โดยใช้ Flag --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 เป็นโค้ดที่สั่งการได้ (สำหรับ CPU x86-64 เท่านั้น)

ตรวจสอบโปรแกรมที่คอมไพล์แล้ว

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 สร้าง 1 รายการต่อคลัสเตอร์ที่คอมไพล์แต่ละคลัสเตอร์ การแนบไฟล์เหล่านี้เมื่อส่งรายงานข้อบกพร่อง XLA จะเป็นประโยชน์อย่างยิ่ง

  • module_XXXX.ir-*.ll ไฟล์ที่สร้างขึ้นในการนำเสนอตัวกลางของ LLVM ที่มี Intrinics ของ 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

นอกเหนือจาก TensorFlow แล้ว โปรแกรม XLA ยังสร้างขึ้นได้ด้วยสิ่งต่อไปนี้

  • JAX: การเปลี่ยนรูปแบบแบบองค์ประกอบต่างๆ ของโปรแกรม Python+NumPy
  • จูเลีย: ภาษาจูเลียสำหรับ การประมวลผลทางวิทยาศาสตร์
  • PyTorch: เฟรมเวิร์กของ PyTorch
  • Nx: ไลบรารีการประมวลผลที่เป็นตัวเลขสำหรับภาษาโปรแกรมของ Elixir

การบรรยาย

กำลังใช้ XLA จาก TF โดยใช้ jit_compile=True

ภาพรวม XLA