Support multimodal in logit checker + match gemma3 logits with HF

[aireenmei]

Description

I found some mismatch in gemma3 logits compared with HF (b/437988753).
Key changes to fill the gap:

• Match the vocab size: HF uses a bigger vocab size which includes special image tokens.
• Special image token ID matches with HF
• Use linear Rope scaling with factor 8 for global attention layers in language model
• This will be a separate PR due to potential conflicts with the ongoing NNX migration: Allow setting gemma3 vision encoder precision through config.matmul_precision

Tests

• add unittest in check_gemma3_layers.py to check RoPE implementation
• logits test:
  Use generate_hf_golden_logits.py to generate HF logits on a m1 CPU, dtype=float32
  Run forward_pass_logit_checker on a v5p TPU, with all the dype and precision set to float32
  Results: see b/437988753

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[hengtaoguo]

Thank you Aireen for digging deep and finding this precision issue!

Is there any chance that you could split out the "precision" changes in gemma3.py into a separate PR? Asking because we are doing the NNX migration, and all of these layers will be rewritten from nn to nnx (example). I guess we can plugin these "precision" changes after the migration? Either way works for me.

[gagika]

why do we need this change here an in other model configs?

doesn't this change impact checkpoint checkpoint conversion?

There embedding lookup and unembed layers depend on the vocab size.