寒武纪MLU上快速入门PyTorch指南-洪萨配资

寒武纪MLU上快速入门PyTorch指南

在国产AI芯片加速落地的今天，越来越多的研究机构和企业开始将深度学习任务从传统GPU平台迁移到信创生态。寒武纪MLU作为国内领先的AI加速器，凭借其高性能、低功耗和良好的软件兼容性，正在成为替代NVIDIA CUDA生态的重要选择之一。

尤其值得注意的是，寒武纪提供的PyTorch-CUDA-v2.7镜像虽然名称中带有“CUDA”，但并不依赖任何NVIDIA硬件或驱动。它是一个完全由寒武纪定制的PyTorch环境，底层已无缝替换为MLU后端支持，开发者只需进行少量代码调整，即可将原本运行在GPU上的PyTorch项目快速部署到MLU设备上。

本文面向已有PyTorch开发经验的用户，以实战为导向，通过简洁的命令与示例，帮助你迅速掌握在MLU平台上运行深度学习任务的核心流程——从设备查看、环境配置到完整训练任务部署，全程对标CUDA开发习惯，实现平滑迁移。

查看MLU设备状态：类比nvidia-smi的cnmon工具

如果你熟悉NVIDIA平台下的nvidia-smi，那么你在MLU平台的第一个得力助手就是cnmon。这是寒武纪提供的系统级监控工具，用于实时查看MLU卡的状态信息。

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

cnmon

输出示例如下：

+-----------------------------------------------------------------------------+ | CNMon | 2.7.0 | Process Sample Interval: 1s | +----------------+--------------------------+-------------------------------+ | Card ID | Name | SPU Util(%) | Memory Used/Total(MB) | Temp(C) | +---------+----------------+-------------+------------------------+---------+ | 0 | MLU370-S4 | 0% | 1024 / 16384 | 45 | | 1 | MLU370-S4 | 0% | 512 / 16384 | 43 | +---------+----------------+-------------+------------------------+---------+

每列含义清晰明了：
-Card ID：设备编号，对应后续编程中的'mlu:0'、'mlu:1'
-SPU Util(%)：算力单元使用率（类似GPU利用率）
-Memory Used/Total(MB)：显存占用情况
-Temp(C)：芯片温度

实用参数组合

查看指定卡状态
若只关心某一张卡（如第0号卡）：

bash cnmon -c 0

开启动态刷新模式
调试训练时想持续观察资源变化？可以用-r参数开启轮询：

bash cnmon -r -t 1000 # 每秒刷新一次，适合监控训练过程

这个功能在调试大模型显存溢出问题时特别有用——你可以一边跑训练脚本，一边开一个终端运行cnmon -r，直观看到内存增长趋势。

📌 小贴士：如果cnmon命令未找到，请确认是否已加载正确的Docker容器或模块环境。通常在官方镜像中该命令默认可用。

启用MLU后端：导入torch_mlu是关键第一步

要在MLU上运行PyTorch程序，最关键的一步不是改模型结构，也不是换数据流，而是加上这一行：

import torch_mlu

这行代码看似简单，实则至关重要。它是连接标准PyTorch框架与寒武纪MLU硬件之间的桥梁。没有它，哪怕其他所有.to('mlu')都写对了，也会因为后端缺失而报错。

完整的最小可运行示例如下：

import torch import torch_mlu # 创建两个张量并移动至MLU a = torch.randn(3, 3).to('mlu') b = torch.randn(3, 3).to('mlu') c = a + b print(c.device) # 输出应为 mlu:0

✅ 成功标志：当输出显示mlu:0或类似的设备标识时，说明整个软硬件链路已经打通。

📌 工程建议：
- 务必将import torch_mlu放在所有.to('mlu')操作之前。
- 推荐将其放在脚本头部，避免遗漏。
- 在Jupyter Notebook中也需每次重新导入（不能跨cell依赖）。

值得一提的是，torch_mlu是一个动态加载模块，无需编译安装，在主流镜像中均已预装，属于“即插即用”型组件。

模型迁移：一句.to(‘mlu’)搞定全部参数转移

对于熟悉CUDA开发的用户来说，将模型搬到MLU几乎没有任何认知负担——语法完全一致。

以ResNet50为例：

from torchvision.models import resnet50 model = resnet50(pretrained=False) model_mlu = model.to('mlu') # 所有参数自动复制到MLU

验证是否成功也很简单：

print(next(model_mlu.parameters()).device) # 应输出 mlu:0

这套机制不仅适用于单卡场景，也天然支持多MLU卡并行计算。

多卡并行方案

如果你有两块MLU卡（如MLU370-S4双卡），可以这样启用数据并行：

model = resnet50().to('mlu') model = torch.nn.DataParallel(model) # 自动分配到多卡

或者显式指定主设备：

model = model.to('mlu:0') # 强制绑定到第一张卡

⚠️ 注意事项：
- 不要混用.cuda()和.mlu()，否则会导致设备不匹配错误。
- DataParallel 会自动处理梯度同步，但目前暂不支持 DistributedDataParallel（DDP）的所有特性，请根据版本文档确认支持范围。

损失函数也要上MLU：别让损失拖后腿

很多人知道要把模型和数据搬到MLU，却容易忽略损失函数本身也需要显式迁移。

错误示范：

