本课说明（2026-11-14 新编）
适用范围：torch_musa（API 与 PyTorch 保持一致）。第 4 模块第 1 课已有一段 AMP 片段，本课给出完整、可检查的写法；FSDP2 中的混合精度策略见第 4 模块第 4 课。 本课依据官方文档与正式 Release 整理，未在真实 MUSA GPU 上实测；标注“官方”的命令与数值来自原文，标注“本站补充”的脚本由本站编写用于检查与教学，使用前请在自己的环境中验证。
一、混合精度在做什么
- autocast：在前向计算中，把适合低精度的算子（如矩阵乘、卷积）自动切到 FP16/BF16，把对精度敏感的算子（如归约、softmax 的部分实现）保留在 FP32。
- 权重仍是 FP32：优化器更新的是 FP32 参数。
- 损失缩放（GradScaler）：FP16 指数范围小，梯度容易下溢为 0，需要先放大损失再反向传播；BF16 指数范围与 FP32 相同，通常不需要缩放。
上述是 PyTorch 的通用机制，见 PyTorch torch.amp 文档。
二、MUSA 上的精度支持边界
| 卡型 | FP16 | BF16 | 依据 |
|---|---|---|---|
| MTT S80 / S3000 | 支持 | 不支持 | torch_musa 历史 Release（v1.1.0 / v1.2.1）（v1.1.0） |
| MTT S4000 | 支持 | 支持 | torch_musa 历史 Release（v1.1.0 / v1.2.1）（v1.1.0） |
| MTT S5000 | — | 官方示例使用 | 官方 LLaMA-Factory 与 vLLM-MUSA 示例均在 S5000 上使用 BF16（见第 7、10 模块） |
说明
S5000 的 FP16 训练没有找到单独的官方说明，本课不作结论。新卡请优先用 BF16。
三、官方写法：torch.musa.amp
torch_musa 快速入门 给出了示例，以下据此整理：补全原文 model(inputs 缺少的右括号、修正缩进，并精简了 forward / __call__ 写法与注释，逻辑不变：
python
import torch
import torch_musa
import torch.nn as nn
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(5, 4)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(4, 3)
def forward(self, x):
return self.fc2(self.relu(self.fc1(x)))
DEVICE = "musa"
def train_in_amp(low_dtype=torch.float16):
model = SimpleModel().to(DEVICE)
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
scaler = torch.musa.amp.GradScaler()
inputs = torch.randn(6, 5).to(DEVICE)
targets = torch.randn(6, 3).to(DEVICE)
for step in range(20):
optimizer.zero_grad()
with torch.musa.amp.autocast(dtype=low_dtype):
outputs = model(inputs)
assert outputs.dtype == low_dtype
loss = criterion(outputs, targets)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
return loss
if __name__ == "__main__":
train_in_amp(torch.float16)四、与 CUDA 代码对齐的写法
据 torch_musa 历史 Release（v1.1.0 / v1.2.1）（v1.2.1），“AMP 用法已与 CUDA 对齐，torch.autocast 会自动启用 torch_musa 的 AMP”。因此设备无关的写法是：
python
with torch.autocast(device_type="musa", dtype=torch.bfloat16):
loss = criterion(model(x), y)迁移现有 CUDA 代码时，如果代码里大量使用 torch.cuda.amp.*，可以借助 torchada README 让它在 MUSA 上运行。注意 PyTorch 已将 torch.cuda.amp.autocast 标记为弃用，新代码请写 torch.amp.autocast(...)（见 PyTorch torch.amp 文档）。GradScaler 本课仍用官方示例中的 torch.musa.amp.GradScaler()；torch.amp.GradScaler("musa") 未见 torch_musa 官方说明，使用前请自行验证。
五、完整训练循环（本站补充）
下面在官方写法基础上加入：按精度决定是否启用缩放、梯度裁剪前先 unscale、跳过非有限值、保存与恢复 scaler 状态。
python
import torch, torch_musa
import torch.nn as nn
def train(model, loader, epochs=1, amp_dtype=torch.bfloat16, max_norm=1.0, ckpt="amp_ckpt.pt"):
device = torch.device("musa")
model = model.to(device)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
use_scaler = (amp_dtype == torch.float16) # BF16 通常不需要损失缩放
scaler = torch.musa.amp.GradScaler(enabled=use_scaler)
for epoch in range(epochs):
for step, (x, y) in enumerate(loader):
x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
opt.zero_grad(set_to_none=True)
with torch.autocast(device_type="musa", dtype=amp_dtype):
loss = criterion(model(x), y)
if not torch.isfinite(loss):
print(f"step {step}: loss 非有限值，跳过"); continue
scaler.scale(loss).backward()
scaler.unscale_(opt) # 裁剪前必须先还原梯度
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
scaler.step(opt)
scaler.update()
if step % 100 == 0:
print(f"epoch {epoch} step {step} loss {loss.item():.4f} scale {scaler.get_scale() if use_scaler else 1.0}")
torch.save({"model": model.state_dict(), "opt": opt.state_dict(),
"scaler": scaler.state_dict()}, ckpt)
# 恢复时：
# state = torch.load(ckpt); model.load_state_dict(state["model"])
# opt.load_state_dict(state["opt"]); scaler.load_state_dict(state["scaler"])六、FP32 输入下的 TF32
torch_musa 快速入门 说明在 S4000 上可通过 TF32 启用 TensorCore，加速 FP32 矩阵运算：
python
torch.backends.mudnn.allow_tf32 = True
# 或者只在局部启用
with torch.backends.mudnn.flags(allow_tf32=True):
c = a @ bTF32 会降低尾数精度，对数值敏感的任务请先对比结果。
七、怎样确认混合精度生效且没有损害精度
- 看数据类型：在 autocast 内断言线性层输出为低精度（如官方示例的
assert outputs.dtype == low_dtype）。 - 对比曲线：同一随机种子、同一小数据集，分别用 FP32 与 AMP 训练若干百步，比较损失曲线是否接近。
- 看缩放因子：FP16 下
scaler.get_scale()若持续下降，说明频繁出现溢出。 - 看显存与吞吐：用
torch.musa.max_memory_allocated()和每秒样本数对比收益。
八、常见问题
| 现象 | 可能原因与处理 |
|---|---|
| 损失变成 NaN / Inf | FP16 溢出：改用 BF16，或降低学习率；检查是否在 autocast 外手动把张量转成了 half |
| FP16 下梯度全为 0 | 未使用 GradScaler |
| 梯度裁剪后效果异常 | 裁剪前没有调用 scaler.unscale_(opt) |
| S80 / S3000 上 BF16 报错 | 这两款卡只支持 FP16 |
| 定位出错的算子 | 设置 MUSA_LAUNCH_BLOCKING=1 同步执行，并用 TORCH_SHOW_CPP_STACKTRACES=1 查看调用栈（官方快速入门中的环境变量） |
九、回退方式
把 autocast 的 enabled 设为 False、GradScaler 的 enabled 设为 False，即可退回纯 FP32 训练，模型和数据代码无需修改。