Initialize vLLM config for Frontier profiling

This commit is contained in:
2026-07-16 22:29:10 +08:00
parent 9b54075e75
commit 885db8527a
5 changed files with 27 additions and 4 deletions

View File

@@ -9,6 +9,7 @@ unchanged and adapt only that call boundary for the profiling experiment.
from __future__ import annotations
import inspect
import os
from typing import Any
import torch
@@ -51,8 +52,20 @@ def main() -> None:
frontier_rope._VLLM_GET_ROPE_IMPORT_ERROR = None
from frontier.profiling.linear_op.main import main as frontier_main
from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
frontier_main()
model_root = os.environ.get("MODEL_ROOT")
if not model_root:
raise RuntimeError("MODEL_ROOT must point to the profiled Qwen checkpoint")
model_config = ModelConfig(
model=model_root,
dtype="bfloat16",
max_model_len=8192,
skip_tokenizer_init=True,
generation_config="vllm",
)
with set_current_vllm_config(VllmConfig(model_config=model_config)):
frontier_main()
if __name__ == "__main__":