From 74f0efcb1a5325b7ca22e2d456899b90e51ba214 Mon Sep 17 00:00:00 2001 From: CarloLucibello Date: Mon, 3 Aug 2026 22:27:50 +0200 Subject: [PATCH] Autocast (scope-based prototype): mixed precision via ScopedValue + dispatch barrier Draft/comparison implementation of Flux.autocast. Layers consult an ambient ScopedValue at forward time through a per-layer dispatch barrier; the machinery is compiled out until the first `autocast` call (a one-time world-age flip of `autocast_active()`), so code that never uses autocast keeps exact inference and zero overhead. Once enabled, forward passes of the affected layers infer as the small union of the Float32/Float16/BFloat16 paths. Also splits the precision casts: f16/bf16 are full casts (like PyTorch model.half()); f16mix/bf16mix keep BatchNorm/InstanceNorm/GroupNorm statistics and affine parameters in Float32. Known: after the flip, active-scope forward passes are a small union rather than a single concrete type; bf16 + AutoEnzyme is blocked upstream (EnzymeAD/Enzyme.jl#3430). This branch is kept for comparison against a wrapper-based implementation. Co-Authored-By: Claude Fable 5 --- NEWS.md | 4 +- Project.toml | 2 + docs/make.jl | 1 + docs/src/guide/training/mixed_precision.md | 153 ++++++++++++++++++ docs/src/reference/training/reference.md | 11 ++ docs/src/reference/utilities.md | 2 + ext/FluxEnzymeExt.jl | 14 +- ext/FluxFiniteDifferencesExt.jl | 16 +- ext/FluxMooncakeExt.jl | 32 ++-- src/Flux.jl | 7 +- src/autocast.jl | 174 +++++++++++++++++++++ src/functor.jl | 92 +++++++---- src/gradient.jl | 35 +++-- src/layers/basic.jl | 26 +-- src/layers/conv.jl | 24 ++- src/layers/normalise.jl | 18 ++- src/layers/recurrent.jl | 67 +++++--- src/losses/Losses.jl | 1 + src/losses/functions.jl | 17 ++ src/train.jl | 34 ++-- test/autocast.jl | 153 ++++++++++++++++++ test/ext_cuda/autocast.jl | 46 ++++++ test/ext_cuda/layers.jl | 9 +- test/ext_enzyme/enzyme.jl | 4 + test/utils.jl | 22 ++- 25 files changed, 829 insertions(+), 135 deletions(-) create mode 100644 docs/src/guide/training/mixed_precision.md create mode 100644 src/autocast.jl create mode 100644 test/autocast.jl create mode 100644 test/ext_cuda/autocast.jl diff --git a/NEWS.md b/NEWS.md index a34cff7de9..e8cb0c9b68 100644 --- a/NEWS.md +++ b/NEWS.md @@ -4,7 +4,9 @@ See also [github's page](https://github.com/FluxML/Flux.jl/releases) for a compl ## Unreleased -- `f16` and `bf16` now convert `BatchNorm`, `InstanceNorm` and `GroupNorm` in **mixed precision**: their statistics and affine parameters are kept in `Float32` while the data flows in half precision. This matches what the NNlib (cuDNN) normalization kernels require for half-precision feature maps, so half-precision `BatchNorm` now works on the GPU. `LayerNorm` is still converted fully (it wraps `NNlib.normalise`, which has no such requirement). Note the behavior change: these layers' parameters are no longer downcast to `Float16`/`BFloat16`. Also adds bf16 `Conv`/pooling GPU coverage. Requires NNlib ≥ 0.9.42 ([#2700](https://github.com/FluxML/Flux.jl/pull/2700)). +- Added `Flux.autocast` for PyTorch-style mixed-precision training: inside an `autocast(f, Float16)` or `autocast(f, BFloat16)` scope — or via the new `autocast` keyword of `gradient`, `withgradient` and `train!` — matmul/convolution-heavy layers cast parameters and inputs to half precision at call time, while normalization layers and loss functions compute in `Float32`. Parameters stay `Float32` ("master weights") and gradients are accumulated back in `Float32`, so the usual optimiser setup works unchanged. Autocast is compiled out until its first use, so code that never calls it pays no overhead and layer forward passes keep their exact inferred return types; the first `autocast` call triggers a one-time recompilation, after which forward passes infer as a small concrete union over the three precision paths ([#2702](https://github.com/FluxML/Flux.jl/pull/2702)). +- Added `f16mix` and `bf16mix`, which convert a model to half precision while keeping the statistics and affine parameters of `BatchNorm`, `InstanceNorm` and `GroupNorm` in `Float32` (the layout the NNlib/cuDNN normalization kernels require for half-precision feature maps, so these models work on the GPU). `f16` and `bf16` remain full casts of every parameter, like PyTorch's `model.half()`/`model.bfloat16()` ([#2702](https://github.com/FluxML/Flux.jl/pull/2702)). +- Half-precision `BatchNorm`/`InstanceNorm`/`GroupNorm` now work on the GPU via the mixed-precision layout described above. Also adds bf16 `Conv`/pooling GPU coverage. Requires NNlib ≥ 0.9.42 ([#2700](https://github.com/FluxML/Flux.jl/pull/2700)). - The normalization layers `BatchNorm`, `InstanceNorm`, `GroupNorm` and `LayerNorm` now delegate their forward pass to the functional operators `NNlib.batchnorm`, `NNlib.instancenorm`, `NNlib.groupnorm` and `NNlib.normalise` (which requires NNlib v0.9.41): the normalization logic now lives in NNlib and is shared across the ecosystem. As a side effect, `LayerNorm` and `Flux.normalise` now add `eps` (rather than `eps^2`) to the variance for numerical stability, matching the other normalization layers ([#2701](https://github.com/FluxML/Flux.jl/pull/2701)). - `Flux.normalise` is now a thin wrapper around `NNlib.normalise`, which should be preferred ([#2701](https://github.com/FluxML/Flux.jl/pull/2701)). - Removed the `FluxCUDAcuDNNExt` extension and the `cuDNN` dependency: cuDNN-accelerated `BatchNorm` on the GPU is now provided by NNlib and selected automatically for `CuArray`s ([#2701](https://github.com/FluxML/Flux.jl/pull/2701)). On AMDGPU the MIOpen `BatchNorm` fast path was likewise removed; `BatchNorm` uses NNlib's generic path there until a MIOpen fast path lands in NNlib ([NNlib.jl#752](https://github.com/FluxML/NNlib.jl/issues/752)). diff --git a/Project.toml b/Project.toml index df77ee6456..c43e7a65eb 100644 --- a/Project.toml +++ b/Project.toml @@ -26,6 +26,7 @@ Preferences = "21216c6a-2e73-6563-6e65-726566657250" ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" +ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63" Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" @@ -76,6 +77,7 @@ Optimisers = "0.4.1" Preferences = "1" ProgressLogging = "0.1" Reexport = "1.0" +ScopedValues = "1.3" Setfield = "1.1" SpecialFunctions = "2.1.2" Statistics = "1" diff --git a/docs/make.jl b/docs/make.jl index f8d7d6cfb0..2758059e34 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -35,6 +35,7 @@ makedocs( "Fitting a Line" => "guide/models/overview.md", "Gradients and Layers" => "guide/models/basics.md", "Training" => "guide/training/training.md", + "Mixed Precision" => "guide/training/mixed_precision.md", "Recurrence" => "guide/models/recurrence.md", "GPU Support" => "guide/gpu.md", "Saving & Loading" => "guide/saving.md", diff --git a/docs/src/guide/training/mixed_precision.md b/docs/src/guide/training/mixed_precision.md new file mode 100644 index 0000000000..9aa3ae1be9 --- /dev/null +++ b/docs/src/guide/training/mixed_precision.md @@ -0,0 +1,153 @@ +# Mixed Precision + +Training in reduced floating point precision (`Float16` or `BFloat16`) can be +substantially faster on modern GPUs and halves the memory taken by activations. +Flux offers two complementary mechanisms, mirroring the two approaches available +in PyTorch: + +1. **Autocast (recommended for training)**: the model's parameters stay in + `Float32`, and a scoped [`autocast`](@ref) context casts values at layer-call + time. This corresponds to PyTorch's `torch.autocast`. +2. **Static casting (recommended for inference)**: [`f16`](@ref) and + [`bf16`](@ref) convert the parameters themselves, like PyTorch's + `model.half()` and `model.bfloat16()`. + +## Autocast + +Wrap the forward pass — or simply pass the `autocast` keyword to +[`Flux.gradient`](@ref), [`Flux.withgradient`](@ref) or [`Flux.train!`](@ref): + +```julia +using Flux + +model = Chain(Conv((3, 3), 3 => 16, relu), BatchNorm(16), + Flux.flatten, Dense(16 * 26 * 26 => 10)) |> gpu +opt_state = Flux.setup(Adam(1e-3), model) + +for (x, y) in dataloader + loss, grad = Flux.withgradient(model; autocast=BFloat16) do m + Flux.logitcrossentropy(m(x), y) + end + Flux.update!(opt_state, model, grad[1]) +end + +# or, equivalently: +Flux.train!((m, x, y) -> Flux.logitcrossentropy(m(x), y), model, dataloader, + opt_state; autocast=BFloat16) +``` + +Inside the scope: + +- Matmul- and convolution-heavy layers (`Dense`, `Conv`, `ConvTranspose`, + `CrossCor`, `Bilinear`, `Embedding` on onehot input, `MultiHeadAttention`, and + the recurrent cells) cast their parameters and inputs to the requested half + precision before computing, so the compute-intensive kernels run fast and the + large activations take half the memory. +- Numerically sensitive operations compute in `Float32`: the normalization + layers (`BatchNorm`, `LayerNorm`, `InstanceNorm`, `GroupNorm`) and the loss + functions in `Flux.Losses` cast their inputs *up*. +- The parameters are never modified; they act as `Float32` "master weights". + The backward pass of each cast accumulates the gradient back in `Float32`, so + parameter gradients — and therefore the optimiser state and update — are + full-precision, with no change to the training loop. + +Things to keep in mind: + +- **`Float16` can underflow.** Its narrow exponent range means small gradients + flush to zero; robust `Float16` training typically needs dynamic loss scaling + (PyTorch's `GradScaler`), which Flux does not provide yet. `BFloat16` has the + same exponent range as `Float32` and trains reliably without it — prefer + `BFloat16` where supported. +- Raw reductions in user code (e.g. `sum(abs2, m(x))`) are not intercepted: they + compute in the half precision that flows into them, just as in PyTorch. Use + the `Flux.Losses` functions, or cast to `Float32` yourself, for the final + reduction. +- The softmax inside `MultiHeadAttention`'s `dot_product_attention` runs in the + half precision (the projections produce half-precision `q`, `k`, `v`). +- Custom layers are not cast automatically unless they are built out of Flux + layers. To opt in, consult [`Flux.autocast_eltype`](@ref) in your forward pass + (see its docstring). +- Weights are re-cast on every forward pass (there is no cast cache). The cast + is cheap next to the matmul/convolution it enables, and compiled backends fuse + it away. +- Autocast is compiled out until its first use: models that never enter an + `autocast` scope pay zero overhead and keep their exact inferred return types. + The first `autocast` call in a session enables the machinery globally, which + triggers a one-time recompilation of affected layer code. +- Autocast works with Zygote (the default), Mooncake, and — for `Float16` — + Enzyme. `BFloat16` autocast is currently not supported with Enzyme. + +## Static casting: `f16`, `bf16` and `f16mix`, `bf16mix` + +For inference, converting the parameters once avoids the per-call casts: + +```julia +model16 = bf16(model) # full cast, like PyTorch's model.bfloat16() +model16 = bf16mix(model) # same, but norm statistics/affine stay Float32 +``` + +[`f16`](@ref)/[`bf16`](@ref) convert *every* parameter. [`f16mix`](@ref)/ +[`bf16mix`](@ref) keep the statistics and affine parameters of `BatchNorm`, +`InstanceNorm` and `GroupNorm` in `Float32`, which the GPU normalization kernels +require for half-precision inputs — prefer the `mix` variants for models +containing those layers. + +For *training* a statically converted model, the gradients and optimiser state +are also half-precision, which loses update accuracy. The +[`Optimisers.MixedPrecision`](https://fluxml.ai/Optimisers.jl/dev/api/#Optimisers.MixedPrecision) rule +compensates by keeping a `Float32` copy of the parameters inside the optimiser +state: + +```julia +model16 = bf16mix(model) +opt_state = Flux.setup(Optimisers.MixedPrecision(Adam(1e-3)), model16) +``` + +Compared to autocast this halves the model's parameter memory (at the price of +the extra copy in the optimiser state) but computes *everything* except the norm +layers in half precision, including the numerically sensitive reductions. + +## Custom layers under autocast + +Flux's built-in layers consult the ambient autocast scope in their forward pass. +A custom layer written in terms of Flux layers (e.g. a struct holding a `Dense`) +inherits this for free. A layer that multiplies its own weight arrays needs one +extra line to participate: + +```julia +struct Affine{W, B} + weight::W + bias::B +end +Flux.@layer Affine + +function (a::Affine)(x) + Flux._autocast_barrier() do T # T is Float16, BFloat16, or nothing + W = Flux._autocast_down(T, a.weight) + b = Flux._autocast_down(T, a.bias) + xT = Flux._autocast_down(T, x) + return W * xT .+ b + end +end +``` + +`Flux._autocast_down(T, x)` casts a floating-point array to `T` and is a no-op +when `T === nothing` (no active scope), when the array already has eltype `T`, +and for non-float arrays (integer or onehot inputs, a `false` bias, ...); its +gradient casts back, so parameter gradients stay `Float32`. The +`Flux._autocast_barrier` wrapper runs the closure with the current scope value +and acts as a dispatch barrier, keeping the layer type-stable when autocast is +not in use. For a numerically sensitive custom layer, use `Flux._autocast_up(x)` +instead, which casts half-precision input *up* to `Float32` inside a scope. + +!!! warning + Inside the `_autocast_barrier` closure, only assign to *fresh* local names. + Assigning to a variable captured from the enclosing function (like `x` or a + destructured argument) boxes it, which silently breaks both type inference + and Zygote gradients. + +These helpers are currently internal (underscore-prefixed): the API may still +evolve, but they are the supported way to make a custom layer autocast-aware. + +Note that when writing a custom layer as a plain Julia function of arrays, an +alternative is to rely on [`Flux.autocast_eltype`](@ref) directly. diff --git a/docs/src/reference/training/reference.md b/docs/src/reference/training/reference.md index d8e2eaef4c..c690745987 100644 --- a/docs/src/reference/training/reference.md +++ b/docs/src/reference/training/reference.md @@ -24,6 +24,17 @@ Optimisers.update! Optimisers.setup ``` +## Mixed Precision + +`train!`, [`gradient`](@ref Flux.gradient) and [`withgradient`](@ref Flux.withgradient) +accept an `autocast` keyword to compute gradients under mixed precision. +See the [Mixed Precision](@ref) guide page. + +```@docs +Flux.autocast +Flux.autocast_eltype +``` + `train!` uses [`@progress`](https://github.com/JuliaLogging/ProgressLogging.jl) which should show a progress bar in VSCode automatically. To see one in a terminal, you will need to install [TerminalLoggers.jl](https://github.com/JuliaLogging/TerminalLoggers.jl) and follow its setup instructions. diff --git a/docs/src/reference/utilities.md b/docs/src/reference/utilities.md index 9aef988ab3..c20fcbaeb6 100644 --- a/docs/src/reference/utilities.md +++ b/docs/src/reference/utilities.md @@ -68,4 +68,6 @@ Flux.f64 Flux.f32 Flux.f16 Flux.bf16 +Flux.f16mix +Flux.bf16mix ``` diff --git a/ext/FluxEnzymeExt.jl b/ext/FluxEnzymeExt.jl index 4c15fbd8d5..a9fcd5dae5 100644 --- a/ext/FluxEnzymeExt.jl +++ b/ext/FluxEnzymeExt.jl @@ -10,13 +10,19 @@ using Enzyme: autodiff_thunk, Reverse, ReverseSplitWithPrimal EnzymeRules.inactive(::typeof(Flux.Losses._check_sizes), args...) = true +# NOTE: BFloat16 autocast under Enzyme is blocked upstream: Enzyme's type analysis +# crashes on `Core.BFloat16` values (missing `typetree_primitive` method, see +# EnzymeAD/Enzyme.jl#3430). Since autocast is compiled out until its first use, plain +# Enzyme differentiation is unaffected; only `autocast` + `AutoEnzyme` in the same +# session requires the upstream fix. + ### gradient & withgradient -function Flux.gradient(f::F, adtype::AutoEnzyme, x::Vararg{Any,N}; zero::Bool=true) where {F,N} - return _enzyme_gradient(f, map(_trymake_duplicated, x)...; zero) +function Flux.gradient(f::F, adtype::AutoEnzyme, x::Vararg{Any,N}; zero::Bool=true, autocast=nothing) where {F,N} + return Flux._with_autocast(() -> _enzyme_gradient(f, map(_trymake_duplicated, x)...; zero), autocast) end -function Flux.withgradient(f::F, adtype::AutoEnzyme, x::Vararg{Any,N}; zero::Bool=true) where {F,N} - return _enzyme_withgradient(f, map(_trymake_duplicated, x)...; zero) +function Flux.withgradient(f::F, adtype::AutoEnzyme, x::Vararg{Any,N}; zero::Bool=true, autocast=nothing) where {F,N} + return Flux._with_autocast(() -> _enzyme_withgradient(f, map(_trymake_duplicated, x)...; zero), autocast) end _trymake_duplicated(x::EnzymeCore.Duplicated) = x diff --git a/ext/FluxFiniteDifferencesExt.jl b/ext/FluxFiniteDifferencesExt.jl index 836f401201..659e1e7dc8 100644 --- a/ext/FluxFiniteDifferencesExt.jl +++ b/ext/FluxFiniteDifferencesExt.jl @@ -4,30 +4,38 @@ using Flux using ADTypes: AutoFiniteDifferences using FiniteDifferences -function Flux.gradient(f::F, adtype::AutoFiniteDifferences, x) where F +function Flux.gradient(f::F, adtype::AutoFiniteDifferences, x; autocast=nothing) where F + Flux._with_autocast(autocast) do ps, re = Flux.destructure(x) gs = FiniteDifferences.grad(adtype.fdm, p -> f(re(p)...), ps)[1] return (re(gs),) + end end -function Flux.gradient(f::F, adtype::AutoFiniteDifferences, x::Vararg{Any,N}) where {F, N} +function Flux.gradient(f::F, adtype::AutoFiniteDifferences, x::Vararg{Any,N}; autocast=nothing) where {F, N} + Flux._with_autocast(autocast) do ps, re = Flux.destructure(x) gs = FiniteDifferences.grad(adtype.fdm, p -> f(re(p)...), ps)[1] return re(gs) + end end -function Flux.withgradient(f::F, adtype::AutoFiniteDifferences, x) where F +function Flux.withgradient(f::F, adtype::AutoFiniteDifferences, x; autocast=nothing) where F + Flux._with_autocast(autocast) do ps, re = Flux.destructure(x) y = f(re(ps)...) gs = FiniteDifferences.grad(adtype.fdm, p -> f(re(p)...), ps)[1] return y, (re(gs),) + end end -function Flux.withgradient(f::F, adtype::AutoFiniteDifferences, x::Vararg{Any,N}) where {F, N} +function Flux.withgradient(f::F, adtype::AutoFiniteDifferences, x::Vararg{Any,N}; autocast=nothing) where {F, N} + Flux._with_autocast(autocast) do ps, re = Flux.destructure(x) y = f(re(ps)...) gs = FiniteDifferences.grad(adtype.fdm, p -> f(re(p)...), ps)[1] return y, re(gs) + end end end # module diff --git a/ext/FluxMooncakeExt.jl b/ext/FluxMooncakeExt.jl index cad7444538..e6d7e4bfc9 100644 --- a/ext/FluxMooncakeExt.jl +++ b/ext/FluxMooncakeExt.jl @@ -4,22 +4,26 @@ using ADTypes: AutoMooncake using Mooncake: Mooncake import Flux -function Flux.gradient(f::F, adtype::AutoMooncake, args::Vararg{Any,N}) where {F,N} - return Flux.withgradient(f, adtype, args...)[2] +Mooncake.@zero_adjoint Mooncake.MinimalCtx Tuple{typeof(Flux.autocast_eltype)} + +function Flux.gradient(f::F, adtype::AutoMooncake, args::Vararg{Any,N}; autocast=nothing) where {F,N} + return Flux.withgradient(f, adtype, args...; autocast)[2] end -function Flux.withgradient(f::F, adtype::AutoMooncake, args::Vararg{Any,N}) where {F,N} - config = Mooncake.Config(friendly_tangents=true) - cache = Mooncake.prepare_pullback_cache(f, args...; config) - # `prepare_pullback_cache` already runs the forward pass once to build the rule, and stores - # a primal-typed buffer of the output `y = f(args...)`. We reuse it to learn the structure - # of the output (without an extra forward pass) and to build the cotangent seed. - yout = cache.y_cache - seed = yout isa Union{Tuple, NamedTuple} ? _loss_seed(yout) : one(yout) - # `value_and_pullback!!` does a single forward + reverse and returns the full output `y`, - # so auxiliary outputs come for free in `val`. - val, grads = Mooncake.value_and_pullback!!(cache, seed, f, args...) - return (val=val, grad=grads[2:end]) +function Flux.withgradient(f::F, adtype::AutoMooncake, args::Vararg{Any,N}; autocast=nothing) where {F,N} + Flux._with_autocast(autocast) do + config = Mooncake.Config(friendly_tangents=true) + cache = Mooncake.prepare_pullback_cache(f, args...; config) + # `prepare_pullback_cache` already runs the forward pass once to build the rule, and stores + # a primal-typed buffer of the output `y = f(args...)`. We reuse it to learn the structure + # of the output (without an extra forward pass) and to build the cotangent seed. + yout = cache.y_cache + seed = yout isa Union{Tuple, NamedTuple} ? _loss_seed(yout) : one(yout) + # `value_and_pullback!!` does a single forward + reverse and returns the full output `y`, + # so auxiliary outputs come for free in `val`. + val, grads = Mooncake.value_and_pullback!!(cache, seed, f, args...) + return (val=val, grad=grads[2:end]) + end end # Auxiliary outputs: `f` returns a Tuple or NamedTuple whose first element is the scalar loss. diff --git a/src/Flux.jl b/src/Flux.jl index 1bfe2aa4a5..46fd712e30 100644 --- a/src/Flux.jl +++ b/src/Flux.jl @@ -24,6 +24,7 @@ using Random: default_rng using Zygote, ChainRulesCore using Zygote: @adjoint, pullback using EnzymeCore: EnzymeCore +using ScopedValues: ScopedValue, with @reexport using ADTypes # AutoZygote, AutoMooncake, etc... using ADTypes: AbstractADType @@ -49,7 +50,7 @@ export Chain, Dense, Embedding, EmbeddingBag, LayerNorm, BatchNorm, InstanceNorm, GroupNorm, WeightNorm, MultiHeadAttention, Upsample, PixelShuffle, - fmap, cpu, gpu, f32, f64, f16, bf16, BFloat16, rand32, randn32, zeros32, ones32, + fmap, cpu, gpu, f32, f64, f16, bf16, f16mix, bf16mix, BFloat16, rand32, randn32, zeros32, ones32, testmode!, trainmode! @compat(public, ( # unexported symbols marked as API, on Julia 1.11 @@ -59,6 +60,7 @@ export Chain, Dense, Embedding, EmbeddingBag, Bilinear, Scale, # utils outputsize, state, create_bias, @layer, initialstates, normalise, loadmodel!, + autocast_eltype, activations, modules, flatten, rng_from_array, nfan, # from OneHotArrays.jl onehot, onehotbatch, onecold, @@ -100,6 +102,9 @@ export Chain, Dense, Embedding, EmbeddingBag, remove_weight_norms, )) +include("autocast.jl") +export autocast + include("gradient.jl") export gradient, withgradient diff --git a/src/autocast.jl b/src/autocast.jl new file mode 100644 index 0000000000..08af28debf --- /dev/null +++ b/src/autocast.jl @@ -0,0 +1,174 @@ + +# The scope stores the concrete `Type{Float16}`/`Type{BFloat16}` (or `nothing`) so that a +# read is a small concrete union: dispatching the per-layer barrier on it specializes each +# branch on a single precision, keeping the forward pass out of `Any`-typed territory. +const AUTOCAST_ELTYPE = ScopedValue{Union{Nothing, Type{Float16}, Type{BFloat16}}}(nothing) + +# Autocast is compiled out entirely until its first use. `autocast_active()` is a +# constant-`false` method that the compiler folds away, so the half-precision branches +# below are dead-stripped from every layer's forward pass: outside of autocast, layers +# infer their exact concrete return type, pay zero overhead, and contain no half-precision +# types in their IR. The first `autocast` call in a session redefines the method to return +# `true` — a one-time world-age flip that invalidates and recompiles the affected layer +# code with the scope checks included (from then on forward passes infer as the small +# union of the three precision paths). +autocast_active() = false + +const AUTOCAST_FLIP_LOCK = ReentrantLock() + +function _ensure_autocast_active() + autocast_active() && return nothing + lock(AUTOCAST_FLIP_LOCK) do + # re-check under the lock in the latest world: another task may have flipped, + # and this specialization was compiled before the flip so it folds `false` + if !Base.invokelatest(autocast_active) + @eval autocast_active() = true + end + end + return nothing +end + +""" + autocast_eltype() + +Return the floating point type of the innermost enclosing [`autocast`](@ref) scope, +or `nothing` when called outside any `autocast` scope. + +Custom layers can use this, together with the (internal) cast helpers used by the +built-in layers, to opt into mixed precision: + +```julia +function (m::MyLayer)(x) + T = Flux.autocast_eltype() + W = Flux._autocast_down(T, m.weight) + xT = Flux._autocast_down(T, x) + return W * xT +end +``` +""" +autocast_eltype() = AUTOCAST_ELTYPE[] + +ChainRulesCore.@non_differentiable autocast_eltype() +EnzymeCore.EnzymeRules.inactive(::typeof(autocast_eltype), args...) = true + +""" + autocast(f, T::Type) + +Run `f()` with mixed precision: while inside `f`, the forward pass of matmul- and +convolution-heavy Flux layers (`Dense`, `Conv`, `ConvTranspose`, `CrossCor`, `Bilinear`, +`Embedding` on onehot input, `MultiHeadAttention`, and the recurrent cells) casts +parameters and inputs to the half-precision type `T` (`Float16` or `BFloat16`) before +computing, while numerically sensitive operations (the normalization layers and the +loss functions) compute in `Float32`. + +The model's parameters are not modified: they act as `Float32` "master weights", +and the gradients returned by [`Flux.gradient`](@ref) and [`Flux.withgradient`](@ref) +are accumulated back in `Float32`, so the usual optimiser setup works unchanged. +This mirrors PyTorch's `torch.autocast` recipe for mixed-precision training, +at layer rather than operator granularity. + +Usually used through the `autocast` keyword of [`Flux.gradient`](@ref), +[`Flux.withgradient`](@ref) and [`Flux.train!`](@ref) rather than directly. + +Note that with `T = Float16` gradients can underflow; robust `Float16` training +typically also needs loss scaling, which Flux does not provide yet. +`BFloat16` has the same exponent range as `Float32` and does not need it. + +Autocast is compiled out until its first use: code that never calls `autocast` pays no +overhead, and layer forward passes keep their exact inferred return types. The first +`autocast` call in a session enables the machinery globally, which triggers a one-time +recompilation of the affected layer code (from then on, forward passes infer as the +small union of the `Float32`/`Float16`/`BFloat16` paths). + +See also [`f16`](@ref) and [`bf16`](@ref) for statically converting a model instead. + +# Examples + +```julia-repl +julia> model = Chain(Dense(3 => 4, relu), BatchNorm(4), Dense(4 => 2)); + +julia> x = randn(Float32, 3, 8); + +julia> y = autocast(BFloat16) do + model(x) + end; + +julia> eltype(y) # the final Dense ran in BFloat16 +BFloat16 + +julia> eltype(model[1].weight) # parameters are untouched +Float32 + +julia> grad = Flux.gradient(m -> sum(abs2, m(x)), model; autocast=BFloat16)[1]; + +julia> eltype(grad.layers[1].weight) # gradients are Float32, like the parameters +Float32 +``` +""" +function autocast(f, ::Type{T}) where {T<:Union{Float16, BFloat16}} + _ensure_autocast_active() + # `invokelatest`: this call may sit in a specialization compiled before the flip + return Base.invokelatest(with, f, AUTOCAST_ELTYPE => T) +end + +autocast(f, ::Type{T}) where T = + throw(ArgumentError("autocast supports Float16 and BFloat16, got $T")) + +_with_autocast(g, ::Nothing) = g() +_with_autocast(g, ::Type{T}) where T = autocast(g, T) + +# Run `f(autocast_eltype())` as a dispatch barrier: the scope value is a concrete type (or +# `nothing`), so each branch of the small union specializes `f` on a single precision. This +# collapses what would otherwise be `Union`-of-`Union` inference (e.g. `W * x` with both cast) +# into a per-precision concrete computation. Until the first `autocast` call flips +# `autocast_active()`, the whole scope consultation folds away to `f(nothing)`. +@inline function _autocast_barrier(f::F) where {F} + if autocast_active() + return f(autocast_eltype()) + else + return f(nothing) + end +end + +# Cast to the autocast eltype, for matmul/conv-family layers. `nothing` (no active scope) and +# non-float arrays (onehot/integer inputs, `false` bias, `Nil`, ...) pass through unchanged. +# BFloat16 goes through `_to_bf16` rather than a native `convert`/broadcast, which can hang +# LLVM codegen on some platforms (JuliaMath/BFloat16s.jl#107). +_autocast_down(::Nothing, x) = x +_autocast_down(::Type{Float16}, x::AbstractArray{Float16}) = x +_autocast_down(::Type{Float16}, x::AbstractArray{<:AbstractFloat}) = Float16.(x) +_autocast_down(::Type{BFloat16}, x::AbstractArray{BFloat16}) = x +_autocast_down(::Type{BFloat16}, x::AbstractArray{<:AbstractFloat}) = _to_bf16(x) +_autocast_down(::Type{<:Union{Float16, BFloat16}}, x) = x + +_autocast_down_pullback(proj) = dx -> (NoTangent(), NoTangent(), proj(unthunk(dx))) +function ChainRulesCore.rrule(::typeof(_autocast_down), ::Type{T}, + x::AbstractArray{<:AbstractFloat}) where T + proj = ChainRulesCore.ProjectTo(x) # widens the cotangent back to eltype(x) + return _autocast_down(T, x), _autocast_down_pullback(proj) +end +function ChainRulesCore.rrule(::typeof(_autocast_down), ::Nothing, x) + return x, dx -> (NoTangent(), NoTangent(), dx) +end + +# Cast half-precision arrays up to Float32 inside an autocast scope, for +# numerically sensitive operations (normalization layers, losses). Outside any +# scope this is a no-op, so statically converted `f16`/`bf16` models are unaffected. +_autocast_up(x) = autocast_active() ? _autocast_up_scoped(x) : x +_autocast_up_scoped(x) = autocast_eltype() === nothing ? x : _cast_f32(x) + +_cast_f32(x::AbstractArray{<:Union{Float16, BFloat16}}) = convert(AbstractArray{Float32}, x) +_cast_f32(x) = x + +function ChainRulesCore.rrule(::typeof(_cast_f32), x::AbstractArray{Float16}) + proj = ChainRulesCore.ProjectTo(x) + cast_f32_pullback(dx) = (NoTangent(), proj(unthunk(dx))) + return _cast_f32(x), cast_f32_pullback +end + +# For BFloat16 the cotangent must be truncated through `_to_bf16`, not ProjectTo, +# again because of JuliaMath/BFloat16s.jl#107. +function ChainRulesCore.rrule(::typeof(_cast_f32), x::AbstractArray{BFloat16}) + cast_f32_bf16_pullback(dx) = (NoTangent(), _to_bf16(unthunk(dx))) + return _cast_f32(x), cast_f32_bf16_pullback +end diff --git a/src/functor.jl b/src/functor.jl index 5a05e915a5..bcbe59310f 100644 --- a/src/functor.jl +++ b/src/functor.jl @@ -124,26 +124,18 @@ Adapt.adapt_storage(::FluxEltypeAdaptor{T}, x::AbstractArray{<:AbstractFloat}) w Adapt.adapt_storage(::FluxEltypeAdaptor{T}, x::AbstractArray{<:Complex{<:AbstractFloat}}) where {T<:AbstractFloat} = convert(AbstractArray{Complex{T}}, x) -# `Float16`/`BFloat16` are "half precision": under these casts, some layers keep part -# of their parameters in `Float32` (mixed precision, see `_keep_f32_under_halfprec`). -_ishalfprec(::Type) = false -_ishalfprec(::Type{Float16}) = true -_ishalfprec(::Type{BFloat16}) = true - -# Layers that override half-precision conversion to keep (some of) their arrays in -# `Float32`. Extended for normalization layers in `layers/normalise.jl`. +# Layers that override the mixed-precision conversions `f16mix`/`bf16mix` to keep (some +# of) their arrays in `Float32`. Extended for normalization layers in `layers/normalise.jl`. _keep_f32_under_halfprec(::Any) = false -function _paramtype(::Type{T}, m) where T - if _ishalfprec(T) - # Stop the walk at layers that manage their own precision and hand them to `f32` - # (so their statistics/affine parameters stay in `Float32`); convert every other - # leaf array to `T`. - fmap(m; exclude = x -> _keep_f32_under_halfprec(x) || Functors.isleaf(x)) do x - _keep_f32_under_halfprec(x) ? f32(x) : adapt(FluxEltypeAdaptor{T}(), x) - end - else - fmap(adapt(FluxEltypeAdaptor{T}()), m) +_paramtype(::Type{T}, m) where T = fmap(adapt(FluxEltypeAdaptor{T}()), m) + +# Mixed-precision conversion (`f16mix`/`bf16mix`): stop the walk at layers that manage +# their own precision and hand them to `f32` (so their statistics/affine parameters stay +# in `Float32`); convert every other leaf array to `T`. +function _paramtype_mixed(::Type{T}, m) where T + fmap(m; exclude = x -> _keep_f32_under_halfprec(x) || Functors.isleaf(x)) do x + _keep_f32_under_halfprec(x) ? f32(x) : adapt(FluxEltypeAdaptor{T}(), x) end end @@ -205,18 +197,20 @@ f64(m) = _paramtype(Float64, m) """ f16(m) -Converts the `eltype` of model's *floating point* parameters to `Float16`. +Converts the `eltype` of model's *floating point* parameters to `Float16`, +like PyTorch's `model.half()`. All parameters are converted, including the +statistics and affine parameters of the normalization layers; note that the GPU +normalization kernels (cuDNN) require `Float32` statistics/affine parameters for +half-precision inputs, so use [`f16mix`](@ref) for models containing `BatchNorm`, +`InstanceNorm` or `GroupNorm`. Recurses into structs marked with [`@layer`](@ref Flux.@layer). Support for `Float16` is limited on many CPUs. Julia may convert to `Float32` for each operation, which is slow. -The normalization layers `BatchNorm`, `InstanceNorm` and `GroupNorm` are converted in -*mixed precision*: their statistics and affine parameters are kept in `Float32` while -the data flowing through stays in `Float16`. This matches the functional operators in -NNlib (and cuDNN), which require `Float32` parameters for half-precision feature maps. +For mixed-precision *training* with `Float32` master weights, see [`autocast`](@ref). -See also [`f32`](@ref), [`f64`](@ref) and [`bf16`](@ref). +See also [`f16mix`](@ref), [`f32`](@ref), [`f64`](@ref) and [`bf16`](@ref). # Example ```jldoctest @@ -249,12 +243,14 @@ is less prone to overflow/underflow, and is well supported on modern GPUs. Support for `BFloat16` is limited on many CPUs, where Julia may convert to `Float32` for each operation. -The normalization layers `BatchNorm`, `InstanceNorm` and `GroupNorm` are converted in -*mixed precision*: their statistics and affine parameters are kept in `Float32` while -the data flowing through stays in `BFloat16`. This matches the functional operators in -NNlib (and cuDNN), which require `Float32` parameters for half-precision feature maps. +All parameters are converted, including the statistics and affine parameters of the +normalization layers; note that the GPU normalization kernels (cuDNN) require `Float32` +statistics/affine parameters for half-precision inputs, so use [`bf16mix`](@ref) for +models containing `BatchNorm`, `InstanceNorm` or `GroupNorm`. -See also [`f16`](@ref), [`f32`](@ref) and [`f64`](@ref). +For mixed-precision *training* with `Float32` master weights, see [`autocast`](@ref). + +See also [`bf16mix`](@ref), [`f16`](@ref), [`f32`](@ref) and [`f64`](@ref). # Example ```jldoctest @@ -272,3 +268,41 @@ Chain( ``` """ bf16(m) = _paramtype(BFloat16, m) + +""" + f16mix(m) + +Converts the model to *mixed* `Float16` precision: like [`f16`](@ref), but the +statistics and affine parameters of `BatchNorm`, `InstanceNorm` and `GroupNorm` are +kept in `Float32` while the data flowing through them stays in `Float16`. This +matches the functional normalization operators in NNlib (and cuDNN), which require +`Float32` parameters for half-precision feature maps — so unlike a full `f16` cast, +a mixed-precision model works on the GPU. `LayerNorm` contains no such parameters +and is converted fully. + +For mixed-precision *training* with `Float32` master weights, see [`autocast`](@ref). +To keep the optimiser state in `Float32` when training a converted model, see +[`Optimisers.MixedPrecision`](https://fluxml.ai/Optimisers.jl/dev/api/#Optimisers.MixedPrecision). + +See also [`bf16mix`](@ref), [`f16`](@ref), [`f32`](@ref) and [`f64`](@ref). +""" +f16mix(m) = _paramtype_mixed(Float16, m) + +""" + bf16mix(m) + +Converts the model to *mixed* `BFloat16` precision: like [`bf16`](@ref), but the +statistics and affine parameters of `BatchNorm`, `InstanceNorm` and `GroupNorm` are +kept in `Float32` while the data flowing through them stays in `BFloat16`. This +matches the functional normalization operators in NNlib (and cuDNN), which require +`Float32` parameters for half-precision feature maps — so unlike a full `bf16` cast, +a mixed-precision model works on the GPU. `LayerNorm` contains no such parameters +and is converted fully. + +For mixed-precision *training* with `Float32` master weights, see [`autocast`](@ref). +To keep the optimiser state in `Float32` when training a converted model, see +[`Optimisers.MixedPrecision`](https://fluxml.ai/Optimisers.jl/dev/api/#Optimisers.MixedPrecision). + +See also [`f16mix`](@ref), [`bf16`](@ref), [`f32`](@ref) and [`f64`](@ref). +""" +bf16mix(m) = _paramtype_mixed(BFloat16, m) diff --git a/src/gradient.jl b/src/gradient.jl index dc46e26ab0..8a124da557 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -19,6 +19,9 @@ The package corresponding to any chosen backend (except Zygote) must be loaded i If no `adtype` is given, then Zygote.jl is used by default, unless at least one argument is of type `Duplicated` from Enzyme.jl, in which case Enzyme.jl is used. +The keyword `autocast` accepts `Float16` or `BFloat16` to compute the forward and backward +pass under mixed precision, see [`autocast`](@ref). + See also [`withgradient`](@ref) to keep the value `f(args...)`. # Examples @@ -48,18 +51,19 @@ julia> Flux.gradient(f, AutoMooncake(), [1.0, 2.0, 3.0]) ([2.0, 2.0, 2.0],) ``` """ -function gradient(f, adtype::AbstractADType, args...) +function gradient(f, adtype::AbstractADType, args...; kws...) error("AD backend has to be loaded to use `gradient(f, AutoXXX(), args...)`. Make sure to `using` the corresponding package, e.g. `using Mooncake` for `AutoMooncake()`. Supported backends are $SUPPORTED_AD_BACKENDS.") end -gradient(f, adtype::AutoZygote, args...) = Zygote.gradient(f, args...) +gradient(f, adtype::AutoZygote, args...; autocast::Union{Nothing,Type}=nothing) = + _with_autocast(() -> Zygote.gradient(f, args...), autocast) # Default gradient using Zygote -function gradient(f, args...; zero::Bool=true) +function gradient(f, args...; zero::Bool=true, autocast::Union{Nothing,Type}=nothing) for a in args - a isa Union{EnzymeCore.Duplicated, EnzymeCore.Const} && return gradient(f, AutoEnzyme(), args...; zero) + a isa Union{EnzymeCore.Duplicated, EnzymeCore.Const} && return gradient(f, AutoEnzyme(), args...; zero, autocast) end for a in args _ensure_noenzyme(a) @@ -70,7 +74,7 @@ function gradient(f, args...; zero::Bool=true) If you are writing new code, then Zygote over Zygote is heavily discouraged. """) end - return Zygote.gradient(f, args...) + return _with_autocast(() -> Zygote.gradient(f, args...), autocast) end # Without any Duplicated, check for no stray Enzyme types before calling Zygote @@ -137,7 +141,8 @@ julia> Flux.gradient(dup_model, [1]; zero=false) do m, x # implicit Const([1]), ((layers = ((weight = [12.0;;], bias = [12.0], σ = nothing),),), nothing) ``` """ -gradient(f, args::Union{EnzymeCore.Const, EnzymeCore.Duplicated}...; zero::Bool=true) = gradient(f, AutoEnzyme(), args...; zero) +gradient(f, args::Union{EnzymeCore.Const, EnzymeCore.Duplicated}...; zero::Bool=true, autocast::Union{Nothing,Type}=nothing) = + gradient(f, AutoEnzyme(), args...; zero, autocast) """ @@ -152,6 +157,9 @@ The package corresponding to the chosen backend must be loaded in advance. If no `adtype` is given, then Zygote.jl is used by default, unless at least one argument is of type `Duplicated` from Enzyme.jl, in which case Enzyme.jl is used. +The keyword `autocast` accepts `Float16` or `BFloat16` to compute the forward and backward +pass under mixed precision, see [`autocast`](@ref). + Se also [`gradient`](@ref) to get just the gradient. # Examples @@ -203,7 +211,7 @@ julia> Flux.withgradient(AutoMooncake(), [1.0, 2.0, 4.0]) do x (val = (1.75, [1.0, 0.5, 0.25]), grad = ([-1.0, -0.25, -0.0625],)) ``` """ -function withgradient(f, adtype::ADTypes.AbstractADType, args...) +function withgradient(f, adtype::ADTypes.AbstractADType, args...; kws...) error("AD backend has to be loaded to use `withgradient(f, AutoXXX(), args...)`. Make sure to `using` the corresponding package, e.g. `using Mooncake` for `AutoMooncake()`. Supported backends are $SUPPORTED_AD_BACKENDS.") @@ -211,9 +219,9 @@ end # Default withgradient using Zygote -function withgradient(f, args...; zero::Bool=true) +function withgradient(f, args...; zero::Bool=true, autocast::Union{Nothing,Type}=nothing) for a in args - a isa Union{EnzymeCore.Duplicated, EnzymeCore.Const} && return withgradient(f, AutoEnzyme(), args...; zero) + a isa Union{EnzymeCore.Duplicated, EnzymeCore.Const} && return withgradient(f, AutoEnzyme(), args...; zero, autocast) end for a in args _ensure_noenzyme(a) @@ -224,12 +232,12 @@ function withgradient(f, args...; zero::Bool=true) If you are writing new code, then Zygote over Zygote is heavily discouraged. """) end - return Zygote.withgradient(f, args...) + return _with_autocast(() -> Zygote.withgradient(f, args...), autocast) end ## Zygote version, supporting aux output too. -function withgradient(f::F, adtype::AutoZygote, x::Vararg{Any,N}) where {F,N} - return Zygote.withgradient(f, x...) +function withgradient(f::F, adtype::AutoZygote, x::Vararg{Any,N}; autocast::Union{Nothing,Type}=nothing) where {F,N} + return _with_autocast(() -> Zygote.withgradient(f, x...), autocast) end """ @@ -262,7 +270,8 @@ julia> Flux.withgradient(m -> m(3), Duplicated(model)) # this uses Enzyme (val = 14.52, grad = ((layers = ((weight = [0.0 0.0 4.4],), (weight = [3.3;;], bias = [1.0], σ = nothing), nothing),),)) ``` """ -withgradient(f, args::Union{EnzymeCore.Const, EnzymeCore.Duplicated}...; zero::Bool=true) = withgradient(f, AutoEnzyme(), args...; zero) +withgradient(f, args::Union{EnzymeCore.Const, EnzymeCore.Duplicated}...; zero::Bool=true, autocast::Union{Nothing,Type}=nothing) = + withgradient(f, AutoEnzyme(), args...; zero, autocast) ## ADD BACK TO withgradient docstring above when AUX is SUPPORTED # The function `f` may return Tuple or NamedTuple, with the loss as the first element. diff --git a/src/layers/basic.jl b/src/layers/basic.jl index b24d37e2b6..2c5314b412 100644 --- a/src/layers/basic.jl +++ b/src/layers/basic.jl @@ -195,8 +195,12 @@ end function (a::Dense)(x::AbstractVecOrMat) _size_check(a, x, 1 => size(a.weight, 2)) - xT = _match_eltype(a, x) # fixes Float64 input, etc. - return NNlib.bias_act!(a.σ, a.weight * xT, a.bias) # does σ.(W*x .+ b), with fast paths + _autocast_barrier() do T # under `autocast`, cast W/b/x to half precision; else a no-op + W = _autocast_down(T, a.weight) + b = _autocast_down(T, a.bias) + xT = _autocast_down(T, _match_eltype(a, x)) # _match_eltype fixes Float64 input, etc. + return NNlib.bias_act!(a.σ, W * xT, b) # does σ.(W*x .+ b), with fast paths + end end function (a::Dense)(x::AbstractArray) @@ -460,21 +464,25 @@ end Bilinear((in12, out)::Pair{<:Integer, <:Integer}, σ = identity; kw...) = Bilinear((in12, in12) => out, σ; kw...) function (a::Bilinear)(x::AbstractMatrix, y::AbstractMatrix) - W, b, σ = a.weight, a.bias, a.σ + _autocast_barrier() do T + # only fresh local names in here: assigning to captured `x`/`y` would box them + W, b, σ = _autocast_down(T, a.weight), _autocast_down(T, a.bias), a.σ + xT, yT = _autocast_down(T, x), _autocast_down(T, y) d_z, d_x, d_y = size(W) - d_x == size(x,1) && d_y == size(y,1) || throw(DimensionMismatch("number of rows in data must match W")) - size(x,2) == size(y,2) || throw(DimensionMismatch("Data inputs must agree on number of columns, got $(size(x,2)) and $(size(y,2))")) + d_x == size(xT,1) && d_y == size(yT,1) || throw(DimensionMismatch("number of rows in data must match W")) + size(xT,2) == size(yT,2) || throw(DimensionMismatch("Data inputs must agree on number of columns, got $(size(xT,2)) and $(size(yT,2))")) # @einsum Wy[o,i,s] := W[o,i,j] * y[j,s] - Wy = reshape(reshape(W, (:, d_y)) * y, (d_z, d_x, :)) + Wy = reshape(reshape(W, (:, d_y)) * yT, (d_z, d_x, :)) # @einsum Z[o,s] := Wy[o,i,s] * x[i,s] - Wyx = batched_mul(Wy, reshape(x, (d_x, 1, :))) + Wyx = batched_mul(Wy, reshape(xT, (d_x, 1, :))) Z = reshape(Wyx, (d_z, :)) # @einsum out[o,s] := σ(Z[o,i] + b[o]) NNlib.bias_act!(σ, Z, b) # σ.(Z .+ b) + end end (a::Bilinear)(x::AbstractVecOrMat) = a(x, x) @@ -774,8 +782,8 @@ Embedding((in, out)::Pair{<:Integer, <:Integer}; init = randn32) = Embedding(ini (m::Embedding)(x::AbstractVector) = NNlib.gather(m.weight, x) (m::Embedding)(x::AbstractArray) = reshape(m(vec(x)), :, size(x)...) -(m::Embedding)(x::AbstractVector{Bool}) = m.weight * x # usually OneHotVector -(m::Embedding)(x::AbstractMatrix{Bool}) = m.weight * x # usually OneHotMatrix +(m::Embedding)(x::AbstractVector{Bool}) = _autocast_down(autocast_eltype(), m.weight) * x # usually OneHotVector +(m::Embedding)(x::AbstractMatrix{Bool}) = _autocast_down(autocast_eltype(), m.weight) * x # usually OneHotMatrix (m::Embedding)(x::AbstractArray{Bool}) = reshape(m(reshape(x, size(x,1), :)), :, size(x)[2:end]...) function Base.show(io::IO, m::Embedding) diff --git a/src/layers/conv.jl b/src/layers/conv.jl index 2a3230f321..1756c330c3 100644 --- a/src/layers/conv.jl +++ b/src/layers/conv.jl @@ -197,8 +197,12 @@ ChainRulesCore.@non_differentiable conv_dims(::Any, ::Any) function (c::Conv)(x::AbstractArray) _conv_size_check(c, x) cdims = conv_dims(c, x) - xT = _match_eltype(c, x) - NNlib.bias_act!(c.σ, conv(xT, c.weight, cdims), conv_reshape_bias(c)) + _autocast_barrier() do T + W = _autocast_down(T, c.weight) + b = _autocast_down(T, conv_reshape_bias(c)) + xT = _autocast_down(T, _match_eltype(c, x)) + NNlib.bias_act!(c.σ, conv(xT, W, cdims), b) + end end _channels_in(l::Conv) = size(l.weight, ndims(l.weight)-1) * l.groups @@ -346,8 +350,12 @@ ChainRulesCore.@non_differentiable conv_transpose_dims(::Any, ::Any) function (c::ConvTranspose)(x::AbstractArray) _conv_size_check(c, x) cdims = conv_transpose_dims(c, x) - xT = _match_eltype(c, x) - NNlib.bias_act!(c.σ, ∇conv_data(xT, c.weight, cdims), conv_reshape_bias(c)) + _autocast_barrier() do T + W = _autocast_down(T, c.weight) + b = _autocast_down(T, conv_reshape_bias(c)) + xT = _autocast_down(T, _match_eltype(c, x)) + NNlib.bias_act!(c.σ, ∇conv_data(xT, W, cdims), b) + end end function Base.show(io::IO, l::ConvTranspose) @@ -484,8 +492,12 @@ ChainRulesCore.@non_differentiable crosscor_dims(::Any, ::Any) function (c::CrossCor)(x::AbstractArray) _conv_size_check(c, x) cdims = crosscor_dims(c, x) - xT = _match_eltype(c, x) - NNlib.bias_act!(c.σ, crosscor(xT, c.weight, cdims), conv_reshape_bias(c)) + _autocast_barrier() do T + W = _autocast_down(T, c.weight) + b = _autocast_down(T, conv_reshape_bias(c)) + xT = _autocast_down(T, _match_eltype(c, x)) + NNlib.bias_act!(c.σ, crosscor(xT, W, cdims), b) + end end function Base.show(io::IO, l::CrossCor) diff --git a/src/layers/normalise.jl b/src/layers/normalise.jl index 6958bf2466..a556273372 100644 --- a/src/layers/normalise.jl +++ b/src/layers/normalise.jl @@ -206,7 +206,10 @@ function (a::LayerNorm)(x::AbstractArray) _size_check(a, x, d => size(a.diag.scale, d)) end end - a.diag(NNlib.normalise(x; dims=1:length(a.size), eps=a.ϵ)) + # fresh name: `x` is captured by the `@ignore_derivatives` closure above, so + # assigning to it would box it, hurting inference and Zygote gradients + xF = _autocast_up(x) # normalization computes in Float32 under `autocast` + a.diag(NNlib.normalise(xF; dims=1:length(a.size), eps=a.ϵ)) end function Base.show(io::IO, l::LayerNorm) @@ -291,7 +294,8 @@ end function (BN::BatchNorm)(x::AbstractArray{T,N}) where {T,N} _size_check(BN, x, N-1 => BN.chs) - y = NNlib.batchnorm(BN.γ, BN.β, x, BN.μ, BN.σ², BN.momentum; + xF = _autocast_up(x) # normalization computes in Float32 under `autocast` + y = NNlib.batchnorm(BN.γ, BN.β, xF, BN.μ, BN.σ², BN.momentum; eps=BN.ϵ, training=_isactive(BN, x), track_stats=BN.track_stats) return BN.λ.(y) end @@ -380,7 +384,8 @@ end function (l::InstanceNorm)(x::AbstractArray{T,N}) where {T,N} _size_check(l, x, N-1 => l.chs) - y = NNlib.instancenorm(l.γ, l.β, x, l.μ, l.σ², l.momentum; + xF = _autocast_up(x) # normalization computes in Float32 under `autocast` + y = NNlib.instancenorm(l.γ, l.β, xF, l.μ, l.σ², l.momentum; eps=l.ϵ, training=_isactive(l, x), track_stats=l.track_stats) return l.λ.(y) end @@ -478,7 +483,8 @@ end function (gn::GroupNorm)(x::AbstractArray) _size_check(gn, x, ndims(x)-1 => gn.chs) - return gn.λ.(NNlib.groupnorm(gn.γ, gn.β, x, gn.G; eps=gn.ϵ)) + xF = _autocast_up(x) # normalization computes in Float32 under `autocast` + return gn.λ.(NNlib.groupnorm(gn.γ, gn.β, xF, gn.G; eps=gn.ϵ)) end testmode!(m::GroupNorm, mode = true) = @@ -503,12 +509,12 @@ See [`BatchNorm`](@ref), [`InstanceNorm`](@ref), [`GroupNorm`](@ref), and [`Laye """ hasaffine(l::Union{BatchNorm, InstanceNorm, LayerNorm, GroupNorm}) = l.affine -# Under `f16`/`bf16`, keep the statistics and affine parameters of `BatchNorm`, +# Under `f16mix`/`bf16mix`, keep the statistics and affine parameters of `BatchNorm`, # `InstanceNorm` and `GroupNorm` in `Float32` (mixed precision). `NNlib.batchnorm`, # `NNlib.instancenorm` and `NNlib.groupnorm` require `Float32` scale/bias/statistics # for half-precision feature maps (as does cuDNN), and half-precision running # statistics are numerically poor. `LayerNorm` is excluded: it wraps `NNlib.normalise`, -# which takes no scale/bias, so it has no such requirement. See `_paramtype`. +# which takes no scale/bias, so it has no such requirement. See `_paramtype_mixed`. _keep_f32_under_halfprec(::Union{BatchNorm, InstanceNorm, GroupNorm}) = true struct WeightNorm{L, G, D} diff --git a/src/layers/recurrent.jl b/src/layers/recurrent.jl index 48e8a98990..d0c0071c28 100644 --- a/src/layers/recurrent.jl +++ b/src/layers/recurrent.jl @@ -192,9 +192,15 @@ end function (m::RNNCell)(x::AbstractVecOrMat, h::AbstractVecOrMat) _size_check(m, x, 1 => size(m.Wi, 2)) - σ = NNlib.fast_act(m.σ, x) - h = σ.(m.Wi * x .+ m.Wh * h .+ m.bias) - return h, h + _autocast_barrier() do T + # NOTE: assigning to captured arguments (`x`, `h`) inside this closure would box them, + # breaking both Zygote gradients and inference — only fresh local names below. + Wi, Wh, b = _autocast_down(T, m.Wi), _autocast_down(T, m.Wh), _autocast_down(T, m.bias) + xT, hT = _autocast_down(T, x), _autocast_down(T, h) + σ = NNlib.fast_act(m.σ, xT) + hnew = σ.(Wi * xT .+ Wh * hT .+ b) + return hnew, hnew + end end function Base.show(io::IO, m::RNNCell) @@ -410,12 +416,16 @@ end function (m::LSTMCell)(x::AbstractVecOrMat, (h, c)) _size_check(m, x, 1 => size(m.Wi, 2)) - b = m.bias - g = m.Wi * x .+ m.Wh * h .+ b - input, forget, cell, output = chunk(g, 4; dims = 1) - c = @. sigmoid_fast(forget) * c + sigmoid_fast(input) * tanh_fast(cell) - h = @. sigmoid_fast(output) * tanh_fast(c) - return h, (h, c) + _autocast_barrier() do T + # only fresh local names in here: assigning to captured `x`/`h`/`c` would box them + Wi, Wh, b = _autocast_down(T, m.Wi), _autocast_down(T, m.Wh), _autocast_down(T, m.bias) + xT, hT, cT = _autocast_down(T, x), _autocast_down(T, h), _autocast_down(T, c) + g = Wi * xT .+ Wh * hT .+ b + input, forget, cell, output = chunk(g, 4; dims = 1) + cnew = @. sigmoid_fast(forget) * cT + sigmoid_fast(input) * tanh_fast(cell) + hnew = @. sigmoid_fast(output) * tanh_fast(cnew) + return hnew, (hnew, cnew) + end end Base.show(io::IO, m::LSTMCell) = @@ -613,18 +623,23 @@ end function (m::GRUCell)(x::AbstractVecOrMat, h) _size_check(m, x, 1 => size(m.Wi, 2)) - gxs = chunk(m.Wi * x, 3, dims = 1) - ghs = chunk(m.Wh * h, 3, dims = 1) - if m.b isa AbstractArray - bs = chunk(m.b, 3, dims = 1) + _autocast_barrier() do T + # only fresh local names in here: assigning to captured `x`/`h` would box them + Wi, Wh, b = _autocast_down(T, m.Wi), _autocast_down(T, m.Wh), _autocast_down(T, m.b) + xT, hT = _autocast_down(T, x), _autocast_down(T, h) + gxs = chunk(Wi * xT, 3, dims = 1) + ghs = chunk(Wh * hT, 3, dims = 1) + if b isa AbstractArray + bs = chunk(b, 3, dims = 1) else # b == false bs = [false, false, false] end r = @. sigmoid_fast(gxs[1] + ghs[1] + bs[1]) z = @. sigmoid_fast(gxs[2] + ghs[2] + bs[2]) h̃ = @. tanh_fast(gxs[3] + r * ghs[3] + bs[3]) - h = @. (1 - z) * h̃ + z * h - return h, h + hnew = @. (1 - z) * h̃ + z * hT + return hnew, hnew + end end Base.show(io::IO, m::GRUCell) = @@ -794,18 +809,24 @@ end function (m::GRUv3Cell)(x::AbstractVecOrMat, h) _size_check(m, x, 1 => size(m.Wi, 2)) - gxs = chunk(m.Wi * x, 3, dims = 1) - ghs = chunk(m.Wh * h, 3, dims = 1) - if m.b isa AbstractArray - bs = chunk(m.b, 3, dims = 1) - else # m.b == false + _autocast_barrier() do T + # only fresh local names in here: assigning to captured `x`/`h` would box them + Wi, Wh, b = _autocast_down(T, m.Wi), _autocast_down(T, m.Wh), _autocast_down(T, m.b) + Wh_h̃ = _autocast_down(T, m.Wh_h̃) + xT, hT = _autocast_down(T, x), _autocast_down(T, h) + gxs = chunk(Wi * xT, 3, dims = 1) + ghs = chunk(Wh * hT, 3, dims = 1) + if b isa AbstractArray + bs = chunk(b, 3, dims = 1) + else # b == false bs = [false, false, false] end r = @. sigmoid_fast(gxs[1] + ghs[1] + bs[1]) z = @. sigmoid_fast(gxs[2] + ghs[2] + bs[2]) - h̃ = tanh_fast.(gxs[3] .+ (m.Wh_h̃ * (r .* h)) .+ bs[3]) - h = @. (1 - z) * h̃ + z * h - return h, h + h̃ = tanh_fast.(gxs[3] .+ (Wh_h̃ * (r .* hT)) .+ bs[3]) + hnew = @. (1 - z) * h̃ + z * hT + return hnew, hnew + end end Base.show(io::IO, m::GRUv3Cell) = diff --git a/src/losses/Losses.jl b/src/losses/Losses.jl index ec5f7ae360..82edbb9f4d 100644 --- a/src/losses/Losses.jl +++ b/src/losses/Losses.jl @@ -7,6 +7,7 @@ using ChainRulesCore # using ..Flux: ofeltype, epseltype ofeltype(x, y) = convert(float(eltype(x)), y) epseltype(x) = eps(float(eltype(x))) +using ..Flux: _autocast_up using NNlib: logsoftmax, logσ, ctc_loss, ctc_alpha, ∇ctc_loss import Base.Broadcast: broadcasted diff --git a/src/losses/functions.jl b/src/losses/functions.jl index 67b9ec51ce..dcc496e608 100644 --- a/src/losses/functions.jl +++ b/src/losses/functions.jl @@ -20,6 +20,7 @@ julia> Flux.mae(y_model, 1:3) """ function mae(ŷ, y; agg = mean) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) agg(abs.(ŷ .- y)) end @@ -44,6 +45,7 @@ julia> Flux.mse(y_model, y_true) """ function mse(ŷ, y; agg = mean) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) agg(abs2.(ŷ .- y)) end @@ -68,6 +70,7 @@ julia> Flux.msle(Float32[0.9, 1.8, 2.7], 1:3) """ function msle(ŷ, y; agg = mean, eps::Real = epseltype(ŷ)) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) agg((log.((ŷ .+ eps) ./ (y .+ eps))) .^2 ) end @@ -103,6 +106,7 @@ julia> Flux.huber_loss(ŷ, 1:3, delta=0.05) # changes behaviour as |ŷ - y| > function huber_loss(ŷ, y; agg = mean, delta::Real = 1) δ = ofeltype(ŷ, delta) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) abs_error = abs.(ŷ .- y) agg(_huber_metric.(abs_error, δ)) @@ -230,6 +234,7 @@ julia> Flux.crossentropy(y_model, y_smooth) """ function crossentropy(ŷ, y; dims = 1, agg = mean, eps::Real = epseltype(ŷ)) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) agg(.-sum(xlogy.(y, ŷ .+ eps); dims = dims)) end @@ -269,6 +274,7 @@ julia> Flux.crossentropy(softmax(y_model), y_label) """ function logitcrossentropy(ŷ, y; dims = 1, agg = mean) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) agg(.-sum(y .* logsoftmax(ŷ; dims = dims); dims = dims)) end @@ -318,6 +324,7 @@ julia> Flux.crossentropy(y_prob, y_hot) """ function binarycrossentropy(ŷ, y; agg = mean, eps::Real = epseltype(ŷ)) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) agg(@.(-xlogy(y, ŷ + eps) - xlogy(1 - y, 1 - ŷ + eps))) end @@ -348,6 +355,7 @@ julia> Flux.binarycrossentropy(sigmoid.(y_model), y_bin) """ function logitbinarycrossentropy(ŷ, y; agg = mean) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) agg(@.((1 - y) * ŷ - logσ(ŷ))) end @@ -388,6 +396,7 @@ Inf """ function kldivergence(ŷ, y; dims = 1, agg = mean, eps::Real = epseltype(ŷ)) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) entropy = agg(sum(xlogx.(y); dims = dims)) cross_entropy = crossentropy(ŷ, y; dims, agg, eps) return entropy + cross_entropy @@ -413,6 +422,7 @@ julia> Flux.poisson_loss(y_model, 1:3) """ function poisson_loss(ŷ, y; agg = mean) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) agg(ŷ .- xlogy.(y, ŷ)) end @@ -448,6 +458,7 @@ true """ function hinge_loss(ŷ, y; agg = mean) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) agg(max.(0, 1 .- ŷ .* y)) end @@ -483,6 +494,7 @@ true """ function squared_hinge_loss(ŷ, y; agg = mean) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) agg((max.(0, 1 .- ŷ .* y)) .^ 2) end @@ -510,6 +522,7 @@ julia> 1 - Flux.dice_coeff_loss(y_pred, 1:3) # ~ F1 score for image segmentatio function dice_coeff_loss(ŷ, y; smooth = 1) s = ofeltype(ŷ, smooth) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) # TODO add agg 1 - (2 * sum(y .* ŷ) + s) / (sum(y .^ 2) + sum(ŷ .^ 2) + s) end @@ -528,6 +541,7 @@ Calculated as: function tversky_loss(ŷ, y; beta::Real = 0.7, β = nothing) β = ofeltype(ŷ, beta) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) #TODO add agg num = sum(y .* ŷ) + 1 den = sum(y .* ŷ + β * (1 .- y) .* ŷ + (1 - β) * y .* (1 .- ŷ)) + 1 @@ -565,6 +579,7 @@ true function binary_focal_loss(ŷ, y; agg=mean, gamma=2, eps::Real=epseltype(ŷ)) γ = gamma isa Integer ? gamma : ofeltype(ŷ, gamma) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) ŷϵ = ŷ .+ eps p_t = y .* ŷϵ + (1 .- y) .* (1 .- ŷϵ) ce = .-log.(p_t) @@ -610,6 +625,7 @@ See also: [`Losses.binary_focal_loss`](@ref) for binary (not one-hot) labels function focal_loss(ŷ, y; dims=1, agg=mean, gamma=2, eps::Real=epseltype(ŷ), ϵ=nothing, γ=nothing) γ = gamma isa Integer ? gamma : ofeltype(ŷ, gamma) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) ŷϵ = ŷ .+ eps agg(sum(@. -y * (1 - ŷϵ)^γ * log(ŷϵ); dims)) end @@ -637,6 +653,7 @@ julia> Flux.siamese_contrastive_loss(ŷ, 1:3, margin = 2) """ function siamese_contrastive_loss(ŷ, y; agg = mean, margin::Real = 1) _check_sizes(ŷ, y) + ŷ, y = _autocast_up(ŷ), _autocast_up(y) margin < 0 && throw(DomainError(margin, "Margin must be non-negative")) return agg(@. (1 - y) * ŷ^2 + y * max(0, margin - ŷ)^2) end diff --git a/src/train.jl b/src/train.jl index a20185ddfb..7e25e5ad5e 100644 --- a/src/train.jl +++ b/src/train.jl @@ -104,9 +104,13 @@ It adds only a few features to the loop above: * Show a progress bar using [`@withprogress`](https://github.com/JuliaLogging/ProgressLogging.jl). * Manage memory. Runs an incremental garbage collection adaptively. + +The keyword `autocast` accepts `Float16` or `BFloat16` to compute each gradient +under mixed precision, see [`autocast`](@ref). """ function train!(loss, adtype::AbstractADType, model, data, opt; cb = nothing, - caching_allocator::Bool = false, gc_interval::Union{Integer, Symbol} = :auto) + caching_allocator::Bool = false, gc_interval::Union{Integer, Symbol} = :auto, + autocast::Union{Nothing, Type} = nothing) isnothing(cb) || error("""train! does not support callback functions. For more control use a loop with `gradient` and `update!`.""") gc_interval isa Symbol && gc_interval !== :auto && @@ -133,11 +137,11 @@ function train!(loss, adtype::AbstractADType, model, data, opt; cb = nothing, # blow past GPU memory. From the second step on the algorithm is fixed, so the cache # only ever sees the real, reusable training buffers. if cache === nothing || i == 1 - opt, model = _train_step!(loss, adtype, model, opt, d_splat, i) + opt, model = _train_step!(loss, adtype, model, opt, d_splat, i; autocast) else # Reuse the memory allocated during the previous step, see issue #2523. GPUArrays.@cached cache begin - opt, model = _train_step!(loss, adtype, model, opt, d_splat, i) + opt, model = _train_step!(loss, adtype, model, opt, d_splat, i; autocast) end end @@ -190,8 +194,8 @@ end # A single training step, factored out so that `train!` can run it with or without the # caching allocator without duplicating the body. -function _train_step!(loss, adtype, model, opt, d_splat, i) - l, gs = Flux.withgradient(m -> loss(m, d_splat...), adtype, model) +function _train_step!(loss, adtype, model, opt, d_splat, i; autocast = nothing) + l, gs = Flux.withgradient(m -> loss(m, d_splat...), adtype, model; autocast) if !isfinite(l) throw(DomainError(lazy"Loss is $l on data item $i, stopping training")) @@ -208,12 +212,14 @@ function _update!(opt_state, model::Duplicated, grad) end -train!(loss, model, data, opt; cb = nothing, caching_allocator::Bool = false, gc_interval::Union{Integer, Symbol} = :auto) = - train!(loss, AutoZygote(), model, data, opt; cb, caching_allocator, gc_interval) +train!(loss, model, data, opt; cb = nothing, caching_allocator::Bool = false, gc_interval::Union{Integer, Symbol} = :auto, + autocast::Union{Nothing, Type} = nothing) = + train!(loss, AutoZygote(), model, data, opt; cb, caching_allocator, gc_interval, autocast) # This method let you use Optimisers.Descent() without setup, when there is no state -function train!(loss, model, data, rule::Optimisers.AbstractRule; cb = nothing, caching_allocator::Bool = false, gc_interval::Union{Integer, Symbol} = :auto) - return train!(loss, model, data, _rule_to_state(model, rule); cb, caching_allocator, gc_interval) +function train!(loss, model, data, rule::Optimisers.AbstractRule; cb = nothing, caching_allocator::Bool = false, gc_interval::Union{Integer, Symbol} = :auto, + autocast::Union{Nothing, Type} = nothing) + return train!(loss, model, data, _rule_to_state(model, rule); cb, caching_allocator, gc_interval, autocast) end function _rule_to_state(model, rule::Optimisers.AbstractRule) @@ -228,12 +234,14 @@ function _rule_to_state(model, rule::Optimisers.AbstractRule) return state end -train!(loss, model::Duplicated, data, opt; cb = nothing, caching_allocator::Bool = false, gc_interval::Union{Integer, Symbol} = :auto) = - train!(loss, AutoEnzyme(), model, data, opt; cb, caching_allocator, gc_interval) +train!(loss, model::Duplicated, data, opt; cb = nothing, caching_allocator::Bool = false, gc_interval::Union{Integer, Symbol} = :auto, + autocast::Union{Nothing, Type} = nothing) = + train!(loss, AutoEnzyme(), model, data, opt; cb, caching_allocator, gc_interval, autocast) # This method let you use Optimisers.Descent() without setup, when there is no state -function train!(loss, model::Duplicated, data, rule::Optimisers.AbstractRule; cb=nothing, caching_allocator::Bool = false, gc_interval::Union{Integer, Symbol} = :auto) - return train!(loss, model, data, _rule_to_state(model, rule); cb, caching_allocator, gc_interval) +function train!(loss, model::Duplicated, data, rule::Optimisers.AbstractRule; cb=nothing, caching_allocator::Bool = false, gc_interval::Union{Integer, Symbol} = :auto, + autocast::Union{Nothing, Type} = nothing) + return train!(loss, model, data, _rule_to_state(model, rule); cb, caching_allocator, gc_interval, autocast) end end # module Train diff --git a/test/autocast.jl b/test/autocast.jl new file mode 100644 index 0000000000..c5882e1e36 --- /dev/null +++ b/test/autocast.jl @@ -0,0 +1,153 @@ +# This testset MUST run before any `autocast` call in this file: it checks that until +# the first use flips `Flux.autocast_active()`, the machinery is compiled out entirely +# and layers infer their exact concrete return types. +@testset "autocast is compiled out before first use" begin + @test Flux.autocast_active() == false + @test @inferred(Dense(3 => 4, relu)(randn(Float32, 3, 8))) isa Matrix{Float32} + @test @inferred(Conv((3,), 2 => 4, relu)(randn(Float32, 10, 2, 5))) isa Array{Float32, 3} + @test @inferred(Chain(Dense(3 => 4, relu), Dense(4 => 2))(randn(Float32, 3, 8))) isa Matrix{Float32} +end + +@testset "autocast eltype flow ($T)" for T in (Float16, BFloat16) + x2 = randn(Float32, 3, 8) # for Dense-like layers + x4 = randn(Float32, 8, 8, 2, 3) # for conv layers + xseq = randn(Float32, 3, 7, 2) # for recurrent layers + + @testset "cast-down layers" begin + for (l, x) in ( + (Dense(3 => 4, relu), x2), + (Flux.Bilinear((3, 3) => 4), x2), + (Conv((3, 3), 2 => 4, relu), x4), + (ConvTranspose((3, 3), 2 => 4), x4), + (CrossCor((3, 3), 2 => 4), x4), + (RNN(3 => 5), xseq), + (LSTM(3 => 5), xseq), + (GRU(3 => 5), xseq), + (GRUv3(3 => 5), xseq), + ) + y = autocast(() -> l(x), T) + @test eltype(y) == T + # parameters are untouched + @test all(p -> eltype(p) == Float32, Flux.trainables(l)) + end + + mha = MultiHeadAttention(16) + xmha = randn(Float32, 16, 5, 2) + y, α = autocast(() -> mha(xmha), T) + @test eltype(y) == T + + e = Embedding(5 => 4) + @test eltype(autocast(() -> e(Flux.onehotbatch([1, 3], 1:5)), T)) == T + @test eltype(autocast(() -> e([1, 3]), T)) == Float32 # gather path stays Float32 + end + + @testset "normalization computes in Float32" begin + for (l, x) in ( + (BatchNorm(3), x2), + (LayerNorm(3), x2), + (InstanceNorm(2; affine=true), x4), + (GroupNorm(2, 2), x4), + ) + y = autocast(() -> l(x), T) + @test eltype(y) == Float32 + end + # half-precision input to a norm layer is upcast inside the scope + xh = T == Float16 ? f16(x2) : bf16(x2) + @test eltype(autocast(() -> LayerNorm(3)(xh), T)) == Float32 + end + + @testset "losses upcast to Float32" begin + half = T == Float16 ? f16 : bf16 + ŷ, y = half(rand(Float32, 4, 8)), half(rand(Float32, 4, 8)) + for loss in (Flux.mse, Flux.mae, Flux.crossentropy, Flux.logitcrossentropy, + Flux.huber_loss) + @test autocast(() -> loss(ŷ, y), T) isa Float32 + @test loss(ŷ, y) isa T # unchanged outside the scope + end + end + + @testset "gradients are Float32 and close to the fp32 reference" begin + model = Chain(Dense(3 => 4, relu), BatchNorm(4), Dense(4 => 2)) + ytarget = randn(Float32, 2, 8) + loss(m) = Flux.mse(m(x2), ytarget) + + val, grad = Flux.withgradient(loss, model; autocast=T) + @test val isa Float32 + gflat = filter(g -> g isa AbstractArray, Functors.fleaves(grad[1])) + @test !isempty(gflat) + @test all(g -> eltype(g) == Float32, gflat) + @test all(g -> all(isfinite, g), gflat) + + val32, grad32 = Flux.withgradient(loss, model) + rtol = T == Float16 ? 0.03 : 0.15 + @test val ≈ val32 rtol=rtol + @test grad[1].layers[1].weight ≈ grad32[1].layers[1].weight rtol=rtol atol=0.05 + + # raw reductions (not Flux losses) stay in half precision, like PyTorch + vraw, _ = Flux.withgradient(m -> sum(abs2, m(x2)), model; autocast=T) + @test vraw isa T + end + + @testset "do-block and keyword forms agree" begin + model = Chain(Dense(3 => 4, tanh), Dense(4 => 2)) + loss(m) = Flux.mse(m(x2), zeros(Float32, 2, 8)) + g1 = autocast(() -> Flux.gradient(loss, model), T) + g2 = Flux.gradient(loss, model; autocast=T) + @test g1[1].layers[1].weight == g2[1].layers[1].weight + wg = Flux.withgradient(loss, AutoZygote(), model; autocast=T) + @test wg.grad[1].layers[1].weight == g2[1].layers[1].weight + end + + @testset "train! with autocast" begin + model = Chain(Dense(3 => 4, relu), Dense(4 => 2)) + w0 = copy(model[1].weight) + opt = Flux.setup(Adam(1e-3), model) + data = [(randn(Float32, 3, 8), randn(Float32, 2, 8)) for _ in 1:3] + Flux.train!((m, x, y) -> Flux.mse(m(x), y), model, data, opt; autocast=T) + @test eltype(model[1].weight) == Float32 + @test model[1].weight != w0 + end +end + +@testset "autocast is a no-op outside the scope" begin + model = Chain(Dense(3 => 4, relu), BatchNorm(4), Dense(4 => 2)) + x = randn(Float32, 3, 8) + y0 = model(x) + autocast(() -> model(x), Float16) # entering and leaving a scope changes nothing + @test model(x) == y0 + @test Flux.mse(y0, zero(y0)) isa Float32 +end + +@testset "forward pass infers as a small union after first use" begin + # The earlier testsets flipped `autocast_active()`. From then on the inferred return + # is at worst the 3-type union over the Float32/Float16/BFloat16 paths (which the + # compiler union-splits) — not `Any` — and the union must not widen through a Chain. + @test Flux.autocast_active() == true + MatUnion3 = Union{Matrix{Float32}, Matrix{Float16}, Matrix{BFloat16}} + @test Base.promote_op(Dense(3 => 4, relu), Matrix{Float32}) <: MatUnion3 + @test Base.promote_op(Conv((3,), 2 => 4, relu), Array{Float32, 3}) <: + Union{Array{Float32, 3}, Array{Float16, 3}, Array{BFloat16, 3}} + @test Base.promote_op(Chain(Dense(3 => 4, relu), Dense(4 => 2)), Matrix{Float32}) <: MatUnion3 +end + +@testset "autocast argument checking" begin + @test_throws ArgumentError autocast(() -> 1, Float32) + @test_throws ArgumentError autocast(() -> 1, Float64) + @test_throws ArgumentError autocast(() -> 1, Int) +end + +@testset "outputsize under autocast" begin + model = Chain(Dense(3 => 4), Conv((3, 3), 1 => 2)) + m2 = Chain(Dense(3 => 7), Dense(7 => 2)) + @test autocast(() -> outputsize(m2, (3, 5)), Float16) == (2, 5) +end + +@testset "autocast with Mooncake" begin + model = Chain(Dense(3 => 4, tanh), Dense(4 => 2)) + x = randn(Float32, 3, 8) + loss(m) = Flux.mse(m(x), zeros(Float32, 2, 8)) + g16 = Flux.gradient(loss, AutoMooncake(config=nothing), model; autocast=Float16) + g32 = Flux.gradient(loss, model) + @test eltype(g16[1].layers[1].weight) == Float32 + @test g16[1].layers[1].weight ≈ g32[1].layers[1].weight rtol=0.03 atol=0.05 +end diff --git a/test/ext_cuda/autocast.jl b/test/ext_cuda/autocast.jl new file mode 100644 index 0000000000..679ae9b5a6 --- /dev/null +++ b/test/ext_cuda/autocast.jl @@ -0,0 +1,46 @@ +@testset "autocast on GPU ($T)" for T in (Float16, BFloat16) + x2 = CUDA.randn(Float32, 3, 8) + x4 = CUDA.randn(Float32, 8, 8, 2, 3) + + @testset "eltype flow" begin + model = Chain(Dense(3 => 4, relu), BatchNorm(4), Dense(4 => 2)) |> gpu + y = autocast(() -> model(x2), T) + @test y isa CuArray{T} + + c = Chain(Conv((3, 3), 2 => 4, relu), BatchNorm(4), MaxPool((2, 2))) |> gpu + yc = autocast(() -> c(x4), T) + @test eltype(yc) == T + end + + @testset "training step" begin + model = Chain(Dense(3 => 4, relu), BatchNorm(4), Dense(4 => 2)) |> gpu + ytarget = CUDA.randn(Float32, 2, 8) + loss(m) = Flux.mse(m(x2), ytarget) + + val, grad = Flux.withgradient(loss, model; autocast=T) + @test val isa Float32 + gW = grad[1].layers[1].weight + @test gW isa CuArray{Float32} + @test all(isfinite, Array(gW)) + + # close to the full-precision gradient + val32, grad32 = Flux.withgradient(loss, model) + rtol = T == Float16 ? 0.03 : 0.15 + @test val ≈ val32 rtol=rtol + @test Array(gW) ≈ Array(grad32[1].layers[1].weight) rtol=rtol atol=0.05 + + # parameters stay Float32 through a train! step + opt = Flux.setup(Adam(1e-3), model) + data = [(CUDA.randn(Float32, 3, 8), CUDA.randn(Float32, 2, 8)) for _ in 1:3] + Flux.train!((m, x, y) -> Flux.mse(m(x), y), model, data, opt; autocast=T) + @test model[1].weight isa CuArray{Float32} + end + + @testset "conv gradient" begin + c = Conv((3, 3), 2 => 4) |> gpu + g = Flux.gradient(m -> sum(abs2, m(x4)), c; autocast=T)[1] + @test g.weight isa CuArray{Float32} + g32 = Flux.gradient(m -> sum(abs2, m(x4)), c)[1] + @test Array(g.weight) ≈ Array(g32.weight) rtol=0.15 atol=0.1 + end +end diff --git a/test/ext_cuda/layers.jl b/test/ext_cuda/layers.jl index 21230a6528..d903792c45 100644 --- a/test/ext_cuda/layers.jl +++ b/test/ext_cuda/layers.jl @@ -305,14 +305,15 @@ end @test eltype(gm2(gx)) == BFloat16 @test Float32.(gm2(gx)) ≈ f32(gm2)(f32(gx)) rtol=0.1 - # BatchNorm, InstanceNorm and GroupNorm are converted in mixed precision: statistics - # and affine parameters stay Float32 while the data flows in bf16, so they dispatch to - # NNlib's (cuDNN) half-precision kernels. + # Under `bf16mix`, BatchNorm, InstanceNorm and GroupNorm are converted in mixed + # precision: statistics and affine parameters stay Float32 while the data flows in + # bf16, so they dispatch to NNlib's (cuDNN) half-precision kernels. (A plain `bf16` + # full cast of these layers is rejected by those kernels.) gx4 = gpu(bf16(randn(Float32, 4, 4, 3, 2))) @testset "$(nameof(typeof(l)))" for l in (BatchNorm(3), InstanceNorm(3; affine=true, track_stats=true), GroupNorm(3, 3)) - gm = bf16(l) |> gpu + gm = bf16mix(l) |> gpu @test eltype(gm.γ) == eltype(gm.β) == Float32 # affine params kept in Float32 y = gm(gx4) @test eltype(y) == BFloat16 # data flow stays bf16 diff --git a/test/ext_enzyme/enzyme.jl b/test/ext_enzyme/enzyme.jl index 36c75d9efc..a7652e80ff 100644 --- a/test/ext_enzyme/enzyme.jl +++ b/test/ext_enzyme/enzyme.jl @@ -77,3 +77,7 @@ end @test Flux.withgradient(sum ∘ LayerNorm(3), z).grad[1] ≈ [0.0, 0.0, 0.0] @test Flux.withgradient(|>, z, _duplicated(sum ∘ LayerNorm(3))).grad[1] ≈ [0.0, 0.0, 0.0] end + +# NOTE: autocast + AutoEnzyme is not tested yet: after the first `autocast` call the +# layer IR contains BFloat16 values, and Enzyme's type analysis crashes on those until +# EnzymeAD/Enzyme.jl#3430 is fixed upstream. diff --git a/test/utils.jl b/test/utils.jl index fcc41297b2..da79da9e84 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -311,10 +311,11 @@ end @test bf16(Float32[1, 1.5, -2.25, 96]) == BFloat16[1, 1.5, -2.25, 96] # exact round-to-nearest-even @test gradient(x -> sum(bf16(x)), x32)[1] isa Vector{Float32} - @testset "mixed-precision normalization ($cast)" for cast in (f16, bf16) - Thalf = cast === f16 ? Float16 : BFloat16 - # BatchNorm/InstanceNorm/GroupNorm keep statistics and affine parameters in Float32 - # under half precision; LayerNorm is converted fully. + @testset "mixed-precision normalization ($cast)" for cast in (f16mix, bf16mix) + Thalf = cast === f16mix ? Float16 : BFloat16 + full = cast === f16mix ? f16 : bf16 + # Under `f16mix`/`bf16mix`, BatchNorm/InstanceNorm/GroupNorm keep statistics and + # affine parameters in Float32; LayerNorm and everything else is converted fully. for layer in (BatchNorm(3), InstanceNorm(3; affine=true, track_stats=true), GroupNorm(4, 2)) @@ -322,17 +323,22 @@ end params = filter(x -> x isa AbstractArray, collect(Functors.fleaves(hl))) @test !isempty(params) @test all(p -> eltype(p) == Float32, params) + # plain `f16`/`bf16` is a full cast, like PyTorch's `model.half()` + @test all(p -> eltype(p) == Thalf, + filter(x -> x isa AbstractArray, collect(Functors.fleaves(full(layer))))) end @test all(p -> eltype(p) == Thalf, filter(x -> x isa AbstractArray, collect(Functors.fleaves(cast(LayerNorm(3)))))) + # non-norm layers convert fully under the mixed cast too + @test eltype(cast(Dense(3 => 4)).weight) == Thalf # Forward is only exercised for Float16: the generic bf16 CPU normalization path can # hang LLVM codegen (JuliaMath/BFloat16s.jl#107), so bf16 norms are tested on the GPU. - if cast === f16 + if cast === f16mix xh = f16(randn(Float32, 4, 4, 3, 2)) - @test eltype(f16(BatchNorm(3))(xh)) == Float16 # mixed: Float32 params, f16 data - @test eltype(f16(InstanceNorm(3; affine=true))(xh)) == Float16 - @test eltype(f16(GroupNorm(4, 2))(f16(randn(Float32, 4, 4, 4, 2)))) == Float16 + @test eltype(f16mix(BatchNorm(3))(xh)) == Float16 # mixed: Float32 params, f16 data + @test eltype(f16mix(InstanceNorm(3; affine=true))(xh)) == Float16 + @test eltype(f16mix(GroupNorm(4, 2))(f16(randn(Float32, 4, 4, 4, 2)))) == Float16 end end