Initial commit: Neural Trace：观察本地 Qwen3.5-4B 计算过程的极简可视化工作台
This commit is contained in:
WpyQwq
@@ -0,0 +1,512 @@
|
||||
"""Lazy local model adapter for observable forward and backward traces.
|
||||
|
||||
The adapter is deliberately conservative: it summarizes tensors in the
|
||||
runtime process instead of sending activations or weights to the browser.
|
||||
Large checkpoints are loaded only after the first stream request and failures
|
||||
fall back to the mock event contract in ``main.py``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import base64
|
||||
import json
|
||||
import queue
|
||||
import re
|
||||
import threading
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
DEFAULT_ARCHITECTURE = {
|
||||
"model_type": "qwen3_5",
|
||||
"num_layers": 32,
|
||||
"pattern": "8 × (3L + 1A)",
|
||||
"hidden_size": 2560,
|
||||
"vocab_size": 248320,
|
||||
"context_length": 262144,
|
||||
"intermediate_size": 9216,
|
||||
"delta_heads_v": 32,
|
||||
"delta_heads_qk": 16,
|
||||
"attention_heads": 16,
|
||||
"attention_kv_heads": 4,
|
||||
"attention_head_dim": 256,
|
||||
"rope_dim": 64,
|
||||
"vision_encoder": True,
|
||||
"sparse_moe": True,
|
||||
"mtp": True,
|
||||
"precision": "BF16 / INT4",
|
||||
}
|
||||
|
||||
|
||||
class ModelRuntime:
|
||||
def __init__(self, model_dir: Path, model_name: str) -> None:
|
||||
self.model_dir = model_dir
|
||||
self.model_name = model_name
|
||||
self.status = "idle"
|
||||
self.error: str | None = None
|
||||
self.model: Any = None
|
||||
self.processor: Any = None
|
||||
self._torch: Any = None
|
||||
self._hooks: list[Any] = []
|
||||
self._load_lock = threading.Lock()
|
||||
self._trace_context = threading.local()
|
||||
|
||||
@property
|
||||
def ready(self) -> bool:
|
||||
return self.status == "ready" and self.model is not None
|
||||
|
||||
@property
|
||||
def checkpoint_path(self) -> Path | None:
|
||||
if not self.model_dir.exists():
|
||||
return None
|
||||
candidates = [self.model_dir]
|
||||
candidates.extend(path for path in sorted(self.model_dir.iterdir()) if path.is_dir())
|
||||
for path in candidates:
|
||||
has_config = (path / "config.json").exists()
|
||||
has_weights = any(path.glob("*.safetensors")) or any(path.glob("*.safetensors.index.json"))
|
||||
if has_config and has_weights:
|
||||
return path
|
||||
return None
|
||||
|
||||
@property
|
||||
def checkpoint_ready(self) -> bool:
|
||||
return self.checkpoint_path is not None
|
||||
|
||||
@property
|
||||
def architecture(self) -> dict[str, Any]:
|
||||
architecture = dict(DEFAULT_ARCHITECTURE)
|
||||
checkpoint_path = self.checkpoint_path
|
||||
if checkpoint_path is None:
|
||||
return architecture
|
||||
try:
|
||||
config = json.loads((checkpoint_path / "config.json").read_text(encoding="utf-8"))
|
||||
except (OSError, ValueError, UnicodeDecodeError):
|
||||
return architecture
|
||||
text_config = config.get("text_config") if isinstance(config.get("text_config"), dict) else config
|
||||
field_map = {
|
||||
"num_hidden_layers": "num_layers",
|
||||
"hidden_size": "hidden_size",
|
||||
"vocab_size": "vocab_size",
|
||||
"max_position_embeddings": "context_length",
|
||||
"intermediate_size": "intermediate_size",
|
||||
"num_attention_heads": "attention_heads",
|
||||
"num_key_value_heads": "attention_kv_heads",
|
||||
"head_dim": "attention_head_dim",
|
||||
"rope_theta_dim": "rope_dim",
|
||||
"model_type": "model_type",
|
||||
}
|
||||
for source_key, target_key in field_map.items():
|
||||
value = text_config.get(source_key)
|
||||
if isinstance(value, (int, float, str)):
|
||||
architecture[target_key] = value
|
||||
if not architecture.get("attention_head_dim") and architecture.get("attention_heads"):
|
||||
architecture["attention_head_dim"] = architecture["hidden_size"] // architecture["attention_heads"]
|
||||
return architecture
|
||||
|
||||
def ensure_loaded(self) -> bool:
|
||||
if self.ready:
|
||||
return True
|
||||
checkpoint_path = self.checkpoint_path
|
||||
if checkpoint_path is None:
|
||||
self.status = "downloading"
|
||||
return False
|
||||
|
||||
with self._load_lock:
|
||||
if self.ready:
|
||||
return True
|
||||
self.status = "loading"
|
||||
self.error = None
|
||||
try:
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoProcessor
|
||||
|
||||
self._torch = torch
|
||||
try:
|
||||
self.processor = AutoProcessor.from_pretrained(checkpoint_path, local_files_only=True, trust_remote_code=True)
|
||||
except Exception:
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
self.processor = AutoTokenizer.from_pretrained(checkpoint_path, local_files_only=True, trust_remote_code=True)
|
||||
|
||||
model_class = AutoModelForCausalLM
|
||||
image_text_model = getattr(__import__("transformers", fromlist=["AutoModelForImageTextToText"]), "AutoModelForImageTextToText", None)
|
||||
if image_text_model is not None:
|
||||
model_class = image_text_model
|
||||
|
||||
load_kwargs = {
|
||||
"local_files_only": True,
|
||||
"trust_remote_code": True,
|
||||
"torch_dtype": "auto",
|
||||
}
|
||||
if importlib.util.find_spec("accelerate") is not None:
|
||||
load_kwargs.update(device_map="auto", low_cpu_mem_usage=True)
|
||||
|
||||
try:
|
||||
self.model = model_class.from_pretrained(
|
||||
checkpoint_path,
|
||||
**load_kwargs,
|
||||
)
|
||||
except Exception:
|
||||
self.model = AutoModelForCausalLM.from_pretrained(
|
||||
checkpoint_path,
|
||||
**load_kwargs,
|
||||
)
|
||||
|
||||
self.model.eval()
|
||||
self._install_hooks()
|
||||
self.status = "ready"
|
||||
return True
|
||||
except Exception as exc: # pragma: no cover - depends on local GPU/runtime
|
||||
self.status = "error"
|
||||
self.error = f"{type(exc).__name__}: {exc}"
|
||||
self.model = None
|
||||
return False
|
||||
|
||||
def start_trace(self, prompt: str, mode: str, image_data: str | None = None) -> queue.Queue[dict[str, Any]] | None:
|
||||
if not self.ensure_loaded():
|
||||
return None
|
||||
events: queue.Queue[dict[str, Any]] = queue.Queue()
|
||||
worker = threading.Thread(target=self._trace_worker, args=(prompt, mode, image_data, events), daemon=True)
|
||||
worker.start()
|
||||
return events
|
||||
|
||||
def _install_hooks(self) -> None:
|
||||
if self._hooks or self.model is None:
|
||||
return
|
||||
|
||||
seen_layers: set[int] = set()
|
||||
for name, module in self.model.named_modules():
|
||||
match = re.search(r"(?:^|\.)(?:language_model\.)?(?:layers|(?:model|transformer|decoder)\.(?:layers|h))\.(\d+)$", name)
|
||||
if match is not None:
|
||||
layer_index = int(match.group(1))
|
||||
if layer_index not in seen_layers:
|
||||
seen_layers.add(layer_index)
|
||||
self._hooks.append(module.register_forward_hook(self._make_forward_hook(layer_index)))
|
||||
self._hooks.append(module.register_full_backward_hook(self._make_backward_hook(layer_index)))
|
||||
elif re.search(r"(?:^|\.)\b(?:visual|vision_model|vision_tower)\b$", name):
|
||||
self._hooks.append(module.register_forward_hook(self._make_vision_hook()))
|
||||
elif re.search(r"(?:^|\.)\b(?:router|routing|gate|experts|expert_gate|mtp)\b$", name, re.IGNORECASE):
|
||||
kind = "router_weights" if re.search(r"router|routing|gate", name, re.IGNORECASE) else "expert_mixture"
|
||||
if re.search(r"mtp", name, re.IGNORECASE):
|
||||
kind = "mtp_logits"
|
||||
self._hooks.append(module.register_forward_hook(self._make_aux_hook(name, kind)))
|
||||
elif re.search(r"(?:^|\.)lm_head$", name):
|
||||
self._hooks.append(module.register_forward_hook(self._make_logits_hook()))
|
||||
|
||||
def _make_forward_hook(self, layer_index: int):
|
||||
def forward_hook(_module: Any, _inputs: Any, output: Any) -> None:
|
||||
context = getattr(self._trace_context, "queue", None)
|
||||
if context is None:
|
||||
return
|
||||
tensor = self._first_tensor(output)
|
||||
summary = self._summarize(tensor)
|
||||
full_attention = (layer_index + 1) % 4 == 0
|
||||
mode = getattr(self._trace_context, "mode", "forward")
|
||||
step = getattr(self._trace_context, "step", 0)
|
||||
self._trace_context.step = step + 1
|
||||
if mode == "attention":
|
||||
kind = "attention_output" if full_attention else "delta_context"
|
||||
elif mode == "state":
|
||||
kind = "delta_state"
|
||||
else:
|
||||
kind = "full_attention" if full_attention else "delta_state"
|
||||
context.put({
|
||||
"step": min(step, 511),
|
||||
"layer": f"L{layer_index + 1:02d}",
|
||||
"kind": kind,
|
||||
"mode": mode,
|
||||
"value": summary["rms"],
|
||||
"mean": summary["mean"],
|
||||
"shape": summary["shape"],
|
||||
"source": "forward_hook",
|
||||
})
|
||||
|
||||
return forward_hook
|
||||
|
||||
def _make_vision_hook(self):
|
||||
def vision_hook(_module: Any, _inputs: Any, output: Any) -> None:
|
||||
context = getattr(self._trace_context, "queue", None)
|
||||
if context is None:
|
||||
return
|
||||
tensor = self._first_tensor(output)
|
||||
summary = self._summarize(tensor)
|
||||
step = getattr(self._trace_context, "step", 0)
|
||||
self._trace_context.step = step + 1
|
||||
context.put({
|
||||
"step": min(step, 511),
|
||||
"kind": "vision_encoder",
|
||||
"mode": getattr(self._trace_context, "mode", "forward"),
|
||||
"value": summary["rms"],
|
||||
"mean": summary["mean"],
|
||||
"shape": summary["shape"],
|
||||
"source": "vision_hook",
|
||||
})
|
||||
|
||||
return vision_hook
|
||||
|
||||
def _make_aux_hook(self, module_name: str, kind: str):
|
||||
def aux_hook(_module: Any, _inputs: Any, output: Any) -> None:
|
||||
context = getattr(self._trace_context, "queue", None)
|
||||
if context is None:
|
||||
return
|
||||
tensor = self._first_tensor(output)
|
||||
summary = self._summarize(tensor)
|
||||
layer_match = re.search(r"layers\.(\d+)", module_name)
|
||||
layer = f"L{int(layer_match.group(1)) + 1:02d}" if layer_match else "AUX"
|
||||
step = getattr(self._trace_context, "step", 0)
|
||||
self._trace_context.step = step + 1
|
||||
context.put({
|
||||
"step": min(step, 511),
|
||||
"layer": layer,
|
||||
"kind": kind,
|
||||
"mode": getattr(self._trace_context, "mode", "forward"),
|
||||
"value": summary["rms"],
|
||||
"mean": summary["mean"],
|
||||
"shape": summary["shape"],
|
||||
"module": module_name,
|
||||
"source": "aux_hook",
|
||||
})
|
||||
|
||||
return aux_hook
|
||||
|
||||
def _make_backward_hook(self, layer_index: int):
|
||||
def backward_hook(_module: Any, _grad_input: Any, grad_output: Any) -> None:
|
||||
context = getattr(self._trace_context, "queue", None)
|
||||
if context is None or getattr(self._trace_context, "mode", "forward") != "backward":
|
||||
return
|
||||
tensor = self._first_tensor(grad_output)
|
||||
summary = self._summarize(tensor)
|
||||
step = getattr(self._trace_context, "step", 0)
|
||||
self._trace_context.step = step + 1
|
||||
context.put({
|
||||
"step": min(step, 511),
|
||||
"layer": f"L{layer_index + 1:02d}",
|
||||
"kind": "gradient",
|
||||
"value": summary["rms"],
|
||||
"mean": summary["mean"],
|
||||
"shape": summary["shape"],
|
||||
"source": "backward_hook",
|
||||
})
|
||||
|
||||
return backward_hook
|
||||
|
||||
def _make_logits_hook(self):
|
||||
def logits_hook(_module: Any, _inputs: Any, output: Any) -> None:
|
||||
context = getattr(self._trace_context, "queue", None)
|
||||
if context is None:
|
||||
return
|
||||
tensor = self._first_tensor(output)
|
||||
if tensor is not None and tensor.ndim >= 2:
|
||||
tensor = tensor[:, -1, :]
|
||||
summary = self._summarize(tensor)
|
||||
step = getattr(self._trace_context, "step", 0)
|
||||
self._trace_context.step = step + 1
|
||||
context.put({
|
||||
"step": min(step, 511),
|
||||
"layer": "OUT",
|
||||
"kind": "logits",
|
||||
"value": summary["rms"],
|
||||
"mean": summary["mean"],
|
||||
"shape": summary["shape"],
|
||||
"source": "logits_hook",
|
||||
})
|
||||
|
||||
return logits_hook
|
||||
|
||||
def _trace_worker(self, prompt: str, mode: str, image_data: str | None, events: queue.Queue[dict[str, Any]]) -> None:
|
||||
self._trace_context.queue = events
|
||||
self._trace_context.mode = mode
|
||||
self._trace_context.step = 0
|
||||
try:
|
||||
inputs = self._encode(prompt, image_data)
|
||||
if image_data and "pixel_values" in inputs:
|
||||
summary = self._summarize(inputs["pixel_values"])
|
||||
events.put({
|
||||
"step": 0,
|
||||
"kind": "vision_input",
|
||||
"mode": mode,
|
||||
"value": summary["rms"],
|
||||
"mean": summary["mean"],
|
||||
"shape": summary["shape"],
|
||||
"source": "processor",
|
||||
})
|
||||
if mode == "backward":
|
||||
self._run_backward(inputs, events)
|
||||
elif mode == "attention":
|
||||
self._run_attention(inputs, events)
|
||||
else:
|
||||
self._run_generation(inputs, events)
|
||||
except Exception as exc: # pragma: no cover - depends on checkpoint/runtime
|
||||
events.put({"kind": "runtime_error", "error": f"{type(exc).__name__}: {exc}", "source": "model_runtime"})
|
||||
finally:
|
||||
events.put({"kind": "done", "source": "model_runtime"})
|
||||
self._trace_context.queue = None
|
||||
|
||||
def _run_attention(self, inputs: dict[str, Any], events: queue.Queue[dict[str, Any]]) -> None:
|
||||
torch = self._torch
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
outputs = self.model(
|
||||
**inputs,
|
||||
use_cache=False,
|
||||
output_attentions=True,
|
||||
return_dict=True,
|
||||
)
|
||||
attentions = getattr(outputs, "attentions", None)
|
||||
emitted = False
|
||||
if attentions:
|
||||
for layer_index, attention in enumerate(attentions):
|
||||
if attention is None or getattr(attention, "ndim", 0) < 4:
|
||||
continue
|
||||
size = min(7, int(attention.shape[-1]))
|
||||
matrix = attention[0, :, -size:, -size:].float().mean(dim=0)
|
||||
values = matrix.detach().cpu().reshape(-1).tolist()
|
||||
summary = self._summarize(matrix)
|
||||
step = getattr(self._trace_context, "step", 0)
|
||||
self._trace_context.step = step + 1
|
||||
events.put({
|
||||
"step": min(step, 511),
|
||||
"layer": f"L{layer_index + 1:02d}",
|
||||
"kind": "attention_weight",
|
||||
"mode": "attention",
|
||||
"value": summary["rms"],
|
||||
"mean": summary["mean"],
|
||||
"shape": list(attention.shape),
|
||||
"attention": [round(float(value), 6) for value in values],
|
||||
"attention_shape": [size, size],
|
||||
"source": "attention_output",
|
||||
})
|
||||
emitted = True
|
||||
if not emitted:
|
||||
events.put({
|
||||
"step": min(getattr(self._trace_context, "step", 0), 511),
|
||||
"kind": "attention_unavailable",
|
||||
"mode": "attention",
|
||||
"source": "model_runtime",
|
||||
})
|
||||
except Exception:
|
||||
events.put({
|
||||
"step": min(getattr(self._trace_context, "step", 0), 511),
|
||||
"kind": "attention_unavailable",
|
||||
"mode": "attention",
|
||||
"source": "model_runtime",
|
||||
})
|
||||
self._run_generation(inputs, events)
|
||||
|
||||
def _encode(self, prompt: str, image_data: str | None = None) -> dict[str, Any]:
|
||||
processor = self.processor
|
||||
rendered_prompt = prompt
|
||||
image = None
|
||||
if image_data:
|
||||
try:
|
||||
from PIL import Image
|
||||
|
||||
encoded_image = image_data.split(",", 1)[-1]
|
||||
image = Image.open(BytesIO(base64.b64decode(encoded_image))).convert("RGB")
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"invalid image input: {type(exc).__name__}") from exc
|
||||
apply_chat_template = getattr(processor, "apply_chat_template", None)
|
||||
if callable(apply_chat_template):
|
||||
try:
|
||||
content = [{"type": "image", "image": image}, {"type": "text", "text": prompt}] if image is not None else prompt
|
||||
rendered_prompt = apply_chat_template(
|
||||
[{"role": "user", "content": content}],
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
)
|
||||
except Exception:
|
||||
rendered_prompt = prompt
|
||||
processor_kwargs: dict[str, Any] = {"text": rendered_prompt, "return_tensors": "pt", "truncation": True, "max_length": 256}
|
||||
if image is not None:
|
||||
processor_kwargs["images"] = image
|
||||
try:
|
||||
encoded = processor(**processor_kwargs)
|
||||
except Exception:
|
||||
processor_kwargs.pop("truncation", None)
|
||||
processor_kwargs.pop("max_length", None)
|
||||
encoded = processor(**processor_kwargs)
|
||||
device = self._input_device()
|
||||
return {key: value.to(device) if hasattr(value, "to") else value for key, value in encoded.items()}
|
||||
|
||||
def _run_generation(self, inputs: dict[str, Any], events: queue.Queue[dict[str, Any]]) -> None:
|
||||
torch = self._torch
|
||||
with torch.inference_mode():
|
||||
generated = self.model.generate(**inputs, max_new_tokens=32, do_sample=False, use_cache=True)
|
||||
if hasattr(generated, "sequences"):
|
||||
generated = generated.sequences
|
||||
input_length = 0
|
||||
if "input_ids" in inputs and getattr(inputs["input_ids"], "ndim", 0) >= 2:
|
||||
input_length = int(inputs["input_ids"].shape[-1])
|
||||
if getattr(generated, "ndim", 0) >= 2 and input_length:
|
||||
generated = generated[:, input_length:]
|
||||
try:
|
||||
text = self.processor.batch_decode(generated, skip_special_tokens=True)[0]
|
||||
except Exception:
|
||||
tokenizer = getattr(self.processor, "tokenizer", None)
|
||||
try:
|
||||
text = tokenizer.batch_decode(generated, skip_special_tokens=True)[0] if tokenizer is not None else ""
|
||||
except Exception:
|
||||
text = ""
|
||||
events.put({
|
||||
"step": min(getattr(self._trace_context, "step", 0), 511),
|
||||
"kind": "generation",
|
||||
"token": text[-1:] if text else "",
|
||||
"text": text,
|
||||
"value": 1.0,
|
||||
"source": "model_runtime",
|
||||
})
|
||||
|
||||
def _run_backward(self, inputs: dict[str, Any], events: queue.Queue[dict[str, Any]]) -> None:
|
||||
torch = self._torch
|
||||
if "input_ids" not in inputs:
|
||||
raise RuntimeError("processor did not return input_ids; backward trace requires text tokens")
|
||||
self.model.zero_grad(set_to_none=True)
|
||||
labels = inputs["input_ids"].clone()
|
||||
with torch.enable_grad():
|
||||
outputs = self.model(**inputs, labels=labels, use_cache=False)
|
||||
loss = outputs.loss
|
||||
loss.backward()
|
||||
events.put({
|
||||
"step": min(getattr(self._trace_context, "step", 0), 511),
|
||||
"kind": "loss",
|
||||
"loss": float(loss.detach().float().item()),
|
||||
"value": float(loss.detach().float().item()),
|
||||
"source": "autograd",
|
||||
})
|
||||
|
||||
def _input_device(self):
|
||||
try:
|
||||
return self.model.device
|
||||
except Exception:
|
||||
for parameter in self.model.parameters():
|
||||
if parameter.device.type != "meta":
|
||||
return parameter.device
|
||||
return self._torch.device("cpu")
|
||||
|
||||
def _first_tensor(self, value: Any):
|
||||
if self._torch is not None and self._torch.is_tensor(value):
|
||||
return value
|
||||
if isinstance(value, (tuple, list)):
|
||||
for item in value:
|
||||
tensor = self._first_tensor(item)
|
||||
if tensor is not None:
|
||||
return tensor
|
||||
if isinstance(value, dict):
|
||||
for item in value.values():
|
||||
tensor = self._first_tensor(item)
|
||||
if tensor is not None:
|
||||
return tensor
|
||||
return None
|
||||
|
||||
def _summarize(self, tensor: Any) -> dict[str, Any]:
|
||||
if tensor is None:
|
||||
return {"mean": 0.0, "rms": 0.0, "shape": []}
|
||||
torch = self._torch
|
||||
sampled = tensor.detach().float().reshape(-1)
|
||||
if sampled.numel() > 4096:
|
||||
stride = max(1, sampled.numel() // 4096)
|
||||
sampled = sampled[::stride]
|
||||
mean = float(sampled.mean().item())
|
||||
rms = float(torch.sqrt(torch.mean(sampled * sampled)).item())
|
||||
return {"mean": round(mean, 6), "rms": round(rms, 6), "shape": list(tensor.shape)}
|
||||
Reference in New Issue
Block a user