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

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

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

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

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 รุ่นต่อๆ ไปได้ทั้งหมด แต่ไม่รับประกันว่าจะมีการจำลองสำหรับรุ่นเก่ากว่า (ซึ่งการแคสต์บางรุ่นอาจมีค่าใช้จ่ายสูงมาก)