หมวดหมู่: รันไทม์: การจัดสรรบัฟเฟอร์ไม่สำเร็จ
ข้อผิดพลาดนี้บ่งบอกว่าตัวจัดสรรหน่วยความจำของรันไทม์ XLA:TPU ไม่พบบล็อกหน่วยความจำที่เหมาะสมใน HBM ของตัวเร่งสำหรับการจัดสรรที่ขอ
ตัวอย่างข้อความแสดงข้อผิดพลาด:
ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)
แบ็กเอนด์ XLA: TPU
ภาพรวม
ข้อผิดพลาดนี้จะเกิดขึ้นเมื่อ
- การจัดสรรบัฟเฟอร์ที่ผู้ใช้เริ่มต้นผ่าน
jax.device_putหรือ - การจัดสรรเอาต์พุตของโปรแกรมที่ผู้ใช้กำหนดเวลาไม่สำเร็จ
โดยการดำเนินการที่ไม่สำเร็จมักเกิดจากสาเหตุ 2 ประการต่อไปนี้
- หน่วยความจำไม่พอ (OOM): ผู้ใช้พยายามจัดสรรหน่วยความจำ ก้อนหนึ่งที่มีขนาดใหญ่กว่าหน่วยความจำว่างทั้งหมดที่พร้อมใช้งานใน HBM ของ TPU
- หน่วยความจำกระจัดกระจาย: การจัดสรรล้มเหลวเนื่องจากไม่มีบล็อกว่างที่ต่อเนื่องกันในพื้นที่หน่วยความจำที่มีขนาดใหญ่พอที่จะตอบสนองขนาดที่ขอ ปริมาณหน่วยความจำว่างทั้งหมดเพียงพอสำหรับการ จัดสรร แต่กระจายอยู่ทั่วพื้นที่หน่วยความจำในบล็อกขนาดเล็กที่ไม่ได้อยู่ติดกัน
รันไทม์ TPU มีกลไกหลายอย่างเพื่อลองจัดสรรอีกครั้งเมื่อจัดสรรไม่สำเร็จ ซึ่งรวมถึงกลไกต่อไปนี้
- หากมีการยกเลิกการจัดสรรที่อยู่ในคิว รันไทม์จะลองจัดสรรที่ล้มเหลวอีกครั้ง
- ใน OOM ที่เกิดจากการกระจายตัว รันไทม์จะทริกเกอร์การจัดเรียงข้อมูลใหม่และการลองอีกครั้งโดยอัตโนมัติ
- รันไทม์ TPU จะให้ความสำคัญกับการจัดสรรบัฟเฟอร์มากกว่าการโหลดโปรแกรมไว้ หากการจัดสรรบัฟเฟอร์ล้มเหลวเนื่องจาก HBM ไม่เพียงพอ ระบบจะนำโปรแกรม TPU ที่โหลดออกจนกว่าจะมีหน่วยความจำเพียงพอสำหรับบัฟเฟอร์
ดังนั้น ข้อผิดพลาดที่พบหลังจากใช้มาตรการลดความเสี่ยงข้างต้นมักจะต้องมีการดำเนินการจากผู้ใช้
การแก้ไขข้อบกพร่อง
- ลดหน่วยความจำที่ใช้ของโมเดล
- ลดขนาดกลุ่ม: การลดขนาดกลุ่มจะช่วยลดการใช้หน่วยความจำได้โดยตรง
- การแบ่งพารามิเตอร์: สำหรับโมเดลขนาดใหญ่มาก ให้ใช้เทคนิคต่างๆ เช่น การขนานกันของโมเดลหรือการแบ่งเพื่อกระจายพารามิเตอร์ใน HBM ของหลายๆ คอร์หรือโฮสต์ TPU
- ลดความยาวของลำดับ/บริบท: สำหรับโมเดลที่ทำงานกับลำดับ (เช่น โมเดลภาษา) การลดความยาวของลำดับอินพุตจะช่วยลดการใช้หน่วยความจำได้อย่างมาก
- การบริจาคบัฟเฟอร์: ใช้ฟีเจอร์ของเฟรมเวิร์ก (เช่น
jax.jit(..., donate_argnums=...)) เพื่อส่งสัญญาณไปยัง XLA ว่าสามารถเขียนทับและนำบัฟเฟอร์อินพุตบางรายการมาใช้ซ้ำสำหรับเอาต์พุตได้ - เพิ่มประสิทธิภาพกลยุทธ์การตรวจสอบ: แทนที่จะบันทึกสถานะโมเดลทั้งหมด ในครั้งเดียว ให้ลองบันทึกเฉพาะน้ำหนักของโมเดลหรือใช้กลยุทธ์การตรวจสอบแบบ Sharded
- เลย์เอาต์และการเพิ่มแพดดิ้งในหน่วยความจำของที่อยู่
- ระบบจะจัดสรรหน่วยความจำ TPU เป็นกลุ่ม และการเพิ่ม Padding อาจเพิ่มขนาดจริงของเทนเซอร์
- ตรวจสอบว่าไม่มีหน่วยความจำรั่วไหล
- ตรวจสอบว่าการอ้างอิงถึงออบเจ็กต์
jax.Arrayไม่ได้นานเกินกว่า ที่ตั้งใจไว้ การเก็บออบเจ็กต์jax.Arrayอาจทำให้ระบบไม่สามารถ ยกเลิกการจัดสรรโดยอัตโนมัติแม้ว่าการคอมไพล์โปรแกรมจะเสร็จสมบูรณ์แล้วก็ตาม
- ตรวจสอบว่าการอ้างอิงถึงออบเจ็กต์
เครื่องมือ
เปิดใช้tpu_log_allocations_on_oomแฟล็กซึ่งตัวจัดสรรจะทิ้งรายงานโดยละเอียดของการจัดสรรปัจจุบันทั้งหมดเมื่อเกิด OOM ซึ่งอาจมีประโยชน์อย่างมากในการแก้ไขข้อบกพร่อง