Merge pull request #575 from LuxDL/ap/ad_benchmarks

More Continuous Benchmarks

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "0.5.32"
+version = "0.5.33"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -70,7 +70,7 @@ LuxAMDGPU = "0.2.2"
 LuxCUDA = "0.3.2"
 LuxCore = "0.1.12"
 LuxDeviceUtils = "0.1.16"
-LuxLib = "0.3.10"
+LuxLib = "0.3.11"
 LuxTestUtils = "0.1.15"
 MacroTools = "0.5.13"
 Markdown = "1.10"

bench/Project.toml

@@ -1,9 +1,18 @@
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Tapir = "07d77754-e150-4737-8c94-cd238a1fb45b"
+Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

bench/helpers.jl

@@ -1,19 +1,131 @@
 # TODO: Special Handling for GPU Arrays with @sync
-function benchmark_forward_pass(tag::String, end_tag::String, model, x, ps_nt::NamedTuple,
- st)
+function benchmark_forward_pass(
+ tag::String, end_tag::String, model, x_dims; simple_chains=nothing,
+ flux_model=nothing)
  SUITE[tag]["cpu"]["forward"]["NamedTuple"][end_tag] = @benchmarkable Lux.apply(
- $model, $x, $ps_nt, $st)
+ $model, x, ps_nt, st) setup=((x, ps_nt, st) = general_setup($model, $x_dims))
 
- ps_ca = ComponentArray(ps_nt)
  SUITE[tag]["cpu"]["forward"]["ComponentArray"][end_tag] = @benchmarkable Lux.apply(
- $model, $x, $ps_ca, $st)
+ $model, x, ps_ca, st) setup=((x, ps_nt, st) = general_setup($model, $x_dims); ps_ca = ComponentArray(ps_nt))
+
+ if simple_chains !== nothing
+ simple_chains_model = simple_chains(model)
+ SUITE[tag]["cpu"]["forward"]["SimpleChains"][end_tag] = @benchmarkable Lux.apply(
+ $simple_chains_model, x, ps_simple_chains, st_simple_chains) setup=((x, ps_simple_chains, st_simple_chains) = general_setup(
+ $simple_chains_model, $x_dims))
+ end
+
+ if flux_model !== nothing
+ SUITE[tag]["cpu"]["forward"]["Flux"][end_tag] = @benchmarkable fmodel(x) setup=(x = randn(
+ StableRNG(0), Float32, $x_dims);
+ fmodel = $(flux_model()))
+ end
+
+ return
+end
+
+function benchmark_reverse_pass(
+ tag::String, end_tag::String, backends, model, x_dims;
+ simple_chains=nothing, flux_model=nothing)
+ for backend in backends
+ __benchmark_reverse_pass(tag, end_tag, backend, model, x_dims)
+ end
+
+ if simple_chains !== nothing
+ simple_chains_model = simple_chains(model)
+ __benchmark_reverse_pass_simple_chains(
+ tag, end_tag, AutoZygote(), simple_chains_model, x_dims)
+ end
+
+ if flux_model !== nothing
+ __benchmark_reverse_pass_flux(tag, end_tag, AutoZygote(), flux_model, x_dims)
+ end
 
  return
 end
 
 function general_setup(model, x_dims)
  rng = StableRNG(0)
  ps, st = Lux.setup(rng, model)
+ x_dims === nothing && return ps, st
  x = randn(rng, Float32, x_dims)
  return x, ps, st
 end
