Biên dịch các mô hình bằng Google Tensor SDK

Google Tensor SDK cho phép bạn tối ưu hoá và chạy các mô hình học máy (ML) ngay trên TPU chuyên dụng của Google Tensor. Ngoài các mô hình ML tiêu chuẩn, bạn có thể biên dịch và tăng tốc các mô hình ngôn ngữ trong quy trình LiteRT.

Đối với một số mô hình được tối ưu hoá trước, chúng tôi cũng cung cấp một tệp nhị phân tuỳ chọn trình biên dịch bổ sung mà bạn có thể đưa vào quy trình làm việc của mình. Để đạt được hiệu suất tối ưu trên Google Tensor, chúng tôi đề xuất sử dụng các cờ tối ưu hoá sau trong quá trình biên dịch:

Cờ biên dịch cho Google Tensor

Tối ưu hoá quy trình biên dịch mô hình bằng cách áp dụng các cờ cấu hình cụ thể để điều chỉnh hiệu suất và mức sử dụng tài nguyên. Bạn có thể sử dụng các tham số này trong quy trình LiteRT Python để điều chỉnh hành vi biên dịch cho cả mô hình PyTorch và TFLite.

Cờ Yêu cầu Mô tả Giá trị mặc định
google_tensor_truncation_type Không bắt buộc Đặt kiểu dữ liệu mục tiêu cho các phép toán dấu phẩy động.
Giá trị được hỗ trợ: auto (mặc định), bfloat16, half, no_truncation		 tự động
google_tensor_sharding_intensity Không bắt buộc Kiểm soát mức độ phân chia mô hình để xử lý song song.
Các lựa chọn: tối thiểu, vừa phải, mở rộng, tối đa.		 minimal
google_tensor_int64_to_int32 Không bắt buộc Đặt thành true để cho phép trình biên dịch chuyển đổi số nguyên 64 bit thành số nguyên 32 bit. Điều này có thể cần thiết đối với một số mô hình. Sai
google_tensor_enable_large_model_support Không bắt buộc Đặt thành true nếu mô hình của bạn lớn hơn 2 GB. Sai
google_tensor_enable_4bit_compilation Không bắt buộc Cho phép biên dịch các mô hình có các hoạt động tích chập chứa các tham số 4 bit. Sai
google_tensor_extra_options_path Không bắt buộc Truyền các lựa chọn trình biên dịch bổ sung dưới dạng tệp nhị phân. "" (Trống)

Bạn có thể dùng các cờ này với quy trình LiteRT Python như minh hoạ trong các đoạn mã sau:

  • Khi biên dịch một mô hình pytorch bằng 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))

  • Khi biên dịch một mô hình tflite

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

Ví dụ về cách sử dụng

Trong ví dụ sau, cờ google_tensor_truncation_type="half" được dùng:

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

Hãy xem Colab AOT LiteRT để biết thêm thông tin.

Biên dịch mô hình ngôn ngữ cho Google Tensor

Để biên dịch các mô hình ngôn ngữ cho Google Tensor, hãy làm theo hướng dẫn trong phần Biên dịch AOT NPU.

Để xuất LLM cho TPU Google Tensor, hãy làm theo ví dụ về các cờ bổ sung cần thiết cho quá trình biên dịch NPU.

Ví dụ:

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}'