torch_musa入门 | 摩尔线程 GPU 学习中心

2026-09-17 月度技术更新

先阅读本节，再使用后面的历史教程。旧章节保留用于理解演进；出现版本、镜像或依赖冲突时，以本节明确范围及官方发布说明为准。本次为文档核验，未在真实 GPU 上执行。

正式发布：torch_musa 2.11.0.post1

GitHub 正式 Release 的发布时间为 2026-08-17；本版本基于 PyTorch 2.11，支持 MUSA SDK 5.1 和 5.2，补充 PyTorch 2.9–2.11 范围内 LibTorch Stable ABI 的 Common/CUDA API 支持，并升级 ProcessGroupMCCL。这里的 ABI 支持不等于任意 Python wheel 可以混装。torch_musa 2.11.0.post1

升级前先理解已知风险

发布说明指出 PyTorch 2.11 中 Python 包装的 Tensor/Storage 的 intrusive_ptr 引用计数在 1↔2 转换时可能获取 GIL，从而影响 kernel launch 吞吐；DataLoader 的 num_workers > 1 与 pin_memory=True 的组合可能加剧问题。该说明没有给出适用于所有模型的固定下降比例。Known Issues

对照组	目的	必须记录
原环境、原参数	建立回归基线	版本、batch、精度、输入长度、吞吐、训练 loss
新环境、原参数	只改变软件栈	GPU 时间、CPU 等待、数据加载时间、显存峰值
新环境、num_workers=0	定位多进程数据流水线影响	端到端速度，而非仅算子微基准
新环境、减少预取或在主线程 pinning	进一步验证 GIL 相关瓶颈	与前组相同负载下的多次测量及波动

后两种措施来自发布说明中的缓解方向，不保证更快。num_workers=0 时不要沿用仅适用于多进程的 prefetch_factor / persistent_workers 配置；应先依据使用的 PyTorch API 调整测试脚本。官方缓解建议

动手练习：做一张可回滚的升级验收单

固定数据集切分、随机种子、模型权重、batch 与精度；保存旧环境包清单及镜像标识。
新建隔离环境，按照正式发行包与配套 PyTorch 的说明安装。不要在现有推理镜像内直接执行无版本限制的 pip upgrade。
先做单卡矩阵计算和训练一步，再跑完整短训练；核对 NaN、loss 趋势、输出精度与显存。
执行上表四组回归，统一预热、测量轮数和同步边界，保留原始日志。
多机作业额外验证 all_to_all：该版本修复了跨节点一侧为零大小张量时可能挂起的问题。为测试设置超时，失败时回滚整套环境而非只降级单个包。

通信修复及性能分析限制均见正式 Release。MUPTI 升级未包含在本次 torch_musa 发布内；不要假定换上 2.11 就自动解决所有 profiler 问题。

历史教程与基础知识

一、什么是torch_musa？

torch_musa是摩尔线程为PyTorch开发的MUSA后端插件，它让你可以像使用CUDA一样在摩尔线程GPU上训练和推理深度学习模型——只需要把 device='cuda' 改为 device='musa'。

torch_musa基于PyTorch的设备后端扩展机制（PrivateUse1）实现，支持PyTorch的绝大部分API和算子。这意味着你现有的PyTorch代码只需极小的修改就能在摩尔线程GPU上运行。

💡

核心思想：最小修改

torch_musa的设计哲学是最小化代码修改。大部分情况下，你只需要三步：1) import torch_musa；2) 将 cuda 替换为 musa；3) 运行。你不需要学习全新的API。

二、安装torch_musa

2.1 使用pip安装

推荐使用pip从摩尔线程官方源安装：

bash

# 确保已安装MUSA SDK（参考课程2.1）
# 确保Python版本为3.8+

# 安装torch_musa（会自动安装匹配版本的PyTorch）
pip install torch_musa -f https://oss.mthreads.com/release-rc/cuda_compatible/torch_musa/

# 或者指定PyTorch版本
pip install torch==2.0.0 torch_musa==1.1.0 -f https://oss.mthreads.com/release-rc/cuda_compatible/torch_musa/

2.2 从源码编译安装

bash

# 克隆torch_musa仓库
git clone https://github.com/MooreThreads/torch_musa.git
cd torch_musa

# 安装依赖
pip install -r requirements.txt

# 编译安装
python setup.py install

# 验证安装
python -c "import torch_musa; print('torch_musa installed successfully')"

2.3 使用Docker（最简单）

bash

