Release a recipe on how to convert Gemma 4 (and other) models to .litertlm

#9
by jlotti - opened

I tried to follow instructions from https://developers.google.com/edge/litert-lm/file_builder and convert a model using litert-torch, but it failed with an error:

litert-torch export_hf --model=/media/user/1B22F52D7210D721/gemma-4-E4B-it --output_dir=/home/user/AI2/gemma-4-e4b-it --bundle_litert_lm=true --externalize_embedder=true

WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1784071713.523266    6404 port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
/home/user/AI2/venv/lib/python3.11/site-packages/torch/cuda/__init__.py:187: UserWarning: CUDA initialization: The NVIDIA driver on your system is too old (found version 12000). Please update your GPU driver by downloading and installing a new version from the URL: http://www.nvidia.com/Download/index.aspx Alternatively, go to: https://pytorch.org to install a PyTorch version that has been compiled with your version of the CUDA driver. (Triggered internally at /pytorch/c10/cuda/CUDAFunctions.cpp:119.)
  return torch._C._cuda_getDeviceCount() > 0
W0715 02:28:39.127000 6404 torch/utils/_pytree.py:630] <enum 'KernelPreference'> is an Enum subclass and is now natively supported by torch.compile as an opaque value type. Calling register_constant() on Enum subclasses is deprecated and will be an error in a future release.
W0715 02:28:41.964000 6404 torch/utils/_pytree.py:630] <enum 'ScaleCalculationMode'> is an Enum subclass and is now natively supported by torch.compile as an opaque value type. Calling register_constant() on Enum subclasses is deprecated and will be an error in a future release.
/home/user/AI2/venv/lib/python3.11/site-packages/torch/cuda/__init__.py:1074: UserWarning: Can't initialize NVML
  raw_cnt = _raw_device_count_nvml()
