Skip to content

Commit 5e2ea70

Browse files
committed
Add unit test support for Llama4ForCausalLM arch due to there's no text only model to be registered
1 parent 716ec8c commit 5e2ea70

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

tests/models/test_initialization.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,11 @@ def test_can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch):
2626
"KimiVLForConditionalGeneration"):
2727
pytest.skip("Avoid OOM")
2828

29+
if model_arch in ("Llama4ForCausalLM", "EagleLlama4ForCausalLM"):
30+
from vllm.model_executor.models.llama4 import Llama4ForCausalLM
31+
from vllm.model_executor.models.registry import ModelRegistry
32+
ModelRegistry.register_model("Llama4ForCausalLM", Llama4ForCausalLM)
33+
2934
# Avoid OOM and reduce initialization time by only using 1 layer
3035
def hf_overrides(hf_config: PretrainedConfig) -> PretrainedConfig:
3136
hf_config.update(model_info.hf_overrides)

0 commit comments

Comments
 (0)