# 使用预装torch_musa的官方Docker镜像
docker run -it --rm --runtime=mt \
  registry.mthreads.com/musa-dev/musa-pytorch:latest \
  python -c "import torch; import torch_musa; print(torch.musa.is_available())"

三、基本API使用

3.1 设备检查

python

import torch
import torch_musa

# ===== 设备信息查询 =====
# 检查MUSA是否可用
print("MUSA available:", torch.musa.is_available())
# 输出: MUSA available: True

# 查看可用GPU数量
print("Device count:", torch.musa.device_count())
# 输出: Device count: 1

# 获取GPU名称
print("Device name:", torch.musa.get_device_name(0))
# 输出: Device name: MTT S80

# 查看当前设备
print("Current device:", torch.musa.current_device())
# 输出: Current device: 0

# 查看显存信息
print("Total memory:", torch.musa.get_device_properties(0).total_memory / 1024**3, "GB")
print("Allocated:", torch.musa.memory_allocated(0) / 1024**2, "MB")
print("Cached:", torch.musa.memory_reserved(0) / 1024**2, "MB")

3.2 张量操作

python

import torch
import torch_musa

# ===== 创建MUSA张量 =====
# 方法1：直接在MUSA设备上创建
a = torch.tensor([1.2, 2.3], dtype=torch.float32, device='musa')
b = torch.tensor([1.8, 1.2], dtype=torch.float32, device='musa')
c = a + b
print(c)  # tensor([3.0000, 3.5000], device='musa:0')

# 方法2：先在CPU创建，再转移到MUSA
x = torch.randn(3, 4)      # CPU上创建
x_musa = x.to('musa')      # 转移到MUSA设备
print(x_musa.device)        # musa:0

# 方法3：使用.musa()方法
y = torch.zeros(5, 5).musa()
print(y.device)              # musa:0

# ===== 张量运算 =====
# 所有PyTorch运算都可以在MUSA上执行
m1 = torch.randn(100, 200, device='musa')
m2 = torch.randn(200, 50, device='musa')
result = torch.matmul(m1, m2)  # 矩阵乘法
print(result.shape)             # torch.Size([100, 50])

# 将结果转回CPU
result_cpu = result.cpu()
# 或
result_cpu = result.to('cpu')

⚠️

设备一致性

参与同一运算的所有张量必须在相同设备上。例如，不能将一个CPU张量和一个MUSA张量直接相加——会报RuntimeError。请确保使用 .to('musa') 统一设备。

3.3 设备上下文管理

python

# 使用with语句管理设备上下文
with torch.musa.device(0):
    # 此block内的所有操作默认在GPU 0上执行
    a = torch.randn(10, 10, device='musa')
    b = torch.randn(10, 10, device='musa')
    c = a @ b

# 多GPU场景下指定设备
if torch.musa.device_count() > 1:
    with torch.musa.device(1):
        d = torch.randn(10, 10, device='musa')  # 在GPU 1上

四、ResNet推理示例

下面展示如何在摩尔线程GPU上进行ResNet-50图像分类推理：

python

import torch
import torch_musa
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image

# ===== 第1步：加载预训练模型 =====
model = models.resnet50(pretrained=True)
model = model.eval()    # 切换到推理模式
model = model.to("musa")  # 将模型转移到MUSA设备
print("模型已加载到MUSA设备")

# ===== 第2步：准备输入数据 =====
# 模拟一张224x224的RGB图像
# 实际使用时可以加载真实图片：
# img = Image.open("cat.jpg")
# x = transform(img).unsqueeze(0)
x = torch.rand((1, 3, 224, 224), device="musa")

# ===== 第3步：执行推理 =====
with torch.no_grad():  # 推理时不需要计算梯度
    y = model(x)

print("输出shape:", y.shape)  # torch.Size([1, 1000])
print("预测类别:", y.argmax(dim=1).item())

# ===== 第4步：性能测量 =====
import time

# 预热
for _ in range(10):
    with torch.no_grad():
        _ = model(x)
torch.musa.synchronize()  # 确保GPU操作完成

# 计时
start = time.time()
num_runs = 100
for _ in range(num_runs):
    with torch.no_grad():
        _ = model(x)
torch.musa.synchronize()
end = time.time()

avg_time = (end - start) / num_runs * 1000
print(f"平均推理时间: {avg_time:.2f} ms")
print(f"吞吐量: {1000/avg_time:.1f} FPS")

五、CIFAR-10训练完整流程

下面是一个完整的CIFAR-10图像分类训练流程，展示了在摩尔线程GPU上训练深度学习模型的全过程。

