Gemma3 to hf mapping error

[jedcheng]

Bug report

After finetuning the model, I used the following command to convert my MaxText Gemma3 Checkpoint to HuggingFace format:

python3 -m MaxText.utils.ckpt_conversion.to_huggingface src/MaxText/configs/base.yml \ model_name='gemma3-4b'     \
hf_access_token={hf token}   \
load_parameters_path={my checkpoint path}    \
base_output_directory=/tmp/gemma3-4B-cpt-hf     \
use_multimodal=false     # tried both false & true \
scan_layers=false   # cannot create scanned checkpoint (another bug)

For testing purposes, I set up a vllm server on the TPU

pip3.11 install vllm-tpu

vllm serve /tmp/gemma3-4B-cpt-hf \
   --disable-log-requests \
    --tensor_parallel_size=4 \
    --api-key {my api key}

Error message from vllm

ValueError: Gemma3 uses `gelu_pytorch_tanh` as the hidden activation function. Please set `hidden_act` and `hidden_activation` to `gelu_pytorch_tanh`.

MaxText Conversion config uses "hidden_activation": "gelu" for text generation

huggingface config for text generation uses gelu_pytorch_tanh

vllm explicitly check for gelu_pytorch_tanh which resulted in the error above

Logs/Output

Environment Information

TPU creation command

export TPU_COUNT=4
gcloud alpha compute tpus tpu-vm create $TPU_ID \
    --zone="${ZONE}" \
    --accelerator-type="v6e-${TPU_COUNT}" \
    --version=v2-alpha-tpuv6e \
    --spot \
    --service-account=${service_account}

Tried installation from source and uv as suggested in the documentation of MaxText.

Additional Context

No response

[jedcheng]

I changed this line to use gelu_pytorch_tanh.

I then converted a new MaxText checkpoint from HF with multimodal on.
Converted it back to HF with multimodal on as well (note: vllm would not load a Gemma model without the vision stuff).
vllm successfully loaded and compiled the Gemma HF model

I ran mmlu before and after MaxText conversion with vllm for the sake of confirmation. The model after MaxText conversion lost 0.73 percentage points. Is it normal?

[shralex]

@hengtaoguo can you take a look

[jedcheng]

Another update:

As suggested in the previous follow-up, I converted an HF model to MaxText and then back to HF with multimodal on.

I did a (continuous) pre-training run today on Gemma3. The Pre-training pipeline doesn't support multi-modal but I didn't expect it to remove the entire vision blocks such that I cannot convert to HF with multi-modal on.

Vllm then refuses to load the model without a vision encoder.
But I guess this is a feature, not a bug.

I came up with a mitigation: read the model with HuggingFace transformers and then save it again. Then manually copy the tokenizer and image generation config to the new model folder