华为云昇腾Atlas推理调优：AOE自动优化与多卡部署实战

文章摘要：PyTorch 模型迁移到昇腾 NPU 后跑了 28ms，但业务要求 15ms 以内。通过 AOE（Ascend Optimization Engine）自动优化 + 算子融合 + 多卡负载均衡，最终把延迟降到了 12ms。本文记录了昇腾推理的完整调优过程，包括 AOE 工具的使用方法、推理精度验证、多卡部署的负载均衡策略，以及那些文档里没写清楚的坑。

从 28ms 到 12ms：调优路径

AOE 自动调优

AOE（Ascend Optimization Engine）是昇腾的自动调优工具，可以对 OM 模型做算子级的性能优化。

# 1. 准备校准数据
mkdir -p calibration_data
python3 prepare_calibration.py  # 生成约 100 条推理请求数据

# 2. 执行 AOE 调优
aoe --model=bert_model.om \
    --job_type=2 \
    --output=bert_model_optimized \
    --calib_data_path=./calibration_data \
    --log=info

# job_type 说明:
# 1: 算子调优（推荐首次使用）
# 2: 子图调优（深度优化，耗时较长）

# 调优完成后
ls -la bert_model_optimized.om
# 比原始 om 小了约 15%（算子融合去掉了冗余计算）

AOE 做了什么？

# 查看 AOE 的调优日志
cat aoe_result.log
[INFO] Fused 32 operators into 12 operator groups
[INFO] Optimized memory allocation: peak memory reduced from 512MB to 384MB
[INFO] Total throughput improvement: 21.3%

1. AOE 调优后精度掉到了 85%

现象：AOE 调优后，推理速度从 28ms 降到了 15ms，但模型 accuracy 从 91.8% 掉到了 85.3%。

原因：AOE 的算子融合优化中，部分 FP32 算子被自动降到了 FP16，导致精度损失。

解决：

# 指定精度模式为 force_fp32，不允许降精度
aoe --model=bert_model.om \
    --job_type=2 \
    --output=bert_model_optimized \
    --precision_mode=force_fp32 \
    --op_precision_lib=/usr/local/Ascend/opp/op_impl/built-in/aicpu/aicpu_kernels

加了这个参数后，精度恢复到 91.5%，延迟 18ms，做了取舍。

2. 不同模型 Batch Size 差异导致显存溢出

现象：同一台推理服务器上同时部署了 BERT 和 GPT 两个模型。BERT 的 batch size 设为 32 时推理正常（延迟 22ms），但 GPT 模型在 batch size 只设了 8 的情况下就报了 OOM：

# GPT 模型推理日志
[ERROR] Ascend310P3 OOM: alloc 156MB failed, device memory used 95.2%
[ERROR] Model inference failed: mem_alloc failed

原因：不同架构模型的显存占用模式差异巨大。GPT 的自回归解码需要缓存完整的 Key-Value Cache，每个 token 的 KV cache 大小 = 2 × hidden_size × num_layers × sizeof(float16)。以 GPT-2 1.5B 为例：

BERT 是 encoder-only 架构，显存主要消耗在中间激活值上，随序列长度增长较慢。GPT 是 decoder-only，KV cache 随 batch size 和序列长度线性增长，显存消耗大得多。

解决：

1. 为每个模型独立配置 batch size，不搞"一刀切"：

# model_config.yaml
models:
  bert_classifier:
    om_path: "/models/bert_model.om"
    max_batch_size: 32
    device_id: 0

  gpt_generator:
    om_path: "/models/gpt_model.om"
    max_batch_size: 4   # 不能超过 4！
    device_id: 1        # 单独分配到另一张卡

2. 将 GPT 单独分配到一张 NPU 卡，避免与 BERT 争抢显存：

docker run -d --name gpt-infer \
    --device=/dev/davinci1 \
    -e NPU_ID=1 \
    -e MAX_BATCH_SIZE=4 \
    ascend-infer:latest