+
+# TODO: Remove these once DifferentiationInterface has been released
+function __benchmark_reverse_pass(
+ tag::String, end_tag::String, ::AutoEnzyme, model, x_dims)
+ # TODO: Enable this. But enzyme doesn't handle closures well it seems...
+ # SUITE[tag]["cpu"]["reverse"]["Enzyme"][end_tag] = @benchmarkable Enzyme.gradient(
+ # $Enzyme.Reverse, $f, $x)
+ return error("Enzyme backend hasn't been implemented yet.")
+end
+function __benchmark_reverse_pass(
+ tag::String, end_tag::String, ::AutoTapir, model, x_dims)
+ SUITE[tag]["cpu"]["reverse"]["Tapir"][end_tag] = @benchmarkable Tapir.value_and_pullback!!(
+ trrule, 1.0f0, f, ps_ca) setup=begin
+ (x, ps, st) = general_setup($model, $x_dims)
+ ps_ca = ComponentArray(ps)
+ f = @closure(p->sum(abs2, first(Lux.apply($model, x, p, st))))
+ trrule = Tapir.build_rrule(f, ps_ca)
+ end
+ return
+end
+function __benchmark_reverse_pass(
+ tag::String, end_tag::String, ::AutoTracker, model, x_dims)
+ SUITE[tag]["cpu"]["reverse"]["Tracker"][end_tag] = @benchmarkable Tracker.gradient(
+ f, ps_ca) setup=begin
+ (x, ps, st) = general_setup($model, $x_dims)
+ ps_ca = ComponentArray(ps)
+ f = @closure(p->sum(abs2, first(Lux.apply($model, x, p, st))))
+ end
+ return
+end
+function __benchmark_reverse_pass(
+ tag::String, end_tag::String, ad::AutoReverseDiff, model, x_dims)
+ if ad.compile
+ SUITE[tag]["cpu"]["reverse"]["ReverseDiff (compiled)"][end_tag] = @benchmarkable ReverseDiff.gradient!(
+ ∂ps, tape, ps_ca) setup=begin
+ (x, ps, st) = general_setup($model, $x_dims)
+ ps_ca = ComponentArray(ps)
+ ∂ps = similar(ps_ca)
+ f = @closure(p->sum(abs2, first(Lux.apply($model, x, p, st))))
+ tape = ReverseDiff.compile(ReverseDiff.GradientTape(f, ps_ca))
+ end
+ else
+ SUITE[tag]["cpu"]["reverse"]["ReverseDiff"][end_tag] = @benchmarkable ReverseDiff.gradient(
+ f, ps_ca) setup=begin
+ (x, ps, st) = general_setup($model, $x_dims)
+ ps_ca = ComponentArray(ps)
+ f = @closure(p->sum(abs2, first(Lux.apply($model, x, p, st))))
+ end
+ end
+end
+function __benchmark_reverse_pass(tag::String, end_tag::String, ::AutoZygote, model, x_dims)
+ SUITE[tag]["cpu"]["reverse"]["Zygote"][end_tag] = @benchmarkable Zygote.gradient(
+ f, ps_ca) setup=begin
+ (x, ps, st) = general_setup($model, $x_dims)
+ ps_ca = ComponentArray(ps)
+ f = @closure(p->sum(abs2, first(Lux.apply($model, x, p, st))))
+ end
+ return
+end
+function __benchmark_reverse_pass_simple_chains(
+ tag::String, end_tag::String, ::AutoZygote, model, x_dims)
+ SUITE[tag]["cpu"]["reverse"]["SimpleChains"][end_tag] = @benchmarkable Zygote.gradient(
+ f, ps) setup=begin
+ (x, ps, st) = general_setup($model, $x_dims)
+ f = @closure(p->sum(abs2, first(Lux.apply($model, x, p, st))))
+ end
+ return
+end
+function __benchmark_reverse_pass_flux(
+ tag::String, end_tag::String, ::AutoZygote, model, x_dims)
+ SUITE[tag]["cpu"]["reverse"]["Flux"][end_tag] = @benchmarkable Zygote.gradient(
+ f, m) setup=begin
+ x = randn(StableRNG(0), Float32, $x_dims)
+ m = $(model)()
+ f = @closure(m->sum(abs2, m(x)))
+ end
+ return
+end

bench/layers.jl

@@ -1,8 +1,14 @@
 function add_dense_benchmarks!()
  for n in (2, 20, 200, 2000)
  layer = Dense(n => n)
- x, ps, st = general_setup(layer, (n, 128))
- benchmark_forward_pass("Dense($n => $n)", "($n, 128)", layer, x, ps, st)
+ simple_chains = n ≤ 200 ? Lux.ToSimpleChainsAdaptor((static(n),)) : nothing
+ flux_model = () -> Flux.Dense(n => n)
+ benchmark_forward_pass(
+ "Dense($n => $n)", "($n, 128)", layer, (n, 128); simple_chains, flux_model)
+ benchmark_reverse_pass(
+ "Dense($n => $n)", "($n, 128)",
+ (AutoTracker(), AutoReverseDiff(), AutoReverseDiff(true), AutoZygote()),
+ layer, (n, 128); simple_chains, flux_model)
  end
 
  return
@@ -11,9 +17,15 @@ end
 function add_conv_benchmarks!()
  for ch in (1, 3, 16, 64)
  layer = Conv((3, 3), ch => ch)
