训练大模型需要海量算力，而自动混合精度（AMP）是提升效率的关键。CANN 原生支持 AMP，可在几乎不损精度的前提下，显著加速训练并降低显存占用。

AMP 原理

AMP 将前向/反向计算转为 FP16，参数仍为 FP32。为避免梯度下溢，采用 Loss Scaling：放大损失 → FP16 反向 → 缩小梯度 → FP32 更新。

CANN 中的使用

启用 AMP 极其简单：

python

import cann

model = cann.nn.Linear(1024, 10).cuda()
optimizer = cann.optim.Adam(model.parameters())
scaler = cann.cuda.amp.GradScaler()

for data, target in loader:
    optimizer.zero_grad()
    
    with cann.cuda.amp.autocast():  # 自动选择数据类型
        output = model(data)
        loss = cann.nn.functional.cross_entropy(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

算子类型策略

CANN 内置三类算子策略：

• 白名单（Conv、MatMul）：默认 FP16
• 黑名单（Softmax、Log）：强制 FP32
• 灰名单（Add、Mul）：依输入决定

开发者也可手动覆盖：

python

with cann.cuda.amp.autocast(dtype=cann.float16):
    y = cann.softmax(x)  # 不推荐，可能数值不稳定

性能收益

以 BERT-base（batch=64）为例：

表格

显存减少 40%，速度提升 61%，精度几乎无损。

分布式训练支持

AMP 与分布式无缝集成：

python

model = cann.nn.parallel.DistributedDataParallel(model)
# AMP 代码不变

CANN 通信库自动处理 FP16 梯度聚合。

故障排查

若训练发散，可尝试：

• 降低初始缩放因子：GradScaler(init_scale=2**10)
• 对敏感层强制 FP32：layer = layer.to(cann.float32)
• 启用梯度检查

小结

CANN 的 AMP 实现兼顾性能与稳定性，是大规模训练的必备技术。合理使用，可在有限资源下训练更大模型，加速 AI 创新。

cann组织链接：https://atomgit.com/cann
ops-nn仓库链接：https://atomgit.com/cann/ops-nn