mizorewww

Laya-MLX

把 Laya 的預訓練權重移植到 MLX,在 Apple 晶片上本地推理,不需要 PyTorch 或雲端 API。獨立埠,非 Convai Innovations 官方釋出。

核驗於 2026-10-05

Laya MLX 玩貪吃蛇 —— 真實決策,原速播放

開放權重的型別化決策,在 Apple Silicon 上原生執行。

英文短型別化決策端到端的中位延遲為 13.4 ms;使用 multilingual checkpoint 時是 7.4 ms。0 個輸出 token。 本地 MLX 推理,不需要 PyTorch、Transformers 執行時或雲端 API。

中文 · 基準 · Snake 演示 · Hugging Face 權重

這段 GIF 是一次真實本地 Snake 執行的原速渲染。每一步都會呼叫 Laya;可見的迴圈安全層可以糾正不安全的提議。上面的延遲數字來自單獨的單問題 API 基準,不是三問題 Snake 迴圈的幀時間。觀看 30 秒 MP4 · Snake 速度與穩定性。

快速開始

pip install laya-mlx
import laya_mlx as laya

agent = laya.load("aac6fef/laya-mlx")
result = agent.predict(
    "I was billed twice. Please refund the duplicate.",
    {
        "department": {
            "type": "choice",
            "instructions": "Who should handle this?",
            "criteria": ["billing", "technical", "sales"],
        }
    },
)
print(result["answers"]["department"])

Apple Silicon,Python 3.11+,macOS 14+。首次載入會下載 checkpoint;之後的推理完全在本地。測量環境是 macOS 27.2、Python 3.12.13 和 MLX 0.32.2。該 MLX 版本提供 macOS 14、15 和 26 的 wheel;本地安裝器選的是 26 的 wheel。這臺機器上沒有測試更早的支援 macOS 版本。

執行終端 demo:

pip install 'laya-mlx[demo]'
hf download aac6fef/laya-multilingual-mlx
laya-snake

離線 demo 前先下載一次。使用至少 104 × 35 格的終端。空格暫停,↑/↓ 改變速度,R 重置,Q 退出。laya-snake --max-speed 讓每一步都重新決策、不做節流。錄製、操作與各指標的準確含義。

laya-snake --optimize --max-speed 啟用經過測試的編譯與字首複用路徑:在成對 M3 Max 測試中,2,400 步達到 75.40 步/秒,零死亡、2 次可見的安全乾預。這比同一輪次中的 eager 對照快約 6.5%。玩法、效能與正確性證據。

在 M3 Max 上的效能

FP16,端到端 Laya 421M Multilingual 322M
單個短問題,P50 13.42 ms 7.39 ms
單個短問題,P95 13.92 ms 7.79 ms
50 問題吞吐 146.8 q/s 395.0 q/s
單個短問題的 MLX 峰值分配 943.6 MiB 687.6 MiB

M3 Max,40 個 GPU 核心,128 GiB 記憶體。計時包含提示準備、分詞、張量、同步推理、校準與結果格式化;模型載入不計入。50 問題的測量使用 batch_size=64;API 預設是 16。不同的長度、問題數量和執行時條件會改變延遲。完整方法與每一個計時樣本。

埠保真度: 三個 checkpoint 在 FP32 和 FP16 下各 63/63 個驗證問題上都與上游選出的答案一致 —— 378/378 次比較。每種配置還通過了 100 次重複的有限、確定性呼叫,測得的活動記憶體增長為零。這衡量的是這些固定樣本上的保真度,不是每一個可能問題上的準確率。機率誤差與驗證。

為什麼用型別化決策?

軟體常常需要一個選擇、一個評分或一個機率。Laya 在一次雙向前向傳播裡回答這些受約束的問題,無需逐 token 解碼,也不生成 JSON。

state + typed question → bidirectional encoder → decision heads → probabilities
  • choice:在具名選項上的機率。
  • score:在有序評分檔位上的機率及其期望分值。
  • noul:某個命題為真的機率 P(true)。

問題行相互獨立地批處理。它們的雙向編碼器表示同時取決於 state 和 question;這個執行時並不聲稱把 state 編碼一次就能在任意問題之間複用其隱藏狀態。

編碼器、決策 Transformer、評分頭與動作頭全部在 MLX 中執行。分詞使用 Hugging Face 的 Rust tokenizer。原始的預訓練權重、問題格式、校準與輸出 schema 都保留。這是獨立的 MLX 埠,不是 Convai Innovations 的官方釋出。

