05.srt

1
00:00:00,000 --> 00:00:04,560
字幕生成: BLACK 字幕校对: 杨绎

2
00:00:05,560 --> 00:00:07,560
哈喽大家好,我是ZOMI

3
00:00:07,560 --> 00:00:12,000
今天回到一个AI编译器里面的PyTorch这个系列

4
00:00:12,000 --> 00:00:15,160
给大家汇报一下AoT Autograde

5
00:00:15,160 --> 00:00:17,560
那AoT就Ahead of Time

6
00:00:18,160 --> 00:00:21,240
这个所谓的Time是指真正执行之前

7
00:00:21,240 --> 00:00:23,960
去做的自动微分的功能

8
00:00:25,160 --> 00:00:29,520
现在来看一下今天要给大家汇报的一个技术点

9
00:00:29,520 --> 00:00:32,240
首先简单的回顾一下

10
00:00:32,240 --> 00:00:34,880
之前讲了PyTorch 2.0的一个新特性

11
00:00:34,880 --> 00:00:39,240
然后又去看了一下Torch Dynamo的一个解读

12
00:00:39,240 --> 00:00:42,920
还有PyTorch关于静态图的一些尝试的方式

13
00:00:42,920 --> 00:00:46,960
今天来到AoT Autograde Ahead of Time Autograde

14
00:00:46,960 --> 00:00:50,320
那Autograde里面主要分开三个内容给大家介绍的

15
00:00:50,320 --> 00:00:53,280
一个就是Autograde的一个具体的实现方式

16
00:00:53,280 --> 00:00:55,280
接着看一下Autograde的效果

17
00:00:55,320 --> 00:00:57,040
那因为Autograde是严重

18
00:00:57,040 --> 00:01:01,240
或者对Torch Dispatch这个机制依赖非常重

19
00:01:01,240 --> 00:01:05,360
所以会单独的去讲一讲Torch Dispatch这个机制

20
00:01:05,360 --> 00:01:07,640
还有它对应的原因

21
00:01:08,760 --> 00:01:10,960
那现在来简单的回顾一下

22
00:01:10,960 --> 00:01:14,760
Torch Dynamo里面的一些具体的Concept

23
00:01:14,760 --> 00:01:16,960
这里面我总结了两条

24
00:01:16,960 --> 00:01:20,240
第一条就是Torch Dynamo里面主要是根据

25
00:01:20,240 --> 00:01:23,640
在Python真正解析实现之前的Cpython里面

26
00:01:23,640 --> 00:01:26,080
去修改了Cpython的Python的Bytecode

27
00:01:26,080 --> 00:01:27,560
就是它的字节码

28
00:01:27,560 --> 00:01:31,080
具体的实现方式是通过Cpython提供的

29
00:01:31,080 --> 00:01:33,600
Frame Evaluation API去实现的

30
00:01:33,600 --> 00:01:37,560
为啥这么辛苦在Cpython解析的时候去实现呢

31
00:01:37,560 --> 00:01:39,920
是因为想把PyTorch的一些操作

32
00:01:39,920 --> 00:01:42,360
后来用Python去写的一些代码了

33
00:01:42,360 --> 00:01:45,200
去把它变成PyTorch的FX的图

34
00:01:45,200 --> 00:01:48,160
从而很好的去捕捉Python的Bytecode

35
00:01:48,160 --> 00:01:51,480
使得PyTorch的动态图跟静态图

36
00:01:51,520 --> 00:01:54,360
都在Python的Bytecode里面去解析

37
00:01:54,360 --> 00:01:57,960
这样的话动态图统一的功能就比较完善

38
00:02:00,120 --> 00:02:03,240
不过有一点值得注意的就是因为在PyTorch

39
00:02:03,240 --> 00:02:06,120
大部分的时间都是用它做训练

40
00:02:06,120 --> 00:02:07,240
就Training的工作

41
00:02:08,640 --> 00:02:13,160
而训练就严重依赖于Automatic Diffusion

42
00:02:13,160 --> 00:02:14,480
自动微分的功能

43
00:02:14,480 --> 00:02:16,200
而PyTorch自动微分的功能

44
00:02:16,200 --> 00:02:18,200
或者自动微分的Engine引擎

45
00:02:18,200 --> 00:02:21,480
其实不是使用Python去实现的

