รหัสข้อผิดพลาด: E1000

หมวดหมู่: เวลาคอมไพล์: HBM OOM

ข้อผิดพลาดนี้บ่งชี้ว่าโปรแกรมต้องการหน่วยความจำแบนด์วิดท์สูง (HBM) มากกว่าที่มีอยู่ในอุปกรณ์ TPU

ตัวอย่างข้อความแสดงข้อผิดพลาด

RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.

แบ็กเอนด์ XLA: TPU

ภาพรวม

XLA จะทำการตรวจสอบเพื่อให้แน่ใจว่าขนาดรวมของการจัดสรรแบบคงที่ที่จำเป็นทั้งหมด พอดีกับ HBM ของอุปกรณ์

คอมไพเลอร์จัดการความจุ HBM แบบคงที่ของ TPU สำหรับการจัดสรรหลายประเภท ดังนี้

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

ข้อผิดพลาดนี้เกิดขึ้นเมื่อคอมไพเลอร์ XLA จัดสรรทั้งหมดข้างต้น ลงใน HBM ของอุปกรณ์ไม่ได้

การแก้ไขข้อบกพร่อง

วิเคราะห์ข้อความแสดงข้อผิดพลาดและบันทึกอย่างละเอียดเพื่อพิจารณาว่าหมวดหมู่ของ HBM OOM ด้านล่างนี้อธิบายข้อผิดพลาดของคุณได้ดีที่สุด

สถานการณ์ที่ 1 รักษาสมดุลการใช้งาน HBM ของ TC และ SC

หากข้อผิดพลาดระบุรายละเอียดการใช้งานอย่างชัดเจน เช่น "การใช้งาน TC Hbm: X, การใช้งาน SC Hbm Y" แสดงว่าการใช้งาน TensorCore (TC) + SparseCore (SC) โดยรวมเกิน ขีดจำกัด HBM เปรียบเทียบค่าทั้ง 2 เพื่อระบุจุดคอขวด

  • การใช้งาน SparseCore สูง
    • เพิ่มประสิทธิภาพการใช้งานสแต็ก HBM: การใช้หน่วยความจำสแต็ก HBM จะปรับขนาดตาม feature_width, max_unique_nz_per_row และ logical_replica_count คุณสามารถลดการใช้งานสแต็กสูงสุดได้โดยการปรับ--xla_sc_num_serialized_tables_to_optimize_hbmแฟล็กซึ่งจะทำให้การประมวลผลตารางเป็นแบบอนุกรม แต่จะทำให้ความขนานลดลง
    • ตรวจสอบค่าใช้จ่ายในการเพิ่ม Padding: SparseCore จะจัดแนวตารางการฝังเป็น 32B (8 โฟลต) ตารางที่มีความกว้างของฟีเจอร์ขนาดเล็ก (เช่น น้อยกว่า 8 โฟลต) จะมีค่าใช้จ่ายในการเพิ่มพื้นที่อย่างมาก ซึ่งทำให้ HBM เสียเปล่า
    • ลดการใช้ฮีป: ค่าสูงสำหรับ maximum_parallel_iterations จะเพิ่มปริมาณข้อมูลอินพุตที่ดึงข้อมูลล่วงหน้าลงในฮีป HBM การลดค่านี้จะช่วยเพิ่มหน่วยความจำได้มาก
    • ยืนยันการแบ่งข้อมูล: ตรวจสอบว่าตารางการฝังได้รับการแบ่งข้อมูลอย่างถูกต้อง ในชิปทั้งหมด ดูวิธีที่ขีดจำกัดเปลี่ยนเป็นตาราง
    • ดูแนวคิดเพิ่มเติมได้ที่ SC: คอขวดด้านประสิทธิภาพและหน่วยความจำ
  • การใช้งาน TensorCore สูง
  • สมดุล
    • หากแต่ละอย่างไม่มากเกินไป แต่เมื่อรวมกันแล้วสูงเกินไป แสดงว่าคุณใช้ความสามารถของชิป จนถึงขีดจำกัดแล้ว คุณต้องลองลดการใช้งานทั้ง 2 องค์ประกอบ ทำตามคำแนะนำในทั้ง 3 ส่วน

สถานการณ์ที่ 2 หน่วยความจำไม่พอเนื่องจากการจัดสรรขนาดใหญ่โดยไม่คาดคิด