支援的 checkpoint

模型 編碼器 參數量 上下文上限 用途
convaiinnovations/laya ModernBERT-large 421M 512 英文
convaiinnovations/laya-multilingual mmBERT-base 322M 1,024 多語言輸入
convaiinnovations/laya-typed-decisions ModernBERT-large 421M 1,024 上游的 typed-decisions 工作流

上下文包含指令、選項和 state。三者都使用原始權重、提示格式、溫度校準和輸出 schema。本倉庫提供推理與轉換;RLCD 訓練與微調仍在上游專案裡。這是獨立埠,不是 Convai Innovations 的官方釋出。

預轉換的 FP16 checkpoint 釋出在 Hugging Face 上:

用 laya.load("aac6fef/laya-mlx") 直接載入它們,或使用上面的原始 checkpoint ID。每個已釋出的 checkpoint 都包含其模型卡、驗證結果、來源、許可和檔案校驗和。全部 36 個已釋出檔案都通過了嚴格的遠端校驗和驗證;固定的 revision 與權重雜湊記錄在 hub-publication.json 中。

開發版安裝

gh repo clone mizorewww/laya-mlx
cd laya-mlx
uv sync --extra demo
uv run --extra demo laya-snake

或者用 pip install 'git+https://github.com/mizorewww/laya-mlx.git' 安裝最新的 GitHub revision。模型權重單獨下載,不包含在 Git 中。

Python API

import laya_mlx as laya

agent = laya.load("aac6fef/laya-mlx", dtype="float16")
result = agent.predict(
    "I was billed twice. Please refund the duplicate today.",
    {
        "department": {
            "type": "choice",
            "instructions": "Which team should handle this request?",
            "criteria": {
                "billing": "invoices, payments, refunds",
                "technical": "bugs and outages",
                "sales": "new purchases",
            },
        },
        "urgency": {
            "type": "score",
            "instructions": "How urgent is this request?",
            "criteria": ["not urgent", "soon", "critical"],
        },
        "refund": {
            "type": "noul",
            "instructions": "Does the customer ask for money back?",
        },
    },
)
print(result["answers"])

system_one 是 predict 的別名。state 可以是文本、JSON 字典或對話列表。choice 接受一個字典或一個唯一標籤列表;score 返回基於零的期望評分檔位;noul 返回 P(true)。結果保留下游的四位小數舍入、action.act_probability 和 token 用量欄位。

預設精度是 FP16。使用 dtype="float32" 可獲得更接近的數值一致性。即使所選標籤一致,不同精度下的機率也可能略有差異;實測誤差見 BENCHMARKS.md。可以請求 BF16,但它不在已釋出的驗證矩陣內。

跟隨上游 v0.3.5,擬合出的校準溫度在使用前被夾緊到 [0.5, 5.0]:已釋出的 choice:11+ 桶是 0.1006,它會將 logits 銳化約 10 倍,把一個擲硬幣報告成接近確定。checkpoint 的原始值仍可通過 agent.temperature_raw 和 agent.temperature_by_options_raw 獲得,載入時會有 RuntimeWarning 命名每一個被夾緊的桶。

batch_size=16 限制每次前向傳播的問題數;更大的請求會分塊處理。記憶體允許時可以提高它。device="gpu" 或 device="cpu" 顯式選擇裝置;否則使用 MLX 的預設裝置。

對於重複性工作負載,載入 Agent 時可選擇啟用 compile=True、pad_to_multiple=16 和 cache_prompts=True。字首快取上限為 128 個問題,並共享 CPU 側的 state 分詞,而每個問題仍各自做一次編碼器計算。編譯有首次使用成本和形狀特化;填充可能讓某些工作負載變慢。這三個選項預設都關閉。實測的 Snake 消融與用法。

agent = laya.load("./models/laya", dtype="float32", batch_size=32)
# Select one checkpoint inside upstream's bundled repository:
multi = laya.load("convaiinnovations/laya", subfolder="multilingual")
# Pin a Hub revision for reproducibility:
agent = laya.load(
    "convaiinnovations/laya",
    revision="c5d78730f3493e4fe16d61507ef4b78eef7318cf",
)