46
00:02:21,480 --> 00:02:23,680
而是使用C++去实现的

47
00:02:24,120 --> 00:02:26,920
这个时候Dynamo之前的一个特性

48
00:02:26,920 --> 00:02:29,680
其实只能获取Python层面的一个Level

49
00:02:29,680 --> 00:02:33,200
也就是只能获取一个静态的正向图

50
00:02:33,200 --> 00:02:34,000
反向图

51
00:02:34,000 --> 00:02:35,880
C++层产生的反向图

52
00:02:35,880 --> 00:02:37,840
其实是没办法去获取的

53
00:02:37,840 --> 00:02:39,000
为了解决这个问题

54
00:02:39,200 --> 00:02:41,440
PyTorch就推出了另外一个新的特性

55
00:02:41,440 --> 00:02:43,080
叫做Ahead of Time

56
00:02:43,080 --> 00:02:44,680
不会去做这个工作的

57
00:02:44,680 --> 00:02:46,400
但是在实现这个工作之前

58
00:02:46,640 --> 00:02:49,560
看一下PyTorch是怎么产生的

59
00:02:49,560 --> 00:02:51,880
首先在写PyTorch的代码的时候

60
00:02:52,200 --> 00:02:54,520
会明确的去声明点Backward

61
00:02:54,520 --> 00:02:56,360
然后去通过Backward这个Path

62
00:02:56,640 --> 00:02:58,800
也就是C++的自动微分的Engine

63
00:02:59,160 --> 00:03:00,600
去具体的实现的

64
00:03:00,600 --> 00:03:03,360
而一般来说会使用eager的模式

65
00:03:03,360 --> 00:03:04,600
就最典型的模式

66
00:03:04,600 --> 00:03:05,840
点Backward来去实现

67
00:03:05,840 --> 00:03:08,680
当然了也可以通过Torch Script来去获取

68
00:03:08,680 --> 00:03:09,840
但是Torch Script

69
00:03:10,120 --> 00:03:12,400
在之前AI编译器的一个系列里面

70
00:03:12,520 --> 00:03:14,240
在自动微分AI编译器

71
00:03:14,360 --> 00:03:17,440
还有计算图去讲了Torch Script这种方式

72
00:03:17,600 --> 00:03:19,400
其实它要不就基于Trace

73
00:03:19,400 --> 00:03:21,080
要不就基于原码解析

74
00:03:21,080 --> 00:03:23,320
这两种方式都不是做得很彻底

75
00:03:24,080 --> 00:03:25,920
所以不能够获取所有的操作

76
00:03:25,920 --> 00:03:27,400
或者所有重取的可能性

77
00:03:27,760 --> 00:03:29,320
这个时候PyTorch就推出了

78
00:03:29,320 --> 00:03:31,480
Ahead of Time Autogrid的功能

79
00:03:32,120 --> 00:03:33,920
利用了PyTorch的Dispatch

80
00:03:33,920 --> 00:03:35,880
很好的去捕捉了整个

81
00:03:36,080 --> 00:03:37,600
反向的自动微分的图

82
00:03:37,920 --> 00:03:39,640
下面我给大家汇报一下

83
00:03:39,640 --> 00:03:42,120
AoT Autogrid具体怎么实现

84
00:03:42,120 --> 00:03:43,480
或者实现了哪些功能

85
00:03:44,040 --> 00:03:46,280
回顾一下整个PyTorch Complier Mode

86
00:03:46,280 --> 00:03:48,520
就是它的编译模式的权占

87
00:03:48,520 --> 00:03:50,080
首先在前端

88
00:03:50,080 --> 00:03:53,240
前端更多的是指面向用户看到的API

89
00:03:53,240 --> 00:03:54,960
而不是单单的指Python层

90
00:03:54,960 --> 00:03:56,000
或者C++层

91
00:03:56,000 --> 00:03:57,680
没有这么严格的定义

92
00:03:57,680 --> 00:03:59,120
前端就只有Python

93
00:03:59,120 --> 00:04:00,440
或者只有C++

94
00:04:00,880 --> 00:04:03,200
但是后端更多的是聚焦于

95
00:04:03,200 --> 00:04:04,680
算子的生成

96
00:04:04,680 --> 00:04:05,680
Kernel的生成