3. 增加显存监控告警，在显存使用超过 80% 时自动触发限流：

# memory_monitor.py
def check_npu_memory():
    result = subprocess.run(
        ["npu-smi", "info", "-t", "memory", "-o", "json"],
        capture_output=True, text=True
    )
    info = json.loads(result.stdout)
    for chip in info["data"]:
        pct = chip["memory_usage"]["used"] / chip["memory_usage"]["total"] * 100
        if pct > 80:
            print(f"[WARN] NPU {chip['device_id']} memory: {pct:.1f}%")
            # 触发限流，降低该卡 batch size

3. 内存碎片导致推理服务运行几天后 OOM

现象：推理服务上线后前 2 天一切正常，第 3 天开始频繁 OOM，重启容器后恢复，但过几天又复现。

# 第 1 天：显存正常
npu-smi info  # 占用 4.2GB / 24GB

# 第 3 天：显存泄漏迹象
npu-smi info  # 占用 18.6GB / 24GB

# OOM 日志
[ERROR] Ascend310 alloc memory failed, error: 100002
[ERROR] Current device memory pool is exhausted

原因：推理请求的输入长度不固定，每次分配不同大小的显存块。昇腾的内存分配器在频繁分配-释放不同大小的显存块时产生内存碎片，就像磁盘碎片一样。总空闲显存还有 5.4GB，但没有一个连续块能满足新请求的分配需求。

# 模拟碎片化过程
# 请求序列：128→释放→512→释放→256→释放→1024→释放→...
# 多轮后空闲块零散分布：[64MB, 128MB, 256MB, ...]
# 一个 512MB 的请求无法分配，尽管总空闲 > 5GB

解决：

1. 启用昇腾内存池预分配，服务启动时一次性申请大块显存：

export ASCEND_MEMORY_POOL_INIT_SIZE=8192   # 预分配 8GB
export ASCEND_MEMORY_POOL_MAX_SIZE=20480   # 最大 20GB

2. 固定推理输入的 shape，避免动态 shape 导致碎片：

MAX_SEQ_LENGTH = 512  # 固定最大长度

def preprocess(text, tokenizer):
    tokens = tokenizer.encode(text)
    padded = np.zeros((MAX_SEQ_LENGTH,), dtype=np.int64)
    padded[:len(tokens)] = tokens[:MAX_SEQ_LENGTH]
    mask = np.zeros((MAX_SEQ_LENGTH,), dtype=np.int64)
    mask[:len(tokens)] = 1
    return padded, mask

3. 设置定时重启策略作为兜底方案：

    healthcheck:
      test: ["CMD", "npu-smi", "info"]
      interval: 30s
      timeout: 10s
      retries: 3
    deploy:
      restart_policy:
        condition: any

加上内存池预分配 + 固定 shape 后，连续运行 30 天未出现 OOM。

算子二进制优化

对于 AOE 优化后仍有瓶颈的特定算子，可以进行手工算子替换。

# 查看推理过程中的算子耗时
msprof --output=./profiling_data

# 解析 profiling 结果
python3 -m mindspore.analyze.profiler --output=./profiling_data

# profiling 输出示例
Operator         Time(us)    %     Calls
BertAttention    3,245      23%   1
MatMul           2,100      15%   8
Softmax          1,800      13%   12
LayerNorm        1,200      9%    12
Gelu             890        6%    6

MatMul 占了 15% 的时间。替换为昇腾的最优化实现：

# 使用 ATC 的 --op_select_implmode 参数
atc --model=bert_model.onnx \
    --framework=5 \
    --output=bert_model \
    --soc_version=Ascend310P3 \
    --op_select_implmode=high_precision  # high_precision / high_performance

精度验证

调优后必须验证精度，这是上线前最关键的一步。

# validate_accuracy.py
import numpy as np
from mindspore import Tensor, load

