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

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

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

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

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.

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

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

ภาพรวม

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

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

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

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

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

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

  • การใช้งาน TensorCore (TC) + SparseCore (SC) HBM เกินขีดจำกัด: หากข้อผิดพลาดระบุรายละเอียดการใช้งานอย่างชัดเจน เช่น "การใช้งาน TC Hbm: X, การใช้งาน SC Hbm Y" → ข้ามไปที่ ส่วนที่ 1 ปรับสมดุลการใช้งาน HBM ของ TC และ SC
  • การจัดสรรขนาดใหญ่โดยไม่คาดคิด: หากข้อผิดพลาดระบุว่า "หน่วยความจำในพื้นที่หน่วยความจำ HBM เต็ม" ให้ตรวจสอบบันทึก เพื่อดูการแจงนับการจัดสรรที่ใหญ่ที่สุดใน HBM ในกรณีที่มีเทนเซอร์ขนาดใหญ่เกินคาดอย่างน้อย 1 รายการ (เช่น มากกว่า 50% ของขีดจำกัด HBM) → ข้ามไปที่ส่วนที่ 2 การจัดสรรขนาดใหญ่โดยไม่คาดคิด
  • การจัดสรรรวมเกินขีดจำกัด HBM: หากข้อผิดพลาดระบุว่า "หน่วยความจำในพื้นที่หน่วยความจำ HBM หมด" แต่ไม่มีเทนเซอร์ขนาดใหญ่ผิดปกติในบันทึก → ไปที่ ส่วนที่ 3 การจัดสรรรวมเกินขีดจำกัด HBM

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

หากข้อผิดพลาดระบุรายละเอียดการใช้งานอย่างชัดเจน เช่น "การใช้งาน TC Hbm: X, การใช้งาน SC Hbm Y" เปรียบเทียบค่าทั้ง 2 เพื่อระบุปัญหาคอขวด

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

ส่วนที่ 2 การจัดสรรที่มีขนาดใหญ่เกินคาด

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

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

ส่วนที่ 3 การจัดสรรรวมเกินขีดจำกัด HBM

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

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

ก. ตรวจสอบการเพิ่มและการจัดแนวเทนเซอร์

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

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

ข. ปรับการกำหนดค่า

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

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

C. เพิ่มประสิทธิภาพสถาปัตยกรรมและการแบ่งข้อมูล

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

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

ง. ปรับแต่งฟีเจอร์สำคัญที่ส่งผลต่อหน่วยความจำของ XLA

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

E. Tune XLA Rematerialization Pass / Manual Checkpointing

หากโมเดลใกล้จะพอดีกับหน่วยความจำ คุณสามารถบังคับให้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 เปิดใช้การส่งผ่านการสร้างใหม่เพิ่มเติมก่อนการส่งผ่านการผสาน ช่วยประหยัดหน่วยความจำได้มากขึ้น แต่จะเพิ่มเวลาในการคอมไพล์และอาจส่งผลต่อความเสถียรของตัวเลข

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

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

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

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