97
00:04:05,680 --> 00:04:08,320
前端指的更多的是指API

98
00:04:08,880 --> 00:04:09,680
在前端的时候

99
00:04:09,800 --> 00:04:12,320
使用Dynamo去获取FX的图

100
00:04:13,800 --> 00:04:14,880
获取到FX的图

101
00:04:15,040 --> 00:04:16,120
这个只有正向

102
00:04:16,360 --> 00:04:18,440
接着使用AoT Autogrid

103
00:04:18,720 --> 00:04:20,120
自动微分的反向的图

104
00:04:20,120 --> 00:04:21,440
反向的图产生完之后

105
00:04:21,560 --> 00:04:24,200
其实现在还是一个FX的Graph

106
00:04:24,200 --> 00:04:26,600
也就是对应的FX的图

107
00:04:27,160 --> 00:04:30,440
不过这个IR已经编成aten或者Prime IR

108
00:04:31,120 --> 00:04:32,040
有了这一层之后

109
00:04:32,200 --> 00:04:33,720
才是真正的走向了

110
00:04:33,720 --> 00:04:35,040
Backend就是后端

111
00:04:35,040 --> 00:04:36,080
算子的生成

112
00:04:36,080 --> 00:04:37,880
或者具体算子的执行了

113
00:04:38,120 --> 00:04:38,680
可以看到

114
00:04:38,680 --> 00:04:40,520
现在还是在这一层里面

115
00:04:41,520 --> 00:04:43,720
根据PyTorch之前做的一些特性

116
00:04:43,880 --> 00:04:46,280
假设想在训练的时候去加速

117
00:04:46,280 --> 00:04:48,160
其实有很多种方法

118
00:04:48,840 --> 00:04:50,240
有基于符号的Script

119
00:04:50,440 --> 00:04:51,200
有LazyTensor

120
00:04:51,560 --> 00:04:54,440
也有自己可能自己去写一个自动微分

121
00:04:54,440 --> 00:04:57,920
但是这些方式都不是说做得非常彻底

122
00:04:58,640 --> 00:04:59,400
AoT Autogrid

123
00:04:59,560 --> 00:05:01,720
它做的比较彻底的有三点

124
00:05:01,880 --> 00:05:04,800
第一点就是可以使用任何可以编译的

125
00:05:04,800 --> 00:05:05,440
编译的后端

126
00:05:05,440 --> 00:05:06,240
或者AI编译器

127
00:05:06,240 --> 00:05:08,280
都可以对接到前端就行了

128
00:05:08,280 --> 00:05:09,400
而AoT Autogrid

129
00:05:09,520 --> 00:05:12,000
主要更多的是获取反向的图

130
00:05:12,240 --> 00:05:14,760
第二个就是可以很好的去使用

131
00:05:14,760 --> 00:05:16,640
PyTorch写的训练的代码

132
00:05:16,640 --> 00:05:18,000
无欠用术的去修改

133
00:05:18,000 --> 00:05:19,320
第三点就是所有代码

134
00:05:19,320 --> 00:05:21,240
都在Python层里面去执行的

135
00:05:21,480 --> 00:05:24,040
这种是非常方便对图进行操作

136
00:05:24,400 --> 00:05:25,520
有了这些基础的概念

137
00:05:25,520 --> 00:05:26,400
看一下Autogrid

138
00:05:26,400 --> 00:05:28,760
其实是严重依赖于TorchDispatch

139
00:05:28,760 --> 00:05:29,560
这个功能的

140
00:05:29,560 --> 00:05:31,000
所以TorchDispatch这个功能

141
00:05:31,240 --> 00:05:33,680
我也会后面详细的去给大家汇报

142
00:05:33,680 --> 00:05:35,680
现在看一下AoT Autogrid

143
00:05:35,680 --> 00:05:37,040
具体的几个实现方式

144
00:05:37,040 --> 00:05:38,760
第一它主要分为三个步骤

145
00:05:38,960 --> 00:05:40,040
第一个步骤就是

146
00:05:40,480 --> 00:05:42,120
使用TorchDispatch这个功能

147
00:05:42,240 --> 00:05:43,200
或者它的调度功能

148
00:05:43,320 --> 00:05:45,560
去追踪正向和反向的图

