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

XLA: การปรับแต่งคอมไพเลอร์เพื่อการเรียนรู้ของเครื่อง

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

ผลลัพธ์คือการปรับปรุงความเร็วและการใช้หน่วยความจำ: มาตรฐานภายในส่วนใหญ่ทำงานได้เร็วขึ้น ~ 1.15x หลังจากเปิดใช้งาน 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 ฟิวชั่นเป็นการเพิ่มประสิทธิภาพที่สำคัญที่สุดเพียงครั้งเดียวของ XLA แบนด์วิดท์หน่วยความจำโดยทั่วไปจะเป็นทรัพยากรที่หายากที่สุดในตัวเร่งความเร็วฮาร์ดแวร์ดังนั้นการลบหน่วยความจำจึงเป็นวิธีที่ดีที่สุดวิธีหนึ่งในการปรับปรุงประสิทธิภาพ

เปิดใช้งาน XLA สำหรับรุ่น TensorFlow

Auto-การจัดกลุ่ม

วิธีที่ง่ายที่สุดในการเริ่มใช้ 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.function

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

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))

API แบบ experimental_compile _compileมีซีแมนทิกส์ที่ ต้องรวบรวม : ฟังก์ชันทั้งหมดถูกคอมไพล์ด้วย 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 และเรียกใช้ซ้ำบนโปรแกรมที่สร้างขึ้น

Frontals XLA

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

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

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