带SHARK的Apple M1 MAX GPU上的PyTorch–比TensorFlow Metal更快

在这个博客中，我们用SHARK演示了PyTorch在Apple M1Max GPU上的训练和推理，只需几行额外的代码，性能优于Apple的Tensorflow metal插件。尽管苹果已经通过现在不受欢迎的Tensorflow macos插件和更新的Tensorflow metal插件发布了对Tensorflow的GPU支持，但最流行的机器学习框架PyTorch在苹果硅上缺乏GPU支持。直到现在。

SHARK基于MLIR和IREE构建，可以无缝瞄准各种硬件。由于SHARK为每个工作负载动态生成内核，因此您可以移植到新的体系结构，如M1Max，而无需供应商提供的手写/手动调整库。

点头ai增加了AMD GPU支持，能够重新定位AMD MI100/MI200类设备的代码生成，并将机器学习工作负载从Nvidia V100/A100无缝转移到AMD MI100。在过去，我们在Intel Alderlake CPU上展示了比Intel MKL和Apache/OctoML TVM更好的codegen，并优于Nvidia的cuDNN/cuBLAS/CUTLASS，后者被Onnxruntime、Pytorch/Torchscript和Tensorflow/XLA等ML框架使用。今天，我们将展示SHARK在M1Max中瞄准苹果的32核GPU，并使用PyTorch模型进行推理和训练。所以，如果你喜欢PyTorch，想在新的Apple Silicon Macbook Pro中使用32个GPU内核，请继续阅读。

在我们的实验中，我们将使用14英寸的MacBook Pro和64GB内存的Apple M1 Max。我们也在16英寸MacBook Pro上运行了相同的基准测试，并注意到了相同的性能，而且在我们的基准测试期间，两者都不会发生热节流。

（基础）anush@MacBook-支持范例%/鲨鱼长凳——模块文件=minilm_jan6_m1max。vmfb——输入函数=预测——函数输入=1x128xi32——函数输入=1x128xi32——函数输入=1x128xi32——基准测试重复次数=102022-02-20T10:17:55-08:00运行/shark benchRun on（10 X 24.1214 MHz CPU）CPU缓存：L1数据64千字节（x10）L1指令128千字节（x10）L2统一4096千字节（x5）平均负载：2.14,1.91，1.81----------------------------------------------------------------基准时间CPU迭代----------------------------------------------------------------------BM_预测/处理时间/实时11.7毫秒1.39毫秒58BM_预测/处理时间/实时11.5毫秒1.34毫秒58BM_预测/处理时间/实时11.7毫秒1.43毫秒58BM_预测/处理时间/实时11.6毫秒1.30毫秒58BM_预测/处理时间/实时11.5毫秒1.33毫秒58BM_预测/处理时间/实时11.7毫秒1.46毫秒58BM_预测/处理时间/实时11.6毫秒1.31毫秒58BM_预测/处理时间/实时11.5毫秒1.33毫秒58BM_预测/处理时间/实时11.7毫秒1.46毫秒58BM_预测/过程时间/实时11.6ms 1.30ms 58BM_预测/过程时间/实时平均11.6ms 1.36ms 10BM_预测/过程时间/实时中位数11.6ms 1.34ms 10BM_预测/过程时间/实时标准差0.074ms 0.063ms 10BM_预测/过程时间/实时cv 0.64%4.65%10（基数）anush@MacBook-专业范例%

