切片原理

通过枚举内部 index 的取值拆分收缩任务,限制每个切片的中间结果大小。

本页目录

把收缩分成多个子任务

选择若干内部 index 后,ArcTN 枚举这些 index 的全部取值组合。每组取值固定相应的输入轴,沿同一条 SSA path 完成一次收缩,最后把全部结果逐元素相加。切片在数学上是精确分解:它不做低秩截断,也不遗漏任何取值组合。实际使用浮点数时,切片求和与未切片收缩可能采用不同的加法顺序,因此可能存在正常的舍入差异,不承诺逐位相同。

切片怎样缩小中间张量

以三张量收缩为例:

Yace=∑b,dAabcBbcdCde.Y_{ace}=\sum_{b,d}A_{abc}B_{bcd}C_{de}.

取 a,b,c,d,ea,b,c,d,e 的维数分别为 2、3、4、5、6。按先收缩 B,CB,C 的顺序,未切片时会生成 ZbceZ_{bce}。固定 b=jb=j 后,每片先计算 Zce(j)=∑dBjcdCdeZ^{(j)}_{ce}=\sum_d B_{jcd}C_{de},再与 AajcA_{ajc} 收缩。

固定内部 index b,把收缩分成三个切片未切片中间张量 Z 有 72 个元素,仅作大小对照。先在输入 A、B 中固定 b 为 0、1、2;每片先计算 24 元素的 Z,再与对应的 A 收缩,得到 48 元素的部分输出 Y。三个部分输出相加得到原结果。未切片中间量 Zbce · 72 个元素在输入中固定 bb = 0Z(0)ce · 24 元素Y(0)ace · 48 元素b = 1Z(1)ce · 24 元素Y(1)ace · 48 元素b = 2Z(2)ce · 24 元素Y(2)ace · 48 元素与 A(j) 收缩
Yace=Yace(0)+Yace(1)+Yace(2).Y_{ace}=Y^{(0)}_{ace}+Y^{(1)}_{ace}+Y^{(2)}_{ace}.

图中每片的中间量有 24 个元素,部分输出有 48 个元素,因此每片最大张量是 48 个元素;72 元素的完整 ZZ 仅作对照,切片计算不会先分配它。Fixed 保持原收缩顺序,代价是可能增加重复计算与结果累加。具体用法见切片示例。

切片数量

若切片 index 的维度为 d1,…,dkd_1,\ldots,d_k,切片数为

Nslice=∏j=1kdj.N_{\mathrm{slice}}=\prod_{j=1}^{k}d_j.

SliceResult.log2_n_slices 保存切片数以 2 为底的对数。index 的维度不必为 2。

中间张量大小上限

target_size 用来控制切片后的中间张量大小,取值为正整数,单位是元素数。每个切片中的二元收缩结果都必须满足这个上限;单张量网络则检查一元处理后的最终输出。它不限制一元预处理产生的全部临时张量,也不是进程内存或 GPU 显存上限。

target_size 约束 target_size 不约束
每个 slice 的最大二元收缩结果元素数;单张量网络则检查最终输出 进程 RSS 或 GPU 显存
按指定路径与切片 index 计算的二元收缩结果 同时保留的所有张量的元素数之和
逻辑元素数 dtype 换算后的字节数
单个 slice 多个并发 slice 的合计驻留内存
结果张量 ∣C∣\lvert C\rvert 单步 ∣A∣+∣B∣+∣C∣\lvert A\rvert+\lvert B\rvert+\lvert C\rvert 或后端临时工作空间

log2_max_contraction_size 报告最大单步逻辑输入加输出元素数的以 2 为底对数,其中原始叶张量按一元求和、迹或对角处理后的逻辑大小计数;log2_peak_size 报告路径模型中存活元素数峰值的以 2 为底对数。它们都不是 target_size 的另一种写法。

单张量网络的 SSA path 为空,但仍可能需要取迹、求和或调整输出轴顺序。ArcTN 会检查完成这些一元处理后的最终输出大小;即使路径没有二元收缩步骤,过小的 target_size 也不能通过检查。

固定路径的切片 index 选择

find_slices_to_size 保持路径不变,把已经选择的 index 维度临时设为 1 并重放路径。若仍有中间结果超过目标,就优先选择在这些结果中出现次数最多的 index;出现次数相同时,选择维度更大的 index;若维度也相同,则选择 ID 更小的 index。

同一路径中,第 3、5、7 步的中间结果超过大小上限。只看候选 index x、y、z:

第 3 步结果
{ x, y }
第 5 步结果
{ x, z }
第 7 步结果
{ x, y, z }

x 出现于 3 个超限结果,y、z 各出现于 2 个,因此先选 x。评估每个切片时,将 x 的有效维度设为 1,沿原路径重新计算大小;若仍超限,就继续选择。

  1. 切片 index 必须在网络和 size_dict 中存在,且维度为正。
  2. 当前不支持对输出 index 切片,最终输出的形状保持不变。
  3. 切片列表中不能重复同一个 index。
  4. 同一输入张量中重复出现的 index 不能切片;这种 index 涉及取迹(trace)或取对角(diagonal)操作。
  5. 没有合法 index 可以继续降低结果规模时,该搜索返回未找到可行方案。

如果最终输出的元素数已经超过 target_size,仅切分内部 index 无法满足这个上限。

rust
use arctn::{find_slices_to_size, slice_result_fits_target_size};

let target_size = 1usize << 24;
let slices = find_slices_to_size(&net, &path, target_size)
    .ok_or("target_size is unreachable for this fixed path")?;
assert!(slice_result_fits_target_size(
    &net, &path, &slices, target_size
)?);

这是固定路径上的贪心选择,不保证切片 index 最少或总 FLOPs 最低。成功返回的方案会通过目标约束检查;搜索失败不等于证明所有可能方案都不可行。

读取 SliceResult

字段 含义
legs 实际选择的内部 index ID
log2_n_slices 全部取值组合数量的以 2 为底的对数
per_slice 将切片 index 的维度设为 1 后计算的 PathStats
log10_flops_total 单片 FLOPs 乘切片数得到的总 FLOPs,以 10 为底取对数

最大中间结果按单片计算,不乘切片数;总 FLOPs 和总读写量则需要乘以切片数。评估一个方案时,可同时查看路径、切片 index、切片数、单片最大结果、总 FLOPs 和实际执行时间。

路径和 SliceResult 必须成对保存。允许改变路径的切片方法会同时返回对应的最终路径。修改路径后,需要重新检查切片方案是否满足大小限制。