แก้ไขข้อผิดพลาด OOM ด้วย XProf

ข้อผิดพลาดเนื่องจากหน่วยความจำไม่เพียงพอ (OOM) จะเกิดขึ้นเมื่อความจุหน่วยความจำแบนด์วิดท์สูง (HBM) ของตัวเร่ง (GPU หรือ TPU) หมด สาเหตุที่พบบ่อยบางประการของปัญหา OOM และเทคนิคการแก้ไขข้อบกพร่องมีรายละเอียดอยู่ใน E1000 - เอกสารประกอบ OOM ของ HBM ในเวลาคอมไพล์และ เอกสารประกอบ JAX เกี่ยวกับการจัดสรรหน่วยความจำ GPU

หน้านี้อธิบายวิธีใช้เครื่องมือดูหน่วยความจำของ XProf เพื่อแสดงภาพการใช้หน่วยความจำของโปรแกรม JAX ระบุอินสแตนซ์การใช้งานสูงสุด และแก้ไขข้อบกพร่องของข้อผิดพลาด OOM ซึ่งเกี่ยวข้องกับขั้นตอนต่อไปนี้

  1. เรียกใช้โปรแกรมด้วย jax.profiler.trace เพื่อบันทึกโปรไฟล์
  2. เริ่ม XProf ในเบื้องหลัง แล้วใช้เครื่องมือดูหน่วยความจำเพื่อดูรายละเอียดการใช้หน่วยความจำ

โปรแกรมตัวอย่าง

โปรแกรม JAX ต่อไปนี้ทำให้เกิดข้อผิดพลาด OOM

import jax
from jax import random
import jax.numpy as jnp


@jax.profiler.trace("/tmp/xprof")
@jax.jit
def oom():
    a = random.normal(random.PRNGKey(1), (327680, 327680), dtype=jnp.bfloat16)
    return a @ a


if __name__ == "__main__":
    oom()

ในเครื่อง TPU โปรแกรมนี้จะล้มเหลวโดยมีข้อความดังนี้

XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=107374182400) would exceed memory (size=17179869184) :: #allocation7 [shape = 'u8[327680,327680]{1,0:T(8,128)(4,1)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of xor_convert_fusion@{}'] :: <no-hlo-instruction>

(ในเครื่อง GPU ข้อผิดพลาดจะมีลักษณะดังนี้ XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 214748364800 bytes.)

เรียกใช้ XProf

ติดตั้ง xprof (pip install xprof) แล้วเริ่มอินสแตนซ์ XProf โดยระบุ ไดเรกทอรีที่จัดเก็บโปรไฟล์

xprof --logdir=/tmp/xprof/ --port=6006

ไปที่อินสแตนซ์ (ในเครื่องที่ใช้ที่ http://localhost:6006) ในเมนูแบบเลื่อนลงเครื่องมือ ให้เลือกโปรแกรมดูหน่วยความจำ แล้วในหน้าต่างเครื่องมือโปรแกรมดูหน่วยความจำ ให้เลือก HBM ในเมนูแบบเลื่อนลงประเภทหน่วยความจำ (โดยปกติจะเลือกไว้โดยค่าเริ่มต้น)

หน้า XProf Memory Viewer สำหรับโปรแกรมตัวอย่างข้างต้น

เอกสารประกอบเครื่องมือ XProf: Memory Viewer อธิบายคอมโพเนนต์ของเครื่องมือและข้อมูลที่แสดง

มุ่งเน้นที่ส่วนการดำเนินการ HLO ที่การจัดสรรหน่วยความจำสูงสุด ซึ่งแสดงแผนภูมิบัฟเฟอร์ 3 รายการที่จุดการใช้งานหน่วยความจำสูงสุด บัฟเฟอร์ประกอบด้วย * อินพุตและเอาต์พุตของโปรแกรม: ชุดการฝึก สถานะของเครื่องมือเพิ่มประสิทธิภาพ ฯลฯ * ข้อมูลชั่วคราวของ TensorCore และ SparseCore: หน่วยความจำแบบไดนามิกที่จำเป็นสำหรับการคำนวณขั้นกลาง (เช่น การเปิดใช้งาน การไล่ระดับสี ฯลฯ)

คุณวางเมาส์เหนือแผนภูมิบัฟเฟอร์เพื่อดูรายละเอียดเพิ่มเติมเกี่ยวกับ Op เช่น ขนาด รูปร่าง ประเภทการจัดสรร และอื่นๆ ซึ่งจะช่วยให้คุณระบุและประเมิน การดำเนินการที่อาจมีอาร์เรย์ชั่วคราวขนาดใหญ่หรือมีอายุการใช้งานยาวนาน รวมถึงเทนเซอร์อินพุต/เทนเซอร์กลาง/เทนเซอร์เอาต์พุตขนาดใหญ่ที่มีการเพิ่มแพดที่ไม่เหมาะสม ฯลฯ ซึ่ง ส่งผลต่อหน่วยความจำสูงสุดและต้องปรับหรือเพิ่มประสิทธิภาพ

ดูเทคนิคการแก้ไขข้อบกพร่องที่เฉพาะเจาะจงได้ใน E1000: การแก้ไขข้อบกพร่อง