该迁移点只要求设置一个NPU上的循环下沉次数,我们称之为npu loop size,您可以通过NPU_LOOP_SIZE的环境变量设置,也可以通过在您的训练脚本中调用npu.set_npu_loop_size接口进行设置,修改的内容很简单,但是需要您理解npu loop size的含义。
npu loop size用于实现NPU训练的极致性能,在介绍npu loop size的由来前,先介绍TF2原生流程中的一些性能损耗点。以GPU为例,GPU环境常规训练方式下的工作时序,如下图所示,脚本侧用户控制执行十次训练,每次在GPU上执行一次训练步,训练步结束后,回到python侧,用户判断当前步数等于10,启动下一次训练,直至训练完十次训练。
观察时序图上的有色区域我们不难发现,此时不论是CPU还是GPU都是间歇性工作的,该模式下的缺陷:
TF2为了省去Python解释器上的额外开销,推荐用户使用While算子来实现训练循环(也就是所谓的循环下沉,循环下沉并非NPU特有的策略),此时判断训练是否达到指定步数的逻辑不再在Python解释器中进行,而是依赖TF2中的While算子,在编码时,使用者应当这样组织自己的训练(下称“循环下沉的编码方式”):
@tf.function def loop_train(iterator, steps): for i in tf.range(steps): train_step(next(iterator))
这样的TF2代码,在编译后,会将训练步嵌套在While算子中执行,时序变为下图所示:
可以看出,采用循环下沉的策略后,在Python解释器上的耗时转移到TF CPU上,耗时更短也更稳定,但是在该形式下,仍然有两部分额外开销:
NPU为了达到极致性能,采取了两个策略来消除这两部分额外耗时:
异步数据传输指TF Adapter的预处理线程主动向NPU发送训练数据,在未使用循环下沉的方式编码时,执行时序如下所示:
此时可以一定程度上消减数据预处理H2D数据传输与CPU调度的耗时(下发训练步的同时,数据传输正在进行)。
当使用了训练循环下沉的编码方式时,NPU上的执行时序图为:
可以看到:
NPU训练循环下沉与异步预处理数据传输方式可以最大程序地减少训练计算无关的耗时,最大化性能收益,但同时对用户训练有额外约束:
由于预处理线程与NPU训练步异步,在使用循环下沉的编码方式时,需要告诉NPU当前的循环下沉次数,所以NPU要求用户在使用循环下沉的编码方式时,额外设置npu loop size,用于指示循环下沉执行的次数。
比如,您使用循环下沉的编码方式组织您的训练:
@tf.function def loop_train(iterator, steps): for i in tf.range(steps): train_step(next(iterator))
当您期望每次loop_train调用会在NPU上训练100个Step时,此时您有两种方式设置npu loop size:
export NPU_LOOP_SIZE=100
npu.set_npu_loop_size(100) loop_train(train_iter, tf.constant(100))
您也可以在训练过程中调用npu.set_npu_loop_size来改变NPU上每次下沉执行的步数,比如您总共训练100个Step,希望每次在NPU上循环执行30Step,显然,最后的91~100Step小于npu loop size,此时你可以在训练完90Step后调用npu.set_npu_loop_size来调整npu loop size的大小:
remaining_steps = 100 # 剩余Steps数 base_loop_size = 30 # 基准npu loop size npu.set_npu_loop_size(base_loop_size) while remaining_steps >= base_loop_size: # 按照基准loop循环下沉训练,直到剩余Step数不足一次loop loop_train(train_iterator, tf.constant(base_loop_size)) remaining_steps -= base_loop_size if remaining_steps > 0: # 如果还有未处理的数据,调整为一个较小的npu loop size处理 npu.set_npu_loop_size(remaining_steps) loop_train(train_iterator, tf.constant(remaining_steps))