Initialize vLLM config for Frontier profiling
This commit is contained in:
@@ -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__":
|
||||
|
||||
Reference in New Issue
Block a user