在寒武纪MLU上快速运行PyTorch指南
在寒武纪MLU上快速运行PyTorch指南
在国产AI芯片加速落地的今天，越来越多深度学习项目正从依赖NVIDIA GPU的CUDA生态向自主可控的异构计算平台迁移。作为国内领先的AI芯片厂商，寒武纪（Cambricon）推出的MLU系列处理器已逐步支持主流框架，其中对PyTorch的适配尤为成熟。
对于长期使用PyTorch-CUDA-v2.7这类镜像进行开发的工程师而言，如何将现有模型平滑迁移到MLU环境，是实际落地中的关键一步。本文不谈理论堆砌，而是以“最小改动、最大兼容”为原则，通过可直接复用的代码片段和真实操作路径，带你快速打通从环境配置到完整训练的全流程。
环境切换：别再用CUDA镜像了
一个常见的误区是：只要PyTorch版本一致（比如v2.7），就能在原有pytorch-cuda镜像中跑MLU任务。事实并非如此——CUDA构建的PyTorch无法调用MLU设备，底层算子库完全不同。
正确的做法是切换至官方提供的Cambricon定制镜像：
cambricon/pytorch:2.7-mlu这个镜像是基于PyTorch 2.7源码重新编译，并启用MLU后端支持的结果。它集成了几个核心组件：
-torch_mlu：MLU设备插件，提供.to('mlu')能力
- MagicMind运行时：用于高性能推理优化
- CNToolkit：包含cnmon、cndev等系统级工具
验证是否就绪很简单：
python -c "import torch; print(torch.__version__)"预期输出类似2.7.x，且不会报任何与CUDA相关的错误。如果看到No module named 'torch_mlu'，说明环境未正确安装或未激活。
查看设备状态：用 cnmon 替代 nvidia-smi
熟悉GPU开发的同学都知道nvidia-smi的重要性。在MLU平台上，对应工具叫cnmon，功能几乎完全一致。
查看当前MLU设备信息：
cnmon典型输出如下：
============ CNMon ============ Driver Version: 5.10.0 Runtime Version: 2.7.0 Device ID: 0 Name: MLU370-S4 Temperature: 58 C Memory Used: 2048 MB / 16384 MB Utilization: 45%你可以像监控GPU一样实时观察显存占用和计算利用率。例如每秒刷新一次：
cnmon -r -t 1000这在调试训练卡顿或内存泄漏时非常有用。尤其当你发现利用率长期低于20%，可能意味着数据加载成了瓶颈，而非算力不足。
启用MLU支持：必须导入 torch_mlu
这是最容易被忽略的一环。即使你写了.to('mlu')，如果不先导入torch_mlu模块，程序仍会抛出：
Invalid device string: 'mlu'正确写法始终是：
import torch import torch_mlu # 必须显式引入！导入之后，就可以像操作CUDA张量一样处理MLU张量了：
a = torch.randn(3, 3).to('mlu') b = torch.randn(3, 3).to('mlu') c = a + b print(c.device) # 输出: mlu:0注意：torch_mlu不会自动注册设备后端，也不能通过torch.backends.mlu访问。它的作用是在导入时动态注册'mlu'设备类型，属于典型的“副作用导入”模式。
模型迁移：只需改一个字符串
得益于良好的API兼容性，模型迁移到MLU几乎不需要重构代码。原来写.to('cuda')的地方，现在改成.to('mlu')即可。
以ResNet为例：
from torchvision import models model = models.resnet50(weights=None) model.to('mlu') # 原地迁移多卡场景下也支持DataParallel（需驱动支持）：
if torch.mlu.device_count() > 1: model = torch.nn.DataParallel(model, device_ids=[0, 1]) model.to('mlu')这里有个小技巧：为了提升代码可移植性，建议封装设备选择逻辑：
device = 'mlu' if torch.mlu.is_available() else 'cpu' model.to(device)这样同一份脚本可以在不同环境中自动适配，避免硬编码带来的维护成本。
损失函数与优化器的设备管理
损失函数本身是nn.Module的子类，因此也需要迁移到MLU设备上：
criterion = nn.CrossEntropyLoss().to('mlu')否则当你执行criterion(output, target)时，若output在MLU而criterion内部参数还在CPU，就会触发跨设备运算异常。
相比之下，优化器更“智能”一些。它不持有可训练参数的副本，而是直接引用模型中的.parameters()。所以只要模型已经to('mlu')，那么：
optimizer = optim.Adam(model.parameters(), lr=0.001)这段代码无需任何修改。优化器会自动跟踪这些参数所在的设备，反向传播时梯度也会自然生成在MLU上。
数据加载：别忘了把输入送进MLU
训练中最常出错的地方，就是只把模型搬到MLU，却忘了数据还在CPU。
标准流程应如下：
for data, target in train_loader: data = data.to('mlu', non_blocking=True) target = target.to('mlu', non_blocking=True) output = model(data) loss = criterion(output, target) # ...两个细节值得注意：
1.non_blocking=True可实现主机内存到设备内存的异步拷贝，尤其在高吞吐数据流中能提升效率。
2. 尽量避免混合使用.cuda()和.mlu()，统一用.to('mlu')形式，便于未来扩展支持更多后端。
另外，虽然data.mlu()是.to('mlu')的快捷方式，但前者不是标准PyTorch API的一部分，建议优先使用通用方法。
完整示例：CIFAR-10上的ResNet训练
下面是一个可在MLU上直接运行的端到端训练脚本，结构清晰，适合初学者参考：
import torch import torch_mlu import torch.nn as nn import torch.optim as optim import torchvision.transforms as transforms from torch.utils.data import DataLoader from torchvision import datasets, models # 1. 数据预处理 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)), ]) train_set = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train) train_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=4) # 2. 模型与设备绑定 device = 'mlu' model = models.resnet18(num_classes=10) model.to(device) criterion = nn.CrossEntropyLoss().to(device) optimizer = optim.Adam(model.parameters(), lr=0.001) # 3. 训练循环 model.train() for epoch in range(5): running_loss = 0.0 for i, (inputs, labels) in enumerate(train_loader): inputs = inputs.to(device) labels = labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() if i % 100 == 99: print(f"Epoch {epoch+1}, Batch {i+1}: Loss {running_loss/100:.4f}") running_loss = 0.0 # 4. 保存模型 torch.save(model.state_dict(), 'resnet18_cifar10_mlu.pth') print("Model saved.")运行该脚本前，请确保：
- 已导入torch_mlu
- 使用的是MLU专用镜像
- 当前机器有可用MLU设备
训练过程中可通过cnmon观察设备利用率变化。若利用率稳定在60%以上，说明计算资源得到了有效利用。
开发接入方式：Jupyter vs SSH
大多数MLU开发服务器提供两种接入方式，各有适用场景。
Jupyter Notebook：交互式调试首选
图形化界面非常适合探索性实验。启动服务后，在浏览器中打开Notebook：
优势明显：
- 支持分块执行，便于调试中间结果
- 可嵌入matplotlib绘图，直观展示训练曲线
- 方便分享完整分析过程给团队成员
但在生产训练中，Jupyter不适合运行长时间任务，容易因网络中断导致进程终止。
SSH远程登录：批量任务的最佳选择
通过终端连接服务器：
ssh username@server_ip -p port配合scp上传代码：
scp train.py username@server_ip:/home/username/长周期任务推荐使用nohup或tmux守护进程：
nohup python train.py > train.log 2>&1 &这种方式更适合自动化流水线、定时训练任务或大规模超参搜索。日志文件持久化存储，方便后续排查问题。
这种高度集成的设计思路，正推动着国产AI基础设施向更高效、更易用的方向发展。随着生态不断完善，未来我们有望看到更多模型无需修改即可在MLU上原生运行。