Support casting to fp32 when word embeddings are tied to lm_head

[pramodith]

What does this PR do?

Per slack discussion https://huggingface.slack.com/archives/C089Q56GPMM/p1761949724954879.

TinyGemma has tie_word_embeddings=True
Qwen3 has tie_word_embeddings=False

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[pramodith]

Profiling with Qwen3-0.6B (tied_weights=True) original model in fp16 with inputs of shape [4, 128]

============================================================
GPU MEMORY USAGE SUMMARY
============================================================
Initial memory:              1.110 GB
After forward pass:          2.789 GB
Peak during forward:         2.789 GB
After backward pass:         2.444 GB
Peak during backward:        3.088 GB
After optimizer step:        4.672 GB
Peak total (overall):        5.782 GB
============================================================

Profiling when word_embeddings and lm_head cast to fp32

============================================================
GPU MEMORY USAGE SUMMARY
============================================================
Initial memory:              1.401 GB
After forward pass:          3.374 GB
Peak during forward:         3.374 GB
After backward pass:         3.172 GB
Peak during backward:        4.332 GB
After optimizer step:        5.978 GB
Peak total (overall):        7.378 GB
============================================================

LM_HEAD and WORD_EMBEDDINGS share the same object so additional memory for FP32
weights at init should be
2 * vocab_size * hidden_dim => 2 * (151936 * 1024)/(1024**2) = 296 MB

This is equal to 1.401GB - 1.110GB = 296 MB (approx)

2 is the factor of additional bytes FP32/FP16

Forward Pass
Additional Memory for FP32 logits should be
vocab_size * seq_len * batch_size => 2 * (151936 * 128 * 4)/(1024 **2) = 148MB

3.374 GB - 2.789 GB should be = (Additional Init Mem + Additional Forward Pass Mem) =  148 MB + 296 MB
0.585 GB !=  444 MB

🤔
The difference is approx 140 MB, I'm wondering if the forward pass stores both the results after the embeddings (fp32) and the post embedding layer hook separately accounting for the extra 140 MB.

[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.

[qgallouedec]

Qwen3 has tied word embedding and Gemma 2 no, correct? If so, I'd just add a small comment so that we remember why we test these two cases

[pramodith]

It's the other way around Qwen3 has untied and Gemma 2 has tied.

[qgallouedec]

for the record, float() is inlace, so in theory, you could just have

target_model.lm_head.float()

it happens that .float() returns self, so target_model.lm_head = target_model.lm_head.float() works as well. I personally prefer the current way, assignment makes it more explicit.

[qgallouedec]

Nice! just a small comment on the test