def validate_model(om_path, test_data, labels):
    model = load(om_path)
    correct = 0
    total = len(test_data)
    
    for i, (inputs, label) in enumerate(zip(test_data, labels)):
        # 昇腾推理
        input_ids = Tensor(inputs["input_ids"])
        attention_mask = Tensor(inputs["attention_mask"])
        
        logits = model(input_ids, attention_mask, None)
        pred = np.argmax(logits.asnumpy(), axis=-1)
        
        if pred == label:
            correct += 1
        
        if i % 100 == 0:
            print(f"Progress: {i}/{total}")
    
    accuracy = correct / total
    print(f"Final accuracy: {accuracy:.4f}")
    return accuracy

# 对比原始 PyTorch 模型的精度
original_accuracy = 0.923
optimized_accuracy = validate_model("bert_model_optimized.om", test_data, labels)
# expected: optimized_accuracy > original_accuracy - 0.01

多卡多实例部署

单卡 310P3 的吞吐不够时，多卡并行部署。

# 查看系统有几张 NPU 卡
npu-smi info
# NPU-ID  Chip
# 0       Ascend310P3
# 1       Ascend310P3
# 2       Ascend310P3
# 3       Ascend310P3

Docker Compose 多卡部署

# docker-compose.yml
version: '3.8'
services:
  infer-0:
    image: ascend-infer:latest
    devices:
      - /dev/davinci0
    environment:
      - NPU_ID=0
    ports:
      - "8501:8501"
      
  infer-1:
    image: ascend-infer:latest
    devices:
      - /dev/davinci1
    environment:
      - NPU_ID=1
    ports:
      - "8502:8501"
      
  infer-2:
    image: ascend-infer:latest
    devices:
      - /dev/davinci2
    environment:
      - NPU_ID=2
    ports:
      - "8503:8501"
      
  infer-3:
    image: ascend-infer:latest
    devices:
      - /dev/davinci3
    environment:
      - NPU_ID=3
    ports:
      - "8504:8501"

  nginx:
    image: nginx:alpine
    ports:
      - "80:80"
    volumes:
      - ./nginx.conf:/etc/nginx/nginx.conf

# nginx.conf — 负载均衡
upstream ascend_servers {
    least_conn;
    server infer-0:8501 max_fails=3 fail_timeout=10s;
    server infer-1:8501 max_fails=3 fail_timeout=10s;
    server infer-2:8501 max_fails=3 fail_timeout=10s;
    server infer-3:8501 max_fails=3 fail_timeout=10s;
}