- x, ps, st = general_setup(layer, (64, 64, ch, 128))
- benchmark_forward_pass(
- "Conv((3, 3), $ch => $ch)", "(64, 64, $ch, 128)", layer, x, ps, st)
+ simple_chains = ch ≤ 16 ?
+ Lux.ToSimpleChainsAdaptor((static(64), static(64), static(ch))) :
+ nothing
+ flux_model = () -> Flux.Conv((3, 3), ch => ch)
+ benchmark_forward_pass("Conv((3, 3), $ch => $ch)", "(64, 64, $ch, 128)",
+ layer, (64, 64, ch, 128); simple_chains, flux_model)
+ benchmark_reverse_pass("Conv((3, 3), $ch => $ch)", "(64, 64, $ch, 128)",
+ (AutoTracker(), AutoReverseDiff(), AutoReverseDiff(true), AutoZygote()),
+ layer, (64, 64, ch, 128); simple_chains, flux_model)
  end
 end
 

bench/runbenchmarks.jl

@@ -1,11 +1,28 @@
+using ADTypes: ADTypes, AutoEnzyme, AutoTracker, AutoReverseDiff, AutoZygote
 using BenchmarkTools: BenchmarkTools, BenchmarkGroup, @btime, @benchmarkable
 using ComponentArrays: ComponentArray
 using InteractiveUtils: versioninfo
+using FastClosures: @closure
+using Flux: Flux
 using Lux: Lux, BatchNorm, Chain, Conv, Dense, Dropout, FlattenLayer, MaxPool
 using NNlib: relu
+using SimpleChains: SimpleChains, static
 using StableRNGs: StableRNG
 using Statistics: median
 
+# AD Backends
+using Enzyme: Enzyme
+using ReverseDiff: ReverseDiff
+using Tapir: Tapir
+using Tracker: Tracker
+using Zygote: Zygote
+
+# BenchmarkTools Parameters
+BenchmarkTools.DEFAULT_PARAMETERS.samples = 100
+BenchmarkTools.DEFAULT_PARAMETERS.seconds = 0.25
+
+struct AutoTapir <: ADTypes.AbstractReverseMode end
+
 @info sprint(versioninfo)
 
 const SUITE = BenchmarkGroup()
@@ -14,7 +31,8 @@ include("helpers.jl")
 include("vgg.jl")
 include("layers.jl")
 
-BenchmarkTools.tune!(SUITE)
+BenchmarkTools.tune!(SUITE; verbose=true)
 results = BenchmarkTools.run(SUITE; verbose=true)
+display(median(results))
 
 BenchmarkTools.save(joinpath(@__DIR__, "benchmark_results.json"), median(results))

bench/vgg.jl

@@ -17,9 +17,33 @@ function add_vgg_benchmarks!()
  BatchNorm(512), MaxPool((2, 2)), FlattenLayer(), Dense(512, 4096, relu),
  Dropout(0.5), Dense(4096, 4096, relu), Dropout(0.5), Dense(4096, 10))
 
+ flux_model = () -> Flux.Chain(
+ Flux.Conv((3, 3), 3 => 64, relu; pad=(1, 1), stride=(1, 1)),
+ Flux.BatchNorm(64), Flux.Conv((3, 3), 64 => 64, relu; pad=(1, 1), stride=(1, 1)),
+ Flux.BatchNorm(64), Flux.MaxPool((2, 2)),
+ Flux.Conv((3, 3), 64 => 128, relu; pad=(1, 1), stride=(1, 1)), Flux.BatchNorm(128),
+ Flux.Conv((3, 3), 128 => 128, relu; pad=(1, 1), stride=(1, 1)), Flux.BatchNorm(128),
+ Flux.MaxPool((2, 2)),
+ Flux.Conv((3, 3), 128 => 256, relu; pad=(1, 1), stride=(1, 1)), Flux.BatchNorm(256),
+ Flux.Conv((3, 3), 256 => 256, relu; pad=(1, 1), stride=(1, 1)), Flux.BatchNorm(256),
+ Flux.Conv((3, 3), 256 => 256, relu; pad=(1, 1), stride=(1, 1)), Flux.BatchNorm(256),
+ Flux.MaxPool((2, 2)),
+ Flux.Conv((3, 3), 256 => 512, relu; pad=(1, 1), stride=(1, 1)), Flux.BatchNorm(512),
+ Flux.Conv((3, 3), 512 => 512, relu; pad=(1, 1), stride=(1, 1)), Flux.BatchNorm(512),
+ Flux.Conv((3, 3), 512 => 512, relu; pad=(1, 1), stride=(1, 1)), Flux.BatchNorm(512),
+ Flux.MaxPool((2, 2)),
+ Flux.Conv((3, 3), 512 => 512, relu; pad=(1, 1), stride=(1, 1)), Flux.BatchNorm(512),
+ Flux.Conv((3, 3), 512 => 512, relu; pad=(1, 1), stride=(1, 1)), Flux.BatchNorm(512),
+ Flux.Conv((3, 3), 512 => 512, relu; pad=(1, 1), stride=(1, 1)), Flux.BatchNorm(512),
+ Flux.MaxPool((2, 2)), Flux.flatten, Flux.Dense(512, 4096, relu), Flux.Dropout(0.5),
+ Flux.Dense(4096, 4096, relu), Flux.Dropout(0.5), Flux.Dense(4096, 10))
+
  for bsize in (1, 16, 64)
