寒武纪MLU上手PyTorch指南

寒武纪MLU上手PyTorch指南

在深度学习国产化加速落地的今天，越来越多开发者开始接触寒武纪MLU平台。面对从NVIDIA GPU生态迁移的需求，如何快速将熟悉的PyTorch训练流程“平移”到MLU设备上，成为关键问题。

好消息是，寒武纪推出的Cambricon PyTorch已实现对主流框架的高度兼容——你不需要重写模型代码，也不必重新学习整套开发范式。只要把'cuda'换成'mlu'，大部分逻辑都能直接跑通。

本文基于PyTorch-MLU-v2.9 镜像环境展开实操讲解。这个镜像是官方PyTorch-CUDA-v2.9的功能移植版本，预装了 Cambricon CATCH 运行时、驱动和扩展模块，支持torch.compile()、AMP混合精度、多卡并行等特性，开箱即用。我们一步步来看怎么真正“跑起来”。

查看MLU设备状态：类比nvidia-smi的mlu-mon命令

如果你用过CUDA，一定熟悉nvidia-smi命令查看GPU资源占用情况。在MLU平台上，对应工具叫mlu-mon。

运行以下命令即可列出所有可用MLU设备：

mlu-mon

输出示例：

+-----------------------------------------------------------------------------+ | MLU Mon V1.8.0 | |-----------------------------------------------------------------------------| | ID Name Temp Util(%) Memory Usage / Total Power(W) | | 0 MLU370-X8 62C 45% 8.2GB / 16.0GB 75W | | 1 MLU370-X8 58C 12% 2.1GB / 16.0GB 68W | +-----------------------------------------------------------------------------+

这和nvidia-smi几乎一模一样：显示设备ID、型号、温度、计算利用率、显存使用和功耗信息。

想看某个设备的详细状态？加-i参数：

mlu-mon -i 0

需要持续监控负载变化？启用轮询模式：

mlu-mon -r -t 1000

其中-r开启自动刷新，-t 1000表示每秒更新一次（单位为毫秒）。适合调试期间观察内存增长或算力波动。

💡 实践建议：长时间训练前先运行mlu-mon确认设备识别正常，避免因驱动异常导致后续操作失败。

启用MLU后端：导入torch_mlu扩展

与CUDA不同，PyTorch原生不支持MLU设备，必须通过导入torch_mlu扩展来注册新后端。

这是整个流程的第一步，也是最关键的一步：

import torch import torch_mlu # 注册MLU为可用设备

一旦导入成功，PyTorch就能识别'mlu'设备，并允许张量创建和运算在其上执行：

a = torch.randn(3, 3, device='mlu') b = torch.ones(3, 3).mlu() # 等价写法 c = a + b print(c.device) # 输出: mlu:0

你可以用下面两行代码验证当前环境是否就绪：

print(torch.mlu.is_available()) # 应返回 True print(torch.mlu.device_count()) # 返回可用设备数量

如果返回False或0，请检查：
- 是否使用的是PyTorch-MLU-v2.9镜像
- MLU驱动和CATCH运行时是否安装正确
- 当前用户是否有/dev/cambricon_ctl等设备节点访问权限

⚠️ 注意：torch_mlu是动态注入式扩展，不导入就不会激活MLU支持，即使硬件存在也无法使用。

模型迁移到MLU：model.to(‘mlu’) 就够了

迁移模型的方式完全对标CUDA习惯，只需替换字符串：

from torchvision import models model = models.resnet50(weights=None) # 默认在CPU mlu_model = model.to('mlu') # 移动到MLU

也可以使用快捷方法.mlu()：

mlu_model = model.mlu()

此时模型的所有参数和缓冲区都会被复制到MLU显存中，后续前向传播、反向传播全部在MLU完成。

对于多卡场景，也支持DataParallel：

if torch.mlu.device_count() > 1: mlu_model = nn.DataParallel(mlu_model, device_ids=[0, 1])

目前暂不支持DistributedDataParallel在MLU上的完整功能（部分版本已实验性支持），建议优先采用DataParallel进行多设备训练。

🔍 经验提示：模型移动应在优化器创建之前完成。否则优化器会绑定CPU参数，导致梯度更新错位。

损失函数也要搬到MLU上

很多人忽略的一点：损失函数虽然是“函数”，但它内部可能包含可学习参数或状态张量（如nn.L1Loss中的 reduction 缓冲区），因此也需要显式移至MLU。

否则会出现经典的设备不匹配错误：

criterion = nn.CrossEntropyLoss() criterion = criterion.to('mlu') # 必须加上这一句！ logits = mlu_model(images) loss = criterion(logits, labels) # 若labels未上MLU，报错！

常见错误信息：

RuntimeError: Expected all tensors to be on the same device...

