332 lines
15 KiB
Markdown
332 lines
15 KiB
Markdown
---
|
||
library_name: transformers
|
||
tags:
|
||
- router
|
||
- orchestrator
|
||
- slm
|
||
- edge-computing
|
||
- mixture-of-experts
|
||
- text-generation
|
||
- TRL
|
||
pipeline_tag: text-generation
|
||
model_type: llama
|
||
datasets:
|
||
- aifeifei798/Prompt-Routing-Dataset-zh-cn
|
||
language:
|
||
- en
|
||
- zh
|
||
base_model:
|
||
- SupraLabs/Supra-Router-51M
|
||
license: apache-2.0
|
||
---
|
||
|
||
<h1 align="center">Supra-Router-51M-zh-cn · Multi-Task Infrastructure Routing Model</h1>
|
||
|
||
你好！请问大熊猫主要生活在中国的哪个省份？
|
||
|
||
Domain: English Language | Complexity: 2 | Math: False | Code: False | Route: small model | Justification: Automated override: Simple text or factual task with low complexity (2), perfectly optimized for an edge SLM.
|
||
|
||
请用 Python 写一段快速排序算法的代码，并简单解释它的时间复杂度。
|
||
|
||
Domain: Programming | Complexity: 2 | Math: False | Code: True | Route: big model | Justification: Automated override: Task complexity is high (2) or involves technical logic (Math: False, Code: True), requiring frontier capabilities.
|
||
|
||
如果一个等差数列的首项是3，公差是5，请问第12项是多少？让我们逐步思考。
|
||
|
||
Domain: Mathematics | Complexity: 2 | Math: True | Code: False | Route: big model | Justification: Automated override: Task complexity is high (2) or involves technical logic (Math: True, Code: False), requiring frontier capabilities.
|
||
|
||
### 使用
|
||
|
||
```python
|
||
import torch
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
||
MODEL_ID = "aifeifei798/Supra-Router-51M-zh-cn"
|
||
|
||
print("[*] Initializing local infrastructure router...")
|
||
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
MODEL_ID, dtype=torch.bfloat16, device_map="auto"
|
||
)
|
||
model.eval()
|
||
|
||
# Example prompt showcasing keyword-trap evasion
|
||
user_prompt = "如果一个等差数列的首项是3，公差是5，请问第12项是多少？让我们逐步思考。"
|
||
print(user_prompt)
|
||
# Format to match internal SFT attention alignment
|
||
formatted_input = f"Task: {user_prompt}\nAnalysis: "
|
||
inputs = tokenizer(formatted_input, return_tensors="pt").to(model.device)
|
||
|
||
with torch.no_grad():
|
||
outputs = model.generate(
|
||
**inputs,
|
||
max_new_tokens=128,
|
||
do_sample=False,
|
||
pad_token_id=tokenizer.pad_token_id,
|
||
eos_token_id=tokenizer.eos_token_id,
|
||
)
|
||
|
||
generated_ids = outputs[0][inputs["input_ids"].shape[1] :]
|
||
print(tokenizer.decode(generated_ids, skip_special_tokens=True).strip())
|
||
|
||
```
|
||
|
||
### TRL 训练脚本
|
||
|
||
```Python
|
||
import torch
|
||
from datasets import load_dataset
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
from trl import SFTTrainer, SFTConfig
|
||
|
||
# ================= 1. 基础配置 =================
|
||
MODEL_ID = "./Supra-Router-51M"
|
||
DATA_PATH = (
|
||
"./Prompt-Routing-Dataset-zh-cn/router_dataset_openai_gpt-oss-120b-zhcn.jsonl"
|
||
)
|
||
OUTPUT_DIR = "./Supra-Router-51M-zh"
|
||
|
||
print("[*] 正在加载 Tokenizer 和 模型...")
|
||
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
||
if tokenizer.pad_token is None:
|
||
tokenizer.pad_token = tokenizer.eos_token
|
||
|
||
# 加载模型
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
MODEL_ID, device_map="auto", torch_dtype=torch.bfloat16
|
||
)
|
||
|
||
print(f"[*] 正在加载中文数据集: {DATA_PATH}")
|
||
dataset = load_dataset("json", data_files=DATA_PATH)["train"]
|
||
|
||
|
||
# ================= 2. 数据格式化与双列拆分 =================
|
||
def format_prompts(batch):
|
||
prompts = []
|
||
completions = []
|
||
for i in range(len(batch["prompt"])):
|
||
raw_prompt = batch["prompt"][i]
|
||
|
||
if len(raw_prompt) > 1500:
|
||
raw_prompt = raw_prompt[:1500] + "\n...[Text Truncated]..."
|
||
|
||
# 1. 输入部分：去掉冒号后面的空格，防止 BPE 分词器边界混淆
|
||
task_input = f"Task: {raw_prompt}\nAnalysis:"
|
||
|
||
# 2. 输出部分：在开头 Domain 前面主动加一个空格，确保无缝对齐
|
||
output_label = (
|
||
f" Domain: {batch['primary_domain'][i]} | "
|
||
f"Complexity: {batch['complexity_score'][i]} | "
|
||
f"Math: {batch['math_task'][i]} | "
|
||
f"Code: {batch['coding_task'][i]} | "
|
||
f"Route: {batch['routing_choice'][i]} | "
|
||
f"Justification: {batch['routing_justification'][i]}"
|
||
f"{tokenizer.eos_token}"
|
||
)
|
||
|
||
prompts.append(task_input)
|
||
completions.append(output_label)
|
||
|
||
return {"prompt": prompts, "completion": completions}
|
||
|
||
|
||
print("[*] 正在格式化数据集并进行双列映射...")
|
||
dataset = dataset.map(format_prompts, batched=True, remove_columns=dataset.column_names)
|
||
|
||
# ================= 3. 训练参数配置 =================
|
||
training_args = SFTConfig(
|
||
output_dir=OUTPUT_DIR,
|
||
max_length=2048,
|
||
completion_only_loss=True, # 开启原生答案专注模式
|
||
per_device_train_batch_size=8,
|
||
gradient_accumulation_steps=4,
|
||
learning_rate=3e-5,
|
||
logging_steps=10,
|
||
num_train_epochs=4,
|
||
save_strategy="epoch",
|
||
bf16=True,
|
||
optim="adamw_torch",
|
||
)
|
||
|
||
trainer = SFTTrainer(
|
||
model=model,
|
||
train_dataset=dataset,
|
||
args=training_args,
|
||
processing_class=tokenizer,
|
||
)
|
||
|
||
# ================= 4. 启动微调 =================
|
||
print("[*] 万事俱备，开始以新版原生【答案专注模式】进行微调！")
|
||
trainer.train()
|
||
|
||
print(f"[*] 训练大功告成！正在保存至 {OUTPUT_DIR}")
|
||
trainer.save_model(OUTPUT_DIR)
|
||
tokenizer.save_pretrained(OUTPUT_DIR)
|
||
print("[[*] 全部任务结束。")
|
||
```
|
||
|
||
## completion_only_loss
|
||
|
||
`completion_only_loss`（仅对答案计算损失）是监督微调（SFT）和对齐训练（Alignment）中**非常核心、甚至可以说是必用**的一项技术。
|
||
|
||
为了让你彻底明白它的底层逻辑，我们从“普通训练”和“答案专注训练”的对比、PyTorch 的底层实现、以及为什么它对小模型至关重要这三个维度来剖析。
|
||
|
||
---
|
||
|
||
### 一、 核心对比：普通 SFT vs Completion-Only SFT
|
||
|
||
在因果语言模型（Causal LM）的日常训练中，模型的任务是**预测下一个 Token**。
|
||
|
||
假设一条完整的训练数据拼接后是：
|
||
> `Task: 1+1等于几？\nAnalysis: 等于2。`
|
||
|
||
#### 1. 普通的 SFT 训练（Full Sequence Loss）
|
||
模型需要对**每一个 Token** 都计算预测损失（Loss），并进行梯度回传：
|
||
* 模型预测 `Task` 后面的第一个字（“1”） -> **计算 Loss**
|
||
* 模型预测 “1” 后面的 “+” -> **计算 Loss**
|
||
* ……
|
||
* 模型预测 `Analysis: ` 后面的 “等” -> **计算 Loss**
|
||
* 模型预测 “等” 后面的 “于” -> **计算 Loss**
|
||
|
||
**弊端**：模型在训练时，不仅要学习怎么回答“等于2”，还要花大量的脑力去学习和记忆“1+1等于几？”这段提示词。但在实际使用时，提示词是用户输入的，模型根本不需要去预测它。
|
||
|
||
#### 2. Completion-Only SFT 训练
|
||
模型依然可以看见完整的 `Task: 1+1等于几？\nAnalysis: 等于2。`（用于提供完整的上下文注意力），但在计算损失时：
|
||
* 提示词（Prompt）部分：**不计算 Loss，完全忽略**。
|
||
* 回答（Completion）部分：**正常计算 Loss**。
|
||
|
||
模型唯一的优化目标，就是在看到前面的输入后，如何精准地吐出后面的答案。
|
||
|
||
---
|
||
|
||
### 二、 底层原理：PyTorch 的 `-100` 魔法
|
||
|
||
在底层代码中，这个机制是通过 PyTorch 交叉熵损失函数（`CrossEntropyLoss`）中的一个特殊参数 **`ignore_index`** 来实现的。
|
||
|
||
在 PyTorch 和 Hugging Face Transformers 中，`ignore_index` 的默认值是 **`-100`**。任何标签值被设为 `-100` 的 Token，在计算 Loss 时都会被**完全忽略，梯度不会回传**。
|
||
|
||
#### 具体的 Token 与 Label 变化：
|
||
假设我们要训练一句话，分词后的 `input_ids` 如下。我们在后台为它准备了对应的 `labels`：
|
||
|
||
| 文本内容 | `Task:` | `大熊猫` | `住哪？` | `\nAnalysis:` | ` Domain:` | ` Mathematics` | ` |` |
|
||
|
||
| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |
|
||
|
||
| **`input_ids` (输入)** | `1203` | `3042` | `9842` | `4322` | `567` | `8831` | `102` |
|
||
|
||
| **`labels` (标签)** | **`-100`** | **`-100`** | **`-100`** | **`-100`** | `567` | `8831` | `102` |
|
||
|
||
1. 在 Forward（前向传播）时，模型能看到完整的 `input_ids`（包含大熊猫、住哪等上下文）。
|
||
2. 在计算 Loss 时，PyTorch 看到前四个 Label 是 `-100`，就会直接跳过，**只对后三个真正的标签值计算 CrossEntropy**。
|
||
3. 这样，模型的所有参数权重，都只会为了“在看到前四个 Token 时，如何更准地预测出后三个 Token”而发生改变。
|
||
|
||
---
|
||
|
||
### 三、 为什么新版 TRL 的“双列自动检测”如此方便？
|
||
|
||
在过去（比如旧版的 `trl` 库中），我们需要使用 `DataCollatorForCompletionOnlyLM`。
|
||
它的工作原理是：
|
||
1. 把整段文本拼成一个大 `text`。
|
||
2. 在分词后的 Token 序列里，拿着一根“放大镜”，去人肉搜索 `\nAnalysis:` 这段字符串对应的 Token ID 边界。
|
||
3. 把边界之前的所有 Label 强行改成 `-100`。
|
||
|
||
**缺点**：非常容易因为分词器（Tokenizer）在特殊字符或空格处的切词变化（比如我们遇到的 Mismatch 问题），导致“放大镜”找不到边界，进而导致整条数据都无法计算 Loss，或者报错。
|
||
|
||
#### 新版 TRL 的“双列原生检测”：
|
||
当你在 `dataset` 里保留 `prompt` 和 `completion` 两个独立的列，并在 `SFTConfig` 中开启 `completion_only_loss=True` 时，TRL 采取了更聪明的做法：
|
||
1. 它先单独对 `prompt` 列进行分词，得到它的 Token 长度（比如是 `L`）。
|
||
2. 它再对 `prompt + completion` 整体进行分词，得到一整个 Token 序列。
|
||
3. 它直接将这个完整序列中，前 `L` 个位置的 `labels` 暴力抹成 `-100`。
|
||
|
||
这是一种纯粹数学/数组切片上的操作，不再依赖于字符串的模式匹配。不仅速度极快，而且只要做好了空格对齐，就绝对不会出现对齐失败的警告。
|
||
|
||
---
|
||
|
||
### 四、 总结：它对 51M 小模型带来的质变
|
||
|
||
这就是为什么你在微调 51M 模型时，开启这个功能后，指标发生了翻天覆地的变化：
|
||
|
||
1. **参数解放**：51M 模型太小了，如果让它去背诵你的 992 条中文 Prompt（其中有些可能长达几千字），它会因为记忆载荷过大而发生“模式坍塌”。开启后，它只需要背诵极其简短、高度重复的英文路由标签。
|
||
2. **收敛极快**：因为预测目标变成了极度简单的分类选择（Domain, Code, Math 等），这相当于把一个复杂的**文本生成任务（Generation）**降维成了简单的**选择分类任务（Classification）**。这也是为什么你的 Loss 在第二轮就降到了 `0.1` 附近，准确率直接高达 `97%+`。
|
||
3. **泛化保留**：因为模型不需要在训练中改变它的“中文语言表征”（因为它不用预测中文），它在微调后，依然完美保留了预训练时学到的常识和逻辑能力，从而在推理时能听懂“大熊猫”和“等差数列”的区别。
|
||
|
||
<h1 align="center">Supra-Router-51M · Multi-Task Infrastructure Routing Model</h1>
|
||
|
||
|
||
|
||
<h2 align="center">About the Model</h2>
|
||
|
||
GGUF model [here](https://huggingface.co/SupraLabs/Supra-Router-51M-gguf)
|
||
|
||
**Supra-Router-51M** is an ultra-lightweight, high-speed infrastructure traffic controller optimized for localized edge orchestration. With only **51.7 million parameters**, this micro-LLM acts as a defensive gateway for multi-model ecosystems, accurately determining when user requests can be processed locally by an Edge SLM or when they must be triaged to a cloud-hosted frontier intelligence layer.
|
||
|
||
The model was built by fine-tuning a pre-trained 51M base on the `SupraLabs/Prompt-Routing-Dataset` (992 rows). Rather than acting as a naive binary classifier, the model uses **Multi-Task Sequence Generation** to map out the underlying properties of a prompt before predicting the final routing token, anchoring its attention heads to robust language and structural logic features.
|
||
|
||
---
|
||
|
||
## Multi-Task Decision Sequence
|
||
|
||
To run inference, wrap your user query inside the structural framing tokens used during training (`Task: [Prompt]\nAnalysis: `). The model will output a deterministic, pipe-separated string containing the full telemetry of the prompt's cognitive requirements:
|
||
|
||
### Expected Output Target Schema:
|
||
```text
|
||
Domain: [Semantic Field] | Complexity: [1-5] | Math: [True/False] | Code: [True/False] | Route: [small model/big model] | Justification: [Rule-driven infrastructure reasoning]
|
||
```
|
||
|
||
## Why this works:
|
||
|
||
By forcing a sub-100M parameter model to calculate the semantic domain, structural complexity, and technical flags before it emits the final Route token, the network effectively runs an internal feature-activation map. This multi-task sequence prevents localized weight collapse and guarantees stable routing boundaries.
|
||
|
||
## Training Telemetry & Optimization
|
||
|
||
- Dataset Source: SupraLabs/Prompt-Routing-Dataset (992 samples)
|
||
- Training Duration: 5 Epochs
|
||
- Checkpoint Selection: Peak generalization was reached during Epoch 3 (eval_loss: 0.1342). To eliminate late-stage micro-model memorization and validation drift, the training state was automatically rewound and saved at this numerical peak.
|
||
- Precision: bfloat16
|
||
- Hardware Footprint: Optimized sequence processing length of 3840 tokens, ensuring rapid inference execution with negligible CPU/GPU overhead (sub-millisecond generation speeds).
|
||
|
||
## Inference & Gateway Implementation
|
||
|
||
Use this direct script to test or wrap the model inside a live production orchestrator or FastAPI gateway. It enforces greedy decoding (do_sample=False) for maximum decision stability.
|
||
|
||
```python
|
||
import torch
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
||
MODEL_ID = "aifeifei798/Supra-Router-51M-zh-cn"
|
||
|
||
print("[*] Initializing local infrastructure router...")
|
||
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
MODEL_ID,
|
||
dtype=torch.bfloat16,
|
||
device_map="auto"
|
||
)
|
||
model.eval()
|
||
|
||
# Example prompt showcasing keyword-trap evasion
|
||
user_prompt = "你好！请问大熊猫主要生活在中国的哪个省份？"
|
||
|
||
# Format to match internal SFT attention alignment
|
||
formatted_input = f"Task: {user_prompt}\nAnalysis: "
|
||
inputs = tokenizer(formatted_input, return_tensors="pt").to(model.device)
|
||
|
||
with torch.no_grad():
|
||
outputs = model.generate(
|
||
**inputs,
|
||
max_new_tokens=128,
|
||
do_sample=False,
|
||
pad_token_id=tokenizer.pad_token_id,
|
||
eos_token_id=tokenizer.eos_token_id
|
||
)
|
||
|
||
generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
|
||
print(tokenizer.decode(generated_ids, skip_special_tokens=True).strip())
|
||
```
|
||
|
||
## Proven Benchmarks & Defensive Boundaries
|
||
|
||
During edge validation testing, Supra-Router-51M demonstrated robust resilience against adversarial prompt strings:
|
||
|
||
- Keyword Trap Evasion: Successfully identifies semantic context rather than matching tokens. Prompts containing words like "script" or "calculus" are correctly parsed as creative writing (not programming/math code) and routed locally to the small model when complexity is low.
|
||
- Complexity-Driven Safety Net: In instances where programming syntax or technical boundaries are ambiguous (e.g., complex regex or architectural database frames), the model naturally scales its evaluation metrics to Complexity: 3, automatically triggering a big model route override.
|
||
- Deterministic Offloading: Safely captures multi-step logic paths, calculus concepts, and code generation scripts, instantly assigning them to cloud-scale frontier endpoints. |