# Parallel Reductions with Zygote

**URL:** <https://discourse.julialang.org/t/parallel-reductions-with-zygote/75969>\
**Category:** Performance\
**Tags:** question, parallel, zygote, threads, sum\
**Created:** [February 7, 2022, 7:42pm UTC](https://discourse.julialang.org/t/parallel-reductions-with-zygote/75969 "2022-02-07T19:42:35Z")\
**Posts on this page:** 12\
**Page:** 1

<div class="post-metadata">

**Author:** ![SEA](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sea/32/8956_2.png) [@SEA](https://discourse.julialang.org/u/SEA)\
**Post date:** [February 7, 2022, 7:42pm UTC](https://discourse.julialang.org/t/parallel-reductions-with-zygote/75969/1 "2022-02-07T19:42:35Z")

</div>

I am trying to speed up the gradient computation of a model of the form

```nohighlight
sum(m->m(x), models),

```

where `models` is an array of constituent models. This is (supposed to be) easily parallelizable, but I am having trouble making it work with Zygote.

First, I tried the parallel `sum` implementation in [ThreadsX.jl](https://github.com/tkf/ThreadsX.jl), but I get worse-than-serial performance for the backward pass on an Intel processor - a 24 core Intel(R) Xeon(R) CPU X5690 @ 3.47GHz - while the forward pass is accelerated. Further, I get a [segfault running natively on an Apple M1 Pro](https://github.com/tkf/ThreadsX.jl/issues/192), which could be related to some threading issues that are still present on ARM.

Second, writing my own parallel version of `sum(f, x)` that is compatible with Zygote also proved more difficult than excepted due to Zygote not supporting array mutation, which precludes [the solutions discussed here](https://discourse.julialang.org/t/parallel-reductions/30180).

This leaves writing a custom adjoint for the custom implementation as a still viable option. I don’t yet know how to do this effectively and would be grateful for any help.

---

<div class="post-metadata">

**Author:** ![SEA](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sea/32/8956_2.png) [@SEA](https://discourse.julialang.org/u/SEA)\
**Post date:** [February 7, 2022, 7:44pm UTC](https://discourse.julialang.org/t/parallel-reductions-with-zygote/75969/2 "2022-02-07T19:44:48Z")

</div>

Also, here is an implementation of a parallel summation that works well for the forward pass, but that Zygote can’t differentiate through (because of a `try-catch` block apparently included in the `@threads` implementation, and secondarily, because of the array mutation happening in the parallel section).

```nohighlight
using Base.Threads
function parallel_sum(f, x::AbstractVector)
    n, k = length(x), nthreads()
    m = n ÷ k
    fx1 = f(x[1]) # to get correct output type
    sum_x = fill(fx1, k)
    @threads for i in 1:k
        y = @view x[(i-1)*m+1 : (i*m)]
        sum_x[i] = sum(f, y)
    end
    result = zero(fx1)
    for i in 1:k
        result += sum_x[i]
    end
    if rem(n, k) != 0
        result += sum(f, @view(x[m*k+1:end]))
    end
    return result
end

```

---

<div class="post-metadata">

**Author:** ![goerch](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/goerch/32/29122_2.png) [@goerch](https://discourse.julialang.org/u/goerch)\
**Post date:** [February 7, 2022, 8:19pm UTC](https://discourse.julialang.org/t/parallel-reductions-with-zygote/75969/3 "2022-02-07T20:19:42Z")

</div>

We recently had some success using `Folds` [here](https://discourse.julialang.org/t/high-precision-summation-functions-not-compatible-with-optim-jl-with-autodiff-forward/75714). Don’t know if this helps with `Zygote` though.

---

<div class="post-metadata">

**Author:** ![SEA](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sea/32/8956_2.png) [@SEA](https://discourse.julialang.org/u/SEA)\
**Post date:** [February 7, 2022, 8:42pm UTC](https://discourse.julialang.org/t/parallel-reductions-with-zygote/75969/4 "2022-02-07T20:42:30Z")

</div>

Thank you for the pointer @goerch!

I just tried out `Folds`, but realized it relies on `Transducers`, the same backend as `ThreadsX` and for this reason, has the same performance (Intel) and segfault (ARM) issues as `ThreadsX` in conjunction with `Zygote`.

---

<div class="post-metadata">

**Author:** ![goerch](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/goerch/32/29122_2.png) [@goerch](https://discourse.julialang.org/u/goerch)\
**Post date:** [February 7, 2022, 8:55pm UTC](https://discourse.julialang.org/t/parallel-reductions-with-zygote/75969/5 "2022-02-07T20:55:36Z")

</div>

> [@SEA](#):
>
> I just tried out `Folds` , but realized it relies on `Transducers` , the same backend as `ThreadsX` and for this reason, has the same performance (Intel) and segfault (ARM) issues as `ThreadsX` in conjunction with `Zygote` .

Naive parallel implementations we discussed [here](https://discourse.julialang.org/t/speed-issue-with-kahansummation/75491).

Edit: sorry, corrected link.  
Edit: just for reference, here my latest parallelized version

```julia
function naive_kbn_parallel(xs) # credit to @tkf from another thread
    len = length(xs)
    nt = min(Threads.nthreads(), len)
    ys = Vector{Float64}(undef, 2 * nt)
    chunk = (len + nt - 1) ÷ nt
    Threads.@threads for i in 1:nt
        s, c = _naive_kbn(@view xs[(i - 1) * chunk + 1:min(i * chunk, len)])
        ys[2 * i - 1] = s
        ys[2 * i] = c
    end
    naive_kbn_serial(ys)
end

```

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [February 7, 2022, 9:00pm UTC](https://discourse.julialang.org/t/parallel-reductions-with-zygote/75969/6 "2022-02-07T21:00:53Z")

</div>

Is it `sum(m, models)` or `sum(m -> m(x), models)`? Either way I presume the gradient with respect to `m` is important.

The rule which handles this is here: [mapreduce.jl#L74-L99](https://github.com/JuliaDiff/ChainRules.jl/blob/main/src/rulesets/Base/mapreduce.jl#L74-L99) and the key lines are:

```julia
     fx_and_pullbacks = map(x->rrule_via_ad(config, f, x), xs)
     y = sum(first, fx_and_pullbacks; dims=dims) # forward

     f̄_and_x̄s = call.(pullbacks, broadcast_ȳ) # gradient
     # f̄_and_x̄s = map(((_, back),) -> back(ȳ), fx_and_pullbacks) for sum without dims

```

These two `map` steps are the work you could parallelise, by hand or with something like `ThreadsX.map`. My guess is that hoping for the gradient of `ThreadsX.sum` to just work, and be efficient, is asking too much, although I’d like to be wrong.

---

<div class="post-metadata">

**Author:** ![SEA](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sea/32/8956_2.png) [@SEA](https://discourse.julialang.org/u/SEA)\
**Post date:** [February 7, 2022, 9:19pm UTC](https://discourse.julialang.org/t/parallel-reductions-with-zygote/75969/7 "2022-02-07T21:19:12Z")

</div>

Thank you for the response @mcabbott!

> Is it `sum(m, models)` or `sum(m -> m(x), models)` ? Either way I presume the gradient with respect to `m` is important.

It is `sum(m -> m(x), models)`, I just edited the original post to correct this.

I’ll give this a deeper look and hope I can parallelize it efficiently!

---

<div class="post-metadata">

**Author:** ![SEA](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sea/32/8956_2.png) [@SEA](https://discourse.julialang.org/u/SEA)\
**Post date:** [February 7, 2022, 9:20pm UTC](https://discourse.julialang.org/t/parallel-reductions-with-zygote/75969/8 "2022-02-07T21:20:29Z")

</div>

> Naive parallel implementations we discussed here.

I am having trouble opening your link. Could the link field be empty?

---

<div class="post-metadata">

**Author:** ![tkf](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tkf/32/17635_2.png) [@tkf](https://discourse.julialang.org/u/tkf)\
**Post date:** [February 8, 2022, 5:11am UTC](https://discourse.julialang.org/t/parallel-reductions-with-zygote/75969/9 "2022-02-08T05:11:46Z")

</div>

I guess we need FoldsChainRules.jl or sth and simply port the rrule for the sequential sum. Many clever things are possible but that looks like a good start. I think it’s better to do this at Folds.jl than ThreadsX.jl so that we get Distributed/Dagger-based definition “for free” as well. Adding AD support to JuliaFolds is rather overdue… But my excuse always has been, “yeah but there are things like Tullio.jl.” @mcabbott Does Tullio.jl handle this case?

---

<div class="post-metadata">

**Author:** ![SEA](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sea/32/8956_2.png) [@SEA](https://discourse.julialang.org/u/SEA)\
**Post date:** [February 8, 2022, 9:31am UTC](https://discourse.julialang.org/t/parallel-reductions-with-zygote/75969/10 "2022-02-08T09:31:22Z")

</div>

I just implemented a custom adjoint for `ThreadsX.sum` by replacing the `map` and `sum` calls in the regular `rrule` with `ThreadsX` versions. I am able to achieve a **speedup of a factor of 6.5** using 10 threads on an Apple M1 Pro in my use case (and no segfault). I am attaching the implementation. Thank you for the helpful pointer @mcabbot!

```nohighlight
using ThreadsX
using ChainRules
using ChainRules: RuleConfig, HasReverseMode, rrule_via_ad, ProjectTo, NoTangent, unthunk
function ChainRules.rrule(
    config::RuleConfig{>:HasReverseMode}, ::typeof(ThreadsX.sum), f, xs::AbstractArray)
    fx_and_pullbacks = ThreadsX.map(x->rrule_via_ad(config, f, x), xs)
    y = ThreadsX.sum(first, fx_and_pullbacks)

    pullbacks = ThreadsX.map(last, fx_and_pullbacks)

    project = ProjectTo(xs)

    function sum_pullback(ȳ)
        call(f, x) = f(x)
        # if dims is :, then need only left-handed only broadcast
        # broadcast_ȳ = dims isa Colon ? (ȳ,) : ȳ
        broadcast_ȳ = ȳ
        f̄_and_x̄s = ThreadsX.map(f->f(ȳ), pullbacks)
        # no point thunking as most of work is in f̄_and_x̄s which we need to compute for both
        f̄ = if fieldcount(typeof(f)) === 0 # Then don't need to worry about derivative wrt f
            NoTangent()
        else
            ThreadsX.sum(first, f̄_and_x̄s)
        end
        x̄s = ThreadsX.map(unthunk ∘ last, f̄_and_x̄s) # project does not support receiving InplaceableThunks
        return NoTangent(), f̄, project(x̄s)
    end
    return y, sum_pullback
end

```

---

<div class="post-metadata">

**Author:** ![vchuravy](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/vchuravy/32/8_2.png) [@vchuravy](https://discourse.julialang.org/u/vchuravy)\
**Post date:** [February 8, 2022, 12:49pm UTC](https://discourse.julialang.org/t/parallel-reductions-with-zygote/75969/11 "2022-02-08T12:49:15Z")

</div>

Just as a general note, if you are using the M1 and are using threading I would strongly recommend using the nightly for Julia 1.8. The spurious seqfaults and hangs with threading got fixed there recently.

---

<div class="post-metadata">

**Author:** ![SEA](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sea/32/8956_2.png) [@SEA](https://discourse.julialang.org/u/SEA)\
**Post date:** [February 9, 2022, 1:05pm UTC](https://discourse.julialang.org/t/parallel-reductions-with-zygote/75969/12 "2022-02-09T13:05:38Z")

</div>

Edit: This issue got resolved and was unrelated to the parallel code, but a bug Zygote currently exhibits with mutable structs. The implementation above works.

An update:

While the custom adjoint accelerates the backward pass and executes reliably on test problems, it leads to a non-deterministic `AssertionError` in Zygote’s code with my more complex model. I opened an [issue about it here](https://github.com/FluxML/Zygote.jl/issues/1159).

Also, the segfaults on ARM are indeed fixed on the nightly for Julia 1.8!