在昇腾(Ascend)全栈 AI 软硬件体系中,MindSpore 框架凭借其全场景协同和极简开发的特性,成为了许多开发者的首选。

对于刚接触 MindSpore 的开发者来说,最常见的一个困惑就是:PyNative 模式和 Graph 模式到底有什么区别?在昇腾 NPU 上我该如何选择?

今天这篇干货文章,我们将剥离复杂的理论,通过代码实例,带你深入理解 MindSpore 的两种运行模式,并掌握如何利用 @jit装饰器实现“开发调试灵活”与“运行性能极致”的完美平衡。

一、 两种模式的本质区别

MindSpore 提供了两种计算图执行模式:

1. PyNative 模式 (动态图)

  • 特点:Eager execution(即时执行)。代码运行到哪一行,算子就执行到哪一行。
  • 优势:调试极其方便。你可以使用 Python 原生的调试工具(如 pdb)或直接 print打印中间 Tensor 的值,逻辑控制完全遵循 Python 语法。
  • 劣势:性能相对较低,无法进行全局的图优化。

2. Graph 模式 (静态图)

  • 特点:先编译,后执行。MindSpore 会先将 Python 代码编译成一张中间表达(IR)的计算图,经过图优化(算子融合、内存复用等)后,下沉到 Ascend 芯片上执行。
  • 优势:性能极致。特别是在昇腾 910/310 NPU 上,静态图能最大化利用芯片算力,显著减少 Host 与 Device 之间的交互开销。
  • 劣势:调试较难,部分 Python 动态语法受限。

二、 代码实战:模式切换与对比

MindSpore 通过 set_context接口来控制全局的运行模式。

场景一:使用 PyNative 进行调试

假设我们正在编写一个自定义的损失函数,逻辑比较复杂,我们需要确保每一步计算都是正确的。

import mindspore as ms
import mindspore.ops as ops
from mindspore import Tensor
import numpy as np

# 设置为 PyNative 模式,方便调试
ms.set_context(mode=ms.PYNATIVE_MODE, device_target="Ascend")

def custom_operation(x, y):
    # 在 PyNative 模式下,这里可以随意 print 调试
    z = ops.add(x, y)
    print(f"Debug Info - Intermediate value z: {z}") 
    res = ops.mul(z, 2.0)
    return res

x = Tensor(np.ones([2, 2]), ms.float32)
y = Tensor(np.ones([2, 2]), ms.float32)

output = custom_operation(x, y)
print(f"Final Output:\n{output}")

运行结果分析: 你会看到 Debug Info被打印出来。这意味着 Python 解释器逐行执行了代码,这对于排查 NaN值或维度不匹配问题至关重要。

场景二:切换到 Graph 模式加速

当调试完成,我们希望代码在昇腾卡上飞快地跑起来。

# 切换为 Graph 模式
ms.set_context(mode=ms.GRAPH_MODE, device_target="Ascend")

# 注意:在 Graph 模式下,编译期间可能不会打印 Python 层面的 print
# 如果需要打印 Tensor 值,通常需要使用 ops.Print 算子
def custom_operation_fast(x, y):
    z = ops.add(x, y)
    res = ops.mul(z, 2.0)
    return res

output = custom_operation_fast(x, y)
print(f"Graph Mode Output:\n{output}")

底层发生了什么?当你运行这段代码时,MindSpore 首先将 custom_operation_fast及其依赖编译成一张静态计算图。在昇腾 NPU 上,这张图会被进一步优化(例如 add和 mul可能会被融合),然后整图下沉执行。

三、 进阶技巧:混合执行(JIT)

在实际的大模型训练中,我们不需要非黑即白地选择全图 PyNative 或全图 Graph。MindSpore 提供了 @jit装饰器(早期版本称为 ms_function),让我们可以在 PyNative 模式下,将核心计算函数编译为静态图运行。

这是目前昇腾开发中最推荐的“黄金搭档”模式:全局 PyNative 控制流程 + 核心函数 JIT 加速。

最佳实践示例

import mindspore as ms
from mindspore import nn, ops, Tensor
import time

ms.set_context(mode=ms.PYNATIVE_MODE, device_target="Ascend")

class SimpleNet(nn.Cell):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc = nn.Dense(1024, 1024)
        self.relu = nn.ReLU()

    def construct(self, x):
        return self.relu(self.fc(x))

net = SimpleNet()
loss_fn = nn.MSELoss()
optimizer = nn.Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)

# 定义前向和反向传播的组合函数
# 使用 @jit 装饰器,将此函数及其调用的子函数编译为静态图
@ms.jit 
def train_step(data, label):
    def forward_fn(data, label):
        logits = net(data)
        loss = loss_fn(logits, label)
        return loss, logits
  
    grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)
    (loss, _), grads = grad_fn(data, label)
    optimizer(grads)
    return loss

# 模拟输入数据
input_data = Tensor(ms.numpy.randn(32, 1024), ms.float32)
label_data = Tensor(ms.numpy.randn(32, 1024), ms.float32)

# 预热(触发图编译)
print("Start warmup (Graph Compilation)...")
train_step(input_data, label_data)
print("Warmup done.")

# 性能测试循环
start_time = time.time()
for _ in range(100):
    loss = train_step(input_data, label_data)
end_time = time.time()

print(f"100 steps finished. Average time per step: {(end_time - start_time) / 100 * 1000:.2f} ms")

为什么这样做是最佳实践?

  1. 灵活的 Python 流程控制:数据加载、预处理、模型保存等非计算密集型操作保留在 PyNative 模式下,方便编写复杂的 Python 逻辑。
  2. 极致的算力释放:核心的 train_step(包括前向计算、梯度求导、权重更新)被 @jit封包,编译成静态图在 NPU 上全速运行。
  3. 算子融合:在 train_step内部,MindSpore 的编译器会自动识别可以融合的算子(例如 Conv + BatchNorm + ReLU),大幅减少内存读写次数。

四、 避坑指南

在使用 Graph 模式或 @jit时,有几点需要特别注意:

  1. 控制流限制:在静态图编译范围内,尽量避免使用复杂的 Python 原生控制流(如依赖 Tensor 值的 if条件判断),除非使用 ops.cond或 ops.while_loop等等效算子。
  2. 第三方库限制:在编译函数内部,尽量不要调用 numpy 等第三方库进行计算,因为编译器无法识别 numpy 的算子,这会导致这些操作回退到 Host 侧 CPU 执行,阻断图的优化。
  3. 副作用处理:如果需要在静态图中修改全局变量(如更新 Parameter),请确保逻辑符合 MindSpore 的副作用处理机制。

五、 总结

在昇腾 NPU 上开发 AI 模型,MindSpore提供了强大的工具链。

  • 刚开始写模型结构、调试 Loss 函数时,请毫不犹豫地使用 ms.set_context(mode=ms.PYNATIVE_MODE)
  • 当模型跑通,需要大规模训练或追求性能时,使用 @jit装饰核心训练步,或者全局切换到 GRAPH_MODE

掌握好动静结合的开发模式,你就能在昇腾社区的大模型开发浪潮中游刃有余!

Logo

作为“人工智能6S店”的官方数字引擎,为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。

更多推荐