import os import torch import torch.nn as nn import torch.nn.functional as F import uvicorn import threading import time import uuid import json from fastapi import FastAPI, Request, HTTPException from fastapi.responses import JSONResponse from fastapi.responses import StreamingResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field, ConfigDict from typing import List, Optional, Union, Dict, Any from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer from transformers import BitsAndBytesConfig from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training from torch.utils.checkpoint import checkpoint from transformers.modeling_outputs import CausalLMOutputWithPast import asyncio # ========================================== # 1. 初始化 FastAPI App 與 跨域(CORS) 設定 # ========================================== app = FastAPI(title="OpenAI-Compatible FrankenQwen Server", version="1.0") # 開啟 CORS，讓前端 UI (如 Open WebUI) 可以順利連線 app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) MAX_CONCURRENT_REQUESTS = 4 gpu_semaphore = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS) def get_vram_status(): import torch if not torch.cuda.is_available(): return {"status": "CUDA 不可用"} # 轉換為 GiB (除以 1024^3) allocated = torch.cuda.memory_allocated() / (1024 ** 3) reserved = torch.cuda.memory_reserved() / (1024 ** 3) peak = torch.cuda.max_memory_allocated() / (1024 ** 3) return { "current_tensor_used_gib": round(allocated, 2), "pytorch_reserved_gib": round(reserved, 2), "peak_vram_used_gib": round(peak, 2) } # ========================================== # 2. 原始模型組件與邏輯 (保持不變) # ========================================== class HiddenStateCompressor(nn.Module): def __init__(self, dim, ratio=4): super().__init__() self.ratio = ratio self.wgate = nn.Linear(dim, 1, bias=False) self.ape = nn.Parameter(torch.empty(ratio, 1)) nn.init.normal_(self.ape, std=0.02) def forward(self, x): B, L, D = x.shape assert L % self.ratio == 0, "輸入長度必須是壓縮倍率的整數倍" orig_dtype = x.dtype target_dtype = self.wgate.weight.dtype x_reshaped = x.view(B, L // self.ratio, self.ratio, D).to(target_dtype) scores = self.wgate(x_reshaped) + self.ape.view(1, 1, self.ratio, 1) scores_fp32 = scores.to(torch.float32) scores_probs = F.softmax(scores_fp32, dim=2).to(target_dtype) compressed_x = (x_reshaped * scores_probs).sum(dim=2) return compressed_x.to(orig_dtype) def get_custom_qwen_forward(compressor_module, window_size=128, ratio=4): from collections.abc import Callable from transformers.models.qwen3_5.modeling_qwen3_5 import ALL_ATTENTION_FUNCTIONS, apply_rotary_pos_emb, eager_attention_forward, FlashAttentionKwargs from transformers.processing_utils import Unpack from transformers.cache_utils import Cache def forward2( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None, past_key_values: Cache | None = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple[torch.Tensor, torch.Tensor | None]: cos, sin = position_embeddings bsz, q_len, _ = hidden_states.size() input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) hidden_states = hidden_states.contiguous() query_states, gate = torch.chunk( self.q_proj(hidden_states).view(*input_shape, -1, self.head_dim * 2), 2, dim=-1 ) gate = gate.contiguous().reshape(*input_shape, -1) query_states = self.q_norm(query_states.contiguous().view(hidden_shape)).transpose(1, 2) key_states, value_states = None, None do_compress = getattr(self, "compress_enabled", True) if do_compress: if q_len > 1: # Prefill window = min(window_size, q_len) global_len = q_len - window cutoff = global_len - (global_len % ratio) compressed = compressor_module(hidden_states[:, :cutoff, :]) if cutoff > 0 else hidden_states[:, :0, :] residual = hidden_states[:, cutoff:global_len, :] if global_len > cutoff else hidden_states[:, :0, :] tail = hidden_states[:, global_len:, :] kv_hidden = torch.cat([compressed, residual, tail], dim=1).contiguous() compressed_pos = torch.arange(ratio - 1, cutoff, ratio, device=hidden_states.device, dtype=torch.long) if cutoff > 0 else torch.arange(0, device=hidden_states.device, dtype=torch.long) residual_pos = torch.arange(cutoff, global_len, device=hidden_states.device, dtype=torch.long) tail_pos = torch.arange(global_len, q_len, device=hidden_states.device, dtype=torch.long) kv_positions = torch.cat([compressed_pos, residual_pos, tail_pos], dim=0) kv_seq_len = kv_positions.size(0) q_pos = torch.arange(q_len, device=hidden_states.device).view(q_len, 1) k_pos = kv_positions.view(1, kv_seq_len) custom_causal_mask = (q_pos >= k_pos) additive_mask = torch.zeros((q_len, kv_seq_len), dtype=hidden_states.dtype, device=hidden_states.device) additive_mask.masked_fill_(~custom_causal_mask, torch.finfo(hidden_states.dtype).min) additive_mask = additive_mask.view(1, 1, q_len, kv_seq_len).expand(bsz, 1, q_len, kv_seq_len) attention_mask = (attention_mask[..., kv_positions] + additive_mask) if attention_mask is not None else additive_mask kv_hidden_shape = (bsz, kv_seq_len, -1, self.head_dim) key_states = self.k_norm(self.k_proj(kv_hidden).view(kv_hidden_shape)).transpose(1, 2).contiguous() value_states = self.v_proj(kv_hidden).view(kv_hidden_shape).transpose(1, 2).contiguous() query_states, _ = apply_rotary_pos_emb(query_states, query_states, cos, sin) k_cos = cos[:, kv_positions, :] if cos.dim() == 3 else cos[kv_positions, :] k_sin = sin[:, kv_positions, :] if sin.dim() == 3 else sin[kv_positions, :] _, key_states = apply_rotary_pos_emb(key_states, key_states, k_cos, k_sin) if past_key_values is not None: key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) else: # Decode new_key = self.k_norm(self.k_proj(hidden_states).view(bsz, 1, -1, self.head_dim)).transpose(1, 2) new_value = self.v_proj(hidden_states).view(bsz, 1, -1, self.head_dim).transpose(1, 2) query_states, _ = apply_rotary_pos_emb(query_states, query_states, cos, sin) _, new_key = apply_rotary_pos_emb(new_key, new_key, cos, sin) target_dtype = hidden_states.dtype new_key = new_key.to(target_dtype).contiguous() new_value = new_value.to(target_dtype).contiguous() if past_key_values is not None: if not hasattr(past_key_values, "uncompressed_buffer_k"): past_key_values.uncompressed_buf_h = {} past_key_values.uncompressed_buffer_k = {} past_key_values.uncompressed_buffer_v = {} layer_idx = self.layer_idx buf_k = past_key_values.uncompressed_buffer_k.get(layer_idx, torch.empty((bsz, self.config.num_key_value_heads, 0, self.head_dim), device=new_key.device, dtype=target_dtype)) buf_v = past_key_values.uncompressed_buffer_v.get(layer_idx, torch.empty((bsz, self.config.num_key_value_heads, 0, self.head_dim), device=new_value.device, dtype=target_dtype)) buf_h = past_key_values.uncompressed_buf_h.get(layer_idx, torch.empty((bsz, 0, hidden_states.size(-1)), device=hidden_states.device, dtype=hidden_states.dtype)) buf_h = torch.cat([buf_h, hidden_states], dim=1) buf_k = torch.cat([buf_k, new_key], dim=2) buf_v = torch.cat([buf_v, new_value], dim=2) buffer_len = buf_k.shape[2] if getattr(self, "compress_enabled", True) and buffer_len >= (window_size + ratio): to_compress_h = buf_h[:, :ratio, :] compressed_h = compressor_module(to_compress_h) compressed_k = self.k_norm(self.k_proj(compressed_h).view(bsz, 1, -1, self.head_dim)).transpose(1, 2) compressed_v = self.v_proj(compressed_h).view(bsz, 1, -1, self.head_dim).transpose(1, 2) past_key_values.update(compressed_k, compressed_v, layer_idx) buf_h = buf_h[:, ratio:, :] buf_k = buf_k[:, :, ratio:, :] buf_v = buf_v[:, :, ratio:, :] past_key_values.uncompressed_buffer_k[layer_idx] = buf_k past_key_values.uncompressed_buffer_v[layer_idx] = buf_v past_key_values.uncompressed_buf_h[layer_idx] = buf_h layer_cache = past_key_values.layers[layer_idx] if hasattr(past_key_values, "layers") else None history_k = layer_cache.keys if layer_cache else past_key_values.key_cache[layer_idx] history_v = layer_cache.values if layer_cache else past_key_values.value_cache[layer_idx] key_states = torch.cat([history_k, buf_k], dim=2) if history_k is not None and history_k.shape[2] > 0 else buf_k value_states = torch.cat([history_v, buf_v], dim=2) if history_v is not None and history_v.shape[2] > 0 else buf_v kv_seq_len = past_key_values.get_seq_length(self.layer_idx) + buf_k.shape[2] else: key_states, value_states = new_key, new_value else: kv_hidden_shape = (bsz, q_len, -1, self.head_dim) key_states = self.k_norm(self.k_proj(hidden_states).view(kv_hidden_shape)).transpose(1, 2).contiguous() value_states = self.v_proj(hidden_states).view(kv_hidden_shape).transpose(1, 2).contiguous() query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_values is not None: key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) kv_seq_len = key_states.shape[2] target_dtype = hidden_states.dtype query_states, key_states, value_states = query_states.to(target_dtype).contiguous(), key_states.to(target_dtype).contiguous(), value_states.to(target_dtype).contiguous() attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(self.config._attn_implementation, eager_attention_forward) attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = attn_output * torch.sigmoid(gate) return self.o_proj(attn_output), attn_weights return forward2 def get_chunked_causal_forward(chunk_size=512): def custom_causal_forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, labels=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None, **kwargs): outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, **kwargs) hidden_states = outputs[0] loss, logits = None, None if labels is not None: shift_hidden_states = hidden_states[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() seq_len = shift_hidden_states.shape[1] total_loss = 0.0 active_tokens = (shift_labels != -100).sum().clamp(min=1) def compute_chunk_loss(h_chunk, lbl_chunk): c_logits = self.lm_head(h_chunk).float() return F.cross_entropy(c_logits.reshape(-1, c_logits.size(-1)), lbl_chunk.reshape(-1), reduction='sum') for i in range(0, seq_len, chunk_size): end_i = min(i + chunk_size, seq_len) c_loss = checkpoint(compute_chunk_loss, shift_hidden_states[:, i:end_i, :], shift_labels[:, i:end_i], use_reentrant=False) total_loss += c_loss loss = total_loss / active_tokens else: logits = self.lm_head(hidden_states) return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions) return custom_causal_forward def load_custom_checkpoint(model, filepath="/vault/franken_qwen2B_compressor_checkpoint.pt"): if not os.path.exists(filepath): print(f"⚠️ 找不到權重檔 {filepath}！將使用原始權重。") return model print(f"\n📂 正在載入客製化權重: {filepath}") trainable_state_dict = torch.load(filepath, map_location="cpu", weights_only=True) clean_state_dict = { (k[7:] if k.startswith('module.') else k): v for k, v in trainable_state_dict.items() } model.load_state_dict(clean_state_dict, strict=False) print("✅ 客製化權重完美掛載！") return model def build_franken_qwen(qwen_id="Qwen/Qwen3.6-27B", device="cuda"): quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16) qwen_model = AutoModelForCausalLM.from_pretrained(qwen_id, dtype=torch.bfloat16, device_map="auto", quantization_config=quantization_config) #qwen_model = prepare_model_for_kbit_training(qwen_model) qwen_model.compressors = nn.ModuleList() for i, layer in enumerate(qwen_model.model.layers): if i % 4 == 3: compressor = HiddenStateCompressor(dim=qwen_model.config.hidden_size, ratio=4).to(device, dtype=torch.bfloat16) qwen_model.compressors.append(compressor) custom_fw = get_custom_qwen_forward(compressor, window_size=128, ratio=4) layer.self_attn.forward = custom_fw.__get__(layer.self_attn, type(layer.self_attn)) qwen_model.forward = get_chunked_causal_forward(chunk_size=512).__get__(qwen_model, type(qwen_model)) for param in qwen_model.parameters(): param.requires_grad = False for compressor in qwen_model.compressors: for param in compressor.parameters(): param.requires_grad = True if qwen_id=="Qwen/Qwen3.6-27B": lora_config = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj"], lora_dropout=0.05, bias="none", task_type="CAUSAL_LM") qwen_model = get_peft_model(qwen_model, lora_config) return load_custom_checkpoint(qwen_model, filepath="./Downloads/franken_qwen27B_compressor_lora.pt") elif qwen_id=="Qwen/Qwen3.5-2B": return load_custom_checkpoint(qwen_model, filepath="./Downloads/franken_qwen2B_compressor_checkpoint.pt") qwen_model.enable_input_require_grads() return load_custom_checkpoint(qwen_model, filepath="/vault/franken_qwen27B_compressor_lora.pt") # 全域模型初始化 torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True MODEL_ID = "Qwen/Qwen3.5-2B" tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) model = build_franken_qwen(qwen_id=MODEL_ID, device="cuda") model=model.get_base_model() if hasattr(model, "get_base_model") else model model=torch.compile(model) model.eval() print(get_vram_status()) # ========================================== # 3. OpenAI 官方相容的 Pydantic 資料模型定義 # ========================================== class ChatMessage(BaseModel): role: str content: str class ChatCompletionRequest(BaseModel): # 🌟 關鍵優化：extra="allow" 允許 lm-eval 傳入額外的 OpenAI 參數 (如 n, logprobs) 而不報錯 model_config = ConfigDict(extra="allow") model: str messages: List[ChatMessage] temperature: Optional[float] = 0.7 top_p: Optional[float] = 0.9 max_tokens: Optional[int] = Field(default=512, alias="max_tokens") stream: Optional[bool] = False stop: Optional[Union[str, List[str]]] = None # 🌟 接收 lm-eval 的停止符號 class CompletionRequest(BaseModel): # 🌟 針對標準文字補全的 Request Schema model_config = ConfigDict(extra="allow") model: str prompt: Union[str, List[str]] temperature: Optional[float] = 0.7 top_p: Optional[float] = 0.9 max_tokens: Optional[int] = Field(default=512, alias="max_tokens") stream: Optional[bool] = False stop: Optional[Union[str, List[str]]] = None # ========================================== # 4. OpenAI 相容核心端點 (/v1/chat/completions) # ========================================== @app.exception_handler(Exception) async def universal_exception_handler(request: Request, exc: Exception): return JSONResponse( status_code=500, content={"error": {"message": f"Internal Server Error: {str(exc)}", "type": "server_error"}} ) def prepare_gen_config(request, inputs): """公用生成參數建立函式""" gen_config = { **inputs, "max_new_tokens": request.max_tokens, "do_sample": True if request.temperature > 0 else False, "temperature": request.temperature if request.temperature > 0 else None, "top_p": request.top_p, "repetition_penalty": 1.05, "use_cache": True, "pad_token_id": tokenizer.eos_token_id, "eos_token_id": [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|im_end|>")] } # 🌟 處理 lm-eval 傳過來的 stop 截斷字串 if request.stop: gen_config["stop_strings"] = [request.stop] if isinstance(request.stop, str) else request.stop gen_config["tokenizer"] = tokenizer return gen_config @app.post("/v1/chat/completions") async def chat_completions_endpoint(request: ChatCompletionRequest): async with gpu_semaphore: # 佇列鎖管控 formatted_messages = [{"role": msg.role, "content": msg.content} for msg in request.messages] input_text = tokenizer.apply_chat_template(formatted_messages, tokenize=False, add_generation_prompt=True) inputs = tokenizer(input_text, return_tensors="pt", add_special_tokens=False).to(model.device) gen_config = prepare_gen_config(request, inputs) streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) gen_config["streamer"] = streamer threading.Thread(target=lambda: torch.no_grad()(model.generate)(**gen_config)).start() if not request.stream: # 🌟 優化：將同步阻塞的 "".join 丟進線程池，解放 Event Loop full_text = await asyncio.to_thread(lambda: "".join([new_text for new_text in streamer])) return { "id": f"chatcmpl-{uuid.uuid4()}", "object": "chat.completion", "created": int(time.time()), "model": request.model, "choices": [{"index": 0, "message": {"role": "assistant", "content": full_text}, "finish_reason": "stop"}], "usage": {"prompt_tokens": inputs["input_ids"].shape[1], "completion_tokens": len(tokenizer.encode(full_text)), "total_tokens": inputs["input_ids"].shape[1] + len(tokenizer.encode(full_text))} } else: async def openai_chat_stream(): chunk_id, created_time = f"chatcmpl-{uuid.uuid4()}", int(time.time()) yield f"data: {json.dumps({'id': chunk_id, 'object': 'chat.completion.chunk', 'created': created_time, 'model': request.model, 'choices': [{'index': 0, 'delta': {'role': 'assistant'}, 'finish_reason': None}]}, ensure_ascii=False)}\n\n" # 🌟 優化：建立安全不阻塞事件循環的取 Token 函式 def get_next_token(): try: return next(streamer), False except StopIteration: return "", True while True: # 🌟 關鍵：每次拿 Token 都讓出主執行緒，讓監控端點可以插隊進來執行 new_text, is_done = await asyncio.to_thread(get_next_token) if is_done: break if new_text: yield f"data: {json.dumps({'id': chunk_id, 'object': 'chat.completion.chunk', 'created': created_time, 'model': request.model, 'choices': [{'index': 0, 'delta': {'content': new_text}, 'finish_reason': None}]}, ensure_ascii=False)}\n\n" yield f"data: {json.dumps({'id': chunk_id, 'object': 'chat.completion.chunk', 'created': created_time, 'model': request.model, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]}, ensure_ascii=False)}\n\n" yield "data: [DONE]\n\n" return StreamingResponse(openai_chat_stream(), media_type="text/event-stream") # ----- 端點 B：文字接龍 (/v1/completions) ----- @app.post("/v1/completions") async def completions_endpoint(request: CompletionRequest): async with gpu_semaphore: # 佇列鎖管控 input_text = request.prompt[0] if isinstance(request.prompt, list) else request.prompt inputs = tokenizer(input_text, return_tensors="pt", add_special_tokens=False).to(model.device) gen_config = prepare_gen_config(request, inputs) streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) gen_config["streamer"] = streamer threading.Thread(target=lambda: torch.no_grad()(model.generate)(**gen_config)).start() if not request.stream: full_text = "".join([new_text for new_text in streamer]) return { "id": f"cmpl-{uuid.uuid4()}", "object": "text_completion", "created": int(time.time()), "model": request.model, "choices": [{"text": full_text, "index": 0, "logprobs": None, "finish_reason": "stop"}], "usage": {"prompt_tokens": inputs["input_ids"].shape[1], "completion_tokens": len(tokenizer.encode(full_text)), "total_tokens": inputs["input_ids"].shape[1] + len(tokenizer.encode(full_text))} } else: async def openai_text_stream(): chunk_id, created_time = f"cmpl-{uuid.uuid4()}", int(time.time()) for new_text in streamer: if new_text: yield f"data: {json.dumps({'id': chunk_id, 'object': 'text_completion.chunk', 'created': created_time, 'model': request.model, 'choices': [{'text': new_text, 'index': 0, 'logprobs': None, 'finish_reason': None}]}, ensure_ascii=False)}\n\n" yield f"data: {json.dumps({'id': chunk_id, 'object': 'text_completion.chunk', 'created': created_time, 'model': request.model, 'choices': [{'text': '', 'index': 0, 'logprobs': None, 'finish_reason': 'stop'}]}, ensure_ascii=False)}\n\n" yield "data: [DONE]\n\n" return StreamingResponse(openai_text_stream(), media_type="text/event-stream") # ========================================== # 5. 補齊 OpenAI 的 /v1/models 查詢端點 (許多 UI 啟動時會先呼叫這個) # ========================================== @app.get("/v1/models") async def list_models(): return { "object": "list", "data": [ { "id": MODEL_ID, "object": "model", "created": int(time.time()), "owned_by": "custom" } ] } @app.get("/v1/debug/vram") async def vram_monitor(): """隨時動態查詢當前 GPU 記憶體狀態""" return { "object": "vram_status", "data": get_vram_status() } if __name__ == "__main__": uvicorn.run(app, host="127.0.0.1", port=8001)