หากคุณเห็นข้อความแสดงข้อผิดพลาด "หน่วยความจำในพื้นที่หน่วยความจำ HBM เต็ม" และ มีการจัดสรรขนาดใหญ่ที่ไม่คาดคิดอย่างน้อย 1 รายการในบันทึก (> 50% ของขีดจำกัด HBM ) แสดงว่าแทบจะไม่ใช่ปัญหาด้านความจุของฮาร์ดแวร์ โดยปกติแล้วจะเป็น ข้อผิดพลาดในการกำหนดค่า ตรวจสอบป้ายกำกับ XLA (หากมี) ของการจัดสรรขนาดใหญ่ เพื่อดูคำแนะนำเกี่ยวกับซอร์สโค้ด JAX

  • นำอาร์ติแฟกต์การแก้ไขข้อบกพร่องออก
    • การใช้ jax.debug.print() ในการเรียกใช้ขนาดใหญ่อาจบังคับให้คอมไพเลอร์สร้างเทนเซอร์ทั้งหมดใน HBM เพื่อโอนไปยัง CPU ซึ่งจะทำให้การผสานหยุดชะงักและเพิ่มการใช้งานหน่วยความจำสูงสุด นำ jax.debug.print() ที่เหลืออยู่ออก
  • แก้ไขรูปร่างตาข่ายหรือการแยกส่วนที่ไม่มีประสิทธิภาพ
    • รูปร่างของเมชที่ไม่ถูกต้องหรือคำอธิบายประกอบการแบ่งพาร์ติชันที่ขาดหายไปอาจทำให้คอมไพเลอร์ใช้เรพลิเคชันเป็นค่าเริ่มต้น ซึ่งบังคับให้คอมไพเลอร์พยายามใส่ Tensor ขนาดใหญ่มากในชิปเดียว
    • ตรวจสอบรูปร่างของการจัดสรรขนาดใหญ่และยืนยันว่า XLA ได้ระบุและเผยแพร่การแยกส่วนอย่างถูกต้อง

สถานการณ์ที่ 3 หน่วยความจำหมดเนื่องจากการจัดสรรแบบรวม

หากคุณเห็นข้อความแสดงข้อผิดพลาด "หน่วยความจำในพื้นที่หน่วยความจำ HBM เต็ม" และ ไม่มีเทนเซอร์ขนาดใหญ่ผิดปกติในบันทึก แสดงว่าโปรแกรมมี ความจุไม่เพียงพอเนื่องจากผลรวมของการจัดสรรเกินขีดจำกัด HBM ในกรณีนี้ การแสดงภาพโปรไฟล์หน่วยความจำมักจะมีประโยชน์ในการระบุบัฟเฟอร์ที่เฉพาะเจาะจงซึ่งทำให้เกิดการใช้งานสูงสุด ดูแก้ไขข้อบกพร่องข้อผิดพลาด OOM ด้วย XProf เพื่อดู คําแนะนําทีละขั้นตอนในการระบุผู้ใช้หน่วยความจําสูงสุด

เมื่อระบุผู้มีส่วนร่วมอันดับต้นๆ ได้แล้ว ให้ทำตามขั้นตอนต่อไปนี้ เพื่อเพิ่มประสิทธิภาพหน่วยความจำที่ใช้

สถานการณ์ที่ 3.A ปรับการกำหนดค่า

คุณมักจะแก้ไข OOM ได้ด้วยการปรับการกำหนดค่าต่อไปนี้

  • ลดขนาดกลุ่ม: หน่วยความจำที่จำเป็นสำหรับการเปิดใช้งานและ การไล่ระดับกลางจะแปรผันตรงกับขนาดกลุ่ม การลดขนาดกลุ่มมักจะช่วยลดการใช้งานหน่วยความจำได้
  • บริจาคบัฟเฟอร์อินพุต: เมื่อใช้ jax.jit ให้ระบุ donate_argnums สำหรับ พารามิเตอร์โมเดล ซึ่งจะช่วยให้ XLA เขียนทับหน่วยความจำอินพุตด้วย เอาต์พุตได้
  • เปิดใช้ความแม่นยำแบบผสม (bfloat16): ใช้ bfloat16 หรือการหาปริมาณ (int8 ฯลฯ) สำหรับเทนเซอร์ที่ใหญ่ที่สุดในโปรแกรม หากสถาปัตยกรรมโมเดลและ ข้อกำหนดด้านคุณภาพอนุญาต โปรดทราบว่าการเปลี่ยนแปลงนี้อาจส่งผลต่อลักษณะการทำงานของโมเดล และควรพิจารณาอย่างรอบคอบ