149
00:05:45,560 --> 00:05:47,520
接着去把正向和反向的图

150
00:05:47,800 --> 00:05:49,040
分成两个图

151
00:05:49,040 --> 00:05:50,000
就两个子图

152
00:05:50,000 --> 00:05:51,280
一个图是正向

153
00:05:51,280 --> 00:05:52,920
一个图是反向

154
00:05:53,320 --> 00:05:55,280
最后一点就是在Ai编辑器

155
00:05:55,280 --> 00:05:57,800
去调用正向的图和反向的图

156
00:05:57,800 --> 00:06:00,520
当它作为一个具体的函数去调用的

157
00:06:00,920 --> 00:06:03,080
这三步就是AoT Autogrid的

158
00:06:03,080 --> 00:06:04,280
具体的执行方式

159
00:06:04,640 --> 00:06:05,520
可能实现的时候

160
00:06:05,680 --> 00:06:07,320
更多的是工程化的问题

161
00:06:07,440 --> 00:06:09,920
但是它的idea还是很outstanding的

162
00:06:10,880 --> 00:06:12,960
首先看一下TorchDispatch的功能

163
00:06:12,960 --> 00:06:13,760
TorchDispatch

164
00:06:13,880 --> 00:06:15,880
会在后面详细的展开的

165
00:06:15,880 --> 00:06:17,560
这里只是简单的过一下

166
00:06:17,920 --> 00:06:20,520
这条红线主要是分开Python的一些Line

167
00:06:20,520 --> 00:06:22,360
还有PyTorch的一个核心的代码

168
00:06:22,640 --> 00:06:23,960
在Python具体实现的时候

169
00:06:24,160 --> 00:06:26,840
更多的是调用TorchDispatch这个功能

170
00:06:27,400 --> 00:06:28,040
另外知道

171
00:06:28,040 --> 00:06:30,040
实际上去写Python的代码的时候

172
00:06:30,160 --> 00:06:31,440
它不是马上执行的

173
00:06:31,440 --> 00:06:34,160
虽然都说PyTorch会马上执行

174
00:06:34,160 --> 00:06:35,880
但实际上在Python框里面

175
00:06:35,920 --> 00:06:37,520
会有很多调度的方式

176
00:06:37,760 --> 00:06:39,080
会走到Aten的算子

177
00:06:39,240 --> 00:06:41,600
Aten的算子它其实做了一个封装

178
00:06:41,800 --> 00:06:44,080
接着再去执行AutoGrid AMP

179
00:06:44,080 --> 00:06:46,560
混合进度相关的一些流程代码

180
00:06:46,800 --> 00:06:48,160
最后才是Kernel Launch

181
00:06:48,480 --> 00:06:49,840
在Kernel Launch就真正的

182
00:06:49,840 --> 00:06:51,320
把算子调起来之前

183
00:06:51,680 --> 00:06:55,040
通过TorchDispatch去捕获反向的图

184
00:06:55,280 --> 00:06:56,080
通过这种方式

185
00:06:56,440 --> 00:06:59,160
直接在Python层面获得反向的图

186
00:06:59,440 --> 00:07:00,840
这个就是整个AoT的

187
00:07:00,840 --> 00:07:02,520
一个具体的架构和逻辑

188
00:07:03,520 --> 00:07:05,160
但是现在回顾一下

189
00:07:05,160 --> 00:07:05,880
AutoGrid

190
00:07:05,880 --> 00:07:07,720
就自动微分怎么去实现的

191
00:07:08,000 --> 00:07:10,240
假设现在有一层forward0

192
00:07:10,240 --> 00:07:11,200
然后forward1

193
00:07:11,200 --> 00:07:13,080
每一层假设它是一个卷积

194
00:07:13,080 --> 00:07:14,640
with loop激活

195
00:07:14,640 --> 00:07:16,440
在真正PyTorch去实现的时候

196
00:07:16,600 --> 00:07:18,680
会去声明Loss.backward

197
00:07:18,680 --> 00:07:21,800
然后才开始真正的构建反向图

198
00:07:22,000 --> 00:07:23,000
构建反向图的时候

199
00:07:23,120 --> 00:07:25,320
大家一开始的概念就是以为

200
00:07:25,320 --> 00:07:26,680
正向图跟反向图

