หน้านี้ได้รับการแปลโดย Cloud Translation API
Switch to English

XLA: การเพิ่มประสิทธิภาพคอมไพเลอร์สำหรับ Machine Learning

XLA (Accelerated Linear Algebra) เป็นคอมไพเลอร์เฉพาะโดเมนสำหรับพีชคณิตเชิงเส้นที่สามารถเร่งโมเดล TensorFlow โดยที่อาจไม่มีการเปลี่ยนแปลงซอร์สโค้ด

ผลลัพธ์คือการปรับปรุงความเร็วและการใช้หน่วยความจำ: เกณฑ์มาตรฐานภายในส่วนใหญ่ทำงานเร็วขึ้น ~ 1.15 เท่าหลังจากเปิดใช้งาน XLA ชุดข้อมูลด้านล่างนี้ได้รับการประเมินจาก NVidia V100 GPU ตัวเดียว:

บทนำ

เมื่อรันโปรแกรม 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

การจัดคลัสเตอร์อัตโนมัติ

วิธีที่ง่ายที่สุดในการเริ่มใช้ 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 บทช่วยสอนการทำคลัสเตอร์อัตโนมัติ

การคอมไพล์อย่างชัดเจนด้วยฟังก์ชัน tf

การจัดคลัสเตอร์อัตโนมัติเป็นเครื่องมือที่ยอดเยี่ยมในการทำให้โมเดลเร็วขึ้นโดยไม่ต้องเปลี่ยนแปลงโค้ด แต่อาจยากที่จะเข้าใจว่ามีการเปลี่ยนแปลงอะไรบ้าง

Explicit compilation API นำเสนอการควบคุมที่ละเอียดยิ่งขึ้นสำหรับการเลือกฟังก์ชันที่ควรคอมไพล์ ตัวอย่างเช่นฟังก์ชัน TensorFlow ต่อไปนี้ซึ่งดำเนินการฝึกอบรม MNIST จะคอมไพล์ด้วย XLA: 

@tf.function(experimental_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))

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

@tf.function
def not_compilable(x):
  return tf.unique(x)

รูปร่างอาจแตกต่างกันไปตามการวิ่งแม้ว่า: 

@tf.function(experimental_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 บทช่วยสอน สำหรับตัวอย่างการใช้งานโดยละเอียดเพิ่มเติม

การรวบรวม 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 intrinsics

  • 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 เดียวโดยใช้ replay_computation และรันซ้ำบนโปรแกรมที่สร้างขึ้น

ส่วนหน้า XLA

นอกเหนือจาก TensorFlow แล้วโปรแกรม XLA ยังสามารถสร้างได้โดย:

  • JAX : การแปลงแบบคอมโพสิตของโปรแกรม Python + NumPy
  • Julia : ภาษา Julia สำหรับการคำนวณทางวิทยาศาสตร์
  • PyTorch : กรอบ PyTorch

อ่านเพิ่มเติม