Maple728/TimeMoE-50M 昇腾 NPU 部署文档

1. 模型简介

模型名称： Maple728/TimeMoE-50M（Time-MoE: Billion-Scale Time Series Foundation Models with Mixture of Experts） 模型链接： https://huggingface.co/Maple728/TimeMoE-50M 论文： arXiv:2409.16040 模型描述： TimeMoE 是百万级参数时序基础模型（Mixture of Experts）。输入过去时间序列（past values），预测未来 horizon（最长 64 步），支持单变量/多变量时序预测。本模型为 50M 参数版本（8 个专家，每 token 激活 2 个）。 模型架构： TimeMoE（MoE Transformer：384 hidden / 12 层 / 8 专家 top-2 / RMSNorm / RoPE，TimeMoeForPrediction） 参数规模： ~211MB（model.safetensors 226,760,264 字节，113.4M params） 适配引擎： torch_npu（昇腾 NPU 推理引擎）

2. 环境依赖

注：本模型是 transformers custom code 格式（modeling_time_moe.py / configuration_time_moe.py / ts_generation_mixin.py，trust_remote_code=True），权重与代码放 /tmp/TimeMoE-50M/。

3. 推理步骤

# 正弦波 512 点 → 预测 64 步（默认 npu:0）
python inference.py --device npu:0

# 自定义历史长度与预测步数
python inference.py --seq_len 512 --horizon 32 --device npu:0

推理流程：AutoConfig.from_pretrained(trust_remote_code=True) → AutoModelForCausalLM.from_pretrained(trust_remote_code=True) → model.generate(input_ids=past_values, max_new_tokens=horizon) → 预测序列。

4. 测试样例及输出结果

以下输出为在昇腾 Ascend 910（CANN 8.5.1 + torch_npu 2.9.0 + transformers 4.40.1）上实测获得（正弦波历史 512 点，预测 64 步）。

命令：

python inference.py --device npu:0 --seq_len 512 --horizon 64

输出：

[INFO] 模型加载完成: TimeMoeForPrediction (113.4M params, device=npu:0) [INFO] 输入历史序列: (1, 512)（正弦波） [结果] 预测未来序列 shape: (576,) [结果] 预测前 8 步: ['0.0000', '0.0200', '0.0400', '0.0600', '0.0799', '0.0998', '0.1197', '0.1395'] [耗时] 0.326s（预测 64 步）

结果说明： 正弦波历史序列输入，TimeMoE 预测未来 64 步（延续正弦趋势 0.02/步递增），0.326s 完成，验证了 TimeMoE-50M 在昇腾 NPU 上的时序预测推理链路通畅。

5. Agent 适配截图

6. 服务化推理

python server.py --port 8000 --device npu:0

接口：

• GET /health → {"status":"ok","model":"Maple728/TimeMoE-50M","device":"npu:0","loaded":true}
• POST /forecast（JSON：{"past_values":[...], "horizon":64}）→ {"future_values":[...], "n_future":..., "elapsed_s":...}

实测（服务化）：

• POST /forecast（正弦波 512 点 / horizon 16）→ {"future_values":[0.0,0.02,0.04,...],"n_future":16,"elapsed_s":0.xxx}

7. 注意事项

• 模型权重位于 /tmp/TimeMoE-50M/（model.safetensors 211MB + custom code + config），避免占用有限的 /opt 空间，不纳入本仓库
• transformers 版本强约束——custom code 声明 transformers_version: 4.40.1：
  • transformers 5.15.1：报 DynamicCache 缺 seen_tokens（可 monkey-patch 但仍遇 _prepare_4d_causal_attention_mask mask 尺寸 511 vs 512 不匹配）
  • transformers 4.44.2：报 GenerationMixin._extract_past_from_model_output() 收到意外参数 standardize_cache_format
  • 4.40.1 正常（部署请固定此版本）
• generate 参数——custom code 的 _greedy_search 接收 input_ids（历史序列），不是 past_values
• 输入：历史序列 (1, seq_len) float32；输出：预测序列（含历史拼接，前 seq_len 为历史、后 horizon 为预测）