201
00:07:26,680 --> 00:07:28,280
每一个算子是对应的

202
00:07:28,280 --> 00:07:30,520
我有一个卷积肯定有一个卷积的

203
00:07:30,520 --> 00:07:31,200
反向的图

204
00:07:31,240 --> 00:07:33,400
我有一个激活肯定有个激活的反向

205
00:07:33,400 --> 00:07:34,320
这是一对应

206
00:07:34,520 --> 00:07:36,160
这只是在概念上面的

207
00:07:36,160 --> 00:07:37,320
但是在实现上面

208
00:07:37,560 --> 00:07:38,560
回顾一下

209
00:07:38,560 --> 00:07:42,080
之前在自动微分的系列里面

210
00:07:42,080 --> 00:07:43,600
讲了PyTorch

211
00:07:43,800 --> 00:07:45,200
它设计于一个Tapebase的

212
00:07:45,200 --> 00:07:46,320
一个面向对象

213
00:07:46,320 --> 00:07:48,200
实现的自动微分的功能

214
00:07:48,440 --> 00:07:51,080
这个就是当时候写的一个伪代码

215
00:07:51,080 --> 00:07:52,640
或者当时候手把手的

216
00:07:52,640 --> 00:07:55,480
带着大家一起去实现PyTorch的

217
00:07:55,480 --> 00:07:57,720
可以看到tangent就是导数

218
00:07:57,920 --> 00:07:58,840
out fn

219
00:07:58,840 --> 00:08:00,400
这是正向的一个执行

220
00:08:00,400 --> 00:08:02,200
backward out就是Tapebase

221
00:08:02,200 --> 00:08:03,960
就是记录整个tab的

222
00:08:04,280 --> 00:08:06,480
最后反馈的就是out和backward out

223
00:08:06,480 --> 00:08:08,800
可以看到这里面主要是基于Tapebase

224
00:08:08,960 --> 00:08:10,280
基于Tapebase这种方式

225
00:08:10,440 --> 00:08:12,800
就不能够去构建一个正向

226
00:08:12,800 --> 00:08:13,880
然后构建一个反向

227
00:08:13,880 --> 00:08:16,480
而是在真正声明backward的时候

228
00:08:16,480 --> 00:08:18,360
才把反向构建出来

229
00:08:18,360 --> 00:08:20,800
所以说正向和反向是分开的

230
00:08:21,000 --> 00:08:22,800
而通过TorchDispatch这个功能

231
00:08:23,040 --> 00:08:25,680
把反向图单独的捕获出来

232
00:08:25,680 --> 00:08:28,200
所以最后就变成两张图

233
00:08:28,200 --> 00:08:29,080
两张子图

234
00:08:29,080 --> 00:08:32,080
第一张就是正向的一个forward图

235
00:08:32,080 --> 00:08:34,080
第二个就是反向的

236
00:08:34,280 --> 00:08:35,440
backward的图

237
00:08:35,920 --> 00:08:37,480
所以说AOT Autogrid

238
00:08:37,600 --> 00:08:39,040
就分开两个图

239
00:08:39,040 --> 00:08:40,600
那最后了解完

240
00:08:40,600 --> 00:08:42,600
AOT Autogrid的一个具体的实现

241
00:08:42,600 --> 00:08:44,440
或者一个简单的实现逻辑之后

242
00:08:44,680 --> 00:08:46,360
看一下AOT Autogrid的

243
00:08:46,360 --> 00:08:48,600
一个具体的效果

244
00:08:49,880 --> 00:08:52,960
现在先看看下面的这个图

245
00:08:52,960 --> 00:08:54,440
那这个图的上半部分

246
00:08:54,440 --> 00:08:55,920
仔细的去看一下

247
00:08:56,040 --> 00:08:58,240
这里面分为前端front end

248
00:08:58,240 --> 00:09:00,600
前端更多的是指正向

249
00:09:00,600 --> 00:09:03,240
然后有一个backward的反向的capture

250
00:09:03,360 --> 00:09:04,640
就反向的获取

251
00:09:04,640 --> 00:09:06,240
然后还有个后端

252
00:09:06,240 --> 00:09:09,160
后端就是真正去执行这些算子的

253
00:09:09,160 --> 00:09:11,520
那实际上最原始的时候