載入時會校驗每一個參數名和形狀。不支援的編碼器和非預設的 RoPE 縮放會顯式報錯。ModernBERT 的全域性/區域性注意力模式、包含邊界的滑動視窗、區域性與全域性不同的 RoPE 基,以及第一層的歸一化行為都被保留。

語言路由與預設

from laya_mlx import Router, triage_questions

router = Router(dtype="float16", max_loaded=2)
result = router.predict({"message": "发票被重复扣款,请退款。"}, triage_questions())
print(result["routing"])  # multilingual

# Choose the specialized checkpoint explicitly:
result = router.predict(state, questions, task="typed_decisions")

路由、語言啟發式、郵件輔助函式和應用預設都改編自上游。Router(preload=True) 讓三個 checkpoint 全部常駐;支援 attach、preload、unload、顯式 lang= 和顯式 model=。模型生命週期由可重入鎖保護,因此併發執行緒共享一個已載入的 Agent,而不是各自構建副本;推理本身並不序列化。typed-decisions 工作流檢測保持 opt-in。這個埠保留了模型的侷限:英文 checkpoint 不能替代 multilingual checkpoint,置信度也不保證準確率。

無法識別的拉丁字母語言(羅馬尼亞語、波蘭語、捷克語、土耳其語……)僅憑其非英文字母就會路由到 multilingual checkpoint,而不是被默默假定為英文。detect_language(state) 報告依據:language_undecided 和 diacritic_rate,以及 language 和 is_english。

大型選項集的初篩

choice 選項共享一份 head_max_len token 預算,所以一個帶數百個標籤的問題只會給每個標籤留下幾個 token。predict_shortlist 會分別嵌入 state 和每個標籤,按餘弦相似度保留前 k 個,然後在縮減後的集合上跑一次 predict。這是 opt-in:Agent.predict 仍會給它的每一個 criterion 打分。

import laya_mlx as laya

agent = laya.load("aac6fef/laya-mlx")
embed_fn = laya.embed_fn_from_agent(agent)  # mean-pools the loaded encoder; no extra weights
result = laya.predict_shortlist(agent, state, questions, embed_fn, k=20)
print(result["shortlist"])  # which labels were kept, with cosine scores

作為 embed_fn 傳入的專用雙編碼器通常比決策 checkpoint 自帶的編碼器篩得更好。初篩後 choice 上的機率只覆蓋被保留的標籤。

命令列

uv run laya-mlx predict \
  --model aac6fef/laya-mlx \
  --state-file examples/state.json \
  --questions examples/questions.json

uv run laya-mlx predict \
  --model aac6fef/laya-multilingual-mlx \
  --state '发票被重复扣款,请退款。' \
  --questions examples/questions.json

v0.3.5 之後選用的上游修復

執行時選擇性地引入上游 4aa6761(v0.3.23 原始碼樹)中的輸入、路由和郵件修復。這不會加入上游的 batch、長文件、hooks 或 server API。神經架構的保真度仍針對 573e5b6 測試。

  • 按時間順序的對話列表在上下文填滿時保留最新的 token;字串和字典保留開頭。字首快取使用相同的規則。
  • noul 的 criteria 只接受 false/true 鍵(包括 Python 布林鍵)。可選的 labels={"false": "no", "true": "yes"} 會改變展示給模型的詞,而答案仍是 P(true)。非法鍵現在會報錯,而不是被忽略。
  • 非字串指令保留 Unicode。空指令、null 評分檔位和 None state 會觸發呼叫方錯誤;問題相關的錯誤會指出是哪個問題。
  • 每個答案都新增 answer_confidence,即校準後的最大選項機率。現有的 confidence 保留其基於熵的含義(對 choice/score)和最大機率(對 noul)。這兩個欄位都不保證新任務上的準確率。
  • usage 新增 state_tokens、state_tokens_dropped(各問題中最大的丟棄量)、truncated 和 truncated_questions。usage.options 只在某些問題的選項 token 跨度發生碰撞時出現,報告 total、distinct 和 tokens_per_option。這報告的是丟失的區分度;它不會找回這些區分度,也不移除位置偏差。
  • 增量的 Router.preload() 保留已常駐的模型;preload([]) 什麼都不做。空白或語言中立的提示會落到檢測環節,無法確定的拉丁文本遵循 Router(default=...)。檢測會檢查巢狀的字串值和混合文本。
  • 郵件清洗會保留提及保密、致謝收件人或以 From: 開頭的普通請求,同時識別多語言郵件頁尾。

