Skip to content

[Bugfix] fix pp for llama4 #16746

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Apr 18, 2025
Merged

[Bugfix] fix pp for llama4 #16746

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fdc4236
fix pp for llama4
luccafong Apr 15, 2025
8706d40
Merge remote-tracking branch 'main/main' into llama4-pp-fix
luccafong Apr 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions vllm/model_executor/models/mllama4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -672,9 +672,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
self.config,
None,
prefix=maybe_prefix(prefix, "multi_modal_projector"))

self.language_model = _initialize_model(
vllm_config=vllm_config.with_hf_config(config.text_config),
vllm_config=vllm_config.with_hf_config(config.text_config,
["LlamaForCausalLM"]),
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we use Llama4ForCausalLM?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Llama4ForCasualLm is not registered architecture, we should avoid using that which requires adding a lot of hacks as in the initial PR

Copy link
Member

@ywang96 ywang96 Apr 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Probably a dumb question, but since _initialize_model is already pointing to model_class=Llama4ForCausalLM, why do we need to override the architectures here to LlamaForCausalLM?

vllm/vllm/model_executor/model_loader/loader.py

Lines 114 to 123 in 8cac35b

def _initialize_model(
vllm_config: VllmConfig,
*,
prefix: str = "",
model_class: Optional[type[nn.Module]] = None,
) -> nn.Module:
"""Initialize a model with the given configurations."""
model_config = vllm_config.model_config
if model_class is None:
model_class, _ = get_model_architecture(model_config)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

during __post__init__ of the hf config when we call replace inside with_hf_config function

vllm/vllm/config.py

Line 3781 in a018e55

return replace(self, model_config=model_config)
, it will check if the architecture is supporting PP or not via

vllm/vllm/config.py

Line 3790 in a018e55

self.model_config.verify_with_parallel_config(self.parallel_config)
, where None architecture will raise an issue as in #16385 in below code

vllm/vllm/model_executor/models/registry.py

Lines 441 to 442 in a018e55

normalized_arch = list(
filter(lambda model: model in self.models, architectures))

prefix=maybe_prefix(prefix, "language_model"),
model_class=Llama4ForCausalLM,
)
Expand Down Expand Up @@ -824,7 +824,7 @@ def load_weights(self, weights: Iterable[Tuple[str,
# language_model is an Llama4ForCausalLM instance. We load it's
# using llama4's load_weights routine.
language_model_weights, other_weights = self.separate_weights(
weights, prefix="language_model.model.")
weights, prefix="language_model.")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wondering why this issue was not triggered before?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

language_model.lm_head can also be loaded by the parent model weight loader if not PP enabled, but for PP, as we split the weights into 2 parts, the lm_head is missing in PP=0 so it will raise an issue weight not found. while in llama4.py model loading we have logic handling is_pp_missing_parameter to avoid the exception

loader = AutoWeightsLoader(self)
loaded_language_model_params = loader.load_weights(
language_model_weights)
Expand Down