สถานการณ์ที่ 3.B เพิ่มประสิทธิภาพสถาปัตยกรรมและการแบ่งข้อมูล

หากการเปลี่ยนแปลงการกำหนดค่าไม่เพียงพอ โทโพโลยีของโมเดลอาจมีขนาดใหญ่เกินไป สำหรับการตั้งค่าฮาร์ดแวร์ปัจจุบัน

  • ใช้ TPU รุ่นใหม่กว่า: โดยทั่วไปแล้ว TPU รุ่นใหม่จะมี HBM ต่อชิปมากกว่า เปลี่ยนไปใช้ TPU รุ่นใหม่กว่าหากมี
  • เรียกใช้ในโทโพโลยีชิปที่ใหญ่ขึ้น: หากน้ำหนักของโมเดลมีขนาดใหญ่เกินไปสำหรับโทโพโลยีที่มีอยู่ คุณสามารถลองแบ่งส่วนน้ำหนักของโมเดลในชิปเพิ่มเติมได้
  • ใช้เทคนิคการแบ่งข้อมูลขั้นสูง:
    • สำรวจแนวทางการขนานข้อมูล เทนเซอร์ หรือไปป์ไลน์ขั้นสูงเพิ่มเติม
    • ระบุ คำแนะนำในการแบ่งข้อมูล สำหรับค่ากลางและเอาต์พุต
  • ใช้การโอนภาระงานของโฮสต์ JAX: เทคนิคการโอนภาระงานของโฮสต์ช่วยให้ผู้ใช้สามารถ โอนภาระงานของเทนเซอร์ขนาดใหญ่ไปยังหน่วยความจำ CPU ของโฮสต์ (เช่น การโอนภาระงานของเลเยอร์การกระตุ้น และ การโอนภาระงานของสถานะเครื่องมือเพิ่มประสิทธิภาพ)

สถานการณ์ที่ 3.C ตรวจสอบการเพิ่มแพดและการจัดแนวของเทนเซอร์

รูปร่างของเทนเซอร์ที่ไม่มีประสิทธิภาพเป็นสาเหตุที่พบบ่อยและไม่แสดงข้อความแจ้งของข้อผิดพลาด OOM ใน TPU เพื่อให้ได้ประสิทธิภาพสูงสุดใน TPU, XLA จะเพิ่มขนาด Tensor โดยปกติจะเพิ่มเป็นหลายเท่าของ 128 สำหรับมิติข้อมูลที่เล็กที่สุด และ 8 สำหรับมิติข้อมูลที่เล็กเป็นอันดับ 2 การเพิ่มพื้นที่นี้ ส่งผลต่อทั้งอาร์เรย์อินพุตและเทนเซอร์กลาง (อาร์เรย์ชั่วคราวของ HLO) ซึ่งอาจทำให้การใช้งานหน่วยความจำเพิ่มขึ้นอย่างมาก โดยเฉพาะอย่างยิ่งเมื่อมีขนาดมิติข้อมูลเล็ก ดู เลย์เอาต์อาร์เรย์

  • ตรวจสอบรูปร่างของบัฟเฟอร์ขนาดใหญ่: (ใน TPU v5 ที่มีเลย์เอาต์เริ่มต้น)
    • การวางเมาส์เหนือบัฟเฟอร์ในXprof Memory Viewer จะแสดงการ์ดรายละเอียดบัฟเฟอร์ซึ่งมีรายละเอียดบัฟเฟอร์ รวมถึงข้อมูลการเพิ่มพื้นที่
    • ตัวอย่าง: รูปร่างของ (129, 1024) อาจได้รับการเพิ่มขนาดเป็น (256, 1024) ซึ่งทำให้สิ้นเปลืองหน่วยความจำเกือบ 50%
    • การแก้ไข: รูปร่างของ (128, 1024) ไม่ต้องมีการเว้นวรรคและทำให้เกิดการสูญเสียหน่วยความจำ 0%
  • จัดแนวขนาด: ตรวจสอบว่าขนาดเทนเซอร์ขนาดใหญ่ทั้งหมด (ขนาดกลุ่ม ขนาดการฝัง ขนาดที่ซ่อน) เป็นทวีคูณของ 128 โปรดทราบว่าการเปลี่ยนแปลงนี้อาจส่งผลต่อลักษณะการทำงานของโมเดลและควรพิจารณาอย่างรอบคอบ

