from functools import lru_cache from pathlib import Path from time import perf_counter import gradio as gr import spaces import torch from transformers import ( AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig, ) DEFAULT_MODEL_ID = "dphn/dolphin-2.9.4-llama3.1-8b" ROOT = Path(__file__).parent TEXT_DIR = ROOT / "text" THEORY_PATH = ROOT / "notes" / "deployment_interview_theory.md" LESSONS = { "00 pipeline 自动配置": "00_pipeline_auto.py", "01 tokenizer + model": "01_auto_tokenizer_model.py", "02 模块分开配置": "02_split_configuration.py", "03 streaming + stopping": "03_streaming_and_stopping.py", "04 forward + logits + KV cache": "04_forward_logits_kv_cache.py", "05 架构 + 中间态输出": "05_architecture_and_intermediates.py", } LESSON_NOTES = { "00 pipeline 自动配置": """ `pipeline("text-generation")` 是最高层封装。它会自动加载 tokenizer、model、generation config，并完成 tokenize、generate、decode。 面试表达：pipeline 适合快速验证；生产部署通常需要把 tokenizer、model、生成参数、设备和流式输出拆开控制。 """, "01 tokenizer + model": """ 这一层显式拆开 `AutoTokenizer` 和 `AutoModelForCausalLM`。 关键链路：messages -> chat template -> token ids -> model.generate -> new token ids -> decode。 """, "02 模块分开配置": """ 这一层把架构配置、分词器配置、模型加载配置、生成配置分开。 面试表达：拆开配置后才能精确控制 dtype、device、max_new_tokens、temperature、top_p、pad/eos token 等部署关键项。 """, "03 streaming + stopping": """ streaming 不是减少总计算量，而是让用户更早看到输出。 停止条件通常来自 `eos_token_id`、stop string、`max_new_tokens` 或自定义 `StoppingCriteria`。 """, "04 forward + logits + KV cache": """ 不调用 `generate()`，直接看一次 forward。 模型输出 logits，解码策略从 logits 中选下一个 token。`past_key_values` 就是 KV cache，是长上下文和高并发部署里的显存大头之一。 """, "05 架构 + 中间态输出": """ 这一层用 `print(model)`、`output_hidden_states=True` 和可选的 `output_attentions=True` 看模型内部。 `hidden_states` 能看到 embedding 和每层 decoder block 后的表示；`attentions` 能看到每层每个 head 对输入 token 的注意力矩阵，但它的显存开销是 seq_len 平方级，长 prompt 时要谨慎开启。 """, } def read_text(path, fallback): if path.exists(): return path.read_text(encoding="utf-8") return fallback TEXT_README = read_text(TEXT_DIR / "README.md", "文本教程文件还没有生成。") THEORY_TEXT = read_text(THEORY_PATH, "理论笔记文件还没有生成。") def device_summary(): if torch.cuda.is_available(): name = torch.cuda.get_device_name(0) return f"cuda:0 ({name})" if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): return "mps" return "cpu" def model_load_kwargs(): if torch.cuda.is_available(): dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 return {"torch_dtype": dtype, "device_map": "auto"} return {"torch_dtype": torch.float32} @lru_cache(maxsize=1) def load_components(model_id): tokenizer = AutoTokenizer.from_pretrained(model_id) config = AutoConfig.from_pretrained(model_id) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( model_id, config=config, low_cpu_mem_usage=True, **model_load_kwargs(), ) if not hasattr(model, "hf_device_map"): if torch.cuda.is_available(): model = model.to("cuda") elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): model = model.to("mps") model.eval() return tokenizer, model, config def model_device(model): return next(model.parameters()).device def build_messages(system_prompt, user_prompt): messages = [] if system_prompt.strip(): messages.append({"role": "system", "content": system_prompt.strip()}) messages.append({"role": "user", "content": user_prompt.strip()}) return messages def build_prompt(tokenizer, messages): if tokenizer.chat_template: return tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, ) rendered = [] for message in messages: rendered.append(f"{message['role']}: {message['content']}") rendered.append("assistant:") return "\n".join(rendered) def generation_config(tokenizer, max_new_tokens, do_sample, temperature, top_p): config = { "max_new_tokens": int(max_new_tokens), "do_sample": bool(do_sample), "repetition_penalty": 1.05, "pad_token_id": tokenizer.pad_token_id, "eos_token_id": tokenizer.eos_token_id, } if do_sample: config["temperature"] = float(temperature) config["top_p"] = float(top_p) return GenerationConfig(**config) def inspect_config(config, tokenizer, model): rows = [ ("model_type", getattr(config, "model_type", "unknown")), ("hidden_size", getattr(config, "hidden_size", "unknown")), ("num_hidden_layers", getattr(config, "num_hidden_layers", "unknown")), ("num_attention_heads", getattr(config, "num_attention_heads", "unknown")), ("num_key_value_heads", getattr(config, "num_key_value_heads", "unknown")), ("vocab_size", getattr(config, "vocab_size", "unknown")), ("pad_token_id", tokenizer.pad_token_id), ("eos_token_id", tokenizer.eos_token_id), ("runtime_device", str(model_device(model))), ] return "\n".join(f"{key}: {value}" for key, value in rows) def tensor_summary(value): if value is None: return "None" if hasattr(value, "shape"): return ( f"shape={tuple(value.shape)}, dtype={value.dtype}, device={value.device}" ) return type(value).__name__ def architecture_summary(config, model): lines = [ f"model class: {type(model).__name__}", f"model_type: {getattr(config, 'model_type', 'unknown')}", f"hidden_size: {getattr(config, 'hidden_size', 'unknown')}", f"num_hidden_layers: {getattr(config, 'num_hidden_layers', 'unknown')}", f"num_attention_heads: {getattr(config, 'num_attention_heads', 'unknown')}", f"num_key_value_heads: {getattr(config, 'num_key_value_heads', 'unknown')}", f"vocab_size: {getattr(config, 'vocab_size', 'unknown')}", "", "top-level modules:", ] for name, module in model.named_children(): lines.append(f"- {name}: {type(module).__name__}") if hasattr(model, "model") and hasattr(model.model, "layers"): layers = model.model.layers lines.extend(["", f"decoder layers: {len(layers)}"]) if layers: first_layer = layers[0] lines.append("first decoder layer modules:") for name, module in first_layer.named_children(): lines.append(f"- {name}: {type(module).__name__}") lines.extend(["", "full model repr:", str(model)]) return "\n".join(lines) def summarize_tensor_sequence(name, values): if values is None: return f"{name}: None" lines = [f"{name} count: {len(values)}"] for index, value in enumerate(values): lines.append(f"{name}[{index:02d}]: {tensor_summary(value)}") return "\n".join(lines) def summarize_cache(cache): if cache is None: return "KV cache: None" lines = [f"KV cache type: {type(cache).__name__}"] try: lines.append(f"KV cache layers: {len(cache)}") except TypeError: lines.append("KV cache layers: unknown") if hasattr(cache, "get_seq_length"): lines.append(f"KV cache sequence length: {cache.get_seq_length()}") if hasattr(cache, "key_cache") and getattr(cache, "key_cache"): lines.append( f"layer0 key: {tensor_summary(cache.key_cache[0])}" ) lines.append( f"layer0 value: {tensor_summary(cache.value_cache[0])}" ) elif isinstance(cache, (tuple, list)) and cache: first_layer = cache[0] if isinstance(first_layer, (tuple, list)) and len(first_layer) >= 2: lines.append(f"layer0 key: {tensor_summary(first_layer[0])}") lines.append(f"layer0 value: {tensor_summary(first_layer[1])}") return "\n".join(lines) def explain_tokens(tokenizer, prompt): encoded = tokenizer(prompt, return_tensors="pt") ids = encoded["input_ids"][0].tolist() tokens = tokenizer.convert_ids_to_tokens(ids) preview = list(zip(ids[:80], tokens[:80])) lines = [f"token_count: {len(ids)}", "", "first_tokens:"] lines.extend(f"{token_id:>8} {token}" for token_id, token in preview) if len(ids) > 80: lines.append(f"... {len(ids) - 80} more tokens") return "\n".join(lines) @spaces.GPU(duration=120) def run_experiment( lesson_name, model_id, system_prompt, user_prompt, max_new_tokens, do_sample, temperature, top_p, output_attentions, ): if not user_prompt.strip(): raise gr.Error("请输入 user prompt。") start = perf_counter() tokenizer, model, config = load_components(model_id.strip() or DEFAULT_MODEL_ID) messages = build_messages(system_prompt, user_prompt) prompt = build_prompt(tokenizer, messages) inputs = tokenizer(prompt, return_tensors="pt").to(model_device(model)) input_len = inputs["input_ids"].shape[-1] if lesson_name == "04 forward + logits + KV cache": with torch.no_grad(): outputs = model(**inputs, use_cache=True) logits = outputs.logits[:, -1, :] next_token_id = logits.argmax(dim=-1) next_token = tokenizer.decode(next_token_id) cache = outputs.past_key_values try: cache_layers = len(cache) except TypeError: cache_layers = "unknown" if hasattr(cache, "get_seq_length"): cache_seq_len = cache.get_seq_length() else: cache_seq_len = input_len elapsed = perf_counter() - start answer = ( "这是一次 forward 观察，不是完整生成。\n\n" f"greedy next token id: {next_token_id.item()}\n" f"greedy next token: {next_token!r}\n" f"logits shape: {tuple(logits.shape)}\n" f"KV cache type: {type(cache).__name__}\n" f"KV cache layers: {cache_layers}\n" f"KV cache sequence length: {cache_seq_len}" ) metrics = { "lesson": lesson_name, "input_tokens": input_len, "new_tokens": 1, "elapsed_seconds": round(elapsed, 3), "device": str(model_device(model)), } return answer, inspect_config(config, tokenizer, model), explain_tokens(tokenizer, prompt), metrics if lesson_name == "05 架构 + 中间态输出": try: with torch.no_grad(): outputs = model( **inputs, use_cache=True, output_hidden_states=True, output_attentions=bool(output_attentions), return_dict=True, ) except Exception as exc: if output_attentions: raise gr.Error( "开启 output_attentions 时 forward 失败。attention 矩阵显存开销较大，" "可以先关闭 output_attentions 再观察 hidden_states 和 KV cache。" ) from exc raise logits = outputs.logits last_token_logits = logits[:, -1, :] next_token_id = last_token_logits.argmax(dim=-1) next_token = tokenizer.decode(next_token_id) elapsed = perf_counter() - start answer = "\n\n".join( [ "这是一次架构和中间态观察，不是完整生成。", "\n".join( [ f"outputs type: {type(outputs).__name__}", f"outputs keys: {list(outputs.keys())}", f"logits: {tensor_summary(logits)}", f"last token logits: {tensor_summary(last_token_logits)}", f"greedy next token id: {next_token_id.item()}", f"greedy next token: {next_token!r}", ] ), summarize_tensor_sequence("hidden_states", outputs.hidden_states), summarize_tensor_sequence("attentions", outputs.attentions), summarize_cache(outputs.past_key_values), ] ) metrics = { "lesson": lesson_name, "input_tokens": input_len, "elapsed_seconds": round(elapsed, 3), "logits_shape": list(logits.shape), "hidden_state_count": len(outputs.hidden_states) if outputs.hidden_states is not None else 0, "attention_count": len(outputs.attentions) if outputs.attentions is not None else 0, "output_attentions": bool(output_attentions), "device": str(model_device(model)), } return answer, architecture_summary(config, model), explain_tokens(tokenizer, prompt), metrics gen_config = generation_config( tokenizer, max_new_tokens=max_new_tokens, do_sample=do_sample, temperature=temperature, top_p=top_p, ) with torch.no_grad(): outputs = model.generate(**inputs, generation_config=gen_config) new_token_ids = outputs[0][input_len:] answer = tokenizer.decode(new_token_ids, skip_special_tokens=True).strip() elapsed = perf_counter() - start metrics = { "lesson": lesson_name, "input_tokens": input_len, "new_tokens": int(new_token_ids.shape[-1]), "elapsed_seconds": round(elapsed, 3), "tokens_per_second": round(float(new_token_ids.shape[-1]) / elapsed, 3) if elapsed > 0 else None, "device": str(model_device(model)), } return answer, inspect_config(config, tokenizer, model), explain_tokens(tokenizer, prompt), metrics def lesson_note(lesson_name): return LESSON_NOTES.get(lesson_name, "") def lesson_code(lesson_name): file_name = LESSONS.get(lesson_name) if not file_name: return "" return read_text(TEXT_DIR / file_name, "代码文件不存在。") def lesson_file_name(lesson_name): return LESSONS.get(lesson_name, "") with gr.Blocks(title="大模型文本部署实验台") as demo: gr.Markdown( """ # 大模型文本部署实验台 从 Hugging Face `pipeline` 自动配置开始，逐步拆到 tokenizer、model、generation config、streaming、logits 和 KV cache。 默认模型：`dphn/dolphin-2.9.4-llama3.1-8b`。这是一个基于 Llama 3.1 8B 的 Dolphin ChatML 模型，适合在 ZeroGPU 上观察 8B 级别文本模型部署链路。 """ ) with gr.Tabs(): with gr.Tab("实验"): with gr.Row(): with gr.Column(scale=4): lesson = gr.Dropdown( choices=list(LESSONS.keys()), value="00 pipeline 自动配置", label="实验层级", ) model_id = gr.Textbox( value=DEFAULT_MODEL_ID, label="模型", ) system_prompt = gr.Textbox( value="你是一个讲解大模型部署的老师，回答要清晰、简洁、面试友好。", label="system prompt", lines=2, ) user_prompt = gr.Textbox( value="用三句话解释 tokenizer 在文本大模型部署中的作用。", label="user prompt", lines=4, ) with gr.Column(scale=3): max_new_tokens = gr.Slider( minimum=16, maximum=256, value=120, step=8, label="max_new_tokens", ) do_sample = gr.Checkbox(value=False, label="do_sample") temperature = gr.Slider( minimum=0.1, maximum=1.5, value=0.7, step=0.1, label="temperature", ) top_p = gr.Slider( minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="top_p", ) output_attentions = gr.Checkbox( value=False, label="output_attentions", ) run_btn = gr.Button("运行", variant="primary") with gr.Row(): lesson_summary = gr.Markdown(value=lesson_note("00 pipeline 自动配置")) runtime = gr.Textbox( value="ZeroGPU 会在点击运行后为 @spaces.GPU 函数临时分配 GPU。", label="运行环境", interactive=False, ) with gr.Row(): answer = gr.Textbox(label="模型输出", lines=12) metrics = gr.JSON(label="指标") with gr.Accordion("token、配置与架构观察", open=False): with gr.Row(): config_view = gr.Textbox(label="模型与配置", lines=10) token_view = gr.Textbox(label="token 预览", lines=10) with gr.Tab("理论"): gr.Markdown(TEXT_README) gr.Markdown(THEORY_TEXT) with gr.Tab("代码"): code_selector = gr.Dropdown( choices=list(LESSONS.keys()), value="00 pipeline 自动配置", label="代码文件", ) code_file = gr.Textbox( value=lesson_file_name("00 pipeline 自动配置"), label="文件名", interactive=False, ) code_view = gr.Code( value=lesson_code("00 pipeline 自动配置"), language="python", label="源码", ) lesson.change(lesson_note, inputs=lesson, outputs=lesson_summary) run_btn.click( run_experiment, inputs=[ lesson, model_id, system_prompt, user_prompt, max_new_tokens, do_sample, temperature, top_p, output_attentions, ], outputs=[answer, config_view, token_view, metrics], ) code_selector.change(lesson_file_name, inputs=code_selector, outputs=code_file) code_selector.change(lesson_code, inputs=code_selector, outputs=code_view) if __name__ == "__main__": demo.queue(max_size=16).launch(ssr_mode=False)