参数：名称空间（models=['；microsoft/MiniLM-L12-H384-uncased'；]，模型#源='；第39页；，model#class=None，引擎=['；tensorflow'；]，cache_dir='/cache#u模型'；，onnx_dir='/onnx#U型号'；，使用gpu=True，精度=<；精确浮动32:'；fp32和#39>；，verbose=False，overwrite=False，optimize_onnx=False，validate_onnx=False，fusion_csv='；融合csv和#39；，细节_csv=#39；细节csv和#39；，结果_csv='；后果csv和#39；，输入计数=[1]，测试次数=1000，批量大小=[1]，序列长度=[128]，禁用端口io绑定=假，线程数=[10]）金属设备设置为：Apple M1 MaxsystemMemory:64.00 GBmaxCacheSize:24.00 GBAll model检查点层在初始化TFModel时使用。TFBertModel的所有层都是从位于microsoft/MiniLM-L12-H384-uncased的模型检查点初始化的。如果您的任务与检查点模型训练的任务类似，则无需进一步训练即可使用TFBertModel进行预测。在microsoft/MiniLM-L12-H384-uncased上使用输入形状[1128]运行Tensorflow39；39；39；39；39；39；39；39；39；39；39；39；39；39；39；版本和39；39；39；39；39；39；39；39；39；版本和39；39；39；39；39；39；39；39；39；39；39；39；39；发动机和发动机和39；39；39；发动机和39；39；发动机和39；39；39；39；39；39；39；39；39；39；39；39；39；39；39；39；39；版本和39；版本和39；39；39；版本和39；39；39；39；版本和39；39；39；39；39；版本和39；39；39；39；39；版本和39；39；39；版本和39；39；39；版本和；39；39；39；39；39；39；39；39；39；版本和与与与39；39；39；39；39；39；39；39；版本和和和和和；#39；fp32和#39；'；io#U绑定和#39；'；'；&#模型#名称和#39；'；microsoft/MiniLM-L12-H384-uncated和#39；&#输入和#39；：1和#39；线程和#10, '批量大小'；：1, '序列长度'；：128, '日期时间'；：'2022-02-21 07:48:04.965312', '测试次数'；：1000, '延迟变化'；：'0.00', '潜伏期90%和39%'19.25', '潜伏期95%和39%'22.10', '潜伏期99_39；：'23.14', ' 平均延迟(毫秒#39；：'16.99', 'QPS'；：'58.85'}

在我们的测试中，我们注意到Tensorflow金属插件似乎无法有效地将反向图形卸载到GPU上。由于Tensorflow metal是一个仅二进制版本，我们无法调试它。同样的Tensorflow实现可以很好地卸载到CUDA GPU上。

（基础）anush@MacBook-支持范例%/shark bench——模块文件=bert_training_feb17。vmfb——函数输入=1x512xi32——函数输入=1x512xi32——函数输入=1x512xi32——函数输入=1xi32——输入函数=学习——基准测试重复次数=102022-02-20T23:04:22-08:00运行/（10 X 24.2416 MHz CPU）CPU缓存上的shark benchRun：L1数据64千字节（x10）L1指令128千字节（x10）L2统一4096千字节（x5）平均负载：2.03,2.53，2.31----------------------------------------------------------------基准时间CPU迭代----------------------------------------------------------------------BM_学习/处理时间/实时104毫秒16.2毫秒5BM_学习/处理时间/实时104毫秒16.1毫秒5BM_学习/处理时间/实时105毫秒15.3毫秒5BM_学习/过程时间/实时104毫秒16.0毫秒5BM_学习/过程时间/实时103毫秒15.2毫秒5BM_学习/过程时间/实时105毫秒16.9毫秒5BM_学习/过程时间/实时104毫秒15.5毫秒5BM_学习/过程时间/实时104毫秒14.9毫秒5BM学习/过程时间/实时105毫秒17.4毫秒5BM_学习/过程时间/实时105毫秒15.7毫秒5BM_学习/过程时间/实时平均104毫秒15.9毫秒10BM_学习/过程时间/实时平均104毫秒15.8毫秒10BM_学习/过程时间/实时标准差0.830毫秒0.774毫秒10BM_学习/过程时间/实时标准差0.80%4.86%10（基数）anush@MacBook-专业范例%

