waylong/NeoQuasar-Kronos-Tokenizer-2k-NPU
模型介绍
文件和版本
Pull Requests
讨论
分析

NeoQuasar/Kronos-Tokenizer-2k + Kronos-mini — Ascend NPU 适配

NPU Ascend910 #NPU

模型简介

NeoQuasar/Kronos-Tokenizer-2k 是 Kronos 金融 K 线基础模型的分层离散 Tokenizer（2k codebook），单独无法完成预测，必须与 NeoQuasar/Kronos-mini 联合使用。Kronos 采用两阶段框架：① Tokenizer 将连续 6 维 OHLCV 量化为分层离散 token（s1_bits=10, s2_bits=10, group_size=5）；② 4 层 Decoder-only Transformer 自回归预测。Kronos-mini 为 4.1M 参数最小版，支持 max_context 2048，适配 5min K 线零样本预测。

  • Hugging Face Tokenizer: NeoQuasar/Kronos-Tokenizer-2k Revision 26966d0035065a0cae0ebad7af8ece35bc1fb51c (config.json 301B, model.safetensors 15.8 MB, 96 tensors)
  • Hugging Face Model: NeoQuasar/Kronos-mini Revision f4e68697d9d5aed55cef5c96aabc3376bcad9f81 (model.safetensors 16.4 MB)
  • 架构: KronosTokenizer (d_in=6, d_model=256, n_enc=4, n_dec=4, 4 heads) + Kronos (d_model=256, n_layers=4, 4 heads, DualHead, DependencyAwareLayer)
  • 适配说明：Tokenizer 与预测模型均显式在 npu:0 运行，无静默 CPU fallback；KronosPredictor 统一预处理、归一化、时序编码与自回归解码；torch.npu.synchronize() 计时

重要约束: Tokenizer-only checkpoint 不是完整预测模型（见 time-series-models.md），本仓以联合推理验证 Tokenizer 在 NPU 的真实可用性，符合 Skill 要求的 tokenizer+model 联合验收对象。

数据契约

  • 任务: 时间序列预测 forecasting (K 线 6 维点预测, 回归类)
  • 输入: DataFrame 6 列 ["open","high","low","close","volume","amount"] + timestamps，示例取 sample_data.csv 前 400 行历史
    • 输入 shape [1, 400, 6] (单序列批), 历史时间戳 2024-06-18 11:15:00 -> 2024-06-28 14:00:00 频率 5min
    • 数值范围: close mean 10.96 std 0.16 (400 行窗口), volume mean 1111, amount mean 1.22M
    • 预处理: per-series mean/std 归一化 (x-mean)/(std+1e-5), clip 到 [-5,5], 时间戳展开为 ["minute","hour","weekday","day","month"] 5 维
    • 轴顺序: [B, T, C] 时间在第1维, 通道在第2维; 严格按 Kronos 官方 KronosPredictor 定义, 禁止转置
  • 输出: pred_df 6 列预测, shape (60, 6) (pred_len=60), 索引为未来 y_timestamp 2024-06-28 14:05:00 -> 2024-07-01 15:00:00, 已逆归一化还原为原始量纲, dtype float32/64
    • 样例 (NPU, T=1.0, top_k=1, top_p=1.0, sample_count=1, seed=42): 
      2024-06-28 14:05:00  open 10.815447 high 10.827702 low 10.806034 close 10.816590 volume 443.13 amount 462779.5
2024-06-28 14:10:00  open 10.816620 high 10.827589 low 10.805984 close 10.815671 volume 443.36 amount 464087.3
2024-06-28 14:15:00  open 10.816676 high 10.827069 low 10.805724 close 10.815144 volume 432.36 amount 454280.6
  • 数据来源: sample_data.csv 599 行 7 列 timestamps,open,high,low,close,volume,amount 来自 Kronos 官方示例 (XSHG 5min); 若本地不存在则自动合成等分布 600 行 K 线 (seed 42, base 11.0)
  • 缩放: 归一化参数按输入窗口计算, 预测后 * (std+1e-5) + mean 逆变换; 报告未 inverse transform 需说明, 本仓已还原

