หมวดหมู่: รันไทม์: การจัดสรรบัฟเฟอร์ไม่สำเร็จ
ข้อผิดพลาดนี้บ่งบอกว่าตัวจัดสรรหน่วยความจำของรันไทม์ 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 ว่าสามารถเขียนทับและนำบัฟเฟอร์อินพุตบางรายการมาใช้ซ้ำสำหรับเอาต์พุตได้ อ่านรายละเอียดเพิ่มเติมได้ใน การบริจาคผ่านบัฟเฟอร์ - เพิ่มประสิทธิภาพกลยุทธ์จุดตรวจสอบ: ลองบันทึกเฉพาะน้ำหนักของโมเดลหรือใช้กลยุทธ์การตรวจสอบจุดที่แยกส่วนแทนการบันทึกสถานะของโมเดลทั้งหมดพร้อมกัน
- เลย์เอาต์และการเพิ่มพื้นที่หน่วยความจำแอดเดรส
- ระบบจะจัดสรรหน่วยความจำ TPU เป็นกลุ่ม และการเพิ่มระยะขอบอาจเพิ่มขนาดจริงของเทนเซอร์
- ตรวจสอบว่าไม่มีหน่วยความจำรั่วไหล
- ตรวจสอบว่าการอ้างอิงถึงออบเจ็กต์
jax.Arrayไม่ได้นานเกินกว่า ที่ตั้งใจไว้ การเก็บออบเจ็กต์jax.Arrayไว้ในหน่วยความจำอาจทำให้ระบบไม่สามารถ ยกเลิกการจัดสรรหน่วยความจำโดยอัตโนมัติได้ แม้ว่าการคอมไพล์โปรแกรมจะเสร็จสมบูรณ์แล้วก็ตาม
- ตรวจสอบว่าการอ้างอิงถึงออบเจ็กต์
ดูรหัสข้อผิดพลาด: E1000 สำหรับ กลยุทธ์อื่นๆ ที่คุณใช้เพื่อลดปริมาณหน่วยความจำที่แต่ละโปรแกรมใช้ได้
เครื่องมือ
เปิดใช้tpu_log_allocations_on_oomแฟล็กซึ่งตัวจัดสรรจะทิ้งรายงานโดยละเอียดของการจัดสรรปัจจุบันทั้งหมดเมื่อเกิด OOM ซึ่งอาจมีประโยชน์อย่างมากในการแก้ไขข้อบกพร่อง