python

import torch
import torch_musa
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

# ===== 第1步：数据准备 =====
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

batch_size = 4

# 下载并加载CIFAR-10训练集
train_set = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform
)
train_loader = torch.utils.data.DataLoader(
    train_set, batch_size=batch_size, shuffle=True, num_workers=2
)

# 下载并加载CIFAR-10测试集
test_set = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform
)
test_loader = torch.utils.data.DataLoader(
    test_set, batch_size=batch_size, shuffle=False, num_workers=2
)

# 类别名称
classes = ('飞机', '汽车', '鸟', '猫', '鹿', '狗', '蛙', '马', '船', '卡车')

# 指定MUSA设备
device = torch.device("musa")
print(f"使用设备: {device} ({torch.musa.get_device_name(0)})")

# ===== 第2步：定义网络模型 =====
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)       # 输入3通道，输出6通道，5x5卷积
        self.pool = nn.MaxPool2d(2, 2)          # 2x2最大池化
        self.conv2 = nn.Conv2d(6, 16, 5)       # 输入6通道，输出16通道
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 全连接层
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)            # 10个分类

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))   # 卷积→激活→池化
        x = self.pool(F.relu(self.conv2(x)))
        x = torch.flatten(x, 1)                # 展平
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 创建模型并转移到MUSA设备
net = Net().to(device)
print("模型结构:", net)

# ===== 第3步：定义损失函数和优化器 =====
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# ===== 第4步：训练循环 =====
print("\n开始训练...")
for epoch in range(2):  # 训练2个epoch
    running_loss = 0.0
    for i, data in enumerate(train_loader, 0):
        # 获取输入数据并转移到MUSA设备
        inputs, labels = data
        inputs = inputs.to(device)
        labels = labels.to(device)

        # 梯度清零
        optimizer.zero_grad()

        # 前向传播
        outputs = net(inputs)

        # 计算损失
        loss = criterion(outputs, labels)

        # 反向传播
        loss.backward()

        # 更新权重
        optimizer.step()

        # 统计损失
        running_loss += loss.item()
        if i % 2000 == 1999:
            print(f'[Epoch {epoch + 1}, Batch {i + 1:5d}] loss: {running_loss / 2000:.3f}')
            running_loss = 0.0

print("训练完成!")

# ===== 第5步：在测试集上评估 =====
correct = 0
total = 0
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        images = images.to(device)
        labels = labels.to(device)
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'\n测试集准确率: {100 * correct / total:.1f}%')

# ===== 第6步：保存模型 =====
torch.save(net.state_dict(), 'cifar10_musa.pth')
print("模型已保存到 cifar10_musa.pth")

✅

关键代码差异：CUDA vs MUSA

对比CUDA代码，使用torch_musa只需修改极少的地方：
import torch_musa（新增一行导入）
device = torch.device("musa")（cuda → musa）
其他所有代码完全相同！.to(device)、DataLoader、模型定义、训练循环——全部不变。

六、常用开发模式

6.1 自动设备选择

python

# 兼容CUDA/MUSA/CPU的设备选择
def get_device():
    if hasattr(torch, 'musa') and torch.musa.is_available():
        return torch.device('musa')
    elif torch.cuda.is_available():
        return torch.device('cuda')
    else:
        return torch.device('cpu')

device = get_device()
print(f"Using device: {device}")

6.2 显存管理

python

# 查看显存使用
print(f"已分配: {torch.musa.memory_allocated()/1024**2:.1f} MB")
print(f"已缓存: {torch.musa.memory_reserved()/1024**2:.1f} MB")

# 手动释放缓存
torch.musa.empty_cache()

# 设置显存分配策略（减少碎片化）
import os
os.environ['PYTORCH_MUSA_ALLOC_CONF'] = 'max_split_size_mb:512'

6.3 混合精度训练

python

# 使用自动混合精度(AMP)加速训练
scaler = torch.musa.amp.GradScaler()

for inputs, labels in train_loader:
    inputs = inputs.to(device)
    labels = labels.to(device)

    optimizer.zero_grad()

    # 自动混合精度前向传播
    with torch.musa.amp.autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels)

    # 缩放梯度并反向传播
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

✅

本课小结

torch_musa是摩尔线程GPU的PyTorch后端，通过简单的 import torch_musa 和 device='musa' 即可将现有PyTorch代码迁移到摩尔线程GPU上。本课覆盖了设备检测、张量操作、ResNet推理和CIFAR-10训练的完整流程。torch_musa的设计保证了与CUDA代码的高度兼容性。

