File size: 3,116 Bytes
954e44f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
import argparse
import os
import sys
PLUGIN_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_DIR = os.path.dirname(PLUGIN_DIR)
sys.path.insert(0, PLUGIN_DIR)
# Register the Quartet II quantization plugin before any vLLM imports
import quartet2_quant # noqa: F401 — triggers @register_quantization_config
from vllm import ModelRegistry
from cloverlm_vllm import CloverLMForCausalLM_vLLM
ModelRegistry.register_model(
"CloverLMForCausalLM", CloverLMForCausalLM_vLLM,
)
def main():
parser = argparse.ArgumentParser(description="Serve CloverLM with vLLM")
parser.add_argument(
"--model", default=MODEL_DIR,
help="Path to CloverLM model directory",
)
parser.add_argument("--api", action="store_true", help="Start OpenAI API server")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--host", default="0.0.0.0")
parser.add_argument("--tp", type=int, default=1, help="Tensor parallel size")
parser.add_argument(
"--max-model-len", type=int, default=1024,
help="Maximum context length",
)
parser.add_argument(
"--gpu-memory-utilization", type=float, default=0.9,
)
args = parser.parse_args()
if args.api:
_serve_api(args)
else:
_offline_inference(args)
def _offline_inference(args):
from vllm import LLM, SamplingParams
llm = LLM(
model=args.model,
quantization="quartet2",
trust_remote_code=True,
dtype="bfloat16",
max_model_len=args.max_model_len,
tensor_parallel_size=args.tp,
gpu_memory_utilization=args.gpu_memory_utilization,
enforce_eager=True,
)
sampling_params = SamplingParams(
temperature=0.8,
top_p=0.95,
max_tokens=128,
)
prompts = [
"The capital of France is",
"Large language models are",
"In the year 2030,",
]
print("=" * 60)
print(" CloverLM — vLLM Offline Inference (Quartet II NVFP4)")
print("=" * 60)
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
prompt = output.prompt
generated = output.outputs[0].text
print(f"\nPrompt: {prompt}")
print(f"Generated: {generated}")
def _serve_api(args):
sys.argv = [
"vllm",
"--model", args.model,
"--quantization", "quartet2",
"--trust-remote-code",
"--dtype", "bfloat16",
"--max-model-len", str(args.max_model_len),
"--tensor-parallel-size", str(args.tp),
"--gpu-memory-utilization", str(args.gpu_memory_utilization),
"--enforce-eager",
"--host", args.host,
"--port", str(args.port),
]
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm.entrypoints.openai.cli_args import make_arg_parser
from vllm.entrypoints.openai.api_server import run_server
import asyncio
vllm_parser = make_arg_parser(FlexibleArgumentParser())
vllm_args = vllm_parser.parse_args()
asyncio.run(run_server(vllm_args))
if __name__ == "__main__":
main()
|