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

หมวดหมู่: เวลาคอมไพล์: ไม่รองรับ RHS DataType ในฮาร์ดแวร์

ข้อผิดพลาดนี้เกิดขึ้นเมื่อ TPU รุ่นที่ใช้ไม่รองรับประเภทข้อมูลที่ใช้สำหรับตัวถูกดำเนินการด้านขวามือ (RHS) ในการคูณเมทริกซ์ (เช่น jax.lax.dot_general, jax.lax.conv, jax.numpy.matmul หรือตัวดำเนินการ @) โดยกำเนิด

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

INTERNAL: Mosaic failed to compile TPU kernel: Unsupported matmul RHS type on target: 'vector<256x256xi8>'
...

The MLIR operation involved:
%13440 = "tpu.matmul"(%13435, %13437, %13439) <dimension_numbers = #tpu.dot_dimension_numbers<...>

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

ภาพรวม

หน่วยคูณเมทริกซ์ (MXU) ของ TPU รองรับการดำเนินการ Float32 โดยกำเนิดในฮาร์ดแวร์ทุกรุ่น

อย่างไรก็ตาม การรองรับ BFloat16 และประเภทข้อมูลอื่นๆ ที่มีการควอนไทซ์ (เช่น Int4, Int8 หรือ Float8) จะแตกต่างกันไปตามรุ่นของฮาร์ดแวร์ ข้อผิดพลาดนี้จะเกิดขึ้นเมื่อเคอร์เนลพยายามแมปการคูณเมทริกซ์กับ MXU โดยใช้ประเภทข้อมูลที่ TPU รุ่นของคุณไม่มีวงจรทางกายภาพเพื่อดำเนินการ

โดยปกติแล้วข้อผิดพลาดนี้แสดงว่าพาสการแปลงรูปแบบมาตรฐานของคอมไพเลอร์ ซึ่งพยายามแปลงประเภทที่ไม่รองรับเป็นประเภทที่รองรับโดยอัตโนมัติ (เช่น ผ่านการจำลองซอฟต์แวร์) ไม่พบกฎการแปลงที่ถูกต้องหรือถูกป้องกันไม่ให้ดำเนินการดังกล่าวเนื่องจากปิดใช้โหมดความเข้ากันได้

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

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

1. แคสต์ไปยังประเภทเนทีฟ

การแก้ไขที่เชื่อถือได้มากที่สุดคือการแปลงโอเปอแรนด์เป็นประเภทข้อมูลที่ฮาร์ดแวร์รองรับด้วยตนเอง (เช่น Float32 หรือ BFloat16 ใน TPU v4 ขึ้นไป) ภายในเคอร์เนลก่อน การดำเนินการ matmul

เหตุผล: Float32 เป็นประเภทข้อมูลสากลที่ MXU รองรับโดยเนทีฟใน TPU ทุกรุ่น
ข้อแลกเปลี่ยน: การดำเนินการนี้มีค่าใช้จ่าย VPU (หน่วยประมวลผลเวกเตอร์) ซึ่งเป็น รอบที่จำเป็นในการดำเนินการแคสต์ แต่รับประกันว่าเคอร์เนลจะทำงานบนฮาร์ดแวร์ปัจจุบันได้

2. ตรวจสอบโหมดความเข้ากันได้

โดยปกติแล้วคอมไพเลอร์จะจัดการปัญหาประเภทที่ไม่ตรงกันเหล่านี้ได้โดยอัตโนมัติในโหมดความเข้ากันได้ ซึ่งเปิดใช้โดยค่าเริ่มต้น ตรวจสอบการกำหนดค่า XLA อีกครั้ง เพื่อให้แน่ใจว่าไม่ได้ตั้งค่า --xla_mosaic_compat_mode เป็น false

ซึ่งทำหน้าที่เป็น "polyfill" โดยจะแทรกลำดับการจำลองซอฟต์แวร์สำหรับการดำเนินการ ที่ฮาร์ดแวร์ไม่รองรับโดยกำเนิด

สิ่งที่โหมดความเข้ากันได้ช่วยให้ทำได้

MatMul แบบความแม่นยำผสม: อนุญาตให้ผสมตัวถูกดำเนินการที่เป็นจำนวนเต็มกับตัวสะสม Float โดยการแทรกการดำเนินการแคสต์โดยอัตโนมัติ (เช่น การขยายจำนวนเต็มเป็น Float32 ก่อน matmul)
การจำลองความแม่นยำต่ำ: ในฮาร์ดแวร์บางรุ่น จะจำลองประเภทที่ไม่รองรับ เช่น 4-bitจุดลอยตัว (4E2M1FN) หรือ8-bitจุดลอยตัว (8E4M3FN) โดยการขยายประเภทเหล่านั้นไปยังประเภทที่รองรับ เช่น BFloat16 หรือ Float32 ก่อนการดำเนินการ

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

3. อัปเกรดฮาร์ดแวร์หรือขอรับการสนับสนุน

หากอัลกอริทึมของคุณต้องการประสิทธิภาพเนทีฟอย่างเคร่งครัดสำหรับประเภทต่างๆ เช่น Int4 หรือ Float8 โดยไม่มีค่าใช้จ่ายในการแคสต์หรือการจำลอง คุณจะต้องเรียกใช้ใน TPU รุ่นใหม่กว่าที่รองรับเนทีฟ

คำขอฟีเจอร์: หากเชื่อว่าฮาร์ดแวร์ของคุณรองรับการดำเนินการนี้ หรือหาก คอมไพเลอร์ไม่มีเส้นทางการจำลองที่ถูกต้องแม้ในโหมดความเข้ากันได้ โปรดส่งคำขอฟีเจอร์ โดยปกติแล้ว เราจะรับประกันว่าการดำเนินการจะเข้ากันได้กับเวอร์ชันที่ใหม่กว่า ดังนั้นหากเคอร์เนลทำงานบน TPU รุ่นใดรุ่นหนึ่ง ก็ควรจะทำงานบน TPU รุ่นต่อๆ ไปได้ทั้งหมด แต่เราไม่รับประกันว่าจะมีการจำลองสำหรับ TPU รุ่นเก่า (เนื่องจากบางรุ่นอาจมีค่าใช้จ่ายในการแคสต์สูงมาก)