匯出 MLX checkpoint

uv run laya-mlx convert \
  --model convaiinnovations/laya \
  --dtype float16 \
  --output models/laya-mlx-fp16

uv run laya-mlx predict \
  --model models/laya-mlx-fp16 \
  --state-file examples/state.json \
  --questions examples/questions.json

匯出內容包含 model.safetensors、編碼器和 agent 配置、tokenizer 檔案以及 mlx_config.json。已有的輸出目錄絕不會被覆蓋。這是參數名/精度轉換,不是量化或重訓練。源 checkpoint 已經儲存 FP16 權重;選擇 FP32 提高的是算術精度,不是源權重的精度。

測試與基準

uv sync --extra dev --extra reference --extra benchmark --extra demo
source .venv/bin/activate
gh repo clone NandhaKishorM/laya .upstream
git -C .upstream checkout 573e5b62696ba441230cd6be71d593331b5d23af
pytest -q
python -m benchmarks.download
python -m benchmarks.validate --repeats 100
python -m benchmarks.run --iterations 50 --warmup 5
python -m benchmarks.accuracy --per-class 64
python -m benchmarks.report

按順序執行 GPU 測量。單元測試使用小型隨機模型,並包含與 Transformers 和固定上游決策頭的直接比較。真實 checkpoint 驗證會測試分詞、logits、校準機率、重複輸出和活動記憶體增長。基準在每次執行時都用全新的程序跑每個後端/checkpoint,並把每一個計時樣本存到 benchmarks/results。完整報告 解釋了計時的邊界和精度差異。

GitHub Actions 在 macOS arm64 runner 上執行小模型 CPU 測試。完整的 checkpoint GPU 基準在本地測量,不屬於託管的 CI。

效能研究

效能調查同時包含數學分析和獨立的本地實驗:

  • 初步效能研究:實現瓶頸、MLX kernel 排程,以及一個受控實驗計劃。
  • 關於再快 10× 的數學調查:算術預算、條件頻寬下界、真實權重譜、精確複用,以及更小的模型設計。
  • 工程調查:實測的編譯、量化、末頭選擇、自定義 Metal kernel,以及有代表性的矩陣乘法。

experiments/ 包含研究指令碼及其原始測量。已釋出執行時的效能與驗證結果在 BENCHMARKS.md 中;每個實驗性變體都有自己的計時與正確性結果。

當前調查不支援用相同 checkpoint 再獲得普遍的 10× 加速。選定的用例顯示約 1.03–1.08× 的成對中位加速;工程報告給出了不確定區間、量化保真度結果和自定義 Metal kernel 測量。

要準備模型卡和經過驗證的匯出以供釋出,安裝 reference extras 並執行:

python -m scripts.prepare_hub --account YOUR_HF_USERNAME
hf upload YOUR_HF_USERNAME/laya-mlx models/hub/laya-mlx . --exclude '.cache/*'

準備指令碼會把每一個匯出的張量與其原始 FP16 源逐一核對。用同樣的方式上傳另外兩個準備好的資料夾,然後用 hf cache verify REPO_ID --local-dir EXPORT_PATH 檢查遠端檔案。

署名與許可

Apache-2.0;見 LICENSE 與 NOTICE。Laya 及其預訓練權重由 Convai Innovations 和上游貢獻者完成。提示構造、輸出格式化、語言路由、郵件工具和預設改編自 NandhaKishorM/laya 的 573e5b62696ba441230cd6be71d593331b5d23af 提交。神經架構在 MLX 中依據 Laya 和 Hugging Face ModernBERT 重新實現。

維護與釋出

這個專案通過原生 MLX 實現,跟隨上游 Laya 的行為。相容上游的修復優先於獨立的模型變體、服務 API 和額外 demo。這仍然是一個選擇性的埠,不聲稱與上游 API 完全對齊。

釋出時,更新 pyproject.toml、laya_mlx/__init__.py 和 uv.lock 中的版本,然後推送匹配的 vX.Y.Z tag。GitHub Actions 會執行 macOS 測試套件、校驗版本一致性、構建並檢查 wheel 和原始碼分發包、用倉庫的 PYPI_API_TOKEN secret 把它們釋出到 PyPI,並建立 GitHub release。測試或構建失敗會阻止釋出。