93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Run Frontier's linear profiler against the vLLM 0.20 RoPE API.
|
|
|
|
Frontier passes the pre-v0.20 ``get_rope`` arguments separately, while vLLM
|
|
0.20 carries the same values in ``rope_parameters``. Keep both repositories
|
|
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
|
|
from vllm import _custom_ops as vllm_ops
|
|
from vllm.model_executor.layers.layernorm import (
|
|
GemmaRMSNorm as VllmGemmaRMSNorm,
|
|
fused_add_rms_norm as vllm_fused_add_rms_norm,
|
|
)
|
|
from vllm.model_executor.layers.rotary_embedding import get_rope as vllm_get_rope
|
|
|
|
from frontier.profiling.common.layers import layernorm as frontier_layernorm
|
|
from frontier.profiling.common.layers import rotary_embedding as frontier_rope
|
|
|
|
|
|
def _vllm020_get_rope_adapter(
|
|
*,
|
|
head_size: int,
|
|
rotary_dim: int,
|
|
max_position: int,
|
|
base: int | float,
|
|
is_neox_style: bool,
|
|
rope_scaling: dict[str, Any] | None,
|
|
dtype: torch.dtype | None = None,
|
|
) -> Any:
|
|
rope_parameters = dict(rope_scaling or {})
|
|
rope_parameters["rope_theta"] = base
|
|
rope_parameters["rope_dim"] = rotary_dim
|
|
return vllm_get_rope(
|
|
head_size=head_size,
|
|
max_position=max_position,
|
|
is_neox_style=is_neox_style,
|
|
rope_parameters=rope_parameters,
|
|
dtype=dtype,
|
|
)
|
|
|
|
|
|
def _vllm020_rms_norm_adapter(
|
|
x: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
variance_epsilon: float,
|
|
) -> torch.Tensor:
|
|
output = torch.empty_like(x)
|
|
vllm_ops.rms_norm(output, x, weight, variance_epsilon)
|
|
return output
|
|
|
|
|
|
def main() -> None:
|
|
parameters = inspect.signature(vllm_get_rope).parameters
|
|
if "rope_parameters" not in parameters or "rotary_dim" in parameters:
|
|
raise RuntimeError(
|
|
"Expected the vLLM 0.20 get_rope API with rope_parameters; "
|
|
f"found {inspect.signature(vllm_get_rope)}"
|
|
)
|
|
|
|
frontier_rope._VLLM_GET_ROPE = _vllm020_get_rope_adapter
|
|
frontier_rope._VLLM_GET_ROPE_IMPORT_ERROR = None
|
|
frontier_layernorm.HAS_VLLM_RMSNORM = True
|
|
frontier_layernorm.VllmGemmaRMSNorm = VllmGemmaRMSNorm
|
|
frontier_layernorm.vllm_rms_norm = _vllm020_rms_norm_adapter
|
|
frontier_layernorm.vllm_fused_add_rms_norm = vllm_fused_add_rms_norm
|
|
|
|
from frontier.profiling.linear_op.main import main as frontier_main
|
|
from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
|
|
|
|
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__":
|
|
main()
|