（基础）anush@MacBook-Pro nlp_模型%python/bert_small_tf_run。pyMetal设备设置为：Apple M1 MaxsystemMemory:64.00 GBmaxCacheSize:24.00 GB2022-02-21 07:58:36.593857:I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory。cc:305]无法识别平台GPU ID 0的NUMA节点，默认为0。您的内核可能没有使用NUMA支持构建。2022-02-21 07:58:36.594010:I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory。cc:271]创建了TensorFlow设备（/job:localhost/replica:0/task:0/device:GPU:0，内存为0 MB）>；物理可插拔设备（设备：0，名称：METAL，pci总线id:<；未定义>；）1个物理GPU，1个逻辑GPU模型："；目前，34日，34日，34日，34日，3日，3日，3日，3日，34日，3日，3日，34日，4日，3日，3日，3日，3日，3日，3日，34日，34日，3日，34日，3日，3日，34日，34日，3日，3月月，UUUU厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄厄_____===========================================================================================================================礼礼礼礼礼礼礼礼礼礼礼礼礼礼礼礼礼礼礼礼礼礼礼礼礼礼元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元15250176['；输入单词ID[0][0]'；，（无768人）'人；输入_掩码[0][0]''输入_type_id[0][0]'；]辍学1（辍学）（无，768）0['；bert#u编码器#1[0][1]'；]句子预测（分类（无，5）3845['；辍学者#1[0][0]'；]1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1月1_。。。时间：1.7728650569915771时间/iter:0.19698500633239746

功率测量需要在非常受控和仪表化的环境中进行。然而，对于在A100和M1MAX上运行相同的伯特训练模型来说，这是一个很好的近似值。

A100在训练过程中消耗131W的峰值，M1MAX GPU消耗的最大功率为15.4W。A100在7ms与104ms（使用SHARK）和196ms（使用TF Metal）之间进行迭代。由于TF Metal不能很好地将训练图卸载到GPU上，我们将其从性能/功率比较中删除。

2022年2月21日星期一00:56:33+---------------------------------------------------------------+|NVIDIA-SMI 470.57.02驱动程序版本：470.57.02 CUDA版本：11.4 | | CUDA版本：11.4 | GPU名称持久化-M |总线Id显示。A |挥发性解甲。ECC | | |风扇温度性能压水堆：使用率/上限|内存使用率| GPU Util Compute M.| | | | | | |====================================================================================================================================================================================================================================================================================。。。Off |00000000:00:04.0 Off | 0 | | | | | | | N/A 36C P0131W/350W | 39341MiB/40536MiB | 53%默认| | | | | | |禁用|+-------------------------------+----------------------+----------------------+

