`getindex` thunk: `add!` disagrees with `val`; `sum(f, xs)` accumulates closure tangents quadratically

未关闭
  • 哪里不对？
#863 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 1 天内回复

还没有人认领这个 Issue。

评估

难度
4/5
预计耗时
3-5 天
新手友好度
25/100
Issue 类型
缺陷
描述清晰度
描述清楚
活跃度
活跃
领域
machine-learning

调研方向

Start with the thunked_∇getindex reverse rule for getindex, comparing what its add! does with unthunk(dy) against what val does through ProjectTo(x). Then look at the general branch of rrule for sum(f, xs) and its sum(first, ...) accumulation, which needs an add!! that recurses into Tangent fields. Done means the reproducers in the issue give the same result from add! as from unthunk, and the quadratic allocation case runs in linear time. The design choices for the fix need maintainer input first.

由索引模型根据 Issue 内容生成。

描述

Two things in the reverse rules keep an AD from accumulating getindex cotangents in place. That makes x[i] in a loop quadratic (FluxML/Zygote.jl#365). Zygote now accumulates in place (FluxML/Zygote.jl#1738), but these two cases stay quadratic because of the rules here. Both are unchanged on master (v1.74.8).

1. thunked_∇getindex: add! adds raw dy, val projects it

val is ∇getindex(x, dy, inds...), which ends in ProjectTo(x)(dx), but add! adds unthunk(dy) as is. So t.add!(dx) can differ from dx + unthunk(t), or throw:

using ChainRules, ChainRulesCore
x = [1.0, 2.0, 3.0]
dx = rrule(getindex, x, 1)[2](1.0 + 2.0im)[2]
unthunk(dx)                    # [1.0, 0.0, 0.0]
dx.add!(zeros(3))              # InexactError
dx.add!(zeros(ComplexF64, 3))  # [1.0 + 2.0im, 0, 0]

dx = rrule(getindex, x, 1:2)[2]([1.0im, 2.0im])[2]
dx.add!(zeros(3))              # InexactError, though unthunk(dx) == zeros(3)

x32 = Float32[1, 2, 3]
dx = rrule(getindex, x32, 2)[2](0.1)[2]
dx.add!(zeros(3))[2] == unthunk(dx)[2]  # false: 0.1 vs Float32(0.1)

So an AD can't call add! without knowing dy, which is hidden in the closure, and has to fall back to unthunk. For a slice that builds a dense array the size of the whole x each time. A loop over the columns of a 100×1000 matrix then allocates 767 MiB in Zygote.

Suggested fix: make add! add the same projected value as val, at least for plain Array{<:Number}, where ProjectTo(x) is elementwise. For example, dx -> ∇getindex!(dx, _project_dy(x, unthunk(dy)), Base.to_indices(x, inds)...), where _project_dy applies ProjectTo(zero(eltype(x))) to dy, elementwise for slices. Structured x could keep dx .+= unthunk(val). With that patch, Zygote calling add! when the inferred eltype of val fits its buffer takes that loop from 767 MiB to 4.4 MiB.

A related cost: in the last example unthunk(dx) is a dense Vector{Float32}, not a OneElement. Projecting the OneElement onto a different eltype materializes it, so mixed-precision scalar indexing is O(length(x)) per access. Projecting dy before building the OneElement would avoid that.

2. sum(f, xs) adds closure tangents with sum(first, df_and_dxs)

In the general branch of rrule(::RuleConfig{>:HasReverseMode}, ::typeof(sum), f, xs), each + of two Tangents whose fields are InplaceableThunks unthunks and copies the whole array. For a closure that captures x and reads x[i], that is O(n²):

using Zygote
x = rand(10_000)
@allocated gradient(x -> sum(i -> x[i]^2, eachindex(x)), x)  # ≈ 766 MiB

Generators hit the same path, since sum(x[i]^2 for i in ...) is lowered to sum(f, iter). Accumulating df into a partial sum the rule owns, instead of sum(first, ...), would make it linear. That needs an add!! that recurses into Tangent fields, since add!! falls back to + for Tangent. Zygote can route around this through its own map adjoint (766 MiB → 3.8 MiB with #1738), but other ADs using this rule would still pay.

🧔‍♂️🤖 Assisted by AI

主要语言
Julia
星标
479
派生
100
平均合并
1 天 4 小时
30 天内合并 PR
14

环境准备

这个项目没有提供开发容器、Dockerfile 或贡献指南，环境需要你自己搭建：先看它的 README，通用步骤见我们的新手贡献指南。

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

JuliaDiff/ChainRules.jl 的其他 Issue

查看 JuliaDiff/ChainRules.jl 的全部 Issue

相似的 Issue

更多 Julia Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。