[Gemma4] Add docstrings for Per-Layer Embeddings (PLE) pipeline

[w4nderlust]

Fixes #45206

What does this PR do?

Adds documentation for the Gemma4 Per-Layer Embeddings (PLE) system, which is currently pretty hard to reverse-engineer from the code alone.

I ran into this while implementing Gemma4 inference from scratch in Rust. The PLE system has several non-obvious aspects that aren't documented anywhere:

1. hidden_size_per_layer_input (256) is the per-layer dimension, but the actual embedding weight is [vocab, num_layers * 256] = [262144, 8960] because all layers are packed
2. The embedding is a Gemma4TextScaledWordEmbedding that silently multiplies by sqrt(256) = 16 - this took me a while to track down
3. The full pipeline has a context-aware projection step (per_layer_model_projection + scale + RMSNorm) that combines with the token lookup before being passed to layers, with specific scale factors (1/sqrt(hidden_size) and 1/sqrt(2))

This PR adds:

• Expanded config docstring for hidden_size_per_layer_input explaining the packed layout, scaling, and full pipeline
• Docstrings for get_per_layer_inputs() and project_per_layer_inputs()
• A comment on the PLE init block pointing to the pipeline methods

Hopefully this saves some pain for others implementing Gemma4 outside of transformers.

[Rocketknight1]

This is a tricky one - I don't want to bloat the docstrings, but I agree Gemma4 is unusual and this bit can be hard to understand. cc @stevhliu WDYT?

[stevhliu]

i agree @Rocketknight1, the config docstring should describe how a parameter works instead of how the model uses it. a better home for this is the gemma4 doc here. the config docstring can say:

hidden_size_per_layer_input (`int`, defaults to 256):
    Per-layer hidden dimension for the PLE system. The actual embedding weight has shape
    `[vocab_size_per_layer_input, num_hidden_layers * hidden_size_per_layer_input]`
    because all layers are packed into a single table. See the [Gemma4](https://huggingface.co/docs/transformers/main/en/model_doc/gemma4#gemma4-vision-model) docs
    for a description of the full PLE pipeline.

the docstrings for get_per_layer_inputs() and project_per_layer_inputs() are fine though :)

[w4nderlust]

Happy to make the changes yo want guys, jsut tell me exactly and I'll modify the PR :)

[stevhliu]

thanks for improving! 🤗

[stevhliu]

Suggested change

	Per-layer hidden dimension for the per-layer input embeddings (PLE). This is the
	dimension of each layer's individual embedding slice. Note: the actual embedding
	weight has shape `[vocab_size_per_layer_input, num_hidden_layers * hidden_size_per_layer_input]`
	(e.g. `[262144, 8960]` for 35 layers x 256 dims) because all layers are packed into
	a single embedding table. The embedding is also a `Gemma4TextScaledWordEmbedding`
	that scales lookups by `sqrt(hidden_size_per_layer_input)`.
	Per-layer hidden dimension for the PLE system. The actual embedding weight has shape
	`[vocab_size_per_layer_input, num_hidden_layers * hidden_size_per_layer_input]`
	because all layers are packed into a single table. See the [Gemma4](https://huggingface.co/docs/transformers/main/en/model_doc/gemma4#gemma4-vision-model) docs
	for a description of the full PLE pipeline.

[stevhliu]

move this to a gemma.md :)

[w4nderlust]

Thanks for the review! I've pushed the changes:

• Shortened the hidden_size_per_layer_input config docstring to just describe the parameter shape, with a link to the model docs
• Moved the full PLE pipeline description to a new section in gemma4.md under "Per-Layer Embeddings (PLE)"
• Kept the method docstrings for get_per_layer_inputs() and project_per_layer_inputs() as-is

Let me know if you'd like any further adjustments!

[stevhliu]

just a few nits! make sure you edit the modular_gemma4.py file instead of the modeling_gemma4.py file to pass the CI :)

[stevhliu]

Suggested change

	Gemma 4 introduces a **Per-Layer Embeddings (PLE)** system that feeds an auxiliary residual signal into each decoder layer, rather than relying solely on a single shared embedding at the input.
	Gemma 4 introduces a Per-Layer Embeddings (PLE) system that feeds an auxiliary residual signal into each decoder layer, rather than relying solely on a single shared embedding at the input.

[stevhliu]

the config params are covered in the docstrings so i don't think we need to repeat them here

Suggested change

	#### Key config parameters
	- `vocab_size_per_layer_input` (default 262144): vocabulary size for PLE.
	- `hidden_size_per_layer_input` (default 256): per-layer embedding dimension. The actual embedding weight has shape `[vocab_size_per_layer_input, num_hidden_layers * hidden_size_per_layer_input]` (e.g. `[262144, 8960]` for 35 layers × 256 dims) because all layers are packed into a single embedding table.
	#### Pipeline

[stevhliu]

Suggested change

	1. **Token-identity** (`get_per_layer_inputs`): looks up `input_ids` in `embed_tokens_per_layer`, a `Gemma4TextScaledWordEmbedding` that multiplies by `√(hidden_size_per_layer_input)`. The packed output is reshaped from `[batch, seq, num_hidden_layers * hidden_size_per_layer_input]` to `[batch, seq, num_hidden_layers, hidden_size_per_layer_input]`.
	1. Token-identity (`get_per_layer_inputs`): looks up `input_ids` in `embed_tokens_per_layer`, a `Gemma4TextScaledWordEmbedding` that multiplies by `√(hidden_size_per_layer_input)`. The packed output is reshaped from `[batch, seq, num_hidden_layers * hidden_size_per_layer_input]` to `[batch, seq, num_hidden_layers, hidden_size_per_layer_input]`.

[stevhliu]

Suggested change

	2. **Context-aware** (`project_per_layer_inputs`): projects `inputs_embeds` through `per_layer_model_projection` (a Linear layer), scales by `1/√(hidden_size)`, reshapes to `[batch, seq, num_layers, ple_dim]`, and normalizes with `per_layer_projection_norm` (RMSNorm).
	2. Context-aware (`project_per_layer_inputs`): projects `inputs_embeds` through `per_layer_model_projection` (a Linear layer), scales by `1/√(hidden_size)`, reshapes to `[batch, seq, num_layers, ple_dim]`, and normalizes with `per_layer_projection_norm` (RMSNorm).

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[stevhliu]

nice, i think the last step is to run make fix-repo to automatically align modeling_gemma4.py with modular_gemma4.py!

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: gemma4

[stevhliu]

🤗