สถานการณ์ที่ 3.D ปรับแต่งหน่วยความจำที่สำคัญซึ่งส่งผลต่อการแจ้ง XLA

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

สถานการณ์ที่ 3.E การส่งผ่านการสร้างใหม่ของ Tune XLA/การตรวจสอบด้วยตนเอง

หากโมเดลใกล้จะพอดีกับหน่วยความจำ คุณสามารถใช้ตัวตกแต่ง jax.checkpoint กับ jax.grad เพื่อควบคุมด้วยตนเองว่าจะบันทึกค่ากลางใดใน การส่งต่อเทียบกับการคำนวณซ้ำในการส่งย้อนกลับ ซึ่งจะแลกเปลี่ยนรอบการคำนวณ กับ HBM

หรือจะบังคับให้XLA::Rematerializationส่งผ่านเพื่อจัดลำดับความสำคัญ การประหยัดหน่วยความจำ ซึ่งอาจทำให้การคอมไพล์ช้าลงก็ได้

ธง คำอธิบาย ผลกระทบ / การแลกเปลี่ยน
--xla_tpu_max_hbm_size_mib กำหนดขีดจำกัดขนาด HBM ที่ใช้โดยการส่งผ่านการสร้างใหม่ด้วยตนเอง บังคับให้คอมไพเลอร์ทำงานหนักขึ้นเพื่อให้โปรแกรมมีขนาดเล็กกว่า HBM จริง
--xla_tpu_rematerialization_algo=PEAK_PRIORITY มุ่งเน้นความพยายามที่จุดที่มีการใช้งานหน่วยความจำสูงสุด มีประสิทธิภาพมากกว่าอัลกอริทึมเริ่มต้นสำหรับการลดหน่วยความจำแบบรุก
--xla_tpu_rematerialization_max_block_size_limit=32 ควบคุมจำนวนคำสั่งสูงสุดในบล็อกที่สามารถสร้างใหม่ได้พร้อมกัน การเพิ่มค่านี้จะช่วยประหยัดหน่วยความจำ แต่จะทำให้เวลาคอมไพล์เพิ่มขึ้นอย่างมาก
--xla_tpu_rematerialization_block_effort_factor=10.0 กำหนดปริมาณความพยายาม (เวลาคอมไพล์) ที่ใช้ในการค้นหาบล็อกเพื่อสร้างใหม่ ค่าที่สูงขึ้นจะช่วยให้ค้นหาการประหยัดหน่วยความจำได้ละเอียดมากขึ้น แต่จะใช้เวลาในการคอมไพล์นานขึ้น
--xla_tpu_pre_fusion_remat=true เปิดใช้การส่งผ่านการสร้างใหม่เพิ่มเติมก่อนการส่งผ่านการผสาน ช่วยประหยัดหน่วยความจำได้มากขึ้น แต่จะเพิ่มเวลาในการคอมไพล์และอาจส่งผลต่อความเสถียรของตัวเลข

โปรดทราบว่าการเปลี่ยนแปลงแฟล็ก XLA ควรใช้เป็นมาตรการสุดท้าย เนื่องจากอาจส่งผลเสียต่อประสิทธิภาพ

สถานการณ์ 3.F ใช้เครื่องมือการจัดโปรไฟล์ขั้นสูง

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

เครื่องมือนี้ช่วยให้คุณเห็นการจัดสรรหน่วยความจำสูงสุดและอายุการใช้งานบัฟเฟอร์ ซึ่ง มีความสำคัญอย่างยิ่งต่อการทำความเข้าใจว่าอะไรที่ใช้ HBM ในจุดที่มีการใช้งานสูงสุด ดูการตั้งค่าการสร้างโปรไฟล์ทั่วไปได้ที่ เริ่มต้นใช้งาน Xprof และ การสร้างโปรไฟล์ TensorBoard