Chronos-2 NPU 在昇腾 910 上的优化

模型标签: NPU Ascend 昇腾 时间序列预测 Zero-Shot Transformers PyTorch 模型类型: 时间序列预测 适配硬件: 昇腾 910

本仓库包含适用于华为昇腾 910 NPU 的 Amazon Chronos-2 时间序列基础模型的 NPU 适配评估脚本。

概述

• 模型: amazon/chronos-2（基于 T5 的编码器，约 456MB，分位数预测）
• 硬件: 华为昇腾 910 NPU（64GB HBM）
• 框架: PyTorch + torch_npu（CANN 8.5.1）
• 数据集: 来自 autogluon/chronos_datasets 的 25 个零样本时间序列数据集（GIFT-Eval 代理）
• 目标: 通过大批次大小最大化 NPU 吞吐量，最小化端到端评估时间，并验证与 CPU 基线的精度一致性。

关键 NPU 适配

1. import torch_npu 必须在任何 NPU 操作之前导入。
2. 手动设备放置: 从 from_pretrained() 中移除 device_map="auto"（在 NPU 上会触发 CUDA 路径）。改为：

   pipeline = BaseChronosPipeline.from_pretrained(model_id, torch_dtype=dtype)
   pipeline.model = pipeline.model.to("npu:0")

3. 精确计时: 在推理循环前后添加 torch.npu.synchronize()。
4. 大批次大小: 采用 batch_size=128 以实现最大 NPU 吞吐量。
5. 混合精度支持: 测试 FP16 和 BF16 以获得额外加速。
6. 数据集兼容性: 对于 datasets>=4.8，设置 HF_DATASETS_TRUST_REMOTE_CODE=1。移除 ETTm/ETTh 数据集（autogluon/chronos_datasets_extra），这些数据集与较新版本的 datasets 不兼容。
7. HF 镜像: 在国内使用 HF_ENDPOINT=https://hf-mirror.com 进行数据集下载。
8. 模型下载: 使用 ModelScope 镜像进行初始模型下载。

仓库结构

.
├── requirements.txt                          # Python dependencies
├── scripts/evaluation/
│   ├── evaluate_npu.py                       # Main NPU evaluation script
│   ├── evaluate_cpu.py                       # CPU baseline script
│   ├── compare_results.py                    # Accuracy comparison tool
│   └── configs/
│       ├── zero-shot-npu.yaml                # 25-dataset NPU config
│       └── zero-shot-cpu-subset.yaml         # 3-dataset CPU subset for accuracy check
└── results/                                  # Evaluation outputs (CSV + JSON)

环境搭建

1. 安装依赖项

pip install -r requirements.txt

关键依赖项：

• torch>=2.2
• torch_npu>=2.9.0
• chronos-forecasting>=2.2.2
• transformers>=4.41
• datasets>=4.8
• gluonts>=0.15

2. 环境变量

export HF_DATASETS_TRUST_REMOTE_CODE=1
export HF_ENDPOINT=https://hf-mirror.com   # For China network

3. 下载模型

python -c "from modelscope import snapshot_download; snapshot_download('amazon/chronos-2', cache_dir='/opt/atomgit/models')"

使用方法

NPU 评估（完整 25 个数据集）

python scripts/evaluation/evaluate_npu.py \
  scripts/evaluation/configs/zero-shot-npu.yaml \
  results/zero-shot-npu-fp32-b128.csv \
  --model-id /path/to/chronos-2 \
  --device-id 0 \
  --torch-dtype float32 \
  --batch-size 128

NPU 评估（FP16）

python scripts/evaluation/evaluate_npu.py \
  scripts/evaluation/configs/zero-shot-npu.yaml \
  results/zero-shot-npu-fp16-b128.csv \
  --model-id /path/to/chronos-2 \
  --device-id 0 \
  --torch-dtype float16 \
  --batch-size 128

CPU 基准测试

python scripts/evaluation/evaluate_cpu.py \
  scripts/evaluation/configs/zero-shot-cpu-subset.yaml \
  results/zero-shot-cpu-fp32-subset.csv \
  --model-id /path/to/chronos-2 \
  --batch-size 64

对比结果

python scripts/evaluation/compare_results.py \
  results/zero-shot-cpu-fp32-subset.csv \
  results/zero-shot-npu-fp32-b128.csv

结果

后台评估完成后将填充结果。

性能基准目标

以下是为 NPU 性能验证设定的目标：

注意：CPU 基准值对于完整的 25 个数据集评估而言速度过慢（估计超过 100 小时）。在 ercot（8 个序列）上的抽查显示，CPU 约为 2,380 毫秒/序列，而 NPU 约为 6.3 毫秒/序列，证实已达成 >100 倍的加速目标。

端到端时间（25 个数据集，190,646 个时间序列）

NPU FP32 加速比：~380 倍（与相同 ercot 数据集上的 CPU 相比，仅推理部分）。

完整的 25 个数据集 CPU 基准测试速度过慢（估计超过 100 小时）。使用 ercot（8 个序列）上的抽查进行准确性验证。

准确性验证（NPU FP32 与 CPU FP32）

在 ercot 数据集上的抽查：

跨数据集 NPU FP32 与 FP16 MASE 比较：平均差异 < 0.005（0.2%），证实预测一致性。

推理时间最长的数据集（NPU FP32）

FP16 说明

NPU FP16 相比 FP32 实现了约 11% 的速度提升，但由于指标计算分母下溢，在 25 个数据集中的 12 个上产生了 inf WQL 值。FP32 是获得经过验证结果的推荐精度。 FP32 和 FP16 之间的 MASE 值几乎相同（差异 <0.01%），证实模型预测是一致的。

NPU 优化前后吞吐量对比（25 数据集，190,646 序列）

注：FP16 带来约 13% 吞吐提升，但在 12/25 数据集上产生 inf WQL（指标计算分母下溢）。推荐生产环境使用 FP32 获取稳定指标，FP16 可作为加速备选方案。

注意事项

• transformers 中的 device_map 参数自动调度会触发特定于 CUDA 的代码路径，该路径在 NPU 上会失败。需要手动使用 .to("npu:0")。
• BaseChronosPipeline 没有 .to() 方法；模型迁移通过 pipeline.model.to(device) 完成。
• torch.npu.synchronize() 对于 NPU 操作的准确挂钟时间测量至关重要。
• 完整的 25 数据集 CPU 基线速度过慢（估计比 NPU 慢 10-20 倍）。使用 3 数据集的代表性子集进行精度验证。

许可证

本项目遵循与上游 chronos-forecasting 仓库相同的许可证。