所以记住一个原则：所有参与运算的张量必须在同一设备上。包括输入数据、标签、模型输出、损失函数本身。

稳妥做法是在构建阶段统一指定设备：

device = 'mlu' model = model.to(device) criterion = nn.CrossEntropyLoss().to(device)

这样可以减少后期排查成本。

数据加载：从CPU拷贝到MLU

训练中最频繁的操作之一就是把批次数据从主机内存搬到MLU显存。

标准写法如下：

for data, target in train_loader: data = data.to('mlu', non_blocking=True) target = target.to('mlu', non_blocking=True)

这里推荐使用non_blocking=True，它表示异步传输，在支持 pinned memory 的情况下能显著提升吞吐效率。

为此，建议在构建DataLoader时开启页锁定内存：

train_loader = DataLoader( dataset, batch_size=128, shuffle=True, num_workers=4, pin_memory=True # 关键！用于加速主机到设备传输 )

pin_memory=True会让数据加载器将张量分配在固定内存中，便于DMA直接搬运，尤其在高吞吐训练中效果明显。

✅ 最佳实践组合：pin_memory=True + non_blocking=True，这是提升训练流水线效率的标准配置。

完整示例：CIFAR-10图像分类训练

下面是一个完整的端到端训练脚本，涵盖数据增强、模型定义、训练循环和测试逻辑，可在MLU上直接运行：

import torch import torch_mlu import torch.nn as nn import torch.optim as optim import torchvision import torchvision.transforms as transforms from torch.utils.data import DataLoader from torchvision.models import resnet18 from torch.optim.lr_scheduler import StepLR # 数据预处理 transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) # 加载数据集 train_set = torchvision.datasets.CIFAR-10(root='./data', train=True, download=True, transform=transform_train) test_set = torchvision.datasets.CIFAR-10(root='./data', train=False, download=True, transform=transform_test) train_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=4, pin_memory=True) test_loader = DataLoader(test_set, batch_size=128, shuffle=False, num_workers=4, pin_memory=True) # 构建模型 device = 'mlu' model = resnet18(num_classes=10) model = model.to(device) # 损失函数与优化器 criterion = nn.CrossEntropyLoss().to(device) optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4) scheduler = StepLR(optimizer, step_size=30, gamma=0.1) # 训练函数 def train_epoch(epoch): model.train() running_loss = 0.0 for i, (inputs, targets) in enumerate(train_loader): inputs = inputs.to(device, non_blocking=True) targets = targets.to(device, non_blocking=True) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, targets) loss.backward() optimizer.step() running_loss += loss.item() if i % 100 == 99: print(f'Epoch {epoch}, Batch {i+1}, Loss: {running_loss / 100:.4f}') running_loss = 0.0 # 测试函数 def test(): model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, targets in test_loader: inputs = inputs.to(device, non_blocking=True) targets = targets.to(device, non_blocking=True) outputs = model(inputs) _, predicted = outputs.max(1) total += targets.size(0) correct += predicted.eq(targets).sum().item() acc = 100. * correct / total print(f'Test Accuracy: {acc:.2f}%') # 主训练流程 if __name__ == '__main__': for epoch in range(90): train_epoch(epoch) scheduler.step() if epoch % 10 == 9: test() test() print("Training completed.")

训练结束后保存模型也很简单：

torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': loss, }, 'cifar10_resnet18_mlu.pth')

加载时注意要指定map_location：

checkpoint = torch.load('cifar10_resnet18_mlu.pth', map_location='mlu') model.load_state_dict(checkpoint['model_state_dict'])

开发方式选择：Jupyter还是SSH？

PyTorch-MLU-v2.9 镜像内置了两种主流开发模式：JupyterLab 和 SSH，可根据需求灵活选择。

使用 JupyterLab 快速调试

容器启动后，JupyterLab 默认监听8888端口：

http://<服务器IP>:8888/?token=<自动生成的token>

登录界面如下：

适合做小规模实验、可视化分析、Notebook形式记录过程。建议将常用脚本放在notebooks/目录下方便管理。

使用 SSH 进行工程化开发

更推荐生产级开发使用SSH接入：

ssh user@<服务器IP> -p <指定端口>

登录后可使用vim、tmux、git、rsync等工具进行远程编码和任务调度。

搭配 VS Code 的Remote-SSH 插件，还能实现本地编辑、远程运行的无缝体验：

🔐 安全建议：关闭密码登录，改用SSH密钥认证，提升系统安全性。

这种高度兼容PyTorch生态的设计思路，极大降低了国产AI芯片的学习门槛。开发者无需重构已有项目，只需微调设备指定逻辑，就能将训练任务部署到MLU平台。随着Cambricon对torch.compile()、分布式训练等功能的持续适配，未来在性能和易用性上有望进一步逼近甚至超越传统GPU方案。