环境

  • 硬件: Ascend 910B x2 (npu:0 Ascend910_9362 Phy-ID 10 Bus 0000:0A:00.0, npu:1 Phy 11); 单卡 64GB HBM, 验证设备 npu:0
  • 驱动/CANN: npu-smi 25.5.5 / CANN 8.5.1, torch.npu.is_available()==True device_count=2 device_name=Ascend910_9362
  • 软件: Python 3.11.14, torch 2.9.0+cpu, torch_npu 2.9.0.post1+gitee7ba04, pandas 3.0.2, numpy 1.26.4, huggingface-hub 0.34.3, safetensors 0.6.2, einops 0.8.1, tqdm 4.67.1
  • 依赖见 requirements.txt, 安装: pip install -r requirements.txt
  • 模型加载: 本地 models/Kronos-Tokenizer-2k 与 models/Kronos-mini 优先, 缺失时 snapshot_download 自动拉取指定 revision (需网络)

安装

pip install -r requirements.txt
# 如已有本地权重可跳过下载，否则首次运行会自动从 HF Hub 拉取
# Tokenizer: NeoQuasar/Kronos-Tokenizer-2k (26966d0...)
# Model: NeoQuasar/Kronos-mini (f4e6869...)

NPU 推理

python inference.py

关键日志（npu:0 同步计时）：

============================================================
Model: NeoQuasar/Kronos-mini (revision f4e68697d9d5aed55cef5c96aabc3376bcad9f81)
Tokenizer: NeoQuasar/Kronos-Tokenizer-2k (revision 26966d0035065a0cae0ebad7af8ece35bc1fb51c)
Device: npu:0 Max context: 512 Lookback: 400 Pred_len: 60
[ENV] torch 2.9.0+cpu, torch_npu 2.9.0.post1+gitee7ba04
[ENV] npu_available=True, device_count=2, device_name=Ascend910_9362
[ENV] npu-smi exit 0
[DEVICE] Set NPU device 0, current device 0
[VERIFY] Dummy tensor device: npu:0 dtype=torch.float32
[LOAD] Loading tokenizer from .../models/Kronos-Tokenizer-2k
[LOAD] Loading model from .../models/Kronos-mini
[VERIFY] Tokenizer first param device: npu:0 dtype=torch.float32
[VERIFY] Model first param device: npu:0 dtype=torch.float32
[INFO] Total params: tokenizer+model weight bytes ~30.77 MB
[DATA] loaded rows=599 columns=['timestamps', 'open', 'high', 'low', 'close', 'volume', 'amount']
[INPUT] lookback=400 pred_len=60
[INPUT] x_timestamp 2024-06-18 11:15:00 -> 2024-06-28 14:00:00
[PREDICTOR] max_context=512 device=npu:0
[BENCH] Warming up and benchmarking...
 100%|██████████| 60/60 [00:00<00:01, 95.73it/s]
[PERF] first run latency 880.31 ms (sync) device=npu:0
[RESULT] pred_df shape (60, 6)
[VERIFY] finite values: True
[PERF] stable avg 461.80 ms p50 463.27 p90 465.65 min 451.97 max 469.05 windows/s 2.17
[OUTPUT] forecast shape (60, 6) dtype float32
[OUTPUT] forecast [0,:5] [10.815447 10.827702 10.806034 10.81659 443.13434]
[VERIFY] Model param device after inference: npu:0
[VERIFY] Tokenizer param device after inference: npu:0
[DONE] NPU inference completed successfully device=npu:0

inference.py 默认完成真实 npu:0 推理, 异常返回非零; 计时前后均 torch.npu.synchronize(), 首轮编译与稳定阶段分开报告。

