Chronos-2-mini — 华为昇腾 NPU 推理优化

amazon/chronos-2-mini 时序预测模型在 Ascend 910B4 NPU 上的推理优化工程。

模型信息

项目	说明
模型	Chronos-2-mini (T5 encoder-decoder, 201M 参数)
硬件	1×Ascend 910B4 (32GB HBM)
框架	PyTorch 2.7.1 + torch_npu 2.7.1
输入	时间序列 context_length=512
输出	预测分布 prediction_length=64, 20 样本
基线	NPU FP32 eager 模式

优化结果

Round	优化方案	延迟 (ms)	吞吐 (sps)	加速比
0	Baseline FP32	1584	0.63	1.00×
1	TQE=2 + TE=1	1163	0.86	1.36×
2	+ tcmalloc LD_PRELOAD	1118	0.89	1.42×
3	torch.compile (npugraph_ex)	636	1.57	2.49×
—	CPU Baseline (1步×1样本) ⚠️	42731	0.02	0.04×

⚠️ CPU 基线：Chronos-2-mini (201M T5) 单次前向 (1步×1样本) 耗时 42,731ms。完整 benchmark (64步×20样本) 预计需 ~15.5 小时，无法在合理时间内完成。NPU 单步加速比约 35,000×。

最佳配置

export TASK_QUEUE_ENABLE=1
export TE_PARALLEL_COMPILATION=1
export LD_PRELOAD=/usr/lib64/libtcmalloc.so.4

python3 scripts/benchmark.py --mode graph --iter 50 --warmup 15

torch.compile(model, backend="npugraph_ex", fullgraph=False) 将 encoder+decoder 编译为 GE 执行图，减少算子调度开销，实现 2.49 倍吞吐提升。

文件结构

scripts/
├── benchmark.py           # 统一性能基准测试
├── inference_baseline.py  # 基线推理
├── npu_patch.py           # Chronos NPU 适配补丁
├── accuracy.py            # 精度验证
├── evaluation.py          # M4 指标评测
├── logger.py              # 实验日志
├── test_cases.py          # 测试数据生成
├── optimization_loop.py   # 自动化优化循环
├── profile_ops.py         # 算子级 Profiling
├── profile_l1.sh          # L1 硬件指标采集
└── analyze_l1.py          # L1 数据分析

experiments/               # 各轮实验结果
├── 00_baseline/           # FP32 基线
├── 01_amp_fp16/           # FP16 试验
├── 02_task_queue/         # TQE 调优
├── 03_graph_mode/         # 图模式 (最优)
└── 04_batch/              # Batch 推理

verification/              # 精度验证
├── baseline/outputs/      # 基线预测
├── optimized/outputs/     # 优化后预测
└── reports/               # 验证报告

results/                   # 汇总指标 JSON
data/                      # 测试数据

关键发现

图模式 (torch.compile) 收益最大 — GE 执行图融合减少算子调度，2.49× 提升
FP16 在 910B4 上反效果 — 部分算子 FP16 kernel 开销更大，比 FP32 慢 30%
TQE=1（非 TQE=2） 与图模式兼容，不能同时使用 TQE=2
tcmalloc 额外带来 ~5% 小幅度收益
Profiling 确认 100% AI Core 时间为 MatMul (MatMulV2 53% + BatchMatMulV2 47%)，瓶颈清晰

数据来源

模型权重

模型	下载地址	说明
amazon/chronos-2-mini	https://huggingface.co/amazon/chronos-2-mini	主模型 (201M, T5 encoder-decoder)
autogluon/chronos-2-small	https://huggingface.co/autogluon/chronos-2-small	精度验证用

国内镜像加速：

export HF_ENDPOINT=https://hf-mirror.com

代码通过 ChronosPipeline.from_pretrained("amazon/chronos-2-mini") 自动下载，默认缓存到 /tmp/chronos-model 或 ~/.cache/huggingface/。

测试数据集

来源	说明	用途
`autogluon/chronos_datasets` (HF)	M4 全部子集 (Yearly/Quarterly/Monthly/Weekly/Daily/Hourly) + Monash	精度评估 (CRPS/MASE/sMAPE)
`scripts/test_cases.py`	人造合成时序 (100×512)	性能基准测试

合成数据生成（无外部依赖）：

python scripts/test_cases.py --num 100 --length 512 --seed 42

精度验证

相对精度 (优化后 vs 基线)

对比同一模型在 NPU 上不同优化模式的输出差异：

python scripts/accuracy.py \
  --baseline verification/baseline_preds.npy \
  --optimized verification/optimized_preds.npy \
  --threshold 1e-11

结果：MSE = 0.0 ✅ — 所有代码级优化（TQE、tcmalloc、torch.compile）均不影响模型精度。

M4 Benchmark 指标评估

在 CPU 上运行 Chronos-2 论文协议评估：

export HF_ENDPOINT=https://hf-mirror.com

python scripts/evaluation.py \
  --model autogluon/chronos-2-small \
  --device cpu \
  --max_series 50

指标	宏观平均	说明
CRPS	86.57	分位数概率评分（越低越好）
Rel-CRPS	0.58	相对朴素基线改进 42%
MASE	2.15	均值缩放误差
sMAPE	—	对称平均绝对百分比误差

详细报告见 verification/reports/final_verification_report.md。

环境要求

# CANN
CANN 8.5.0+

# Python 依赖
torch>=2.7.1
torch_npu>=2.7.1
transformers>=5.8.0
chronos-forecasting>=2.0.0
numpy