十、v2.9.0 新特性更新（2026-04-20 追加）

torch_musa 于 2026-03-17 发布 v2.9.0，对应上游 PyTorch 2.7.1。这是一个功能密度极高的版本，将开源社区上有的大模型训练能力成批带入了国产 GPU 生态。主要亮点如下。

10.1 四大新特性速览

特性	价值
FSDP2 Context Parallel（Ulysses）	长序列训练，单次前向可处理 128K+ token
torch.compile reduce-overhead	通过 MUSA Graph 捕获，小 batch 在训练循环中调度开销几乎归零
Sparse operators	原生支持 CSR / COO 稀疏张量计算，图神经网络等场景受益
DLPack zero-copy	与 CuPy / Triton / JAX 等库无复制共享显存

表：torch_musa v2.9.0 四大新特性

10.2 硬性要求：MUSA SDK ≥ 4.3.2

v2.9.0 不再兼容 4.2.x 以前的 SDK。升级前先用下面的命令确认环境：

bash

mcc --version                  # 期望：>= 4.3.2
python -c "import torch_musa; print(torch_musa.__version__)"  # 2.9.0
python -c "import torch; print(torch.__version__)"             # 2.7.1

🚫

版本不匹配的典型报错

如果你在 import torch_musa 阶段看到 MUSA runtime version mismatch 或 undefined symbol: musaLaunchKernelExC，都意味着 SDK 版本低于 4.3.2。升级驱动 + SDK 后再试。

10.3 FSDP2 Context Parallel（Ulysses）最小示例

长序列场景下，把 GPU 划分成 dp × cp 两个 mesh 维度：cp 维度在 attention 内部做 all-to-all，将 head 维与 seq 维互换，从而让每张卡只承担一段序列的计算。

python

import torch, torch_musa
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
from torch.distributed.device_mesh import init_device_mesh

# 举例：8 张 GPU → 4 dp × 2 cp
mesh = init_device_mesh("musa", (4, 2), mesh_dim_names=("dp", "cp"))
mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16,
                          reduce_dtype=torch.float32)

for block in model.blocks:
    fully_shard(block, mesh=mesh["dp"], mp_policy=mp)

# attention 中使用 Ulysses all-to-all（参照模块四的 m4-l4 课程）

10.4 torch.compile 的 reduce-overhead 模式

当单次 step 的纯计算时间很短时，Python 端的调度开销（dispatcher、autograd、引用计数等）会占据不小比重。reduce-overhead 通过 MUSA Graph 捕获函数调用序列，省掉重复开销。

python

import torch, torch_musa

model = MyModel().musa()
opt = torch.optim.AdamW(model.parameters())

compiled = torch.compile(model, mode="reduce-overhead")

for step, (x, y) in enumerate(loader):
    x, y = x.musa(), y.musa()
    out = compiled(x)
    loss = F.cross_entropy(out, y)
    opt.zero_grad(); loss.backward(); opt.step()

10.5 GEMM 默认 FP32，TF32 需显式开启

为了对齐上游 PyTorch 的数值语义，v2.9.0 起，matmul 默认以 FP32 精度累加，不再默认启用 TF32。如果你的模型在之前依赖 TF32 带来的加速，需手动开：

python

import torch
torch.backends.cuda.matmul.allow_tf32 = True      # 注意：这里 API 名字依然是 cuda。
                                                  # 加上 torchada 后会被转发到 MUSA
# 或直接设置算子精度策略
torch.set_float32_matmul_precision("high")        # 允许 TF32 加速

10.6 Sparse 与 DLPack

python

# 稀疏张量
import torch
indices = torch.tensor([[0,1,2],[1,2,0]], device="musa")
values  = torch.tensor([3.0, 4.0, 5.0], device="musa")
sp = torch.sparse_coo_tensor(indices, values, (3, 3))
print(sp.to_dense())

# DLPack 零拷贝与 CuPy/Triton/JAX 互通
import cupy as cp
dl = torch.utils.dlpack.to_dlpack(torch.randn(4, device="musa"))
arr = cp.from_dlpack(dl)   # 无复制，两者共享同一块显存

✨

升级路径推荐

1. 先确认 MUSA SDK ≥ 4.3.2，升级驱动。
2. pip install -U torch==2.7.1 torch_musa==2.9.0 torchada。
3. 训练脚本中，若依赖 TF32，手动开启；若使用 FSDP，升级成 FSDP2 的 fully_shard。
4. 详细版本笔记见 Release Notes。