github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge 88031f1ac1 into f659b55761

pull/881/merge
LiberServer 2026-09-29 17:05:14 +08:00 committed by GitHub
parent f659b55761 88031f1ac1
commit 04203cc1ba
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 808 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
.github/workflows/build-vllm-native-ghcr.yml vendored Normal file
View File

@ -0,0 +1,52 @@
name: Build and publish MiniMind vLLM image
on:
workflow_dispatch:
push:
branches:
- master
- codex/ghcr-clean
paths:
- Dockerfile.vllm-native
- vllm_native/**
- .github/workflows/build-vllm-native-ghcr.yml
permissions:
contents: read
packages: write
env:
REGISTRY: ghcr.io
IMAGE_NAME: liberserver/minimind
jobs:
build-and-push:
runs-on: ubuntu-latest
steps:
- name: Check out repository
uses: actions/checkout@v6
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3
- name: Log in to GitHub Container Registry
uses: docker/login-action@v3
with:
registry: ${{ env.REGISTRY }}
username: ${{ github.actor }}
password: ${{ secrets.GITHUB_TOKEN }}
- name: Build and push image
uses: docker/build-push-action@v6
with:
context: .
file: ./Dockerfile.vllm-native
platforms: linux/amd64
push: true
tags: |
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:vllm-native-0.17.0
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:sha-${{ github.sha }}
labels: |
org.opencontainers.image.source=https://github.com/LiberServer/minimind
provenance: false
sbom: false

34
.gitignore vendored
View File

@ -4,3 +4,37 @@ __pycache__/
out
website/
docs-minimind/
# Model weights and exported checkpoints (keep large binaries out of Git)
*.bin
*.safetensors
*.pt
*.pth
*.ckpt
*.gguf
*.ggml
*.onnx
*.onnx_data
*.tflite
*.engine
*.plan
*.h5
*.hdf5
*.npy
*.npz
*.parquet
*.arrow
# Training datasets and local experiment output
dataset/*.jsonl
dataset/**/*.jsonl
wandb/
swanlog/
# Downloaded archives and model bundles
*.tar
*.tar.gz
*.tgz
*.zip
*.7z
*.rar

10
Dockerfile.vllm-native Normal file
View File

@ -0,0 +1,10 @@
FROM vllm/vllm-openai:v0.17.0
COPY vllm_native /opt/minimind-vllm
RUN uv pip install --system --no-deps /opt/minimind-vllm
ENV VLLM_PLUGINS=minimind_vllm
EXPOSE 8000
ENTRYPOINT ["vllm", "serve"]
CMD ["/models/minimind", "--host", "0.0.0.0", "--port", "8000", "--served-model-name", "minimind", "--trust-remote-code", "--chat-template", "/models/minimind/chat_template.jinja", "--max-model-len", "4096", "--gpu-memory-utilization", "0.5", "--mamba-cache-mode", "align", "--enforce-eager"]

36
vllm_native/README.md Normal file
View File

@ -0,0 +1,36 @@
# MiniMind native vLLM adapter
This out-of-tree plugin targets the MiniMind hybrid checkpoint format and
vLLM 0.17.0. It keeps the checkpoint and the repository's existing model
implementation unchanged.
## Start
From the repository root in PowerShell:
```powershell
.\run_vllm_native.ps1
```
The default checkpoint directory is
`D:\Tools\code\TinyLLM\out\grpo_768_hybrid_hf`. Override it when needed:
```powershell
.\run_vllm_native.ps1 -ModelPath 'D:\path\to\your\converted\checkpoint'
```
The script builds `minimind:vllm-native-0.17.0`, mounts the checkpoint read-only,
and starts the OpenAI-compatible API on `http://127.0.0.1:8000/v1`. If local
Python is on PATH, it also starts a Chinese web chat at
`http://127.0.0.1:8080`. The API and web chat bind to loopback only.
## Stop
Stop the model container with:
```powershell
docker stop minimind-vllm-native
```
The launch script prints the web chat helper's process ID; stop that specific
process with `Stop-Process -Id <PID>` when finished.