samples/ModelCompiler/nlp_models$python bert_small_tf_run。皮耶。。型号："；bert#u分类器"__________________________________________________________________________________________________====================================================================================================================================================================================礼礼礼礼礼礼礼元（类型））输出形状（类型）输出（类型）输出形状（类型）层（类型）输出（类型）形状）层（类型）输出（类型）形状）形形形形参数参数参数（类型）层（类型）元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元元UUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUU0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 UUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUU10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 UUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUU卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢卢_________________________________1.0[0 0[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8月月UU\UUUUUUU\\UUUUUUUUUUUU\\UUUUUUUUUUUUU\\\UUUUUUUUUUUUUUUUUU\\\\UUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUU（无，无，768）无）0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 8 8 8 8 8 8 8 8 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8__________________;句子预测（Classifica（None，5）3845辍学者_1[0][0]================================================================================================================================================================================================================================================================================================================================================================================================================UUUUUUU（UUUUUU（据据据据据据据据据据据据据据据据据据据据据据据据据据据据据据据据据据据据UUUUUU月月月月月月月月月月月月月月月月（UUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUU。。时间：6.907362699508667时间/iter:0.006977134039907744

封装功率：2.48W（平均功率：2.34W峰值：37.78W）油门：无CPU：1.37W（平均功率：1.03W峰值：32.32W）GPU：0.00W（平均功率：0.00W峰值：15.41W）

SHARK运行时以pip包的形式提供，需要在Pytorch/Python文件中更改几行代码。下面是一个使用SHARK从Python运行Resnet50和BERT的示例，我们计划在这里添加更多Python示例。对于本文中的实验，我们将使用一个使用谷歌基准测试支持构建的基准二进制文件来运行一个受控实验，尽管从Python调用会增加几微秒。我们测试了一个BERT MiniLM模型（用于HuggingFace的无限演示）和一个BERT训练模型。您还可以尝试bert base uncased、Resnet50、Mobilenetv3等，所有型号都会通过torch mlir自动卸载到GPU。

BERT训练图的火炬MLIR降低正在该阶段分支中集成。所有重新创建测试的代码都在这里和这里。但是，由于torchvision还不支持M1，您需要从源代码处安装PyTorch torchvision。您还需要从SHARK存储库的apple-m1-max-support分支构建SHARK。

（基础）anush@MacBook-Pro示例%pip列表| grep TensorFlow TensorFlow估计器2.6.0传感器流量-macos 2.8.0传感器流量-metal 0.3.0（基本）anush@MacBook-专业示例%pip list | grep onnx onnx 1.10.1非通用1.8.1（基本）anush@MacBook-专业示例%pip list | grep tensortensortensorboard 2.8.0tensorboard-data-server 0.6.0tensorboard-plugin-wit 1.8.0tensorflow-estimator 2.6.0tensorflow-macos 2.8.0tensorflow-metal 0.3.0（基本）anush@MacBook-专业示例%pip列表| grep火炬1.10.0或视觉0.9.0a0（基本）anush@MacBook-专业示例%pip列表| grep onnxy（基本）anush@MacBook-专业示例%pip list | grep onnx onnx 1.10.1非通用1.8.1（基本）anush@MacBook-专业范例%

我们使用asitop监控电源使用情况、GPU使用情况、内核节流等，这些都可以通过pip安装。

苹果的CoreML不仅能够针对CPU或GPU，还能够针对苹果的神经引擎，尽管只是用于推理。我们确实试图让CoreML进行推理比较，但在这里遇到了模型转换问题，并将其排除在测试之外。最新的coremltools似乎需要旧的Tensorflow版本2.5.0，这不再是使用conda安装Tensorflow时的默认版本。

当所有GPU被最大限度地使用时，OSX窗口服务器崩溃。我们已经就这个错误向苹果公司提出了反馈问题，希望它能很快得到解决，但在重新创建结果时，我们建议在MacBook Pro中使用ssh，不要打开显示屏。

这些存储库中的所有代码都是正在进行的工作，并以技术预览的形式显示，这需要一定程度的润色，然后才能供非技术用户使用。我们计划继续使其用户友好，并为Torch MLIR添加渴望模式支持，因此M1Max GPU上的PyTorch支持为无缝开发和部署提供了条件。如果你不想把所有的开源部分都放在自己的位置上，或者如果你有在M1设备上部署支持GPU的PyTorch和专业解决方案的商业案例，请在这里注册，我们的解决方案团队将伸出援助之手。

使用SHARK Runtime，我们在Apple M1Max GPU上演示了高性能的PyTorch模型。它比Tensorflow Metal的推理能力强1.5倍，训练BERT模型能力强2倍。在不久的将来，我们计划增强最终用户体验，并添加“渴望”模式支持，以便在任何硬件上实现从开发到部署的无缝连接。

如果您想访问SHARK Runtime的商业版，请在此处注册访问，如果您在重新创建结果时遇到问题，请打开一个问题。

SHARK建立在开源软件包Torch MLIR、LLVM/MLIR和Google iRE之上，我们感谢所有开发者和社区的支持。我们特别希望呼吁本·瓦尼克、张磊、斯特拉·劳伦佐和托马斯·拉乌克斯在MLIR/IREE GPU代码生成路径上提供帮助、支持和指导。