MindSpore编译构建后Vmap模块的RuntimeError问题

技术背景

这篇文章来源于MindSpore仓库中的一个Issue，简单描述问题就是，如果你用MindSpore开发了一个python软件供别人使用，那么很有可能涉及到编译构建的问题。但是如果直接使用编译好的whl包去运行的话，就有可能出现一个跟Jit即时编译有关的报错，这里Jit在其他的一些模块中也会被使用到，比如Vmap函数和Grad函数等。

问题复现

我们通过一个简单的测试案例来复现一下这个RuntimeError。

软件环境

-- MindSpore version: 2.2.11
-- Python version: 3.9.16
-- OS platform and distribution: Linux Ubuntu 20.04

执行模式

GRAPH_MODE

复现代码结构

debug/
├── debug
│   ├── __init__.py
│   └── test_vmap.py
├── example
│   └── test.py
└── setup.py

代码内容

setup.py

from setuptools import setup, find_packages
from pathlib import Path
this_directory = Path(__file__).parent

setup(
    name="debugs",
    version="1.0",
    description="Test Vmap",
    license="Apache 2.0 Licence",
    packages=find_packages(),
    platforms="any",
    scripts=[],
    include_package_data=True
)

test_vmap.py

from mindspore import ops

func = lambda ts: ts
batch_func = ops.vmap(func, in_axes=(0, ), out_axes=(0, ))

test.py

# import sys
# sys.path.insert(0, '../')

from debug.test_vmap import batch_func
from mindspore import Tensor

a = Tensor([0, 1])
batch_a = batch_func(a)
print (batch_a)

复现步骤

1. 按照目录树结构构建好本地测试复现环境；
2. 取消test.py文件中的两行注释，运行python3 test.py；
3. 通过setup.py构建安装python3 setup.py install，然后把test.py文件的头两行注释掉，再次运行python3 test.py。

预期结果

两次运行test.py都输出[0 1]

报错信息

第一次测试，为了方便代码阅读，使用了lambda函数，后来发现这个报错可能是即时编译跟lambda函数不兼容所导致的：

后来修改成了普通的函数，才复现了一个RuntimeError：

test_vmap.py

from mindspore import ops, jit

@jit
def func(ts):
    return ts

batch_func = ops.vmap(func, in_axes=(0, ), out_axes=(0, ))

得到的报错信息为：

Traceback (most recent call last):
  File "test.py", line 8, in <module>
    batch_a = batch_func(a)
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/common/api.py", line 718, in staging_specialize
    out = _MindsporeFunctionExecutor(func, hash_obj, input_signature, process_obj, jit_config)(*args, **kwargs)
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/common/api.py", line 121, in wrapper
    results = fn(*arg, **kwargs)
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/common/api.py", line 350, in __call__
    raise err
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/common/api.py", line 344, in __call__
    phase = self.compile(self.fn.__name__, *args_list, **kwargs)
  File "/home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/common/api.py", line 435, in compile
    is_compile = self._graph_executor.compile(self.fn, compile_args, kwargs, phase, True)
RuntimeError: 'VmapOperation' arg0 must be a 'Function' or 'Cell', but got AbstractScalar(Type: External, Value: PythonObject(type: <class 'function'>, value: <function func at 0x7f75c0375ef0>), Shape: NoShape).

----------------------------------------------------
- C++ Call Stack: (For framework developers)
----------------------------------------------------
mindspore/ccsrc/frontend/operator/composite/composite.cc:1420 CheckVmapFunc

----------------------------------------------------
- The Traceback of Net Construct Code:
----------------------------------------------------
# 0 In file /home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages/mindspore/ops/composite/base.py:686
            return vmap_(fn, in_axes, out_axes)(*args, **kwargs)
                   ^
 (See file '/tmp/debug/example/rank_0/om/analyze_fail.ir' for more details. Get instructions about `analyze_fail.ir` at https://www.mindspore.cn/search?inputValue=analyze_fail.ir)

问题原因

按照开发人员的回复，这个问题是因为Vmap会使用到静态图模式。但是编译成whl包之后，静态图模式下会把whl包的内容识别为第三方仓库。从第三方仓库来的函数，有可能出现无法入图的问题：

解决的方法就是，把相应的仓库添加到环境变量中去：

export MS_JIT_MODULES=debug

然后再次运行test.py，问题就解决了。

总结概要

在MindSpore编译计算图的过程中，会把从编译构建好的whl包中引入的模块视为第三方库，也就没有办法在即时编译的阶段入图。普通的math和numpy等第三方库不入图也不会影响计算。但如果是基于MindSpore本身开发的一些函数，如果用到了Jit、Grad和Vmap，那么有可能出现无法入图的问题，就会出现RuntimeError报错。解决方法就是设置一个跟即时编译有关的环境变量，把相关的第三方包引用修改为内部引用。

版权声明

本文首发链接为：https://www.cnblogs.com/dechinphy/p/ms_jit_module.html

作者ID：DechinPhy

更多原著文章：https://www.cnblogs.com/dechinphy/

请博主喝咖啡：https://www.cnblogs.com/dechinphy/gallery/image/379634.html