criterion = nn.CrossEntropyLoss() # 默认在CPU上 loss = criterion(output_mlu, target_mlu) # output在MLU，但criterion还在CPU → 可能出错！

正确做法是：

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

这样，损失计算就会在MLU设备上完成，避免不必要的设备间拷贝开销。

更完整的前向流程如下：

output = model(data) # data已在MLU，output自然也在MLU target = target.to('mlu') # 确保标签也在MLU loss = criterion(output, target) # 全部在同一设备，高效执行

这种“全栈统一设备”的思想是提升性能的关键。一旦出现跨设备操作，PyTorch会自动插入隐式拷贝，轻则降低效率，重则引发异常。

数据搬运：non_blocking=True才是高吞吐秘诀

数据从CPU到MLU的传输是训练流水线中的关键一环。处理不当，很容易变成性能瓶颈。

常规方式：

data_mlu = data.to('mlu')

推荐方式：

data_mlu = data.to('mlu', non_blocking=True)

区别在哪？

non_blocking=True表示异步传输：主线程不必等待数据搬完就能继续执行下一步（比如启动下一个batch的预处理），从而实现计算与通信重叠。
特别配合DataLoader(num_workers>0)使用时，效果尤为显著。

典型的高效数据加载循环如下：

for batch_idx, (data, target) in enumerate(train_loader): data = data.to('mlu', non_blocking=True) target = target.to('mlu', non_blocking=True) optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step()

📌 经验法则：
- 只要数据是从CPU张量转到MLU，一律加上non_blocking=True
- 如果原始数据已经是MLU张量（如持久化缓存），则无需重复搬运

完整训练案例：CIFAR-10图像分类实战

下面是一个端到端的CIFAR-10图像分类训练脚本，涵盖了从数据加载、模型构建到训练测试的全过程，充分展示如何在MLU上运行标准PyTorch流程。

import torch import torch_mlu import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms from torchvision.models import resnet18 from torch.optim.lr_scheduler import StepLR # 超参数设置 BATCH_SIZE = 128 EPOCHS = 5 LR = 0.01 # 数据增强与标准化 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_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train) test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test) train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4) test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4) # 构建模型并移至MLU model = resnet18(num_classes=10) model = model.to('mlu') # 损失函数和优化器 criterion = nn.CrossEntropyLoss().to('mlu') optimizer = optim.SGD(model.parameters(), lr=LR, momentum=0.9, weight_decay=5e-4) scheduler = StepLR(optimizer, step_size=1, gamma=0.9) # 训练函数 def train(epoch): model.train() for batch_idx, (data, target) in enumerate(train_loader): data = data.to('mlu', non_blocking=True) target = target.to('mlu', non_blocking=True) optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() if batch_idx % 100 == 0: print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ' f'({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}') # 测试函数 def test(): model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data, target in test_loader: data = data.to('mlu', non_blocking=True) target = target.to('mlu', non_blocking=True) output = model(data) test_loss += criterion(output, target).item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() test_loss /= len(test_loader) accuracy = 100. * correct / len(test_loader.dataset) print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n') # 主循环 if __name__ == '__main__': for epoch in range(1, EPOCHS + 1): train(epoch) test() scheduler.step() # 保存模型（权重仍在MLU上，保存前会自动回传CPU） torch.save(model.state_dict(), "resnet18_cifar10_mlu.pth") print("Model saved.")

💡 关键点解析：
-import torch_mlu是唯一新增语句
- 所有.to('mlu')替代了原有的.cuda()
-non_blocking=True提升数据吞吐
- 模型保存无需额外操作，PyTorch会自动处理设备转换

该脚本可直接在PyTorch-CUDA-v2.7镜像中运行，无需修改任何依赖。

开发接入方式：Jupyter还是SSH？按需选择

实际工作中，我们通常通过两种方式访问搭载MLU的服务器：Jupyter Notebook和SSH远程终端。各有优势，适用不同场景。

Jupyter Notebook：交互式开发首选

适合算法调优、教学演示、可视化分析等需要频繁试错的场景。

假设服务地址为http://<server_ip>:8888，浏览器打开后输入Token即可进入：

进入后可创建.ipynb文件，逐段运行代码，即时查看中间结果：

✅ 推荐搭配：
-%matplotlib inline：内联绘图
-torchinfo.summary()：模型结构查看
-wandb或tensorboard：训练日志追踪

缺点是不适合长期运行任务，且安全性较低，建议仅用于开发阶段。

SSH远程终端：生产部署利器

对于批量训练、自动化脚本、CI/CD集成等场景，SSH才是王道。

连接命令：

ssh user@<mlu_server_ip> -p 22

上传代码：

scp train.py user@<mlu_server_ip>:/workspace/

登录后可在后台运行训练任务，并结合screen或tmux实现断线不中断：

tmux new -s mlutrain python train.py # Ctrl+B, D 脱离会话

同时可另开终端监控资源：

cnmon -r -t 1000

这种方式稳定性强、权限可控，是工业级部署的标准操作。

这种高度集成且兼容主流生态的设计思路，正推动着国产AI基础设施走向成熟。未来随着寒武纪对FSDP、Torch.compile等新特性的逐步支持，MLU平台将进一步缩小与国际先进水平的差距，为自主可控的AI研发提供坚实底座。

寒武纪MLU上快速入门PyTorch指南