MegEngine Windows Python wheel 包减肥之路
作者:张浩龙 | 旷视科技 MegEngine 架构师
写在之前
本文的目的
通过讲述在支持 MegEngine Windows Python wheel 过程中遇到的问题以及解决问题的流程,此文最后的解决方法可能不是最优,欢迎留言指正。
过程中顺便科普一些关于 MegEngine 的构建以及构建时用到的基础东西,当然这些基础知识我相信是工程之道经常会用到的,包括但不限于:
编译、链接、符号隐藏,符号 export 等。此处先推荐一本 “老书”《程序员的自我修养》,自然它没有 xxx 四库全书让人绞尽脑汁,但是它里面的基础知识依然是目前我们和计算机“交流”中经常遇到的。
Python wheel 包构建
MegEngine 各平台支持情况
cpp 推理支持情况:
TEE:https://en.wikipedia.org/wiki/Trusted_execution_environment
训练:
Python 侧:
目前官方发布的 wheel 包,只有 Windows-X64-CPU-CUDA,many Linux 64bit -X64-CPU-CUDA,MacOS-X64-CPU,其他的可自己编译,或者社区提单索取。
cpp 侧训练支持情况和上面的 cpp 推理情况一致。
从上面的情况,可看见 MegEngine 无论训练还是推理,还是各种硬件,还是各种 OS 都支持的非常全面,如有需求,不妨试用!!!!
遇到的问题
为了全面的支持上面提到的 MegEngine 各个平台,各个 OS,期间或多或少会遇到一些问题,比如 Windows 平台上 Python wheel 包体积过大。
先看一下目前 MegEngine wheel 包体积大小,摘自 1.7 版本 pypi
其中因为 Linux 和 Windows 支持了 CUDA,所以包体积在 900MB 左右,这是一个正常的 size。
在之前 Windows CUDA 包体积在 1.7G 左右:这就是后面尝试分析和修复的问题。
先对问题 MECE 一下
MECE 是 Mutually Exclusive Collectively Exhaustive 缩写,中文意思是“相互独立,完全穷尽”。也就是对于一个问题的议题,能够做到不重叠、不遗漏的分类,而且能够藉此有效把握问题的核心,并解决问题的方法。强调两点:
各部分之间相互独立(MutuallyExclusive)
化简后, 感觉就是分析问题时,可能的方法要尽可能的独立,尽量不要有交集
所有部分完全穷尽(CollectivelyExhaustive)
化简后, 感觉就是分析问题时,尽可能的要把方法想全,尽量不要有遗漏
为啥会选择一条鱼呢:鱼头附近一般都比较大(胖),而越往鱼尾走,会越来越小(瘦),从而希望通过这个流程“Windows wheel 减肥”之路,达到减肥的目的。
问题的影响
这样的体积会有什么问题呢(毕竟人太胖也会有些副作用 xxx etc.)
首先就是体积超过了我们申请的 pypi 上单个文件最大体积限制
给用户体验不好, 为什么相同的版本, Window 比 Linux 大这么多呢
Windows 上显存比 Linux 占用大很多(估计提到这点,大家已经猜测到问题所在)
解决这个问题可能需要的相关知识
问题给人的第一印象是:
编译构建相关的
Python wheel 打包相关的
Windows OS 独有的
先就 MegEngine 如下基础知识做一些基础补充(减肥前总得有一些科普吧,到底是吃药还是锻炼,或者具体到吃什么药吧)
MegEngine CMake 构建流程
其中 CMake 描述主要在:
此文件包含很多的 option,主要用于控制是否编译一些模块,比如是否打开 MGE_BUILD_WITH_ASAN 用于调试内存问题
此文件包含对各种 ARCH 的适配控制, 比如编译 X86_64 还是 AARCH64 等
此文件包含对各种 OS 的适配, 比如是编译 Linux,Android 还是 Windows 等
以及一些杂项配置,比如有优化等级,比如 CUDA SM 的配置等
此文件包含了 MegEngine 核心代码 MegBrain 层所有源代码的编译管理
此文件包含了 MegEngine 核心代码 dnn (主要实现各种 backends)层所有源代码的编译管理
一些杂项 CMakelist
各种 example,比如 lite example
各种 test,比如 megbrain test
各种 helper,见 helper module
其中 Ninja 提供丰富的可视化调试功能,下面列举如何通过 Ninja debug server 来看 MegEngine 部分模块的构建依赖
执行 host_build.sh 来进行 host 编译,同时它会在 build dir 生成整个构建依赖描述文件 build.ninja
更多的编译支持请参考 BUILD_README.md
build.ninja 文件功能类似 GNU Makefile
有了 build.ninja 后,便可进行一些调试
MegEngine wheel 构建流程
有了上面 Ninja 编译出来的各种库后,我们就可以将它们和 MegEngine 中的 py src 一起进行打包,最后生成可安装,可分发的 Python wheel 包了
构建流程主要说明在 BUILD_PYTHON_WHL_README
描述了目前 MegEngine python wheel 的支持状态
自己本地构建需要的一些 env 准备
一些使用说明
MegEngine wheel 包遵从 pep-0571
包 setup 入口在 python wheel setup
调用 setup 前各个 OS 的准备差异化在 wheel scripts
有人可能会问,为什么不使用 auditwheel 来自动管理 wheel 包中的 so 依赖,有两个原因
auditwheel 不支持所有的操作系统,比如 Windows
auditwheel 不支持依赖库使用 dlopen 的情况
auditwheel 不支持 subpackage 的 wheel 包
当你执行 python3 -m pip install megengine -f https://megengine.org.cn/whl/mge.html 后,可以 import megengine,也可以 import megenginelite,是因为 megengine 和 megenginelite 均会存在安装的包中,且他们会复用 megengine_shared 这一体积超大的库
MegEngine 构建上如何适配 Windows
上面介绍了 MegEngine 基于 CMake 的构建基础和使用 Ninja 自带的调试功能以及帮我们从宏观了解了一下 project 的编译依赖和进行一些常规调试,下面再介绍一下 MegEngine 是如何适配 Windows 平台的。
首先 MegEngine 大部分的源代码都是 c++,且 cpp 推理要求是
c++14
,编译 Python 训练要求是c++17
各家编译器其实对这些标准实现不是完全一致的,抛开和系统相关的,比如 POSIX 外,其实还有比较多基本的上层用法各家编译器其实时不太兼容的,特别是明显的是 gcc 和 clang 能编译过的代码,Windows cl.exe 其实是编译不过的。
为了解决上面提到的两个问题
尽可能的抛弃 cl.exe,Windows 上使用 llvm-clang-cl 进行构建
再加上,上面提到的 CMake, Ninja 本身是跨平台的,这样一组合,MegEngine 便原生支持了 Windows,注意不是基于 WSL 的哦
问题简单的分析
在上面“MegEngine CMake 构建流程”小节中,我们提到了 Ninja debug server 能够帮忙可视化整个构建组件的依赖关系,下面我们补充一下在问题修复前 Windows 和 Linux 下 imperative 依赖的可视化结果。
Linux 下
Windows 下
Windows 和 Linux 下最大的差异化如下:
Linux 下 imperative 是依赖的 libmegengine_shared.so
Windows 下 imperative 是依赖的 megbrain 和 megdnn,又因为 megbrain 和 dnn 在 CMake 这边其实一个 OBJECT,所以相当于直接依赖他们的 .obj 了
初步结论:
MegEngine wheel 包,有两个 Python module 接口
MegEngine 用于 python 侧训练的基础接口
imperative: 当你 安装 完成 MegEngine 时,在 Python 中输入 import megengine 时。加载的就是它,我们提供了一些入门的教程供您快速上手 MegEngine
MegEngineLite:易用的 cpp,Python 推理接口
当你使用 MegEngine 完成训练模型后,可参考 部署 文档使用 MegEngineLite,快速将你的模型部署落地。
因为有这两个顶层构建目标的存在,且他们在 Windows 和其他 OS 上,依赖底层的目标不同,导致了问题的产生
为什么不同的依赖关系,会产生这么大的体积区别呢,先看一张 MegEngine 的架构图
从下往上依次是:
MegEngine 不同 backends 的差异化实现被封装到了 dnn(对应到上图的“硬件层”,对应到 CMakeList 中的 megdnn 模块), 而其中 CUDA backends 因为有大量的 kernel 以及对较多的 SM 支持,会对整个库或者可执行程序体积产生大量的体积贡献
图中“硬件抽象层”,部分“核心组件层”,对应 CMakeList 中的 megbrain 模块
在往上“接口层”,对应 CMakeList 中的 imperative 和 MegEngineLite 模块。
由于上述的原因,在 Windows 平台, imperative 模块和 MegEngineLite 模块会同时静态依赖 dnn 和 megbrain 代码,导致体积几乎翻倍。问题修复前的依赖图:
可能的解决方案
通过上面的分析,问题原因已经找到,再来猜想一下
为什么 Windows 平台上和其他平台目标依赖有差异
Windows class member 不能隐式的被 export,需要显式的使用 dllexport 和 dllimport,详细见 Microsoft Specific
dnn,megbrain 层有大量的 data 数据访问并没有抽象成函数,而是需要直接访问数据成员下面举一个栗子来说明跨 dll 动态库访问数据成员方式的差异,主要包含三支文件 api.h 和 api.c 实现函数 func_a 和 定义一个变量 a,被编译成动态库 dllclient.c 会调用上面 api.c 实现方法和访问变量 a
跨 dll 动态库直接访问数据成员方式
跨 dll 动态库通过函数访问数据成员方式
改造上面的 example 代码, 把其中变量 a 的访问封装到一个函数
可以看见在 Windows 上函数符号和数据成员符号 export 是等价的,但是 import data 要求要严格的多
在 Linux, MacOS 下, 函数符号和数据成员符号 export 属性是等价的
由于上面提到种种限制,导致在最初支持 Windows 平台时,所有的上层目标(MegEngine,MegEngineLite)都必须静态依赖 megbrain 和 dnn。
既然问题原因已经找到,需要修复这个问题的目标就变的非常清晰了:让 megengine_shared 动态库 (dll) 在 Windows 平台上可用。
列举一下可能的方案:
方案一:CMake 自带的 WINDOWS_EXPORT_ALL_SYMBOLS
结论:不太适用 MegEngine 这类“大”工程
原因:MegEngine 符号太多,超过了 link.exe max symbols 65536 的限制 (使能 CUDA 时,大约有 1.7W 个符号)
分析 CMake WINDOWS_EXPORT_ALL_SYMBOLS 的原理,能否中间加一些 hook 来过滤不需要 export 的符号,以达到类似 gcc/clang -Wl,--version-script 的效果,cmake 对他的处理逻辑:
(stage a): 生成 CMakeFiles/megengine_export.dir/exports.def.objs,本质是 obj 的集合
(stage b): 插入 PRE_LINK stage 生成 CMakeFiles/megengine_export.dir/exports.def (此文件类似 gcc/clang -Wl,--version-script)
(stage c): LINK_FLAG 自动插入 /DEF exports.def
CMake 提供了对 stage a output 的 hook,意思是可以修改 exports.def.objs,但是没有机会修改 exports.def
加入 hook command,把 exports.def.objs 中所有 DNN 的 obj 删除,想象中应该可以了
但是 imperive 和 megenginelite,不仅仅是和 megbrain 打交到,很多直接使用 dnn 的接口和数据成员
方案二:“优化”版本的 WINDOWS_EXPORT_ALL_SYMBOLS
如上面分析 WINDOWS_EXPORT_ALL_SYMBOLS 有一定的缺陷,会把所有的 obj 的符号全部 export,那能不能手动修改 WINDOWS_EXPORT_ALL_SYMBOLS 生成的 exports.def
保留必要的 symbols
CMakeList 中目标依赖修改过后的 exports.def (让其符号不超过 65536)
结论:不可行
Windows cl linker.exe 不支持 * 通配符,不支持存放一个不存在的符号,导致一旦放了固定的 exports.def,稍微更改一个编译参数,或者加点代码,都会编译不过
方案三:到最后发现没有一个“偷懒的”方式来解决这个问题,回退到最 naive 的方式
把 megbrain、dnn、megenginelite 对外暴露的 API 依赖的成员符号全部显式的加上 declspec(dllexport) 和 declspec(dllimport) 属性描述
修复示例, 完整修改见 PR
想象未来更好的解决方法:
修改 CMake 本身源代码,让 flag WINDOWS_EXPORT_ALL_SYMBOLS 支持用户自定义 filter,让其生成的 exports.def 本身就是带用户过滤参数的
当然因为 Windows 数据成员在 import 部分处还必须显式的加上 dllimport,对这块似乎 CMake 也无能为力
可以考虑工程一开始设计时,API 尽可能的不要存在隐式的数据成员之间的访问,尽可能的将其转换成一个函数 API
更多 MegEngine 信息获取,您可以:查看文档、和 GitHub 项目。欢迎参与 MegEngine 社区贡献,成为 Awesome MegEngineer,荣誉证书、定制礼品享不停。
评论