Huxley
6738855fb1
Add: basic script for OLMoE inference and saving expert logits. Add: OLMoE vLLM module with logits saving. Add: vLLM model registering. Fix: update DataLogger module for multiprocessing.
32 lines
1.2 KiB
Python
32 lines
1.2 KiB
Python
from vllm import ModelRegistry
|
|
|
|
|
|
def register_vllm_logit_logging_models():
|
|
models_mapper = {"OlmoeForCausalLM": "models.log_expert.olmoe:OlmoeForCausalLM"}
|
|
for model_name, model_path in models_mapper.items():
|
|
ModelRegistry.register_model(model_name, model_path)
|
|
|
|
# ModelRegistry.register_model(
|
|
# "Qwen3MoeForCausalLM", "src.modeling_vllm_save.qwen3_moe:Qwen3MoeForCausalLM"
|
|
# )
|
|
# ModelRegistry.register_model(
|
|
# "MixtralForCausalLM", "src.modeling_vllm_save.mixtral:MixtralForCausalLM"
|
|
# )
|
|
# ModelRegistry.register_model(
|
|
# "OlmoeForCausalLM", "models.log_expert.olmoe:OlmoeForCausalLM"
|
|
# )
|
|
# ModelRegistry.register_model(
|
|
# "DeepseekV2ForCausalLM",
|
|
# "src.modeling_vllm_save.deepseek_v2:DeepseekV2ForCausalLM",
|
|
# )
|
|
# ModelRegistry.register_model(
|
|
# "Llama4ForConditionalGeneration",
|
|
# "src.modeling_vllm_save.mllama4:Llama4ForConditionalGeneration",
|
|
# )
|
|
# ModelRegistry.register_model(
|
|
# "GptOssForCausalLM", "src.modeling_vllm_save.gpt_oss:GptOssForCausalLM"
|
|
# )
|
|
# ModelRegistry.register_model(
|
|
# "PhiMoEForCausalLM", "src.modeling_vllm_save.phimoe:PhiMoEForCausalLM"
|
|
# )
|