真实结果

  • 输入: 400 步历史 (open/high/low/close/volume/amount), 预测 60 步
  • 输出: 6 维点预测, 有限值 True, 无 NaN/Inf
  • 预测统计 (NPU, 60 步):
    • open mean 10.9004 std 0.0439 min 10.812 max 10.936
    • close mean 10.8997 std 0.0437
    • volume mean 573.88 std 805.44 min 137.46 max 5623.34
    • amount mean 615611 max 6,038,692
  • 性能: 首轮含编译 880.31 ms, 稳定 10 次平均 461.80 ms (p50 463.27, p90 465.65), 吞吐 2.17 windows/s (batch=1, 每窗 400->60, 约 868 rows/s), 峰值 HBM 约 3105 MB (npu-smi)
  • 可复现: SEED=42, T=1.0, top_k=1 (greedy), top_p=1.0, sample_count=1, determinism 受浮点误差影响但 top_k=1 保证最大一致性

一致性

同一权重、输入、预处理、dtype=FP32、eval 模式、seed=42、固定窗口, CPU vs NPU:0 对比原始 forecast (60,6):

python scripts/compare_outputs.py --cpu cpu_forecast.npy --npu npu_forecast.npy --task regression --atol 1e-4 --rtol 1e-3

结果:

{
  "task": "regression",
  "cpu_shape": [60, 6],
  "npu_shape": [60, 6],
  "finite": true,
  "atol": 0.0001,
  "rtol": 0.001,
  "max_abs_error": 2.0,
  "mean_abs_error": 0.04049,
  "passed": true
}
  • 按列分解: open/high/low/close max diff <1e-6, volume max 0.0014, amount max 2.0 (相对误差 3.7e-06, 属 1e-3 容限内, 因量纲大 6e5, 绝对差 2 相对可忽略)
  • 结论: 全部 np.allclose 通过, 无 CPU fallback, 日志已打印 Model param device: npu:0

单窗口 smoke consistency, 非全数据集 benchmark; 量化阈值结合 FP32 与自回归采样误差确定

性能

  • 计时含 torch.npu.synchronize(), 首轮=编译+推理, 稳定=10 次平均, 预热已包含在首轮
  • 环境: 单卡 npu:0 Ascend910_9362, batch=1, context 400, horizon 60, 6 通道, FP32
  • 结果: 首轮 880.31 ms, 稳定 avg 461.80 ms p50 463.27 p90 465.65 min 451.97 max 469.05, 2.17 windows/s
  • 显存: HBM 3105/65536 MB (npu-smi), 权重 30.77 MB
  • 预处理 (CSV 读取+归一化+时间编码) 计时包含在端到端内, 约 <10 ms, 逆缩放已包含

证据图

图片由 xterm.js 根据本次真实日志生成 (非手工绘图), 固定提示符 atomgit@pod-a94f8701860f4700b161b00e290de466:~$, 展示标签不证明服务器身份; 已脱敏 token 认证头。

agent_workflow 侦察、数据检查、下载、加载、NPU 推理、一致性、benchmark 与提交校验的完整工作流日志

npu_device_call npu-smi info、NPU 可用性、设备名、模型参数 device/dtype、输入输出 Tensor device 的设备调用证据

model_result 默认 python inference.py 的真实完整输出, 包括 schema/窗口摘要、真实预测、同步耗时和状态

限制

  • Tokenizer 单独无法预测, 必须与 Kronos-mini (或 small/base) 联合; 单独提交 tokenizer 为 component-only, 本仓已联合验收
  • 仅验证单窗口 (400->60) 点预测, 未覆盖长上下文 2048、批量并行 predict_batch、概率采样 (T>1, top_p<1, sample_count>1) 及量纲外数据
  • 精度为 FP32; BF16/FP16 未验证, 阈值需重新校准; 未做量化 (msModelSlim) 与图优化
  • 自定义 op 无需改写, 全部为标准 PyTorch; 若遇不支持 op 需改写为等价实现并记录, 本次未触发
  • 数据频率固定 5min, 跨市场/品种零样本泛化未评估; 需 fit() 的场景已区分 CPU 训练与 NPU 推理
  • 单卡资源下 Kronos-small/base (24M/102M) 未在本环境验证, 显存与编译时间将显著增加