Support for Float32

[seabbs]

I see in various places issues using Float32 (i.e. https://discourse.julialang.org/t/issue-with-float32-precision-in-turing-model-sampling/108160) with the problem being that areas of the code hard code Float64.

I've looked but can't see an issue for this. Is there one? If not what work would need to be done to make this possible and could it be itemised to make it easier to chunk out into small parts?

In principle this seems like something that could be tackled by newer contributors and so making that easy seems like it would be worth the effort? That being said I have no idea and it could be far more complicated.

[seabbs]

Update: I just did a quick scan across the ecosystem in the TuringLang org and I see a couple of issues/PRs attacking this (TuringLang/Bijectors.jl#266) so maybe I am just looking in the wrong place?

[torfjelde]

I don't think there is a unifying issue unfortunately, but we can make this the one:)

It is definitively possible to achieve. For example, it's already possible to evaluate a model using, say, Float32 if one is willing to do some work.
But it requires quite some effort to make consistent across code-bases to the point where it's easy to just run the full inference procedure using Float32.

There are a few aspects to this:

• The user needs to write evertyhing in a way that's compatible with whatever type we want. Specifically, that means something like @model demo(::Type{T}=Float64) where T and make sure that all the numbers inside of the model are of type T, e.g. call zeros(T, n) and so on.
• We need to enable ways of specifying that Turing.jl should use some time T everywhere internally, e.g. through either the sampler or the sample method itself.
• We need to replace all constructions of VarInfo (used internally in Turing.jl to keep track of realizations) with a corresponding one that takes the desired type into account.

[penelopeysm]

This actually works to some extent now, 'out of the box'. You have to write the model with f32 things, e.g. the Normal(0.0f0, 1.0f0) below. Unfortunately, there are still some operations which return Float64 (notably, logjoint(model, params) and LogDensityProblems.logdensity(ldf, params)).

julia> using DynamicPPL, LogDensityProblems, ForwardDiff, Distributions

julia> @model function f()
           x ~ Normal(0.0f0, 1.0f0)
       end
f (generic function with 2 methods)

julia> typeof(rand(f())[@varname(x)])
Float32

julia> typeof(logjoint(f(), (; x=0.0f0)))
Float64

julia> ldf = LogDensityFunction(f(), getlogjoint_internal, LinkAll(); adtype=AutoForwardDiff());

julia> typeof(rand(ldf))
Vector{Float32} (alias for Array{Float32, 1})

julia> typeof(LogDensityProblems.logdensity(ldf, [0.0f0]))
Float64

julia> typeof(LogDensityProblems.logdensity_and_gradient(ldf, [0.0f0]))
Tuple{Float32, Vector{Float32}}

Bad news:

The logjoint one is harder to fix. The issue with logjoint is that there is no clean way of knowing that the user wants a logpdf that has Float32. You could argue that we can infer it from the type passed into the NamedTuple (; x = 0.0f0). The catch is that when you have a trivial NamedTuple like that, it's very easy to infer, but once you start mixing and matching different values with different types, it's a nightmare to reverse-engineer. There used to be this function infer_nested_eltype, which attempted to use some heuristics to do this, but I don't want to bring it back because it just wasn't the correct solution.

To correctly fix that, there would need to be some flag that the user can set. Right now in DynamicPPL we declare const LogProbType = float(Real), which usually ends up being Float64 (https://github.com/TuringLang/DynamicPPL.jl/blob/a31e3e396c1640b3344cf86041b2a7d9a52a74cf/src/utils.jl#L5-L10). An obvious initial idea would be to just change that to const LogProbType = Ref(float(Real)), which then lets people control it by setting DynamicPPL.LogProbType[] = Float32. However, I'm not sure what if any implications this might have on performance-sensitive code. It is definitely worth a shot, though.

Good news:

The LogDensityProblems.logdensity one is I suspect a lot easier to fix, though, because the necessary information is right there in the eltype of the vector argument. In fact, we already have an issue that's essentially the same thing: TuringLang/DynamicPPL.jl#906.

The thing is that, many bits of Turing are already using only LogDensityFunction. Indeed, I imagine that if we can make this work for HMC/NUTS, that would satisfy a large number of use cases. So to get this to work correctly, we would need to:

1. Fix the DynamicPPL issue (IMO, very easy, Claude could do it, but it's so easy that it'd probably be faster for me to fix it than to explain it to Claude).
2. Check whether it works 'out of the box'. Note that, since rand(ldf) already returns Vector{Float32}, I have a suspicion that it will already work out of the box. The point here is that because the distribution has Float32 params, if you forward-sample the model, you will get Float32 params. Once you have a vector of Float32, it should just keep on using Float32 forever.

So I think that it wouldn't even be necessary to introduce an extra argument to something like sample(). If you write the model to use F32, it should ideally just use F32 throughout.

3. If my hunch about it working out of the box is wrong, then fix it.

Good news 2:

The earlier bad news was wrong.

Bad news 2:

Making AdvancedHMC work with float32 is surprisingly painful, and the quality of that codebase is questionable.

[penelopeysm]

One potential API for the logjoint thing could be to add an extra type param to InitContext which specifies the log-prob type to use. Then one could for example, hook into this by wrapping the NamedTuple of params in a struct: logjoint(model, WithLogp{Float32}(params)) -> Float32

This actually isn't needed, I found a full solution to this.

[penelopeysm]

The final "blocker" to this, is that certain things in Distributions.jl are also hardcoded to return Float64:

julia> x = DiscreteUniform(1, 2) # Ints only...
DiscreteUniform(a=1, b=2)

julia> typeof(logpdf(x, 1))
Float64

see https://github.com/JuliaStats/Distributions.jl/blob/0421b1890c49417de993b255f6b177a1cb93fca3/src/univariate/discrete/discreteuniform.jl#L24-L27.

This is obviously not something I can control (although I think it should actually be easy enough to make an upstream PR to add a type parameter to DiscreteUniform). In the absence of that, the workaround for this is to define and use a custom version of DiscreteUniform that uses Float32 (or whatever type the user wants).

[penelopeysm]

Now having carefully considered all of this, I think implementing this might actually be a fairly quick win!

[penelopeysm]

@nsiccha Do you have a NUTS implementation that doesn't break down when handling Float32? I got as far as fixing literally everything in DPPL and Turing, but I'm running into so many AdvancedHMC issues 😭

[penelopeysm]

With TuringLang/DynamicPPL.jl#1328 + #2793 + tpapp/DynamicHMC.jl#199 + a trivial compat bump to FlexiChains to account for the DPPL/Turing PRs, one can do this:

julia> using Turing, FlexiChains, DynamicHMC

julia> @model function f()
           x ~ Normal(0.0f0, 1.0f0)
       end
f (generic function with 2 methods)

julia> chn = sample(f(), externalsampler(DynamicHMC.NUTS()), 100; chain_type=VNChain)
Sampling 100%|████████████████████████████████████████████| Time: 0:00:02
FlexiChain (100 iterations, 1 chain)
↓ iter=1:100 | → chain=1:1

Parameter type   VarName
Parameters       x
Extra keys       :logprior, :loglikelihood, :logjoint


julia> eltype(chn[@varname(x)])
Float32

julia> eltype(chn[:logjoint])
Float32

So, that's good. Unfortunately, this messes with type stability for some DynamicPPL models, with concomitant losses in performance, which is less good.

julia> using DynamicPPL, Distributions

julia> @model function f()
           x = zeros(2)
           x .~ Normal()
       end
f (generic function with 2 methods)

julia> @inferred logjoint(f(), (; x = zeros(2)))
ERROR: return type Float64 does not match inferred return type Union{DynamicPPL.NoLogProb, Float64}
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] top-level scope
   @ REPL[4]:1

So, it might well be that we do actually need to force the user to explicitly specify the type they want...

[seabbs]

Nice!! This all looks very promising. Note for awareness on top of the DynamicHMC PR @SamuelBrand1 (from memory, he also explored AdvancedHMC and found the changes needed would be non-trivial) has also been rolling out a F32 wrapper for Distributions.jl (I am not sure if he has explored submitting PRs to fix i.e DiscreteUniform (https://github.com/SamuelBrand1/DistributionsF32).

[SamuelBrand1]

One issue to keep in mind is that if the goal of using f32 float types is to either:

• lift various compute operations to other devices e.g. GPU
• avoid various type system gotcha in Julia (like BigInt suddenly rearing its ugly head)
• Get rid of check_args

then the whole operation has to be f32 and you need to avoid accidentally promoting inside a function call and having a f32 conversion at the end.

I did a rewrite of Normal here https://github.com/SamuelBrand1/DistributionsF32/blob/main/src/univariate/continuous/normal32.jl which is basically a copy and paste to inject this constant https://github.com/SamuelBrand1/DistributionsF32/blob/377c41713de46246010e2fbd0b2d337fcaf0f283/src/DistributionsF32.jl#L18

Which may well have been overkill!

Lots of the Distributions can probably be done minimally e.g. https://github.com/SamuelBrand1/DistributionsF32/blob/main/src/univariate/continuous/gamma32.jl

[SamuelBrand1]

So, that's good. Unfortunately, this messes with type stability for some DynamicPPL models, with concomitant losses in performance, which is less good.

Thats the tricky bit. But thanks for looking at this!

[penelopeysm]

From the perspective of Turing / DynamicPPL, I think my standpoint is that if the user writes the model in a f32-compliant way and feed in f32 inputs, nothing inside Turing should automatically promote it. That ties in very much with

then the whole operation has to be f32 and you need to avoid accidentally promoting inside a function call and having a f32 conversion at the end.

Just from a software design perspective too, I'm not keen to make the user specify something like rand(model; Tlogp=Float32) and then have every intermediate calculation be down-promoted to f32, I think that's an antipattern in terms of software design and that lines up with the fact that it'd be pointless to do so if the target is e.g. GPUs.

The DynamicPPL PR I linked above (TuringLang/DynamicPPL.jl#1328) essentially obeys this -- you can see in the top level comment that I'm not doing anything special, it 'just works'. For the most part, I am very keen to merge it, but I don't want to compromise performance, and if I can find a way around the perf loss I will probably merge it.

This is of course a distinct consideration from whether one should have tools to make it easier to write f32-compliant models. For example, distributions that default to f32 would definitely make it easier. But that's, in my view, not something that DynamicPPL has to take care of, or be responsible for. DPPL's job is to be as unintrusive as possible and not override whatever the user wants.