如果说强大的算力是驱动AI发展的动能，那么完善的生态就是让算力在AI发展中自由流动的通途。

今天，我们很高兴地宣布IPU生态又增添了一位“新成员”——OpenMMLab现已支持Graphcore IPU作为其训练芯片。OpenMMLab系算法框架一直以开箱即用和高度可扩展化享誉AI社区，而Graphcore聚焦于提供通用的高性能算力。借助OpenMMLab在视觉领域多元化的生态，Graphcore将能为广大的客户快速提供不同场景的解决方案。仅需加入一个简单的参数，OpenMMLab的用户就可以轻松将模型训练部署到IPU机器上，充分利用IPU所提供的高效算力，加速推进AI的边界。

本文将简要介绍现已完成的IPU上OpenMMLab的支持工作，并展开说明在OpenMMLab上调用IPU的具体操作。

OpenMMLab

OpenMMLab是深度学习时代最完整的计算机视觉开源算法体系，迄今OpenMMLab已开源了超过20个算法库，包括图像分类、目标检测、语义分割、动作识别等，共开源了300+种算法，提供了2000+预训练模型。

从2018年10月启动开源至今，OpenMMLab在GitHub上总计获得超过5万star，用户遍及全球110个国家和地区，吸引了超过1000位来自国内外的代码贡献者，在国内人工智能算法开源领域遥遥领先，受到了学术界和工业界的广泛好评，在多个国际学术竞赛上被冠军团队所采用，并且在多家顶尖互联网企业和众多中小企业的实际业务中得到广泛使用。在中国开源AI技术梯队里遥遥领先于其他单个垂直领域的算法和框架。

MMCls X IPU

目前，IPU已经接入OpenMMLab各算法库的底层基础库MMCV。基于MMCV对IPU的支持，分类模型框架MMClassification（MMCls）等上层算法库现可调用IPU作为训练芯片。用户只需要加入一个简单的命令行参数“—ipu-replicas”便可毫不费力地将模型训练部署到IPU机器上。

以ResNet的训练为例，在按照get_started文档完成MMCV x IPU环境部署后，只需在MMCls文档下输入指令：

“python3 tools/train.py configs/resnet/resnet18_8xb16_cifar10.py –ipu-replicas 1”

便可轻松使用IPU加速完成ResNet在cifar10数据集上的训练。

设计方案

PyTorch在IPU上的部署主要涉及以下4个方面：

1. IPU options的配置
2. PyTorch模型的修改（混精、模型切分、定义输入输出）转换以及编译
3. 优化器的转换以及参数同步
4. DataLoader的转换（optional）

由于这些转换涉及到一个模型训练的不同部分，在一般模型训练/训练框架的IPU部署上，这些实现难免会分散在工程的不同地方导致代码耦合，需要处处判断GPU/IPU/CPU的环境问题。容易导致大量冗余代码出现，还需要用户介入以进行一些额外的配置处理和判断。这一直是训练代码在不同硬件平台上进行适配的一个痛点。

MMCV创造性地设计了名为Runner的核心组件，可以较好地解决这个问题。该核心组件负责整个模型训练过程的调度，目的是让用户能够以更少的代码进行灵活的、可配置的训练。它具备如下核心特性：

• 支持以EpochBasedRunner和IterBasedRunner为单位的迭代模式以满足不同场景
• 支持定制工作流以满足训练过程中各状态自由切换，目前支持训练和验证两个工作流。工作流可以简单理解为完成训练和验证的迭代过程
• 配合默认和自定义HOOK，对外提供灵活扩展能力
• 在Runner内部可以轻松访问模型训练需要的各种资源，如dataloader、model、optimizer等

由以上特性可以发现，如果使用Runner来统一判断硬件环境，管理调度IPU的各项配置部署，所有的修改都可以放在一处，实现了代码层面的解耦去冗。同时，原本需要用户管理的各项配置，也可以通过Runner对其内部资源和状况的判断自动化处理。如此，模型训练在不同硬件平台之间的切换便可以由框架自动完成，达到用户对此几乎没有感知的效果。

具体而言，我们对Runner进行扩展，额外实现了ipu_runner，实现代码在mmcv/device/ipu/runner.py中。

其中，IPU适配需要4个操作，其中前3个将在ipu_runner的初始化中完成，第4项dataloader的转换将在runner.run函数中判断完成以加速IPU上训练数据的加载。

我们在前端（MMCls/MMDet）只需要加入少量的参数在前端（MMCls/MMDet），我们只需要加入少量的参数解析，就可以开启ipu_runner，进行以上4个部分的处理以及模型训练。具体流程见下图：

如上如上图所示，相比于CPU模式，IPU模式将会多出4个部分的具体实现（加粗部分），它们分别对应了前文所述的4个方面。

同当涉及IPU与host主机之间的通信时，如参数同步或者参数修改，也会由Runner自动加入额外的HOOK来实现。

IPU性能榨汁，以ViT训练为例

除了通用用法以外，我们还增加了专属于IPU AI训练加速的配置。以Vision Transformer（ViT）为例，我们将ViT模型切分成4个部分，分配在4块IPU上，并充分发挥IPU的可扩展和混合精度训练的能力，在IPU-POD₁₆上开启4个replicas进行数据并行训练，成功地将批尺寸扩展到了2176，只用了短短半小时便完成了ViT的微调。

配置文件详见支持了IPU的MMCls版本中的“ configs/vision_transformer/vit-base-p16_ft-4xb544_in1k-224_ipu.py”。纯模型训练的Tput更是达到了6000+，远高于8张A100的训练速度。下表是对比了8张A100在ViT微调的测试结果，8个epoch，224输入的尺寸。

device	size	精度	acc	Tput/images（纯模型训练）	实际训练时间/minutes
IPU	224	混精	81.2	6250	31
GPU	224	TF32	81.2	3686	45

在部署完IPU环境（教程见后文）以后，只需运行如下代码便可以复现该实验。

bash demo/ipu_train_example.sh

log 如下：

更多的细节（log，权重，配置）请参见mmcls/model zoo

（https://mmclassification.readthedocs.io/en/latest/papers/vision_transformer.html）

环境部署教程

1. 向Graphcore申请获取IPU机器和SDK包
2. 在机器上，安装mmcv-full(增加了IPU feature的最新版本)
3. 解压SDK包，得到路径<sdk_path>
4. 安装pytorch的IPU开发套件poptorch，pip install <sdk_path>/poptorch_x.x.x.whl
5. 激活IPU基础环境，source <path to poplar installation>/enable.sh & source <path to popart installation>/enable.sh
6. 验证是否安装完成，输入 python3 -c “import mmcv;import popart;import poptorch;”