คอมไพล์โมเดลด้วย Google Tensor SDK

Google Tensor SDK ช่วยให้คุณเพิ่มประสิทธิภาพและเรียกใช้โมเดลแมชชีนเลิร์นนิง (ML) ได้โดยตรงบน TPU เฉพาะของ Google Tensor นอกจากโมเดล ML มาตรฐานแล้ว คุณยังคอมไพล์และเร่งความเร็วโมเดลภาษาภายในเวิร์กโฟลว์ LiteRT ได้ด้วย

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

แฟล็กการคอมไพล์สำหรับ Google Tensor

เพิ่มประสิทธิภาพกระบวนการคอมไพล์โมเดลโดยใช้แฟล็กการกำหนดค่าที่เฉพาะเจาะจงเพื่อปรับประสิทธิภาพและการใช้ทรัพยากร คุณสามารถใช้พารามิเตอร์เหล่านี้ภายในเวิร์กโฟลว์ LiteRT Python เพื่อปรับลักษณะการทำงานของการคอมไพล์สำหรับทั้งโมเดล PyTorch และ TFLite

แฟล็ก ข้อกำหนด คำอธิบาย ค่าเริ่มต้น
google_tensor_truncation_type ไม่บังคับ กำหนดประเภทข้อมูลเป้าหมายสำหรับการดำเนินการแบบจุดลอยตัว
ค่าที่รองรับ: auto (ค่าเริ่มต้น), bfloat16, half, no_truncation
อัตโนมัติ
google_tensor_sharding_intensity ไม่บังคับ ควบคุมความเข้มข้นของการแยกโมเดลสำหรับการประมวลผลแบบขนาน
ตัวเลือก: minimal, moderate, extensive, maximum
minimal
google_tensor_int64_to_int32 ไม่บังคับ ตั้งค่าเป็น "จริง" เพื่ออนุญาตให้คอมไพเลอร์แปลงจำนวนเต็ม 64 บิตเป็นจำนวนเต็ม 32 บิต ซึ่งอาจจำเป็นสำหรับบางโมเดล เท็จ
google_tensor_enable_large_model_support ไม่บังคับ ตั้งค่าเป็น "จริง" หากโมเดลมีขนาดใหญ่กว่า 2 GB เท็จ
google_tensor_enable_4bit_compilation ไม่บังคับ เปิดใช้การคอมไพล์โมเดลที่มีการดำเนินการแบบ Convolution ซึ่งมีพารามิเตอร์ 4 บิต เท็จ
google_tensor_extra_options_path ไม่บังคับ ส่งตัวเลือกคอมไพเลอร์เพิ่มเติมเป็นไฟล์ไบนารี "" (ว่างเปล่า)

คุณสามารถใช้แฟล็กเหล่านี้กับโฟลว์ LiteRT Python ได้ตามที่แสดงในข้อมูลโค้ดต่อไปนี้

  • เมื่อคอมไพล์โมเดล pytorch โดยใช้ ai_edge_torch

    compiled_models = (
      ai_edge_torch.experimental_add_compilation_backend(tensor_g5_target,
      flag_name1=value1,
      flag_name2=value2, ...).convert(
        channel_last_selfie_segmentation,
        sample_input))
    
  • เมื่อคอมไพล์โมเดล tflite

    compiled_models = aot_lib.aot_compile(
        tflite_model_path,
        target=[tensor_g5_target],
        flag_name1=value1,
        flag_name2=value2, ...)
    

ตัวอย่างการใช้งาน

ตัวอย่างต่อไปนี้ใช้แฟล็ก google_tensor_truncation_type="half"

  compiled_models = aot_lib.aot_compile(
      tflite_model_path,
      target=[tensor_g5_target],
      keep_going=False,
      google_tensor_truncation_type="half"
  )

ดูข้อมูลเพิ่มเติมได้ที่ LiteRT AOT Colab

คอมไพล์โมเดลภาษาสำหรับ Google Tensor

หากต้องการคอมไพล์โมเดลภาษาสำหรับ Google Tensor ให้ทำตามวิธีการใน NPU AOT compilation

หากต้องการส่งออก LLM สำหรับ Google Tensor TPU ให้ทำตามตัวอย่างสำหรับแฟล็กเพิ่มเติมที่จำเป็นสำหรับการคอมไพล์ NPU

ตัวอย่าง:

litert-torch export-hf \
  --model=google/gemma-3-270m-it \
  --output_dir=/tmp/gemma3-270m-google-tensor-g5 \
  --split_cache \
  --externalize_embedder \
  --prefill_lengths=128, \
  --cache_length=1280 \
  --quantization_recipe="weight_only_wi8_afp32" \
  --aot_backend=GOOGLE \
  --aot_soc_model=Tensor_G5 \
  --aot_compilation_config_dict='{"google_tensor_enable_large_model_support": True}'