หมวดหมู่: เวลาคอมไพล์: Scoped Vmem OOM
ข้อผิดพลาดนี้บ่งชี้ว่าโปรแกรมต้องการ Scoped Vector Memory (Vmem) มากกว่าที่จัดสรรไว้
ตัวอย่างข้อความแสดงข้อผิดพลาด:
RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...
แบ็กเอนด์ XLA: TPU
ภาพรวม
TPU มีหน่วยความจำเวกเตอร์ (VMEM) ซึ่งเป็นหน่วยความจำแบบกระดาษทดในเครื่องที่ใช้ เฉพาะ TensorCore (TC) คอมไพเลอร์จัดการ Vmem สำหรับการจัดสรรประเภทต่างๆ ดังนี้
- การจัดสรรขอบเขตคำสั่ง: พื้นที่เก็บข้อมูลชั่วคราวใน Vmem ขณะดำเนินการ คำสั่ง HLO รายการเดียว ซึ่งรวมถึงบัฟเฟอร์ช่วงตัวถูกดำเนินการ (เช่น สำหรับการบัฟเฟอร์คู่) และการล้นรีจิสเตอร์
- การจัดสรรขอบเขตโปรแกรม: การจัดสรรที่อยู่นอกขอบเขตของ คำสั่ง HLO เดียว โดยปกติแล้วจะเป็นผลลัพธ์ชั่วคราวและผลลัพธ์ระดับกลางของ HLO ซึ่งเป็นอินพุตและ/หรือเอาต์พุตของคำสั่ง HLO
ข้อผิดพลาด OOM ของ Vmem ที่กำหนดขอบเขตเวลาคอมไพล์จะเกิดขึ้นเมื่อการจัดสรรที่กำหนดขอบเขตคำสั่ง เกินขีดจำกัดการจัดสรรสำหรับคำสั่งนั้น โดยเราจะควบคุมขีดจำกัดนี้
- ทั่วโลกสำหรับทั้งโปรแกรมผ่านแฟล็ก --xla_tpu_scoped_vmem_limit_kib
และ
- ต่อเคอร์เนลที่กำหนดเองผ่าน พารามิเตอร์ vmem_limit_bytes
โดยปกติแล้ว ข้อผิดพลาดเหล่านี้เกิดจากข้อบกพร่องของคอมไพเลอร์ภายในหรือเคอร์เนลที่กำหนดเองซึ่งเกินขีดจำกัดการจัดสรร
การแก้ไขข้อบกพร่อง
วิเคราะห์ข้อความแสดงข้อผิดพลาดอย่างละเอียดเพื่อระบุว่าข้อผิดพลาดเกิดจากเคอร์เนลที่กำหนดเองหรือ HLO มาตรฐาน ข้อผิดพลาดเนื่องจากเคอร์เนลที่กำหนดเองควรมี ลายเซ็นต่อไปนี้
Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
- Vmem OOM ที่กำหนดขอบเขตเคอร์เนลที่กำหนดเอง: หากข้อผิดพลาดชี้ไปที่เคอร์เนลที่กำหนดเอง → ไปที่ปรับแต่งเคอร์เนลอีกครั้ง
- ปัญหา Vmem ที่ไม่ใช่เคอร์เนล: หาก Vmem OOM เกิดขึ้นเนื่องจาก การดำเนินการที่ไม่ใช่เคอร์เนลที่กำหนดเอง แสดงว่าอาจเป็นข้อบกพร่องของคอมไพเลอร์ภายใน โปรดส่งรายงานข้อบกพร่องใน XLA พร้อมกับ การทิ้ง HLO
ปรับแต่งเคอร์เนล
หากข้อผิดพลาดเกิดจากเคอร์เนลที่กำหนดเอง ให้ใช้เทคนิคต่อไปนี้เพื่อ ลดข้อกำหนดด้านหน่วยความจำของเคอร์เนล
- ปรับขนาดบล็อก: ลดขนาดบล็อก (ขนาดไทล์) ในการกำหนดค่าเคอร์เนล เพื่อลดการใช้ Vmem ที่กำหนดขอบเขต
- ตั้งค่าขีดจำกัด Vmem ที่กำหนดขอบเขตต่อเคอร์เนล: ขอจำนวนหน่วยความจำที่จำเป็นอย่างชัดเจน สำหรับเคอร์เนลที่เฉพาะเจาะจงนั้นโดยใช้ พารามิเตอร์ vmem_limit_bytes
- แก้ไขการระบายสีหน่วยความจำ: ระบายสี/จำกัดอินพุต/เอาต์พุตของเคอร์เนล เป็น Vmem อย่างชัดเจนโดยใช้ pallas.tpu.with_memory_space_constraint ระวังอย่าใส่สีอินพุตเอาต์พุตมากเกินไปใน Vmem เพราะอาจทำให้เกิด OOM โดยรวมของ Vmem
- ปรับขีดจำกัด Vmem: หากการปรับแต่งเฉพาะเคอร์เนลทำได้ยากหรือปัญหา ส่งผลต่อเคอร์เนลจำนวนมาก คุณสามารถปรับขีดจำกัด Vmem ทั่วโลกได้โดยใช้แฟล็ก --xla_tpu_scoped_vmem_limit_kib