[Model] Enable LoRA support for Pixtral

[A1c0r-Z]

Purpose

Enable LoRA adapters on the vision tower and connector components of Pixtral models (Part of #31479).

Technical Details

• Implement get_mm_mapping(), get_num_mm_encoder_tokens(), and get_num_mm_connector_tokens().
• Add SupportsLoRA to PixtralForConditionalGeneration inheritance.
• Patch Merge Logic: Account for the $S \times S$ (typically $2 \times 2$) patch merging in Pixtral.
  • get_num_mm_encoder_tokens scales up the LLM budget by $S^2$ to match the Vision Tower's raw output.
  • get_num_mm_connector_tokens scales down the raw count by $S^2$ to match the Connector's input after merging.

Part of #31479

FIX #8802

[mergify]

Documentation preview: https://vllm--31724.org.readthedocs.build/en/31724/

[gemini-code-assist]

Code Review

This pull request enables LoRA support for the vision tower and connector components of Pixtral models. The changes include adding the SupportsLoRA interface to PixtralForConditionalGeneration and implementing the necessary methods for LoRA integration with multi-modal components, such as get_mm_mapping, get_num_mm_encoder_tokens, and get_num_mm_connector_tokens. My review found a critical issue in the implementation of get_mm_mapping that would prevent LoRA from being applied correctly. I've provided a suggestion to fix it.

[mergify]

Hi @A1c0r-Z, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[jeejeelee]

Thank you for contribution, have you tested it locally?

[A1c0r-Z]

Thank you for contribution, have you tested it locally?

Yes, I have tested it locally using a "probe test" method.

I inserted debug print statements into the Pixtral model definition to verify that the token counts calculated by the new get_num_mm_* methods strictly match the actual tensor shapes observed during the forward pass (_process_image_input).

Here is the test detail:

1. Methodology

I instrumented the code to print:

• Calculated Values: The output of get_num_mm_encoder_tokens and get_num_mm_connector_tokens.
• Runtime Shapes: The actual tensor shapes before and after the patch_merger layer in _process_image_input.

# Instrumentation Logic
def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
    # ... logic ...
    print(f"[DEBUG PROBE encoder]: {calculated_val}")
    return calculated_val

def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
    # ... logic ...
    print(f"[DEBUG PROBE connector]: {calculated_val}")
    return calculated_val

def _process_image_input(self, ...):
    # ...
    image_features = self.vision_encoder(images)
    print(f"[DEBUG PROBE 1] Vision Encoder Output Shape: {image_features[0].shape}")
    
    if self.patch_merger is not None:
         image_features = self.patch_merger(...)
         
    print(f"[DEBUG PROBE 2] Shape BEFORE connector: {image_features.shape}")
    # ...

2. Test Case A: mistralai/Pixtral-12B-2409 (No Patch Merger)

This model does not use a patch merger, so the mapping should be 1:1.

Startup Profile:

(EngineCore_DP0) [DEBUG PROBE encoder]: 4096
(EngineCore_DP0) [DEBUG PROBE connector]: 4096
(EngineCore_DP0) [DEBUG PROBE 1] Vision Encoder Output Shape: torch.Size([4096, 1024])
(EngineCore_DP0) [DEBUG PROBE 2] Shape BEFORE connector: torch.Size([4096, 1024])

Inference (Small Image):

(EngineCore_DP0) [DEBUG PROBE encoder]: 16
(EngineCore_DP0) [DEBUG PROBE connector]: 16
(EngineCore_DP0) [DEBUG PROBE 1] Vision Encoder Output Shape: torch.Size([16, 1024])
(EngineCore_DP0) [DEBUG PROBE 2] Shape BEFORE connector: torch.Size([16, 1024])

Result: The calculated budget matches the runtime tensor shape perfectly (Identity).

3. Test Case B: mistralai/Ministral-3-3B-Instruct-2512 (With Patch Merger)

This model uses a 2x2 patch merger, so the encoder tokens should be the connector tokens.

Startup Profile:

(EngineCore_DP0) [DEBUG PROBE encoder]: 12100
(EngineCore_DP0) [DEBUG PROBE connector]: 3025
(EngineCore_DP0) [DEBUG PROBE 1] Vision Encoder Output Shape: torch.Size([12100, 1024])
(EngineCore_DP0) [DEBUG PROBE 2] Shape BEFORE connector: torch.Size([3025, 1024])

Note: . The calculation is correct.

Inference (Small Image):

(EngineCore_DP0) [DEBUG PROBE encoder]: 16
(EngineCore_DP0) [DEBUG PROBE connector]: 4
(EngineCore_DP0) [DEBUG PROBE 1] Vision Encoder Output Shape: torch.Size([16, 1024])
(EngineCore_DP0) [DEBUG PROBE 2] Shape BEFORE connector: torch.Size([4, 1024])

[jeejeelee]

Thank you