蚂蚁 inclusionAI 开源 SingProbe 训练代码
GitHub - inclusionAI/SingProbe: The official training code of SingProbe. · GitHub
蚂蚁 inclusionAI 在 GitHub 发布 SingProbe 官方训练代码，这是一个基于冻结 Base Model（如 Ling-3.0-flash）多层隐藏状态的轻量级 token 级探针训练框架，每个 token 输出 10 维 logits。Base Model 完全 requires_grad=False 仅做前向推理，由打过 token-probe 补丁的 SGLang 构建实时将指定层隐藏状态写入 tmpfs，训练侧以流式方式学习。框架覆盖三个任务、提供两种探针架构，训练完成后自动转换为标准 HuggingFace 模型目录。
正文
当前显示原文，中文全文译文尚未就绪。
inclusionAI; SingProbe
Repository navigation
NotificationsYou must be signed in to change notification settings; Fork4 (4); Star31 (31)You must be signed in to star a repository
The official training code of SingProbe.
Repository files navigation
README
English | 中文
A lightweight token-level probe training framework built on the multi-layer hidden states of a frozen Base Model (e.g. Ling-3.0-flash). The frozen Base Model is served by a token-probe-patched SGLang build for forward inference, dumping the per-token hidden states of the configured layers to tmpfs in real time; the training-side SingProbe model trains on these hidden states in a streaming fashion and finally outputs 10-dim logits per token.
The Base Model is entirely requires_grad=False and only runs inference to extract hidden states; only the probe model is trained.
Overview
Three tasks (10-dim output, fixed)
Each dataset only labels its own dims; the rest is unsupervised (label -100). See data/README.md for the format details.
Two probe architectures (singprobe_model.arch)
Both architectures share the same input contract (the concatenated hidden states); train.py dispatches on arch. See models/README.md for architecture details and self-test commands.
Environment Setup
The scripts in this repo do not install anything — prepare the environment once following this section before running.
1. Python dependencies
Python 3.10+ recommended, with a CUDA environment matching your hardware.
2. Token-probe-patched SGLang (required)
Hidden-state extraction requires the token-probe-patched SGLang build (upstream SGLang has no such dump capability). Build and install it from the token-probe branch:
3. flash-linear-attention (required for Ling-3.0)
Ling-3.0 contains linear-attention layers; the SGLang side needs the fla kernels:
4. spaCy model (optional)
Only needed when training.use_hallu_assertion_mask: true:
Model Preparation
The two configs under configs/all_models/ use the HuggingFace repos inclusionAI/Ling-3.0-flash / inclusionAI/Ling-3.0-tiny: both the SGLang server launch and the tokenizer load auto-download them from the Hub on first use. For offline environments, download the models in advance and point base_model.name at the local path.
Data Preparation
Training needs two dataset types (format details in data/README.md):
Safety (JSONL): Query / Response / Query_Label (A–H chars) / Response_Label (Safe|Unsafe) — labels dims 0–8; Hallucination (JSONL): character spans inside the response, mapped to token-level labels — labels dim 9
Before running, replace the four path/to/... placeholders under data: in the YAML with your real paths; train.py / the pipeline pre-check reports exactly which paths are missing.
Quickstart: Training
The pipeline scripts/run_train_pipeline.sh runs in order:
[1/3] Start the SGLang hidden-state server in the background (GPUs from sglang_gpus) and wait until it is ready; [2/3] Launch train.py via torchrun (training GPUs pinned by TRAIN_GPUS, default 0,1,2,3); [3/3] Convert the final checkpoint into a HuggingFace model directory (written to the checkpoint directory's safetensors/)
On any exit path (success / failure / interruption) the SGLang process group is always torn down so the GPUs release. The SGLang-side and train-side GPU sets must be disjoint.
Handy env-var knobs: TP DP GPUS TRAIN_GPUS SGLANG_PORT SGLANG_SGLOG PROBE_CKPT SAVE_DIR MEM_FRACTION — see the script header (bash scripts/run_train_pipeline.sh --help).
Resuming: bash scripts/all/ling-3.0-flash.sh --resume path/to/checkpoint-N
Artifacts and Inference
When training finishes, the pipeline automatically converts the final checkpoint:
The output dir is a standard HuggingFace model directory — load it directly with trust_remote_code (no SingProbe repo needed):
Any training checkpoint can also be converted manually:
Configuration Overview
Config parsing lives in config.py (Config.from_yaml); top-level sections:
Repository Layout
Legal
Please read LEGAL.md before use.
Citation
For more information, please refer to the technical report.
The official training code of SingProbe.