117
vllm_native/chat.html Normal file
View File

@ -0,0 +1,117 @@
<!doctype html>
<html lang="zh-CN">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>MiniMind 聊天</title>
<style>
:root { color-scheme: light; font-family: "Segoe UI", "Microsoft YaHei", sans-serif; }
* { box-sizing: border-box; }
body { margin: 0; background: #f4f6fa; color: #182230; }
main { width: min(900px, 100%); height: 100dvh; margin: auto; display: flex; flex-direction: column; padding: 20px; }
header { display: flex; align-items: center; justify-content: space-between; gap: 12px; padding: 4px 2px 16px; }
h1 { margin: 0; font-size: 20px; font-weight: 650; }
#status { color: #667085; font-size: 13px; }
#status.ready { color: #16803c; }
#messages { flex: 1; overflow: auto; display: flex; flex-direction: column; gap: 14px; padding: 18px 4px; }
.message { max-width: min(78%, 720px); padding: 12px 15px; border-radius: 16px; line-height: 1.65; white-space: pre-wrap; overflow-wrap: anywhere; }
.user { align-self: flex-end; color: #fff; background: #3159c9; border-bottom-right-radius: 5px; }
.assistant { align-self: flex-start; background: #fff; border: 1px solid #e4e7ec; border-bottom-left-radius: 5px; }
.system { align-self: center; color: #667085; font-size: 13px; }
form { display: flex; align-items: flex-end; gap: 10px; padding: 12px; border: 1px solid #e4e7ec; border-radius: 18px; background: #fff; box-shadow: 0 5px 20px #1822300d; }
textarea { flex: 1; resize: none; max-height: 180px; min-height: 48px; border: 0; outline: 0; padding: 12px 4px; font: inherit; }
button { border: 0; border-radius: 12px; padding: 12px 18px; background: #3159c9; color: #fff; font: inherit; cursor: pointer; }
button:disabled { opacity: .55; cursor: wait; }
.toolbar { display: flex; align-items: center; gap: 10px; }
.clear { padding: 8px 12px; color: #475467; background: #eaecf0; font-size: 13px; }
@media (max-width: 600px) { main { padding: 12px; } .message { max-width: 90%; } }
</style>
</head>
<body>
<main>
<header>
<h1>MiniMind 聊天</h1>
<div class="toolbar"><span id="status">正在连接模型…</span><button class="clear" id="clear" type="button">新对话</button></div>
</header>
<section id="messages" aria-live="polite">
<div class="message system">你好！发送消息开始与本地 MiniMind 对话。</div>
</section>
<form id="form">
<textarea id="prompt" placeholder="输入消息…（Enter 发送，Shift+Enter 换行）" rows="1" required></textarea>
<button id="send" type="submit">发送</button>
</form>
</main>
<script>
const form = document.querySelector('#form');
const promptBox = document.querySelector('#prompt');
const sendButton = document.querySelector('#send');
const messagesBox = document.querySelector('#messages');
const statusLabel = document.querySelector('#status');
const history = [];
function addMessage(role, content) {
const item = document.createElement('div');
item.className = `message ${role}`;
item.textContent = content;
messagesBox.append(item);
messagesBox.scrollTop = messagesBox.scrollHeight;
return item;
}
async function checkService() {
try {
const response = await fetch('/health');
statusLabel.textContent = response.ok ? '模型服务已连接' : '模型服务启动中';
statusLabel.classList.toggle('ready', response.ok);
} catch {
statusLabel.textContent = '暂时无法连接模型';
}
}
form.addEventListener('submit', async (event) => {
event.preventDefault();
const content = promptBox.value.trim();
if (!content) return;
history.push({ role: 'user', content });
addMessage('user', content);
promptBox.value = '';
sendButton.disabled = true;
const waiting = addMessage('assistant', '正在思考…');
try {
const response = await fetch('/api/chat', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ model: 'minimind', messages: history, max_tokens: 512, temperature: 0.7 })
});
const data = await response.json();
if (!response.ok) throw new Error(data.error?.message || data.error || `HTTP ${response.status}`);
const answer = data.choices?.[0]?.message?.content ?? '(模型没有返回文本)';
waiting.textContent = answer;
history.push({ role: 'assistant', content: answer });
} catch (error) {
waiting.textContent = `请求失败：${error.message}`;
history.pop();
} finally {
sendButton.disabled = false;
promptBox.focus();
messagesBox.scrollTop = messagesBox.scrollHeight;
}
});
promptBox.addEventListener('keydown', (event) => {
if (event.key === 'Enter' && !event.shiftKey) {
event.preventDefault();
form.requestSubmit();
}
});
document.querySelector('#clear').addEventListener('click', () => {
history.length = 0;
messagesBox.replaceChildren();
addMessage('system', '新对话已开始。');
promptBox.focus();
});
checkService();
setInterval(checkService, 10000);
</script>
</body>
</html>

93
vllm_native/chat_server.py Normal file
View File

@ -0,0 +1,93 @@
"""Loopback-only web chat and same-origin proxy for the local vLLM API."""
from __future__ import annotations
import argparse
import json
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen
CHAT_PAGE = Path(__file__).with_name("chat.html")
class ChatHandler(BaseHTTPRequestHandler):
backend = "http://127.0.0.1:8000"
def log_message(self, fmt: str, *args: object) -> None:
print(f"[{self.log_date_time_string()}] {fmt % args}")
def _respond(self, status: int, body: bytes, content_type: str) -> None:
self.send_response(status)
self.send_header("Content-Type", content_type)
self.send_header("Content-Length", str(len(body)))
self.send_header("Cache-Control", "no-store")
self.end_headers()
self.wfile.write(body)
def _proxy(self, path: str, body: bytes | None = None) -> None:
headers = {"Accept": "application/json"}
if body is not None:
headers["Content-Type"] = "application/json"
request = Request(
f"{self.backend}{path}",
data=body,
headers=headers,
method="POST" if body is not None else "GET",
)
try:
with urlopen(request, timeout=600) as response:
self._respond(
response.status,
response.read(),
response.headers.get("Content-Type", "application/json"),
)
except HTTPError as error:
self._respond(error.code, error.read(), "application/json; charset=utf-8")
except (TimeoutError, URLError, OSError) as error:
message = json.dumps(
{"error": f"无法连接到 vLLM 服务：{error}"}, ensure_ascii=False
).encode("utf-8")
self._respond(502, message, "application/json; charset=utf-8")
def do_GET(self) -> None:
if self.path == "/":
try:
self._respond(200, CHAT_PAGE.read_bytes(), "text/html; charset=utf-8")
except OSError as error:
self._respond(500, str(error).encode(), "text/plain; charset=utf-8")
elif self.path == "/health":
self._proxy("/health")
else:
self._respond(404, b"Not found", "text/plain; charset=utf-8")
def do_POST(self) -> None:
if self.path != "/api/chat":
self._respond(404, b"Not found", "text/plain; charset=utf-8")
return
try:
length = int(self.headers.get("Content-Length", "0"))
except ValueError:
length = 0
if not 0 < length <= 1_048_576:
self._respond(400, b'{"error":"Invalid request size"}', "application/json")
return
self._proxy("/v1/chat/completions", self.rfile.read(length))
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=8080)
parser.add_argument("--backend", default="http://127.0.0.1:8000")
args = parser.parse_args()
ChatHandler.backend = args.backend.rstrip("/")
server = ThreadingHTTPServer((args.host, args.port), ChatHandler)
print(f"MiniMind web chat: http://{args.host}:{args.port}", flush=True)
server.serve_forever()
if __name__ == "__main__":
main()

12
vllm_native/minimind_vllm/__init__.py Normal file
View File

@ -0,0 +1,12 @@
"""Out-of-tree vLLM model registration for MiniMind."""
def register() -> None:
from vllm import ModelRegistry
architecture = "MiniMindForCausalLM"
if architecture not in ModelRegistry.get_supported_archs():
ModelRegistry.register_model(
architecture,
"minimind_vllm.model:MiniMindForCausalLM",
)

437
vllm_native/minimind_vllm/model.py Normal file
View File

@ -0,0 +1,437 @@
"""Inference-only MiniMind implementation for vLLM 0.17.x.
Parameter names match the exported Transformers checkpoint. Gated Delta Net
recurrent state is allocated and scheduled by vLLM's hybrid Mamba cache.
"""
from __future__ import annotations
import math
from itertools import islice
from typing import Iterable
import torch
import torch.nn.functional as F
from torch import nn
from transformers.activations import ACT2FN
from vllm.config import CacheConfig, ModelConfig, VllmConfig
from vllm.distributed import get_pp_group
from vllm.forward_context import get_forward_context
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.fla.ops import (
chunk_gated_delta_rule,
fused_recurrent_gated_delta_rule,
)
from vllm.model_executor.layers.mamba.abstract import MambaBase
from vllm.model_executor.layers.mamba.mamba_utils import (
MambaStateCopyFunc,
MambaStateCopyFuncCalculator,
MambaStateDtypeCalculator,
MambaStateShapeCalculator,
)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.interfaces import HasInnerState, IsHybrid, SupportsPP
from vllm.model_executor.models.utils import (
PPMissingLayer,
extract_layer_index,
make_empty_intermediate_tensors_factory,
make_layers,
maybe_prefix,
)
from vllm.sequence import IntermediateTensors
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
class MiniMindRMSNorm(nn.Module):
def __init__(self, dim: int, eps: float):
super().__init__()
self.weight = nn.Parameter(torch.ones(dim))
self.eps = eps
def forward(self, x: torch.Tensor) -> torch.Tensor:
x_fp32 = x.float()
normalized = x_fp32 * torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + self.eps)
return (self.weight.float() * normalized).to(x.dtype)
def _precompute_rope(config) -> tuple[torch.Tensor, torch.Tensor]:
dim = config.head_dim
base = config.rope_theta
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
scaling = getattr(config, "rope_scaling", None)
attention_factor = 1.0
if scaling is not None:
original = scaling.get("original_max_position_embeddings", 2048)
factor = scaling.get("factor", 16)
beta_fast = scaling.get("beta_fast", 32.0)
beta_slow = scaling.get("beta_slow", 1.0)
def correction_dim(beta: float) -> float:
return (dim * math.log(original / (beta * 2 * math.pi))) / (2 * math.log(base))
low = max(math.floor(correction_dim(beta_fast)), 0)
high = min(math.ceil(correction_dim(beta_slow)), dim // 2 - 1)
ramp = torch.clamp(
(torch.arange(dim // 2, dtype=torch.float32) - low)
/ max(high - low, 0.001),
0,
1,
)
inv_freq = inv_freq * (1 - ramp + ramp / factor)
attention_factor = scaling.get("attention_factor", 1.0)
positions = torch.arange(config.max_position_embeddings, dtype=torch.float32)
phase = torch.outer(positions, inv_freq)
phase = torch.cat((phase, phase), dim=-1)
return phase.cos() * attention_factor, phase.sin() * attention_factor
def _apply_rope(x: torch.Tensor, positions: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
pos = positions.reshape(-1).long()
cos_pos = cos.index_select(0, pos).unsqueeze(1)
sin_pos = sin.index_select(0, pos).unsqueeze(1)
half = x.shape[-1] // 2
rotated = torch.cat((-x[..., half:], x[..., :half]), dim=-1)
return (x * cos_pos + rotated * sin_pos).to(x.dtype)
class MiniMindFeedForward(nn.Module):
def __init__(self, config):
super().__init__()
self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
self.act_fn = ACT2FN[config.hidden_act]
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
class MiniMindFullAttention(nn.Module):
def __init__(self, config, cache_config: CacheConfig, prefix: str):
super().__init__()
self.num_heads = config.num_attention_heads
self.num_kv_heads = config.num_key_value_heads
self.head_dim = config.head_dim
self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
self.q_norm = MiniMindRMSNorm(self.head_dim, config.rms_norm_eps)
self.k_norm = MiniMindRMSNorm(self.head_dim, config.rms_norm_eps)
self.attn = Attention(
self.num_heads,
self.head_dim,
self.head_dim**-0.5,
num_kv_heads=self.num_kv_heads,
cache_config=cache_config,
prefix=f"{prefix}.attn",
)
def forward(self, x, positions, cos, sin):
n = x.shape[0]
q = self.q_proj(x).view(n, self.num_heads, self.head_dim)
k = self.k_proj(x).view(n, self.num_kv_heads, self.head_dim)
v = self.v_proj(x).view(n, self.num_kv_heads * self.head_dim)
q = _apply_rope(self.q_norm(q), positions, cos, sin)
k = _apply_rope(self.k_norm(k), positions, cos, sin)
y = self.attn(q.flatten(1), k.flatten(1), v)
return self.o_proj(y)
class MiniMindGatedDeltaNet(MambaBase):
"""MiniMind GDN layer using vLLM-managed per-sequence recurrent state."""
def __init__(self, config, vllm_config: VllmConfig, prefix: str):
super().__init__()
self.prefix = prefix
self.model_config: ModelConfig = vllm_config.model_config
self.cache_config: CacheConfig = vllm_config.cache_config
self.num_heads = config.num_attention_heads
self.num_kv_heads = config.num_key_value_heads
self.head_dim = config.head_dim
self.n_rep = self.num_heads // self.num_kv_heads
self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
self.beta_proj = nn.Linear(config.hidden_size, self.num_heads, bias=False)
self.decay_proj = nn.Linear(config.hidden_size, self.num_heads, bias=False)
self.decay_bias = nn.Parameter(torch.full((self.num_heads,), float(config.linear_decay_bias)))
self.q_norm = MiniMindRMSNorm(self.head_dim, config.rms_norm_eps)
self.k_norm = MiniMindRMSNorm(self.head_dim, config.rms_norm_eps)
self.out_norm = MiniMindRMSNorm(self.head_dim, config.rms_norm_eps)
self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
self.kv_cache: tuple[torch.Tensor, ...] = ()
@property
def mamba_type(self):
return MambaAttentionBackendEnum.GDN_ATTN
def get_state_shape(self) -> tuple[tuple[int, ...], ...]:
# MiniMind has no short convolution state. A zero-sized conv entry keeps
# the GDN backend's standard two-state layout; the second is [H, D, D].
return MambaStateShapeCalculator.gated_delta_net_state_shape(
1, self.num_heads, self.num_heads, self.head_dim, self.head_dim, 1, 0
)
def get_state_dtype(self) -> tuple[torch.dtype, ...]:
return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
self.model_config.dtype,
self.cache_config.mamba_cache_dtype,
self.cache_config.mamba_ssm_cache_dtype,
)
def forward(self, x, positions, cos, sin):
n = x.shape[0]
q = self.q_proj(x).view(n, self.num_heads, self.head_dim)
k = self.k_proj(x).view(n, self.num_kv_heads, self.head_dim)
v = self.v_proj(x).view(n, self.num_kv_heads, self.head_dim)
q = _apply_rope(self.q_norm(q), positions, cos, sin)
k = _apply_rope(self.k_norm(k), positions, cos, sin)
k = k.repeat_interleave(self.n_rep, dim=1)
v = v.repeat_interleave(self.n_rep, dim=1)
q = F.normalize(q.float(), p=2, dim=-1).to(x.dtype)
k = F.normalize(k.float(), p=2, dim=-1).to(x.dtype)
beta = torch.sigmoid(self.beta_proj(x))
log_decay = -F.softplus(
self.decay_proj(x).float() + self.decay_bias.float()
).to(x.dtype)
forward_context = get_forward_context()
raw_metadata = forward_context.attn_metadata
if raw_metadata is None:
# The engine runs a metadata-free profile pass before cache binding.
return x.new_zeros((n, self.num_heads * self.head_dim))
metadata = raw_metadata[self.prefix] if isinstance(raw_metadata, dict) else raw_metadata
if not isinstance(metadata, GDNAttentionMetadata):
raise TypeError(f"Expected GDN metadata for {self.prefix}, got {type(metadata)!r}")
if metadata.spec_sequence_masks is not None:
raise RuntimeError("Disable speculative decoding for the MiniMind adapter")
if not self.kv_cache:
raise RuntimeError(f"The vLLM state cache was not bound for {self.prefix}")
# vLLM may expose state tensors directly or grouped by virtual engine.
cache = self.kv_cache
if not isinstance(cache[0], torch.Tensor):
cache = cache[forward_context.virtual_engine]
actual = metadata.num_actual_tokens
q, k, v = (t[:actual].unsqueeze(0) for t in (q, k, v))
beta = beta[:actual].unsqueeze(0)
log_decay = log_decay[:actual].unsqueeze(0)
ssm_state = cache[1]
if metadata.num_prefills > 0:
state_indices = metadata.non_spec_state_indices_tensor
initial_state = ssm_state[state_indices].contiguous()
initial_state[~metadata.has_initial_state, ...] = 0
recurrent, final_state = chunk_gated_delta_rule(
q=q,
k=k,
v=v,
g=log_decay,
beta=beta,
initial_state=initial_state,
output_final_state=True,
cu_seqlens=metadata.non_spec_query_start_loc,
use_qk_l2norm_in_kernel=False,
)
ssm_state[state_indices] = final_state.to(ssm_state.dtype)
elif metadata.num_decodes > 0:
state_indices = metadata.non_spec_state_indices_tensor
recurrent, _ = fused_recurrent_gated_delta_rule(
q=q,
k=k,
v=v,
g=log_decay,
beta=beta,
initial_state=ssm_state,
inplace_final_state=True,
cu_seqlens=metadata.non_spec_query_start_loc[: metadata.num_decodes + 1],
ssm_state_indices=state_indices,
use_qk_l2norm_in_kernel=False,
)
else:
recurrent = q.new_zeros(q.shape)
recurrent = self.out_norm(recurrent.squeeze(0)).reshape(actual, -1)
output = self.o_proj(recurrent)
if actual < n:
output = F.pad(output, (0, 0, 0, n - actual))
return output
class MiniMindDecoderLayer(nn.Module):
def __init__(self, config, vllm_config: VllmConfig, prefix: str, layer_type: str):
super().__init__()
if layer_type == "linear":
self.self_attn = MiniMindGatedDeltaNet(
config, vllm_config, prefix=f"{prefix}.self_attn"
)
elif layer_type == "full":
self.self_attn = MiniMindFullAttention(
config, vllm_config.cache_config, prefix=f"{prefix}.self_attn"
)
else:
raise ValueError(f"Unsupported MiniMind layer type {layer_type!r}")
self.input_layernorm = MiniMindRMSNorm(config.hidden_size, config.rms_norm_eps)
self.post_attention_layernorm = MiniMindRMSNorm(config.hidden_size, config.rms_norm_eps)
self.mlp = MiniMindFeedForward(config)
def forward(self, hidden_states, positions, cos, sin):
attn_out = self.self_attn(self.input_layernorm(hidden_states), positions, cos, sin)
hidden_states = hidden_states + attn_out
return hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
class MiniMindModel(nn.Module):
def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
config = vllm_config.model_config.hf_text_config
self.config = config
self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size)
def make_layer(layer_prefix: str):
index = extract_layer_index(layer_prefix)
return MiniMindDecoderLayer(
config, vllm_config, layer_prefix, config.layer_types[index]
)
self.start_layer, self.end_layer, self.layers = make_layers(
config.num_hidden_layers,
make_layer,
prefix=maybe_prefix(prefix, "layers"),
)
self.norm = (
MiniMindRMSNorm(config.hidden_size, config.rms_norm_eps)
if get_pp_group().is_last_rank
else PPMissingLayer()
)
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
["hidden_states"], config.hidden_size
)
cos, sin = _precompute_rope(config)
self.register_buffer("freqs_cos", cos, persistent=False)
self.register_buffer("freqs_sin", sin, persistent=False)
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.embed_tokens(input_ids)
def forward(self, input_ids, positions, intermediate_tensors=None):
pp = get_pp_group()
if pp.is_first_rank:
if input_ids is None:
raise ValueError("input_ids is required on the first pipeline rank")
hidden_states = self.embed_input_ids(input_ids)
else:
if intermediate_tensors is None:
raise ValueError("intermediate_tensors is required on non-first ranks")
hidden_states = intermediate_tensors["hidden_states"]
cos = self.freqs_cos.to(hidden_states.device)
sin = self.freqs_sin.to(hidden_states.device)
for layer in islice(self.layers, self.start_layer, self.end_layer):
hidden_states = layer(hidden_states, positions, cos, sin)
if not pp.is_last_rank:
return IntermediateTensors({"hidden_states": hidden_states})
return self.norm(hidden_states)
class MiniMindForCausalLM(nn.Module, HasInnerState, SupportsPP, IsHybrid):
has_inner_state = True
is_hybrid = True
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
self.vllm_config = vllm_config
self.model_config = vllm_config.model_config
self.config = vllm_config.model_config.hf_text_config
self.model = MiniMindModel(vllm_config, prefix=maybe_prefix(prefix, "model"))
self.lm_head = ParallelLMHead(
self.config.vocab_size,
self.config.hidden_size,
prefix=maybe_prefix(prefix, "lm_head"),
)
self.lm_head.weight = self.model.embed_tokens.weight
self.logits_processor = LogitsProcessor(self.config.vocab_size)
self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.embed_input_ids(input_ids)
def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
intermediate_tensors: IntermediateTensors | None = None,
inputs_embeds: torch.Tensor | None = None,
**kwargs,
):
if inputs_embeds is not None:
raise NotImplementedError("This text-only MiniMind adapter does not accept inputs_embeds")
return self.model(input_ids, positions, intermediate_tensors)
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
return self.logits_processor(self.lm_head, hidden_states)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
params = dict(self.named_parameters())
loaded: set[str] = set()
for name, tensor in weights:
if name == "lm_head.weight":
name = "model.embed_tokens.weight"
param = params.get(name)
if param is None:
continue
loader = getattr(param, "weight_loader", default_weight_loader)
loader(param, tensor)
loaded.add(name)
missing = sorted(set(params) - loaded)
if missing:
preview = ", ".join(missing[:8])
suffix = " ..." if len(missing) > 8 else ""
raise ValueError(f"Checkpoint misses {len(missing)} parameters: {preview}{suffix}")
return loaded
@classmethod
def get_mamba_state_dtype_from_config(
cls, vllm_config: VllmConfig
) -> tuple[torch.dtype, torch.dtype]:
return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
vllm_config.model_config.dtype,
vllm_config.cache_config.mamba_cache_dtype,
vllm_config.cache_config.mamba_ssm_cache_dtype,
)
@classmethod
def get_mamba_state_shape_from_config(
cls, vllm_config: VllmConfig
) -> tuple[tuple[int, ...], tuple[int, ...]]:
config = vllm_config.model_config.hf_text_config
tp_size = vllm_config.parallel_config.tensor_parallel_size
num_spec = (
vllm_config.speculative_config.num_speculative_tokens
if vllm_config.speculative_config
else 0
)
return MambaStateShapeCalculator.gated_delta_net_state_shape(
tp_size,
config.num_attention_heads,
config.num_attention_heads,
config.head_dim,
config.head_dim,
1,
num_spec,
)
@classmethod
def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, ...]:
return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()

16
vllm_native/pyproject.toml Normal file
View File

@ -0,0 +1,16 @@
[build-system]
requires = ["setuptools>=68"]
build-backend = "setuptools.build_meta"
[project]
name = "minimind-vllm-native"
version = "0.1.0"
description = "Native vLLM plugin for MiniMind hybrid checkpoints"
requires-python = ">=3.10"
dependencies = []
[project.entry-points."vllm.general_plugins"]
minimind_vllm = "minimind_vllm:register"
[tool.setuptools]
packages = ["minimind_vllm"]