已开启
activation swap/重计算对持有共享函数属性的同级模块独立加 checkpoint 时误报区域重叠 #189
xuxinglei创建于  6月5日
xuxinglei成员
6月5日 创建

该问题是怎么引起的?

activation_swap.py(torch 与 mindspore 两个平台都有)在给模块施加 activation swap / 重计算 checkpoint 时,_check_and_mark_wrapped 会遍历该模块及其所有子模块,对每个“可包装的公开 callable 实例属性”做两件事:先校验它没有被标记过 _is_wrapped(重叠检测),再给它打上 _is_wrapped = True

判定“可包装 callable 属性”的 _iter_wrappable_callable_attrs 之前会把实例上任何非 Cell/Module 的公开 callable 都算进来,其中包含大量以实例属性形式持有的模块级共享函数,例如每层都写的 self.act = F.geluself.reshape = mint.reshapeself.cast = ops.cast

这些是无状态、被多个模块以引用方式共享的同一个对象,把它纳入重叠跟踪有两个问题:

  1. 给共享函数对象写 _is_wrapped 会污染全局对象(Python 普通函数允许写属性),影响程序里其它引用同一函数的位置。
  2. 同一个函数对象被多个同级层共享:A 层 wrap 后把 self.act(= F.gelu)标成 wrapped,再独立 wrap 同级的 B 层时,会看到 self.act._is_wrapped == True,于是误判为区域重叠并抛出异常——而这两个区域其实并不重叠,只是共享了同一个无状态工具函数。

该问题在重计算独立调度(对多个同级模块分别加 checkpoint)的场景下暴露。

重现步骤

  1. 构造若干同级(兄弟)子模块,每个都以实例属性形式持有同一个共享函数,如 self.act = F.geluself.reshape = mint.reshape
  2. 对这些同级模块分别、独立地施加 activation swap / 重计算 checkpoint(彼此不嵌套、不重叠)。
  3. 第二个模块进入 _check_and_mark_wrapped_check_callable_attr_not_wrapped 时,发现共享函数已被第一个模块标记,触发误报。

报错信息

ValueError: Callable 'function' is already wrapped. Wrapping overlapping module regions is not allowed.
likedislike
Xxuxinglei成员
6月5日 关联了pull request:fix(activation_checkpoint): exclude shared functions from callable overlap tracking