SDK Google Tensor позволяет оптимизировать и запускать модели машинного обучения (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 | Необязательный | Определяет, насколько агрессивно модель разделяется для параллельной обработки. Варианты: минимальный, умеренный, обширный, максимальный. | минимальный |
google_tensor_int64_to_int32 | Необязательный | Установите значение true, чтобы разрешить компилятору преобразовывать 64-битные целые числа в 32-битные, что может быть необходимо для некоторых моделей. | ЛОЖЬ |
google_tensor_enable_large_model_support | Необязательный | Установите значение true, если размер вашей модели превышает 2 ГБ. | ЛОЖЬ |
google_tensor_enable_4bit_compilation | Необязательный | Позволяет компилировать модели с операциями свертки, содержащими 4-битные параметры. | ЛОЖЬ |
google_tensor_extra_options_path | Необязательный | Передает дополнительные параметры компилятора в виде бинарного файла. | "" (Пустой) |
Эти флаги можно использовать с потоком LiteRT Python, как показано в следующих фрагментах кода:
При компиляции модели PyTorch с использованием
ai_edge_torchcompiled_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 .
Для экспорта LLM-модулей для TPU Google Tensor следуйте примеру, чтобы узнать о дополнительных флагах, необходимых для компиляции 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}'