============== Export Configuration ==============
aot_backend            : None
aot_compilation_config_dict : None
aot_soc_model          : None
auto_model_override    : None
batch_size             : 1
bundle_litert_lm       : 'true'
cache_implementation   : 'LiteRTLMCache'
cache_length           : 4096
cache_length_dim       : None
enable_dynamic_shape   : False
experimental_lightweight_conversion : False
experimental_use_mixed_precision : False
export_vision_encoder  : False
externalize_embedder   : 'true'
externalize_rope       : False
extra_kwargs           : {}
jinja_chat_template_override : None
k_ts_idx               : 2
keep_temporary_files   : False
litert_lm_llm_metadata_override : None
litert_lm_model_type_override : None
model                  : '/media/user/1B22F52D7210D721/gemma-4-E4B-it'
output_dir             : '/home/user/AI2/gemma-4-e4b-it'
prefill_length_dim     : None
prefill_lengths        : [128]
quantization_recipe    : 'dynamic_wi8_afp32'
single_token_embedder  : False
split_cache            : False
task                   : <ExportTask.TEXT_GENERATION: 'text_generation'>
trust_remote_code      : False
use_jinja_template     : True
v_ts_idx               : 3
vision_encoder_quantization_recipe : 'dynamic_wi8_afp32'
work_dir               : '/home/user/AI2/gemma-4-e4b-it/tmp9qbrtkkt'
==================================================
(00:00) [START] LiteRT GenAI Export
(00:00) [START] LiteRT GenAI Export > Load source model
Gemma4 patch applied.
Loading weights: 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 2076/2076 [05:14<00:00,  6.59it/s]
(05:18) [ DONE] LiteRT GenAI Export > Load source model (+05:18)
(05:18) [START] LiteRT GenAI Export > Export text prefill-decode model
Using Gemma4 exportables.
(05:18) [ FAIL] LiteRT GenAI Export > Export text prefill-decode model
(05:18) [ FAIL] LiteRT GenAI Export
Traceback (most recent call last):
  File "/home/user/AI2/venv/bin/litert-torch", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/home/user/AI2/venv/lib/python3.11/site-packages/litert_torch/cli.py", line 30, in main
    fire.Fire(CLI())
  File "/home/user/AI2/venv/lib/python3.11/site-packages/fire/core.py", line 135, in Fire
    component_trace = _Fire(component, args, parsed_flag_args, context, name)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI2/venv/lib/python3.11/site-packages/fire/core.py", line 468, in _Fire
    component, remaining_args = _CallAndUpdateTrace(
                                ^^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI2/venv/lib/python3.11/site-packages/fire/core.py", line 684, in _CallAndUpdateTrace
    component = fn(*varargs, **kwargs)
                ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI2/venv/lib/python3.11/site-packages/litert_torch/generative/export_hf/export.py", line 194, in export
    exported_model_artifacts = run_export_tasks(
                               ^^^^^^^^^^^^^^^^^
  File "/home/user/.pyenv/versions/3.11.15/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI2/venv/lib/python3.11/site-packages/litert_torch/generative/export_hf/export.py", line 67, in run_export_tasks
    exported_model_artifacts = export_task(
                               ^^^^^^^^^^^^
  File "/home/user/.pyenv/versions/3.11.15/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI2/venv/lib/python3.11/site-packages/litert_torch/generative/export_hf/core/export_lib.py", line 270, in export_text_prefill_decode_model
    sample_prefill_inputs = prefill_module.get_sample_inputs(text_model_config)
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI2/venv/lib/python3.11/site-packages/litert_torch/generative/export_hf/core/exportable_module.py", line 179, in get_sample_inputs
    kv_cache_inputs, kv_cache_dynamic_shapes = self.get_sample_kv_cache(
                                               ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI2/venv/lib/python3.11/site-packages/litert_torch/generative/export_hf/core/exportable_module.py", line 123, in get_sample_kv_cache
    ].create_from_config(
      ^^^^^^^^^^^^^^^^^^^
  File "/home/user/AI2/venv/lib/python3.11/site-packages/litert_torch/generative/export_hf/model_ext/gemma4/cache.py", line 107, in create_from_config
    LiteRTLMCacheLayerForGemma4.create_from_config(
  File "/home/user/AI2/venv/lib/python3.11/site-packages/litert_torch/generative/export_hf/core/cache.py", line 293, in create_from_config
    return cls(
           ^^^^
TypeError: Can't instantiate abstract class LiteRTLMCacheLayerForGemma4 with abstract method get_max_length
LiteRT Community (FKA TFLite) org

@jlotti This isn't Gemma 4-specific it's a transformers version incompatibility with stable litert-torch.
transformers 5.13.0 renamed the abstract method on CacheLayerMixin from get_max_cache_shape to get_max_length. Stable litert-torch 0.9.1 only implements get_max_cache_shape, and its transformers dependency is unpinned, so a fresh venv pulls 5.13+ and every cache layer subclass (including LiteRTLMCacheLayerForGemma4) ends up with an unimplemented abstract method hence the TypeError.
Two fixes:
Pin transformers: pip install "transformers<5.13" (5.12.1 works with 0.9.1)
Or use the nightly, which already implements both methods and is what the official Gemma 4 conversion docs use: pip install --pre litert-torch-nightly (https://developers.google.com/edge/litert-lm/models/gemma-4)
The docs command also passes --jinja_chat_template_override=litert-community/gemma-4-E2B-it-litert-lm (swap in the E4B repo for your model) so the bundled chat template matches the official releases. The CUDA driver warning in your log is unrelated export runs on CPU.
A pinned transformers requirement in the package would prevent this whole class of breakage.

LiteRT Community (FKA TFLite) org
โ€ข
edited 3 days ago

This issue is caused by an incompatibility between the LiteRT-LM runtime's template parser and the Jinja chat template included in the Gemma 4 model on Hugging Face. Specifically, the template uses the map.get() method, a syntax that is not supported by the mobile-side parser.
Official documentation recommends using the --jinja_chat_template_override parameter to resolve this issue.

litert-torch export_hf \
--model=/media/user/1B22F52D7210D721/gemma-4-E4B-it \
--output_dir=/home/user/AI2/gemma-4-e4b-it \
--bundle_litert_lm=true \
--externalize_embedder=true \
--jinja_chat_template_override=litert-community/gemma-4-E2B-it-litert-lm

doc:https://developers.google.com/edge/litert/conversion/pytorch/genai#jinja-template-override


Google็ฟป่จณใ‚’ไฝฟ็”จ

LiteRT-LMใƒฉใƒณใ‚ฟใ‚คใƒ ใฎใƒ†ใƒณใƒ—ใƒฌใƒผใƒˆใƒ‘ใƒผใ‚ตใƒผใจใ€Hugging FaceไธŠใฎGemma 4ใƒขใƒ‡ใƒซใซๅซใพใ‚Œใ‚‹Jinjaใƒใƒฃใƒƒใƒˆใƒ†ใƒณใƒ—ใƒฌใƒผใƒˆใจใฎ้–“ใฎ้žไบ’ๆ›ๆ€งใŒๅŽŸๅ› ใงใ™ใ€‚ๅ…ทไฝ“็š„ใซใฏใ€ใใฎใƒ†ใƒณใƒ—ใƒฌใƒผใƒˆใŒใƒขใƒใ‚คใƒซๅดใฎใƒ‘ใƒผใ‚ตใƒผใงใ‚ตใƒใƒผใƒˆใ•ใ‚Œใฆใ„ใชใ„ๆง‹ๆ–‡ใงใ‚ใ‚‹ map.get() ใƒกใ‚ฝใƒƒใƒ‰ใ‚’ไฝฟ็”จใ—ใฆใ„ใ‚‹ใ“ใจใŒๅŽŸๅ› ใงใ™ใ€‚
ใ“ใฎๅ•้กŒใ‚’่งฃๆฑบใ™ใ‚‹ใŸใ‚ใซ --jinja_chat_template_override ใƒ‘ใƒฉใƒกใƒผใ‚ฟใ‚’ไฝฟ็”จใ™ใ‚‹ใ“ใจใŒๆŽจๅฅจใ•ใ‚Œใฆใ„ใพใ™ใ€‚

litert-torch export_hf \
--model=/media/user/1B22F52D7210D721/gemma-4-E4B-it \
--output_dir=/home/user/AI2/gemma-4-e4b-it \
--bundle_litert_lm=true \
--externalize_embedder=true \
--jinja_chat_template_override=litert-community/gemma-4-E2B-it-litert-lm

Sign up or log in to comment