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

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

บทนำ

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

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

เปิดใช้ XLA สําหรับโมเดล TensorFlow

การคอมไพล์ที่อาจไม่เหมาะสมด้วย tf.function(jit_compile=True)

Explicit compilation 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 ไม่สามารถคอมไพล์ฟังก์ชันที่มิติข้อมูลไม่สามารถอนุมานได้ กล่าวคือ ไม่สามารถอนุมานมิติข้อมูลของเทนเซอร์ทั้งหมดได้โดยไม่ต้องเรียกใช้การคํานวณทั้งหมด ตัวอย่างเช่น ฟังก์ชันต่อไปนี้จะคอมไพล์ไม่ได้

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

ดูตัวอย่างการใช้งานโดยละเอียดและวิดีโอแนะนำการใช้งาน jit_compile=True ได้ที่ tutorial colab

การใช้งานกับ 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 พร้อมอินทริเนต 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
  • Julia: ภาษา Julia สําหรับการประมวลผลทางวิทยาศาสตร์
  • PyTorch: เฟรมเวิร์ก PyTorch
  • Nx: ไลบรารีการคำนวณเชิงตัวเลขสำหรับภาษาโปรแกรม Elixir

การบรรยาย

ใช้ XLA จาก TF โดยใช้ jit_compile=True

ภาพรวม XLA