一套从架构层原生支持 SSD+DRAM 混合加载推理的 MoE 框架,配套完整的 6 阶段训练链路。
主流 MoE 模型(Mixtral、DeepSeek-MoE、Qwen-MoE 等)在消费级硬件上运行时,路由决策是即时、被动、逐 token 的——每一步 decode 才去检查所需专家是否在 VRAM 中。当 VRAM 不足时,系统在 decode 流中同步等待 SSD→RAM→VRAM 的搬运,生成速度跌至 < 5 tokens/s。
这不是参数调优问题,而是架构层的根本缺陷:这些模型设计时假设全量 VRAM 驻留,offload 是事后打的补丁。
sequenceDiagram
autonumber
participant Prompt as Prompt tokens
participant Intent as IntentClassifier
participant SSD as 聚簇 SSD 缓存
participant RAM as Pinned RAM
participant H2D as H2D stream
participant GPU as Decode compute stream
participant LR as LookaheadRouter
rect rgb(236, 248, 255)
Prompt->>Intent: prefill 读取完整 prompt
Intent->>SSD: 预测整次生成的热点专家集合
SSD->>RAM: mmap cluster-packed .ctsr 文件
RAM->>H2D: 首个 decode token 前排队加载热点专家
H2D-->>GPU: per-expert ready event
end
loop 每个 decode token
GPU->>LR: Block 0 hidden state
LR->>H2D: 预取 t+1..t+K 的未来专家
GPU->>GPU: 使用已驻留专家计算 token t
GPU->>H2D: 只等待当前所需专家事件
end
两层关键改进:
- Prefill 时机:
PrefillScheduler+IntentClassifier在第一个 decode token 之前批量预加载专家。 - 事件级同步(M3):
promote_to_vram(blocking=False)在_h2d_stream上记录torch.cuda.Event;compute 流通过current_stream.wait_event(evt)只等所需专家,不再全局stream.synchronize()。在 30ms 模拟 SSD 延迟下,新路径比旧路径有 35ms+/token 的流水线富余空间。
flowchart TB
subgraph GPU["VRAM / Metal 统一内存"]
Dense["Dense layers"]
Shared["共享专家<br/>始终驻留"]
Hot["预测热点专家<br/>使用前提升到 VRAM"]
Events["per-expert CUDA event<br/>无全局 stream sync"]
end
subgraph RAM["Pinned RAM 中转层"]
Buffer["已预取专家权重"]
LRU["RAM LRU cache"]
end
subgraph Disk["NVMe SSD"]
C1["cluster_000.ctsr"]
C2["cluster_001.ctsr"]
Manifest["cluster_manifest.json<br/>expert -> cluster 映射"]
end
Manifest --> C1
Manifest --> C2
C1 -->|"顺序 mmap 读取"| Buffer
C2 -->|"顺序 mmap 读取"| Buffer
Buffer --> LRU
LRU -->|"非阻塞 H2D"| Hot
Hot --> Events
Shared --> Events
Dense --> Events
cluster_manifest.json 与 .ctsr 文件由离线 Louvain 聚类生成;运行时一次 safetensors.safe_open(...).mmap 把整个簇拉进 RAM,把随机读改写成顺序读。
即使最坏情况也不会卡顿——柔性降级(Soft Gating):
# 纯张量乘法,无 Python 分支,torch.compile 不会图断裂
output = avail[i] * expert_output + (1.0 - avail[i]) * shared_expert_output共享专家(常驻 VRAM)按比例混入,生成流从不中断,精度平滑降级,专家后台加载完毕后自动恢复。
flowchart LR
Prompt["Prompt / prefill context"] --> IC["IntentClassifier<br/>整次生成专家先验"]
IC --> Budget["ExpertPredictor<br/>带预算的热点专家集合"]
Budget --> Prefill["PrefillScheduler<br/>decode 前批量预加载"]
Token["Decode token t"] --> Block0["Transformer Block 0"]
Block0 --> LR["LookaheadRouter<br/>Q_t^(1..K)"]
LR --> Future["未来专家预测<br/>t+1 ... t+K"]
Future --> Async["AsyncPrefetcher<br/>预取队列"]
Prefill --> Cache["CacheManager / ExpertStore"]
Async --> Cache
Cache --> MoE["ChronosMOEFeedForward<br/>驻留专家或共享专家兜底"]
| IntentClassifier(第一层) | LookaheadRouter(第二层) | |
|---|---|---|
| 触发时机 | Prefill 阶段一次 | Decode 每个 token |
| 输入 | 完整 Prompt(最多 512 token) | Block 0 的 hidden state |
| 输出 | 整次生成的专家集合 | t+1、t+2 步的专家 ID |
| 训练目标 | 真实激活日志监督(Phase 5) | L_lookahead — 真实路由 t+k 作为 stop-grad teacher(M2) |
| 参数量 | ~10–15M(单独训练) | ~2M(随主模型训练) |
M2 之前 LookaheadRouter 没有任何监督——只是个未训练的 head。M2 引入:
让前瞻头真正学习预测未来 K 步的路由分布。
每个阶段都是独立的 entry script,全部继承 Chronos 损失混合器(lookahead + temporal + balance),并在对齐阶段引入路由 KL 锚定防止缓存命中率被 RL/DPO 梯度毁掉。
flowchart LR
P1["Stage 1<br/>Pretrain<br/>CE + Chronos mix"] --> P2["Stage 2<br/>SFT<br/>assistant-token CE + mix"]
P2 --> P3["Stage 3<br/>DPO<br/>preference loss + anchor"]
P2 --> P4["Stage 4<br/>ORPO<br/>NLL + odds-ratio + anchor"]
P4 --> P5["Stage 5<br/>GRPO<br/>rollout reward + KL + anchor"]
P5 --> P6["Stage 6<br/>Distill<br/>KD + CE + anchor"]
subgraph Shared["共享 Chronos 训练项"]
B["raw load-balance aux"]
T["temporal locality"]
LA["supervised lookahead"]
A["router KL anchor<br/>对齐阶段"]
end
Shared -. 应用于 .-> P1
Shared -. 应用于 .-> P2
Shared -. 应用于 .-> P3
Shared -. 应用于 .-> P4
Shared -. 应用于 .-> P5
Shared -. 应用于 .-> P6
| Stage | 脚本 | 核心损失 | Router KL 锚定 (默认 λ) |
|---|---|---|---|
| 1 Pretrain | train_chronos.py |
CE + balance + temporal + lookahead | 0.0 (off) |
| 2 SFT | train_chronos_sft.py |
+ 上述 mix | 0.01 (weak) |
| 3 DPO | train_chronos_dpo.py |
DPO log-σ(β·logits) + mix | 0.10 (strong) |
| 4 ORPO | train_chronos_orpo.py |
NLL + λ·OR(无 ref model) | 0.10 |
| 5 GRPO | train_chronos_grpo.py |
PG·A − β·KL(含 ToyReward / 可插 LMRewardModel) | 0.10 |
| 6 Distill | train_chronos_distill.py |
α·T²·KL(s‖t) + (1−α)·CE | 0.05 |
完整 6 阶段端到端对比见 tools/compare_minimind_chronos_v3.py。
flowchart TD
Request["用户请求<br/>训练 / 推理 / WebUI / CLI"] --> Dispatcher["BackendDispatcher"]
Dispatcher --> Probe["能力探针"]
Probe --> CUDA["CUDA<br/>训练 + 推理"]
Probe --> MPS["MPS<br/>训练 + 推理"]
Probe --> MLX["MLX<br/>Apple Silicon 原生路径"]
Probe --> CPU["CPU<br/>可移植 fallback"]
Probe --> EXT["扩展钩子<br/>Vulkan / OpenCL"]
Dispatcher --> Choice["优先级 + 可用性决策"]
Choice --> Runtime["Chronos runtime / trainer"]
from chronos.backend import BackendDispatcher
d = BackendDispatcher()
d.available() # ['mlx', 'mps', 'cpu'] on Apple Silicon
# ['cuda', 'cpu'] on NVIDIA host
d.select() # 自动选最佳;可被 CHRONOS_BACKEND 环境变量覆盖
d.describe() # 人类可读的能力总览- 一等公民(训练 + 推理):
cpu、mps、cuda、mlx - 推理仅 / 实验性:
vulkan(仅当 PyTorchUSE_VULKAN=ON自定义构建时存在) - 第三方插件钩子:
opencl(替换chronos/backend/ext/opencl.py:PROBE())
诚实声明:上游 PyTorch 没有 OpenCL 后端、Vulkan 也仅在自定义构建中可用。Chronos 提供 dispatcher 接缝,使第三方插件无需改核心代码即可接入。
-
ChronosForCausalLM继承PreTrainedModel,已注册AutoConfig/AutoModelForCausalLM,无需trust_remote_code:from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained("./out_dir")
-
chronos.model.hf_io.save_chronos_pretrained/load_chronos_pretrained输出标准model.safetensors+config.json,并把cluster_manifest.json+.ctsr一起带过去;roundtrip logits 0.00e+00 偏差。 -
chronos.serving.register_chronos_with_vllm()在已安装 vLLM 时把 Chronos 注册到ModelRegistry;未安装时打印安装提示,不报错。worker 侧 mask 注入 hook 见docs/vllm_integration.md。
| 特性 | llama.cpp offload | vLLM offload | Project Chronos |
|---|---|---|---|
| 专家预测 | 无(被动) | 无(被动) | 主动预测(IntentCLF + LookaheadRouter) |
| Lookahead 训练 | n/a | n/a | L_lookahead 真实监督(M2) |
| IO 时机 | Decode 期间(阻塞) | Decode 期间(阻塞) | Prefill 期间(异步,前置) |
| Decode 流水线 | 同步 | 同步 | 双流 + per-expert event(M3) |
| Cache miss 行为 | 硬阻塞 | 硬阻塞 | Soft Gating(零阻塞) |
| 磁盘格式 | gguf | safetensors | cluster-packed safetensors(.ctsr) |
| 训练集成 | 事后补丁 | 事后补丁 | 6 阶段全栈 + 路由 KL 锚定 |
| 后端调度 | 编译期固定 | CUDA-only | cpu / mps / cuda / mlx 自动 + vulkan/opencl 钩子 |
| Apple Silicon 原生 | 部分 | 无 | 完整 MLX 后端 |
| HuggingFace 兼容 | 仅 GGUF | ✓ | ✓ + 携带专家缓存 |
| vLLM 兼容 | n/a | 原生 | 可选 adapter(按需注册) |
L_base:阶段相关目标(CE / DPO / ORPO / GRPO / KD)。L_aux:未缩放的 MoE load-balance 辅助项;Chronos 在chronos_loss_term中只乘一次lambda_bal。L_temporal:约束相邻 token 的路由分布不要剧烈跳变,提高专家复用和缓存局部性。L_lookahead:未来真实路由分布到前瞻预测的 soft-target cross entropy。sg(...)表示 stop-gradient。L_routerKL:对齐阶段锚定 stage 开始时捕获的参考路由分布,防止 RL/DPO/ORPO/GRPO 梯度破坏聚簇布局。
λ 全部支持 Optuna TPE 自动搜索(包括 hidden_size / num_experts / kv_latent_dim 等结构超参)。
pip install project-chronos或从源码:
git clone https://github.com/FonaTech/Project_Chronos
cd Project_Chronos
pip install -e ".[dev]"MLX(Apple Silicon):
pip install "project-chronos[mlx]"vLLM 服务(可选,仅 Linux+CUDA):
pip install vllmminimind 依赖:Project Chronos 使用 minimind 作为 MoE 内核。 若本地未找到,首次 import 时自动克隆至
~/.cache/chronos/minimind-master/。 minimind 采用 Apache-2.0 授权,完整归属见 THIRD_PARTY_NOTICES.md。
环境要求:Python 3.10+,PyTorch 2.4+
chronos-ui
# 或
python chronos_app.py包含:⚙️ Config(含右侧实时参数估算面板,合并了 Designer)/ 🏋️ Train(拥有 data_path)/ 🧪 6-Stage Pipeline(每阶段独立数据路径)/ 💬 Inference / 📊 Benchmark(Markdown 表 + BarPlot)/ 🔬 Auto-Tune(持久化日志 + 一键 Apply Best → Config)/ 📡 IO Monitor。i18n 支持 zh-Hans / zh-Hant / en / ja。
python train_chronos.py \
--data_path ./tests/fixtures/tiny_pretrain.jsonl \
--hidden_size 256 --num_hidden_layers 4 --num_experts 4 \
--epochs 1 --device cpu --save_dir ./outpython train_chronos_sft.py --data_path ./tests/fixtures/tiny_sft.jsonl --from_weight chronos --save_dir ./out --device cpu
python train_chronos_dpo.py --data_path ./tests/fixtures/tiny_dpo.jsonl --from_weight sft --save_dir ./out --device cpu
python train_chronos_orpo.py --data_path ./tests/fixtures/tiny_dpo.jsonl --from_weight sft --save_dir ./out --device cpu
python train_chronos_grpo.py --data_path ./tests/fixtures/tiny_grpo.jsonl --from_weight orpo --save_dir ./out --device cpu \
--reward toy # 或 lm:/path/to/reward-modelpython train_chronos_distill.py \
--data_path ./tests/fixtures/tiny_sft.jsonl \
--teacher_path ./out/sft_192_moe.pth \
--from_weight grpo --save_dir ./out --device cpu \
--alpha 0.7 --temperature 4.0python tools/compare_minimind_chronos_v3.py \
--pretrain_steps 150 --align_steps 30 --distill_steps 30 \
--simulated_ssd_ms 30 --device cpu \
--output results/compare_results_v3.json输出包括:每阶段 loss、HF roundtrip Δlogit、tokens/sec、激活专家比例、常驻专家字节、M3 流水线富余空间、当前主机后端列表。
from chronos.io.cluster_layout import (
collect_activation_log, build_cooccurrence_matrix,
try_louvain_clustering, repack_expert_weights_safetensors,
)
log = collect_activation_log(model, calib_loader, "cpu", max_batches=50)
clusters = try_louvain_clustering(build_cooccurrence_matrix(log, num_experts))
repack_expert_weights_safetensors(model, clusters, "./expert_cache_clustered")from chronos.tuning.chronos_auto_tuner import ChronosAutoTuner, ChronosSearchSpaceConfig
tuner = ChronosAutoTuner()
tuner.start(
model_id="./out/chronos_256_moe.pth",
dataset_path="./dataset/train.jsonl",
search_space=ChronosSearchSpaceConfig(
tune_lambda_balance=True, tune_lambda_temporal=True,
tune_lambda_lookahead=True, tune_lookahead_steps=True,
tune_hidden_size=True, tune_num_experts=True,
tune_num_shared_experts=True, tune_kv_latent_dim=True,
),
n_trials=20,
)Project_Chronos/
├── chronos/
│ ├── deps.py # 自动下载 minimind(若本地未找到)
│ ├── __init__.py # 注册 AutoConfig / AutoModelForCausalLM
│ ├── model/
│ │ ├── config.py # ChronosConfig(lookahead/temporal/anchor/storage_format 等)
│ │ ├── hybrid_attention.py # MLAAttention + SlidingWindowAttention
│ │ ├── lookahead_router.py # 逐 token 前瞻预测器(第二层)
│ │ ├── moe_chronos.py # ChronosMOEFeedForward + 共享专家 + Soft Gating
│ │ ├── model_chronos.py # ChronosForCausalLM(_tied_weights_keys 已修补)
│ │ ├── temporal_loss.py # 时间局部性 + lookahead 监督损失
│ │ └── hf_io.py # save/load_chronos_pretrained + AutoModel 注册
│ ├── io/
│ │ ├── expert_store.py # 三级存储 + per-expert Event + 非阻塞 promote
│ │ ├── async_prefetcher.py # 异步预取(prefetch_only / promote_current 已分离)
│ │ ├── storage.py # ClusterStorage:.ctsr safetensors + manifest
│ │ ├── cluster_layout.py # 共现聚簇 + safetensors 重排
│ │ └── io_simulator.py # CHRONOS_SIM_SSD_MS 测试钩子(M3)
│ ├── router/
│ │ ├── intent_classifier.py # Prompt 级专家预测器(第一层,~10M 参数)
│ │ ├── expert_predictor.py # IntentVector → ExpertSet(含预算上限)
│ │ └── prefill_scheduler.py # 编排 prefill 阶段批量预加载
│ ├── mlx/ # Apple Silicon 原生后端
│ │ ├── attention.py / moe.py / model.py / expert_store.py / inference.py
│ ├── runtime/
│ │ ├── cache_manager.py # prefetch_for_next_step / ensure_resident(M3)
│ │ ├── inference_engine.py # 端到端推理引擎(重排为 H2D-compute 重叠)
│ │ └── metrics.py # MetricsBus(IO Monitor 数据源)
│ ├── trainer/
│ │ ├── loss_mixin.py # chronos_loss_term + router_kl_anchor + capture_reference_routing
│ │ ├── chronos_trainer.py # Pretrain
│ │ ├── sft_trainer.py # Stage 2
│ │ ├── dpo_trainer.py # Stage 3
│ │ ├── orpo_trainer.py # Stage 4(无 ref model)
│ │ ├── grpo_trainer.py # Stage 5(含 self-contained rollout)
│ │ ├── distill_trainer.py # Stage 6(KL/T² + α 混合)
│ │ └── reward.py # ToyReward / LMRewardModel / build_reward_fn
│ ├── tuning/
│ │ └── chronos_auto_tuner.py # Optuna λ + 结构超参搜索
│ ├── eval/
│ │ ├── io_profiler.py # Phase 1 验证(前瞻准确率)
│ │ └── benchmark.py # 端到端基准测试
│ ├── data/
│ │ └── flexible_dataset.py # 自动识别任意 JSONL 字段格式
│ ├── backend/
│ │ ├── __init__.py # BackendDispatcher(cpu/mps/cuda/mlx)
│ │ ├── dispatcher.py # 探针 + 优先级 + 训练能力声明
│ │ └── ext/opencl.py # 第三方 OpenCL 插件钩子(stub)
│ ├── _backend_legacy.py # 向后兼容 build_model() 等旧 API
│ ├── serving/
│ │ ├── __init__.py
│ │ └── vllm_adapter.py # 可选 vLLM 注册(无 vLLM 时优雅降级)
│ └── cli.py # 统一 CLI
├── ui/ # Gradio Web UI(i18n: zh-Hans/zh-Hant/en/ja)
│ ├── i18n.py
│ ├── estimator.py # 实时参数量/内存估算(与 Config 同步)
│ └── tabs/
│ ├── config_tab.py # 已合并 Designer,右侧实时估算面板
│ ├── train_tab.py # 拥有 data_path(不再属于 Config)
│ ├── pipeline_tab.py # 6 阶段,每段独立 data_path
│ ├── inference_tab.py
│ ├── benchmark_tab.py # Markdown 表 + gr.BarPlot
│ ├── autotune_tab.py # 持久化日志 + Apply Best → Config
│ └── iomon_tab.py # MetricsBus 实时仪表盘
├── chronos_app.py # Web UI 入口
├── train_chronos.py # Stage 1 入口
├── train_chronos_sft.py # Stage 2 入口
├── train_chronos_dpo.py # Stage 3 入口
├── train_chronos_orpo.py # Stage 4 入口
├── train_chronos_grpo.py # Stage 5 入口
├── train_chronos_distill.py # Stage 6 入口
├── tools/
│ ├── compare_minimind_chronos.py # v1 (M1+M2)
│ ├── compare_minimind_chronos_v2.py # v2 (M3+M4)
│ └── compare_minimind_chronos_v3.py # v3 (含 6 阶段 + HF roundtrip + 后端报告)
├── tests/
│ ├── test_smoke.py # 18 个单元测试
│ ├── test_smoke_cuda.py # 仅 CUDA 主机执行
│ └── fixtures/ # tiny_pretrain / tiny_sft / tiny_dpo / tiny_grpo
├── docs/
│ └── vllm_integration.md
├── pyproject.toml
└── README.md / README_zh.md / THIRD_PARTY_NOTICES.md
timeline
title Project Chronos 交付路线
Phase 1
: LookaheadRouter
: 时间局部性正则
: 路由概率采集路径
Phase 2
: 异步 IO 引擎
: SSD / RAM / VRAM 三级存储
: 共激活专家聚簇
Phase 3
: Hybrid MLA + SlidingWindow attention
: PrefillScheduler
: 双层路由
Phase 4
: MLX 原生后端
: Web UI 和 CLI
: Optuna 搜索
: 开源发布
M1-M3
: cluster-aware safetensors 存储
: 真实监督 lookahead loss
: per-expert event 双流解码
M4-M6
: SFT / DPO / ORPO / GRPO 训练器
: Router KL anchor
: HF IO、vLLM adapter、多后端调度
: Stage 6 蒸馏和可插拔 reward
: Web UI v2、Benchmark 图表、IO Monitor
Next
: 大规模激活语料训练 IntentClassifier
: 7B+ checkpoint 基准测试
: vLLM worker path mask 注入
: 真正的 Vulkan / OpenCL kernel
mindmap
root((Chronos 架构创新面))
预测式路由
IntentClassifier
Prompt 级热点专家先验
带预算的专家集合预测
LookaheadRouter
每 token 未来路由预测
真实未来路由 soft-target CE
存储感知 MoE
Clustered safetensors
.ctsr 专家聚簇打包
manifest 驱动 mmap
三级缓存
NVMe SSD
Pinned RAM
VRAM / 统一内存
柔性兜底
共享专家始终驻留
cache miss 不硬阻塞
Decode 流水线
Prefill-time expert loading
AsyncPrefetcher 队列
H2D stream
per-expert CUDA event
训练栈
Pretrain
SFT
DPO
ORPO
GRPO
Distill
Router KL anchor
部署
HF safetensors IO
AutoModel 注册
vLLM adapter
CPU / CUDA / MPS / MLX 调度
Web UI 和 CLI
@misc{chronos2026,
title = {Project Chronos: Prefill-Time Expert Loading and Dual-Layer Routing
for Zero-Stall On-Device MoE Inference},
author = {Fona and Project Chronos Contributors},
year = {2026},
url = {https://github.com/FonaTech/Project_Chronos}
}Project Chronos 基于 jingyaogong 的 minimind(Apache-2.0)构建。完整归属见 THIRD_PARTY_NOTICES.md。
Apache 2.0 — see LICENSE