Tiktok-Talent-Info/pipeline_setup.py

27 lines
745 B
Python

import os
from lmdeploy import pipeline, TurbomindEngineConfig, GenerationConfig
IMAGE_TOKEN = "[IMAGE_TOKEN]"
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
# os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
# Model initialization
# model = "OpenGVLab/InternVL2-26B-AWQ"
# model = "OpenGVLab/InternVL2_5-4B-AWQ"
model = "OpenGVLab/InternVL2_5-8B-MPO-AWQ"
pipe = pipeline(
model,
backend_config=TurbomindEngineConfig(
model_format="awq",
tp=2,
# tp=4,
session_len=16384, # 4096, 8192, 16384, 32768
max_batch_size=1,
cache_max_entry_count=0.15, # 0.05
cache_block_seq_len=16384, # 8192, 16384, 32768
# quant_policy=8,
# precision="fp16",
),
# log_level='DEBUG'
)