Autocast (scope-based prototype) — mixed precision

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)).

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"

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",

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.

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.

docs/src/reference/utilities.md

@@ -68,4 +68,6 @@ Flux.f64
 Flux.f32
 Flux.f16
 Flux.bf16
+Flux.f16mix
+Flux.bf16mix
 ```

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

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

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.

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