server {
    listen 80;
    location /infer {
        proxy_pass http://ascend_servers;
        proxy_set_header Host $host;
        proxy_read_timeout 30s;

## 最终调优效果

| 优化步骤 | 延迟 | 吞吐 | 精度 |
|:---------|:----:|:----:|:----:|
| 初始（OM 直接推理）| 28ms | 960 qps | 91.8% |
| + AOE 自动调优 | 18ms | 1,500 qps | 91.5% |
| + 算子优化 | 15ms | 1,800 qps | 91.5% |
| + 内存池预分配 | 15ms | 1,850 qps | 91.5% |
| + 4 卡负载均衡 | 12ms | 6,800 qps | 91.5% |


## 常见问题（FAQ）

### Q1: AOE 调优需要多久？会影响线上服务吗？

AOE 调优时间取决于模型大小和 `job_type`：

| job_type | 说明 | BERT-Base 耗时 | GPT-2 1.5B 耗时 |
|:---------|:-----|:-------------:|:---------------:|
| 1 | 算子调优 | 20-40 分钟 | 1-2 小时 |
| 2 | 子图调优 | 1-2 小时 | 3-6 小时 |
| 3 | 全流程调优 | 3-4 小时 | 8-12 小时 |

**建议**：首次使用 job_type=1 快速见效，后续再跑 job_type=2 深度优化。AOE 调优会占用 NPU 资源，**不要在线上服务节点上执行**，建议在测试环境或灰度节点上运行。

### Q2: 多卡部署时，如何避免单卡故障影响整体服务？

从三个层面做容错：

1. **Nginx 层**：配置 `max_fails=3 fail_timeout=30s`，连续 3 次健康检查失败自动摘除节点

2. **应用层**：实现优雅降级，单卡故障时自动切换到剩余卡：

```python
class FallbackManager:
    def __init__(self, devices):
        self.devices = devices
        self.healthy = {d: True for d in devices}

    def infer(self, inputs):
        for device in self._healthy_devices():
            try:
                return self._infer_on_device(inputs, device)
            except Exception:
                self.healthy[device] = False
        raise RuntimeError("All devices failed")

    def _healthy_devices(self):
        return [d for d in self.devices if self.healthy.get(d, True)]

3. 监控层：对接告警系统，单卡故障 1 分钟内通知值班人员

推理调优 Checklist

部署前

• [ ] ONNX 模型输入输出 shape 已确认，与原始 PyTorch 一致
• [ ] AOE 调优已在测试环境完成，job_type 至少跑过 1 和 2
• [ ] 精度验证通过：优化后 accuracy drop < 0.5%（或业务可接受范围）
• [ ] 已确认 NPU 型号和 soc_version 参数正确（npu-smi info 查看）
• [ ] 内存池预分配参数 ASCEND_MEMORY_POOL_INIT_SIZE 已配置
• [ ] 最大 batch size 已根据模型显存占用计算确认
• [ ] 静态 shape 已启用，所有输入做 padding 到固定长度

部署中

• [ ] Docker Compose 中每个推理实例绑定到指定 NPU 设备（/dev/davinciN）
• [ ] Nginx 健康检查配置正确，max_fails 和 fail_timeout 已设置
• [ ] 每个实例端口映射不冲突（8501-8504）
• [ ] 负载均衡策略选择合理（least_conn 优于 round_robin）
• [ ] 环境变量 NPU_ID 和 ASCEND_MEMORY_POOL_INIT_SIZE 已设置

部署后

• [ ] 压测验证：P99 延迟满足业务 SLA，吞吐达到预期
• [ ] 显存监控已配置：超过 80% 触发告警
• [ ] 健康检查接口 /health 可正常访问
• [ ] 单卡故障测试：手动停掉一个实例，观察自动摘除和恢复
• [ ] 长期稳定性测试：运行 48 小时以上，无 OOM 和显存泄漏
• [ ] 精度回归测试：上线后用 1,000 条样本验证精度一致

总结与经验教训

核心路径回顾

关键教训

1. 不要盲目相信自动优化：AOE 很强大，但自动降精度可能导致 accuracy 大幅下跌。先跑精度验证，再上线，这个顺序不能反。

2. 不同模型差异巨大：BERT 和 GPT 的显存占用模式完全不同，batch size 不能一刀切。每个模型独立调参是必须的。

3. 内存碎片是隐形杀手：动态 shape 带来的内存碎片不会立即暴露，但运行几天后就会显现。固定 shape + 内存池预分配是最有效的预防手段。

4. 多卡不是银弹：4 卡接近线性扩展的前提是负载均衡策略正确、健康检查配置到位。配不好反而会引入单点故障。

性能数据汇总

后续方向

• 动态 batch：根据实时流量动态调整 batch size，低峰期减小 batch 降延迟，高峰期增大 batch 提吞吐
• 模型量化：INT8 量化可进一步降低显存占用和延迟，精度损失通常 < 0.5%
• 芯片选型：下一篇文章做国产 AI 芯片选型对比 —— 昇腾 vs 寒武纪 vs 天数智芯，真实数据说话

你在做国产 AI 芯片的推理调优吗？遇到过什么奇怪的精度问题？欢迎在评论区分享踩坑经历！