254
00:09:11,520 --> 00:09:13,120
可以使用Torch Script

255
00:09:13,120 --> 00:09:15,360
去获取正反向的图

256
00:09:15,920 --> 00:09:17,400
从结果可以看到

257
00:09:17,400 --> 00:09:18,440
获取正反向图

258
00:09:18,560 --> 00:09:20,200
其实有很多是红的

259
00:09:20,200 --> 00:09:22,520
红的代表没有办法去执行获取

260
00:09:22,520 --> 00:09:24,120
这个图是失败的

261
00:09:25,560 --> 00:09:26,600
接下来看一下

262
00:09:26,640 --> 00:09:28,320
使用2.0Torch Dynamo

263
00:09:28,320 --> 00:09:29,680
作为一个前端

264
00:09:29,680 --> 00:09:31,080
去获取正向的图

265
00:09:31,080 --> 00:09:32,280
然后获取反向图

266
00:09:32,400 --> 00:09:33,840
有两种方式去比较

267
00:09:33,840 --> 00:09:35,400
第一种是使用Torch Script

268
00:09:35,400 --> 00:09:36,520
演讲使用Torch Script

269
00:09:36,520 --> 00:09:38,200
去获取反向的图

270
00:09:38,360 --> 00:09:40,080
另外一种是使用

271
00:09:40,080 --> 00:09:41,920
Ahead of Time Autogrid的这种方式

272
00:09:41,920 --> 00:09:43,000
就刚才说的

273
00:09:43,000 --> 00:09:45,480
使用dispatch的功能去获取的

274
00:09:45,680 --> 00:09:48,520
后端有NCC和MVFusion都是相同的

275
00:09:48,520 --> 00:09:49,840
从结果可以看到

276
00:09:49,840 --> 00:09:52,000
使用AOT Autogrid这种方式

277
00:09:52,160 --> 00:09:54,000
其实大部分的网络模型

278
00:09:54,120 --> 00:09:55,440
都是可以工作的

279
00:09:55,440 --> 00:09:56,360
就大部分网络模型

280
00:09:56,360 --> 00:09:57,480
都是正常工作的

281
00:09:57,480 --> 00:09:59,160
这种获取图的方式

282
00:09:59,160 --> 00:10:01,360
或者获取反向图的方式

283
00:10:01,360 --> 00:10:02,440
会更好

284
00:10:02,600 --> 00:10:05,000
而使用Torch Script去获取反向图

285
00:10:05,000 --> 00:10:06,280
有很多的反向图

286
00:10:06,280 --> 00:10:07,840
是没有办法去获取的

287
00:10:07,880 --> 00:10:09,120
也就是标红

288
00:10:09,120 --> 00:10:10,560
或者标黄的地方

289
00:10:11,000 --> 00:10:12,880
而从真正的性能来看

290
00:10:13,160 --> 00:10:14,640
蓝色的这些性能比较

291
00:10:14,720 --> 00:10:16,320
就效果特别好

292
00:10:16,320 --> 00:10:19,160
而绿色就是有比较好的性能的超越

293
00:10:19,160 --> 00:10:21,160
可以看到使用AOT Autogrid

294
00:10:21,160 --> 00:10:22,840
然后加上Torch Dynamo

295
00:10:22,880 --> 00:10:24,080
对PyTorch 2.0

296
00:10:24,080 --> 00:10:25,480
这个特性的提升来说

297
00:10:25,480 --> 00:10:27,120
还是非常之好的

298
00:10:27,120 --> 00:10:27,760
好了

299
00:10:28,520 --> 00:10:28,960
好了

300
00:10:28,960 --> 00:10:30,800
今天的内容就到这里为止

301
00:10:30,800 --> 00:10:31,480
谢谢各位

302
00:10:31,480 --> 00:10:32,480
拜了个拜

303
00:10:33,120 --> 00:10:33,920
卷的不行了

304
00:10:33,920 --> 00:10:34,800
卷的不行了

305
00:10:34,800 --> 00:10:36,240
记得一键三连加关注

306
00:10:36,600 --> 00:10:37,960
所有的内容都会开源

307
00:10:37,960 --> 00:10:39,800
在下面这条链接里面

308
00:10:40,160 --> 00:10:41,160
拜了个拜