Skip to content

feat(gemma3): add embed_tokens() and forward_embeds() - #3903

Open
sidd-27 wants to merge 1 commit into
huggingface:mainfrom
sidd-27:feat/gemma3-forward-embeds
Open

sidd-27 wants to merge 1 commit into
huggingface:mainfrom
sidd-27:feat/gemma3-forward-embeds

Conversation

@sidd-27

@sidd-27 sidd-27 commented Aug 17, 2026

Copy link
Copy Markdown

gemma3::Model only accepts token ids, and keeps embed_tokens private. A multimodal caller needs to build a mixed sequence — text embeddings with image embeddings spliced in — which has no representation as token ids, so Gemma 3's vision variants can't be wired up from outside the crate.

gemma.rs already exposes both entry points (embed_tokens() at line 371, forward_embeds() at line 412) and paligemma.rs is built on them. This adds the same pair to gemma3.

forward now delegates to forward_embeds, so there's no behaviour change. Unlike gemma, forward_embeds builds the attention masks itself rather than taking them as an argument, because gemma3 uses two (full and sliding) and picks per layer — having the caller supply them would leak that detail.

Verified the two paths agree exactly (tiny randomly-initialised model, CPU, no downloads):

let a = model.forward(&ids, 0).unwrap();
model.clear_kv_cache();
let embeds = model.embed_tokens().forward(&ids).unwrap();
let b = model.forward_embeds(&embeds, 0).unwrap();
logits len              = 64
max |forward - embeds|  = 0.000000000
OK: forward and forward_embeds are identical

Note that running that check at all requires #3902 — gemma3::Model::forward currently fails for every input with slice-set only supports contiguous tensors. This PR compiles and is independent of it, but can't be exercised at runtime until that one lands.

Multimodal callers need to build a mixed embedding sequence - text
embeddings with image embeddings spliced in - which cannot be expressed
as token ids. gemma already exposes both entry points and paligemma is
built on them; gemma3 exposed neither, so Gemma 3's vision variants
cannot be wired up outside the crate.

forward now delegates to forward_embeds, so there is no behaviour change:
the two produce bit-identical logits.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant