Chronos-2 on Ascend NPU
Amazon Chronos-2 时间序列预测模型在华为昇腾 NPU 上的推理适配与性能优化。
快速开始
pip install chronos-forecasting
# 推理
python3 -c "
from chronos import Chronos2Pipeline
import torch, numpy as np
pipe = Chronos2Pipeline.from_pretrained('model_files/chronos-2', local_files_only=True)
pipe.model = pipe.model.to('npu:0').to(torch.bfloat16)
data = torch.tensor(np.random.randn(32, 1, 512), dtype=torch.bfloat16).to('npu:0')
forecasts = pipe.predict(data, prediction_length=64)
print(forecasts[0].shape) # (21 quantiles, 64 horizon)
"性能
基线要求：NPU 吞吐量 ≥ 500 series/s，对比 CPU 基线 (6.9 series/s) 加速 ≥ 72×
测试条件：context_length=512，prediction_length=64，Ascend910_9362
NPU vs CPU (bs=32, ctx=512)
| 配置 | 延迟 | 吞吐量 | 加速比 |
|---|---|---|---|
| CPU FP32 (baseline) | 1.163s | 6.9 series/s | 1× |
| NPU BF16 (LN融合+dropout=0, bs=32) | 0.0269s | 1,189 series/s | 173× |
| NPU BF16 (LN融合+dropout=0, bs=512) | 0.1571s | 3,259 series/s | 474× ← 峰值 |
Batch 摸高 (NPU BF16 优化配置 — 最佳结果)
| Batch Size | 延迟 | 吞吐量 | 加速比 vs CPU |
|---|---|---|---|
| 32 | 0.0269s | 1,189 /s | 173× |
| 64 | 0.0286s | 2,238 /s | 334× |
| 128 | 0.0403s | 3,174 /s | 474× |
| 256 | 0.0773s | 3,310 /s | 494× |
| 512 | 0.1571s | 3,259 /s ← 峰值 | 474× |
| 1024 | 0.3145s | 3,256 /s | 473× |
NPU 峰值吞吐：3,259 series/s (batch=512) = 474× vs CPU (单核 6.9/s)
Context 摸高 (NPU BF16 优化配置, bs=32)
| Context Length | 延迟 | 吞吐量 |
|---|---|---|
| 128 | 0.0325s | 983.6 series/s |
| 256 | 0.0322s | 993.7 series/s |
| 512 | 0.0329s | 971.8 series/s |
| 1024 | 0.0330s | 969.5 series/s |
| 2048 | 0.0369s | 866.8 series/s |
| 4096 | 0.0720s | 444.5 series/s |
优化记录
| 优化手段 | 效果 |
|---|---|
| BF16 混合精度 | 基础加速 |
| TASK_QUEUE_ENABLE=1 | 流并行下发 |
| Chronos2LayerNorm → npu_rms_norm (37x) | +6% 唯一有效算子替换 |
| 推理时 Dropout 禁用 (p=0) | 消除无用计算 |
| NPU_ASD_ENABLE=0 | 消除诊断开销 |
| PYTORCH_NPU_ALLOC_CONF | 内存分配优化 |
| PER_STREAM_QUEUE=1 | 无额外收益 |
| CPU_AFFINITY_CONF=1 | 无额外收益 |
| TASK_QUEUE_ENABLE=2 | 无额外收益 (1.01x) |
| NPUGraph | ❌ 流水线中 host-device 拷贝不支持 |
| RoPE→npu_apply_rotary_pos_emb | ❌ -15% (BNSD↔BSH转置开销) |
| FFN→npu_ffn | ❌ -17% (小层无收益) |
| AUTO_TUNE_MODE=1 | ❌ -17% (编译开销) |
| 全组合 | ❌ -19% (累积开销) |
验证结论
| 检查项 | 结果 |
|---|---|
| NPU 推理跑通 | ✅ 通过 |
| 精度误差 < 1% | ✅ 通过 (0.64%) |
| 性能基线 > 500 series/s | ✅ 通过 (3,259 series/s) |
| 整体状态 | ✅ 全部通过 |
精度
NPU BF16 与 CPU FP32 对比：平均相对误差 0.64%（修剪后 0.39%，< 1% ✅）
结构
chronos-2/
├── model_files/chronos-2/ # 模型权重 (LFS)
│ ├── config.json
│ └── model.safetensors
├── benchmark.py # 全面基准测试脚本
├── deep_optimize.py # 深度算子优化测试
├── deep_optimize_round2.py # 第二轮摸高测试 (LN融合最优配置)
├── deep_optimize_round3.py # 第三轮摸高测试 (NPU亲和算子替换)
├── deep_optimize_round4.py # 第四轮摸高测试 (Pipeline bypass)
├── results/ # 全部结果 JSON
├── SKILL.md # 优化技能文档
└── README.md