[BUGFIX] Fix Pixtral consolidated format vision weight loading

[juliendenize]

Purpose

#36963 replaced the Pixtral vision encoder's nn.Linear layers (wq/wk/wv/wo/w1/w2/w3) with QKVParallelLinear and MergedColumnParallelLinear (qkv_proj/o_proj/gate_up_proj/down_proj) to support LoRA. However, the weight loading stacked_params only mapped HF-style names (q_proj, k_proj, etc.), not Mistral native names (wq, wk, etc.), causing vision encoder weights to be silently dropped when loading consolidated-format checkpoints.

Test Plan

Added a ministral test that is run for small GPUs instead of relying on Pixtral.

Test Result

Passing.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request adds support for Mistral native (consolidated) weight formats to the Pixtral model and introduces a new test case for consolidated loading. Feedback indicates that the added test uses a text-only model, which fails to exercise the vision encoder weight loading logic. Additionally, the weight loading implementation may fail to match keys containing a '.weight' suffix, and the remapping logic for native parameter names is inefficiently located and may lead to dropped weights.

[afurm]

The wo and w2 parameters are handled via _vision_encoder_name_remap rather than through _vision_encoder_stacked_params. Since they're not sharded across TP ranks like qkv/w1/w3, they don't appear in the stacked params list. Is there a reason they couldn't be added to the stacked params list with their shard_id, or is the remap approach more robust to variations in how these keys appear across different checkpoint formats?

[NickLucche]

does the ministral test you added actually excercise the fix? 🤔

[juliendenize]

does the ministral test you added actually excercise the fix? 🤔

Hey thanks for the merge and to answer your question @NickLucche It does by making sure the output is not garbaged with a smaller gpu size than needed to Pixtral. I think I could even lower the size of the GPU but AFAIK it should be always used by the CI for 16 GB right ? So now we should catch whenever a regression happens !