[Model] Refactor JambaForCausalLM #21394
Conversation
Signed-off-by: Jee Jee Li <[email protected]>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the Jamba model implementation, primarily by replacing the custom JambaMLP with the more generic LlamaMLP and by simplifying the weight loading logic in JambaForCausalLM to use the AutoWeightsLoader. These changes improve code clarity and maintainability. I've found two high-severity issues related to weight loading that could cause problems.
| ("qkv_proj", "q_proj", "q"), | ||
| ("qkv_proj", "k_proj", "k"), | ||
| ("qkv_proj", "v_proj", "v"), |
There was a problem hiding this comment.
The stacked_params_mapping entries for qkv_proj are missing leading dots1. This can cause ambiguous or incorrect weight name matching, as q_proj could match unrelated parameter names. Using .q_proj ensures matching a specific component, improving safety and explicitness. This inconsistency with gate_up_proj mappings and implementations in models like LlamaModel further emphasizes the need for correction.
Style Guide References
| ("qkv_proj", "q_proj", "q"), | |
| ("qkv_proj", "k_proj", "k"), | |
| ("qkv_proj", "v_proj", "v"), | |
| (".qkv_proj", ".q_proj", "q"), | |
| (".qkv_proj", ".k_proj", "k"), | |
| (".qkv_proj", ".v_proj", "v"), |
Footnotes
| return hidden_states.view(orig_shape) | ||
|
|
||
|
|
||
| class JambaMLP(JambaMoE): |
jeejeelee
changed the title
Signed-off-by: Jee Jee Li <[email protected]>
| if num_experts > 1: | ||
| self.feed_forward = JambaMoE( | ||
| config, | ||
| quant_config=quant_config, | ||
| prefix=f"{prefix}.feed_forward", | ||
| ) | ||
| else: | ||
| self.feed_forward = JambaMLP( | ||
| config.hidden_size, | ||
| config.intermediate_size, | ||
| config.hidden_act, | ||
| quant_config=quant_config, | ||
| prefix=f"{prefix}.feed_forward", | ||
| ) |
| (".gate_up_proj", ".gate_proj", 0), | ||
| (".gate_up_proj", ".up_proj", 1), |
There was a problem hiding this comment.
why do we have to handle these now? (the old code didn't mention these)
There was a problem hiding this comment.
These exist in the old version of the code https://github.com/vllm-project/vllm/blob/v0.5.4/vllm/model_executor/models/jamba.py#L857, this PR is just a revert
tlrmchlsmth
left a comment
tlrmchlsmth
left a comment
There was a problem hiding this comment.
looks good, thanks - just a couple of nits
Signed-off-by: Jee Jee Li <[email protected]>
Signed-off-by: Jee Jee Li <[email protected]>
Signed-off-by: Jee Jee Li <[email protected]> Signed-off-by: x22x22 <[email protected]>
Signed-off-by: Jee Jee Li <[email protected]>
Signed-off-by: Jee Jee Li <[email protected]>
Signed-off-by: Jee Jee Li <[email protected]> Signed-off-by: Jinzhen Lin <[email protected]>
Signed-off-by: Jee Jee Li <[email protected]> Signed-off-by: Noam Gat <[email protected]>
Signed-off-by: Jee Jee Li <[email protected]> Signed-off-by: Paul Pak <[email protected]>
Signed-off-by: Jee Jee Li <[email protected]> Signed-off-by: Diego-Castan <[email protected]>
Signed-off-by: Jee Jee Li <[email protected]>
Signed-off-by: Jee Jee Li <[email protected]>
2 participants
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.Purpose
The main motivation is to standardize the model implementation and support BNB quantization
Test Plan
Test Result
(Optional) Documentation Update