หมวดหมู่: รันไทม์: จัดสรรโปรแกรมไม่สำเร็จ
ข้อผิดพลาดนี้บ่งชี้ว่ารันไทม์ XLA ในอุปกรณ์ TPU โหลดไฟล์ที่เรียกใช้งานได้ของโปรแกรม XLA ที่คอมไพล์แล้วลงใน HBM ของ TPU ไม่สำเร็จ
ตัวอย่างข้อความแสดงข้อผิดพลาด:
XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).
แบ็กเอนด์ XLA: TPU
ภาพรวม
โดยข้อผิดพลาดนี้มักจะเกิดจากสาเหตุต่อไปนี้
- ขนาดโปรแกรมเกิน HBM ที่พร้อมใช้งาน: โปรแกรม XLA ที่คอมไพล์แล้ว รวมถึง คำสั่ง ข้อมูลแบบคงที่ และค่าคงที่ที่ฝังไว้ มีขนาดใหญ่กว่าจำนวน HBM ที่ว่างทั้งหมด ซึ่งพร้อมใช้งานในขณะนี้ในแกน TPU ที่เฉพาะเจาะจงซึ่งมีการโหลดโปรแกรม
- การกระจายตัวของ HBM: แม้ว่า HBM ที่ว่างทั้งหมดในอุปกรณ์อาจเพียงพอ โดยรวม แต่ก็ไม่พร้อมใช้งานในบล็อกเดียวที่ต่อเนื่องซึ่งมีขนาดใหญ่พอที่จะ รองรับโปรแกรมทั้งหมด
คุณควรทำความเข้าใจว่ารันไทม์ TPU จัดลำดับความสำคัญของหน่วยความจำอย่างไร การจัดสรร Buffer มีสิทธิ์เหนือกว่าโปรแกรมที่โหลด หากการจัดสรรบัฟเฟอร์ล้มเหลว รันไทม์จะนำโปรแกรมที่โหลดแล้วออกจาก HBM เพื่อเพิ่มพื้นที่ว่าง ซึ่งอาจทำให้โปรแกรมที่โหลดได้สำเร็จก่อนหน้านี้ ล้มเหลวด้วยข้อผิดพลาด OOM เนื่องจากตอนนี้ HBM มีบัฟเฟอร์ข้อมูลมากขึ้น
การแก้ไขข้อบกพร่อง
- ลดร่องรอยของหน่วยความจำบัฟเฟอร์: การเพิ่มหน่วยความจำที่ใช้โดยบัฟเฟอร์ข้อมูลจะ
ทำให้มีพื้นที่เหลือสำหรับโปรแกรมมากขึ้น
- ลดขนาดกลุ่ม: นี่เป็นวิธีที่มีประสิทธิภาพที่สุดวิธีหนึ่งในการลด ปริมาณหน่วยความจำที่ใช้สำหรับการเปิดใช้งาน
- การแบ่งพารามิเตอร์: สำหรับโมเดลขนาดใหญ่มาก ให้ใช้การขนานกันของโมเดลหรือ เทคนิคการแบ่ง (เช่น FSDP หรือ Megascale) เพื่อกระจาย พารามิเตอร์และการคำนวณของโมเดลไปยังแกนหรือโฮสต์ TPU หลายรายการ
- ลดความยาวของลำดับ/บริบท: สำหรับโมเดลที่ประมวลผลข้อมูลตามลำดับ (เช่น โมเดล NLP) การลดความยาวของลำดับจะช่วยลดการใช้หน่วยความจำได้อย่างมาก
- การบริจาคบัฟเฟอร์: ใช้ฟีเจอร์เฟรมเวิร์ก (เช่น
jax.jit(..., donate_argnums=...)) เพื่ออนุญาตให้ XLA นำหน่วยความจำของบัฟเฟอร์อินพุตมาใช้ซ้ำ เพื่อจัดเก็บเอาต์พุต ซึ่งจะช่วยลดการใช้หน่วยความจำสูงสุด
- ลดข้อกำหนดด้านหน่วยความจำของโปรแกรมสำหรับไฟล์ชั่วคราว
- ลดการใช้หน่วยความจำของโปรแกรมสำหรับไฟล์ชั่วคราวโดยใช้แฟล็ก
tpu_shared_memory_percentโปรดทราบว่าการดำเนินการนี้อาจส่งผลเสียต่อ ประสิทธิภาพ
- ลดการใช้หน่วยความจำของโปรแกรมสำหรับไฟล์ชั่วคราวโดยใช้แฟล็ก
- เพิ่มประสิทธิภาพกลยุทธ์การดำเนินการ/ลดภาระการแสดงโฆษณา
- จัดการการโหลดโปรแกรม: หากคุณกำลังคอมไพล์ฟังก์ชันหลายรายการแบบ JIT โปรดทราบว่าแต่ละฟังก์ชันอาจส่งผลให้มีการโหลดโปรแกรม ลอง จัดโครงสร้างภาระงานเพื่อลดจำนวนโปรแกรมที่โหลดพร้อมกัน
- ตรวจสอบว่าไม่มีหน่วยความจำรั่วไหล
- ตรวจสอบว่าการอ้างอิงถึงออบเจ็กต์
jax.Arrayไม่ได้นานเกินกว่า ที่ตั้งใจไว้ การเก็บออบเจ็กต์jax.Arrayอาจทำให้ระบบไม่สามารถ ยกเลิกการจัดสรรโดยอัตโนมัติแม้ว่าการคอมไพล์โปรแกรมจะเสร็จสมบูรณ์แล้วก็ตาม
- ตรวจสอบว่าการอ้างอิงถึงออบเจ็กต์
เครื่องมือ
- เปิดใช้แฟล็ก
tpu_log_allocations_on_oomซึ่งตัวจัดสรรจะ ส่งรายงานโดยละเอียดของการจัดสรรปัจจุบันทั้งหมดเมื่อเกิด OOM ซึ่ง มีประโยชน์อย่างมากในการแก้ไขข้อบกพร่อง - สร้างโปรไฟล์โปรแกรม: ใช้โปรไฟล์หน่วยความจำ JAX หรือโปรไฟล์ TensorFlow เพื่อดูรายละเอียดการใช้หน่วยความจำของโปรแกรมเมื่อเวลาผ่านไป ซึ่งจะช่วยระบุการใช้งานหน่วยความจำที่เพิ่มขึ้นอย่างไม่คาดคิดได้