- x, ps, st = general_setup(vgg16, (32, 32, 3, bsize))
- benchmark_forward_pass("vgg16", "(32, 32, 3, $bsize)", vgg16, x, ps, st)
+ benchmark_forward_pass(
+ "vgg16", "(32, 32, 3, $bsize)", vgg16, (32, 32, 3, bsize); flux_model)
+ benchmark_reverse_pass(
+ "vgg16", "(32, 32, 3, $bsize)", (AutoTracker(), AutoZygote()),
+ vgg16, (32, 32, 3, bsize); flux_model)
  end
 
  return

docs/src/.vitepress/config.mts

@@ -62,9 +62,9 @@ export default defineConfig({
  },
  nav: [
  { text: 'Home', link: '/' },
- { text: 'Getting Started', link: '/introduction/index' },
+ { text: 'Getting Started', link: '/introduction' },
  { text: 'Benchmarks', link: 'https://lux.csail.mit.edu/benchmarks/' },
- { text: 'Tutorials', link: '/tutorials/index' },
+ { text: 'Tutorials', link: '/tutorials' },
  { text: 'Manual', link: '/manual/interface' },
  {
  text: 'API', items: [
@@ -104,22 +104,22 @@ export default defineConfig({
  },
  {
  text: 'Versions', items: [
- { text: 'Stable', link: 'https://lux.csail.mit.edu/stable/' },
- { text: 'Dev', link: 'https://lux.csail.mit.edu/dev/' }
+ { text: 'Stable', link: 'https://lux.csail.mit.edu/stable' },
+ { text: 'Dev', link: 'https://lux.csail.mit.edu/dev' }
  ]
  }
  ],
  sidebar: {
  "/introduction/": {
  text: 'Getting Started', collapsed: false, items: [
- { text: 'Introduction', link: '/introduction/index' },
+ { text: 'Introduction', link: '/introduction' },
  { text: 'Overview', link: '/introduction/overview' },
  { text: 'Resources', link: '/introduction/resources' },
  { text: 'Citation', link: '/introduction/citation' }]
  },
  "/tutorials/": {
  text: 'Tutorials', collapsed: false, items: [
- { text: 'Overview', link: '/tutorials/index' },
+ { text: 'Overview', link: '/tutorials' },
  {
  text: 'Beginner', collapsed: false, items: [
  { text: 'Julia & Lux for the Uninitiated', link: '/tutorials/beginner/1_Basics' },

docs/src/index.md

@@ -10,7 +10,7 @@ hero:
  actions:
  - theme: brand
  text: Tutorials
- link: /tutorials/
+ link: /tutorials
  - theme: alt
  text: Ecosystem
  link: /ecosystem
@@ -28,7 +28,7 @@ features:
  - icon: 🚀
  title: Fast & Extendible
  details: Lux.jl is written in Julia itself, making it extremely extendible. CUDA and AMDGPU are supported first-class, with experimental support for Metal Hardware.
- link: /introduction/
+ link: /introduction
 
  - icon: 🧑‍🔬
  title: SciML ❤️ Lux

examples/SimpleChains/main.jl

@@ -78,6 +78,11 @@ function train(model; rng=Xoshiro(0), kwargs...)
  train_state = Lux.Experimental.TrainState(
  rng, model, Adam(3.0f-4); transform_variables=identity)
 
+ ### Warmup the model
+ x_proto = randn(rng, Float32, 28, 28, 1, 1)
+ y_proto = onehotbatch([1], 0:9)
+ Lux.Experimental.compute_gradients(AutoZygote(), loss, (x_proto, y_proto), train_state)
+
  ### Lets train the model
  nepochs = 10
  for epoch in 1:nepochs