207 lines
6.7 KiB
Python
Executable File
207 lines
6.7 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Uncensored chat LLM API — loads an instruct model once, serves an
|
|
OpenAI-compatible /v1/chat/completions endpoint. Runs on the AMD MI50 (gfx906)
|
|
via ROCm 5.7 + torch 2.3.1, mirroring the joycaption service pattern.
|
|
|
|
Model fits in fp16 inside the 32GB VRAM (≈12B ceiling). Override with env MODEL_ID.
|
|
|
|
Endpoints:
|
|
POST /v1/chat/completions — OpenAI-compatible chat (used by gen_poses.py)
|
|
GET /v1/models — list the loaded model
|
|
GET /health — health check + GPU info
|
|
|
|
Env:
|
|
MODEL_ID HuggingFace repo id (default below)
|
|
PORT listen port (default 8001)
|
|
HSA_OVERRIDE_GFX_VERSION=9.0.6 set by start script for gfx906
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
import subprocess
|
|
import time
|
|
import uuid
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextlib import asynccontextmanager
|
|
|
|
import torch
|
|
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
torch.set_num_threads(1)
|
|
torch.set_num_interop_threads(1)
|
|
|
|
gpu_executor = ThreadPoolExecutor(max_workers=1)
|
|
|
|
# Mistral [INST] template — fallback for models without chat_template in tokenizer_config
|
|
_MISTRAL_TEMPLATE = (
|
|
"{{ bos_token }}"
|
|
"{% for message in messages %}"
|
|
"{% if message['role'] == 'system' %}"
|
|
"{{ '[INST] ' + message['content'] + '\n\n' }}"
|
|
"{% elif message['role'] == 'user' %}"
|
|
"{% if not loop.first or messages[0]['role'] != 'system' %}{{ '[INST] ' }}{% endif %}"
|
|
"{{ message['content'] + ' [/INST]' }}"
|
|
"{% elif message['role'] == 'assistant' %}"
|
|
"{{ ' ' + message['content'] + eos_token }}"
|
|
"{% endif %}"
|
|
"{% endfor %}"
|
|
)
|
|
|
|
MODEL_ID = os.environ.get("MODEL_ID", "dphn/Dolphin3.0-Mistral-24B")
|
|
# For 24B fp16 (~48GB): split across 32GB VRAM + CPU RAM via device_map="auto".
|
|
# Override with MAX_GPU_MEM / MAX_CPU_MEM env vars if needed.
|
|
MAX_GPU_MEM = os.environ.get("MAX_GPU_MEM", "30GiB") # leave ~2GB headroom
|
|
MAX_CPU_MEM = os.environ.get("MAX_CPU_MEM", "100GiB") # 113GB available
|
|
|
|
state: dict = {}
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
# Boost GPU to high performance mode (avoids power-saving clock throttle)
|
|
subprocess.run(["/opt/rocm/bin/rocm-smi", "--setperflevel", "high"], capture_output=True)
|
|
print(f"Loading model {MODEL_ID} (gpu≤{MAX_GPU_MEM} cpu≤{MAX_CPU_MEM})...", flush=True)
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
|
if not tokenizer.chat_template:
|
|
print("No chat_template found — applying Mistral [INST] fallback.", flush=True)
|
|
tokenizer.chat_template = _MISTRAL_TEMPLATE
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID,
|
|
torch_dtype=torch.float16,
|
|
device_map="auto",
|
|
max_memory={0: MAX_GPU_MEM, "cpu": MAX_CPU_MEM},
|
|
)
|
|
model.eval()
|
|
state["tokenizer"] = tokenizer
|
|
state["model"] = model
|
|
|
|
# Warm up: tiny generation to trigger kernel compilation
|
|
print("Warming up...", flush=True)
|
|
_run_chat([{"role": "user", "content": "hi"}], max_tokens=4, temperature=0.0, top_p=1.0, stop=None)
|
|
print(f"Model ready on: {next(model.parameters()).device}", flush=True)
|
|
yield
|
|
gpu_executor.shutdown(wait=False, cancel_futures=True)
|
|
state.clear()
|
|
|
|
|
|
app = FastAPI(title="Uncensored Chat LLM API", lifespan=lifespan)
|
|
|
|
|
|
def _run_chat(messages, max_tokens, temperature, top_p, stop):
|
|
tokenizer = state["tokenizer"]
|
|
model = state["model"]
|
|
|
|
input_ids = tokenizer.apply_chat_template(
|
|
messages, add_generation_prompt=True, return_tensors="pt"
|
|
).to(model.device)
|
|
prompt_tokens = input_ids.shape[1]
|
|
|
|
do_sample = temperature is not None and temperature > 0
|
|
gen_kwargs = dict(
|
|
max_new_tokens=max_tokens,
|
|
do_sample=do_sample,
|
|
use_cache=True,
|
|
pad_token_id=tokenizer.eos_token_id,
|
|
)
|
|
if do_sample:
|
|
gen_kwargs["temperature"] = temperature
|
|
gen_kwargs["top_p"] = top_p
|
|
|
|
t0 = time.perf_counter()
|
|
with torch.inference_mode():
|
|
output_ids = model.generate(input_ids, **gen_kwargs)
|
|
dt = time.perf_counter() - t0
|
|
|
|
new_ids = output_ids[0][prompt_tokens:]
|
|
text = tokenizer.decode(new_ids, skip_special_tokens=True)
|
|
|
|
finish = "stop"
|
|
if stop:
|
|
stops = [stop] if isinstance(stop, str) else list(stop)
|
|
cut = min((text.find(s) for s in stops if s and s in text), default=-1)
|
|
if cut != -1:
|
|
text = text[:cut]
|
|
completion_tokens = int(new_ids.shape[0])
|
|
if completion_tokens >= max_tokens:
|
|
finish = "length"
|
|
|
|
del input_ids, output_ids
|
|
return {
|
|
"text": text.strip(),
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"finish_reason": finish,
|
|
"generate_s": round(dt, 2),
|
|
}
|
|
|
|
|
|
async def run_chat(*args):
|
|
return await asyncio.get_running_loop().run_in_executor(gpu_executor, _run_chat, *args)
|
|
|
|
|
|
class ChatMessage(BaseModel):
|
|
role: str
|
|
content: str
|
|
|
|
|
|
class ChatRequest(BaseModel):
|
|
model: str | None = None
|
|
messages: list[ChatMessage]
|
|
max_tokens: int = 512
|
|
temperature: float = 0.8
|
|
top_p: float = 0.95
|
|
stop: list[str] | str | None = None
|
|
|
|
|
|
@app.post("/v1/chat/completions")
|
|
async def chat_completions(req: ChatRequest):
|
|
if not req.messages:
|
|
raise HTTPException(400, "messages must not be empty")
|
|
msgs = [{"role": m.role, "content": m.content} for m in req.messages]
|
|
r = await run_chat(msgs, req.max_tokens, req.temperature, req.top_p, req.stop)
|
|
return {
|
|
"id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
|
|
"object": "chat.completion",
|
|
"created": int(time.time()),
|
|
"model": MODEL_ID,
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": r["text"]},
|
|
"finish_reason": r["finish_reason"],
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": r["prompt_tokens"],
|
|
"completion_tokens": r["completion_tokens"],
|
|
"total_tokens": r["prompt_tokens"] + r["completion_tokens"],
|
|
},
|
|
"timing": {"generate_s": r["generate_s"]},
|
|
}
|
|
|
|
|
|
@app.get("/v1/models")
|
|
async def list_models():
|
|
return {"object": "list", "data": [{"id": MODEL_ID, "object": "model"}]}
|
|
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
gpu = {}
|
|
if torch.cuda.is_available():
|
|
gpu = {
|
|
"name": torch.cuda.get_device_name(0),
|
|
"vram_total_gb": round(torch.cuda.get_device_properties(0).total_memory / 1e9, 1),
|
|
"vram_used_gb": round(torch.cuda.memory_allocated(0) / 1e9, 1),
|
|
}
|
|
return {"status": "ok", "model": MODEL_ID, "loaded": "model" in state, "gpu": gpu}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 8001)))
|