Google I/O สำเร็จแล้ว! ติดตามเซสชัน TensorFlow ดูเซสชัน

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

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 ฟิวชั่นคือการเพิ่มประสิทธิภาพที่สำคัญที่สุดเพียงอย่างเดียวของ 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 ไม่สามารถรวบรวมฟังก์ชั่นที่มีขนาดไม่ inferrable: นั่นคือถ้ามันไม่ได้เป็นไปได้ที่จะอนุมานขนาดของเทนเซอร์ทั้งหมดโดยไม่ต้องใช้การคำนวณทั้งหมด ตัวอย่างเช่น ฟังก์ชันต่อไปนี้จะไม่คอมไพล์:

@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 การใช้งาน

การจัดกลุ่มอัตโนมัติ

วิธีง่ายๆในการเริ่มต้นใช้งานในรูปแบบ XLA TensorFlow โดยไม่มีการเปลี่ยนแปลงใด ๆ ที่จะช่วยให้การจัดกลุ่มโดยอัตโนมัติซึ่งจะพบกลุ่ม (subgraphs ที่เกี่ยวโยงกัน) ภายในฟังก์ชั่น TensorFlow ที่สามารถรวบรวมและดำเนินการโดยใช้ XLA Auto-การจัดกลุ่มใน GPU สามารถเปิดใช้งานโดยการตั้งค่า TF_XLA_FLAGS ตัวแปรสภาพแวดล้อม:

$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program

Auto-การจัดกลุ่มมีการเพิ่มประสิทธิภาพในขณะนี้สำหรับเวิร์กโหลด 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

ทอท (ข้างหน้าของเวลา) รวบรวมสำหรับ 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 ต่อหนึ่งแต่ละคลัสเตอร์รวบรวม การแนบสิ่งเหล่านั้นเมื่อส่งรายงานข้อผิดพลาด 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 Frontends

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

  • JAX แปลง Composable ของโปรแกรมหลาม + NumPy:
  • จูเลีย : จูเลียภาษาสำหรับการคำนวณทางวิทยาศาสตร์
  • PyTorch กรอบ PyTorch:
  • Nx ห้องสมุดการคำนวณเชิงตัวเลขสำหรับการเขียนโปรแกรมภาษา Elixir:

เสวนา

ใช้ XLA จาก TF ใช้ jit_compile=True

ภาพรวม XLA