Getting JAX to use bfloat16 everywhere

[cool-RR]

I have a large JAX codebase which I usually run on my RTX 5080. Recently I ran it on an A100 and it ran slowly. I consulted ChatGPT and it told me I should use bfloat16 on the A100 and it'll run fast.

Is there a way to get my JAX project to use bfloat16 everywhere? Right now I'm going through all the code to put dtype=jnp.bfloat16 on everything, but that's a bad solution. Besides adding lots of code, it also makes it difficult to quickly switch between bfloat16 and the float32 I've been using so far.

[jakevdp]

A similar question was asked here: #30106

I don't know of any way to force JAX to default to bfloat16, short of explicitly adding dtype='bfloat16' wherever arrays are created or loaded.

My suggestion would be to write your function such that it always enforces the dtype of the input: most JAX APIs will already do this, and it would allow you to switch between bfloat16 and float32 at runtime without changing the source code.

[cool-RR]

Thanks for the quick reply.

My suggestion would be to write your function such that

More like my hundred functions... Which is what I'm trying to avoid.

Should I submit a feature request for this? Do you agree it's a good feature?

[jakevdp]

I'm not sure what the concrete ask is: do you want to make all APIs produce bfloat16 rather than float32/float64 by default? Do you want JAX to error if 32-bit values are introduced somewhere? Do you want X64-style semantics where we silently squash all 32-bit floats to bfloat16 in your computation? Do you want this even if the implementation explicitly requests dtype='float32'?

[cool-RR]

I think that making bfloat16 the default dtype for floating point numbers would achieve most of what I want. I'm not sure whether I'd want an exception or a silent squashing in case float32 is explictly specified.

The fact that you don't recognize my need is hinting that maybe I'm going about this wrong. I didn't completely understand what ChatGPT explained to me. I do know that for the little benchmark script I did, once I switched it to bfloat16, the runtime on the A100 was 2-3x faster than the 5080, which is the kind of speedup I expected.

ChatGPT also suggested export NVIDIA_TF32_OVERRIDE=1 which didn't make my program faster, and config.update("jax_default_dtype_bits", 16) which apparently doesn't work anymore.

If you understand what I'm not understanding and can clue me in, that'll be helpful.

[jakevdp]

I think that making bfloat16 the default dtype for floating point numbers would achieve most of what I want.

Good to know – that's something we've been thinking about and may add soon.

and config.update("jax_default_dtype_bits", 16) which apparently doesn't work anymore.

This would never have worked – it's just an LLM hallucination.

once I switched it to bfloat16, the runtime on the A100 was 2-3x faster than the 5080, which is the kind of speedup I expected.

This makes sense, but in general it's only certain operations (typically matmuls) that will benefit from computation in bfloat16, and the precise speedups will depend on the hardware you're using. You may be able to achieve the speedup by globally setting the jax_default_matmul_precision configuration to "bfloat16": with this, even if your values are stored as float32, the matmuls will be computed at bfloat16 precision at the hardware level. Just add this at the top of your script:

jax.config.update('jax_default_matmul_precision', 'bfloat16')

[cool-RR]

I think that making bfloat16 the default dtype for floating point numbers would achieve most of what I want.

Good to know – that's something we've been thinking about and may add soon.

I'm happy to hear that.

jax.config.update('jax_default_matmul_precision', 'bfloat16')

I tried that yesterday, it didn't speed up my code. It was still slower on the A100 than it is on the 5080.

[cool-RR]

I finally finished manually converting my entire project to bfloat16, and it's still 3x slower on the A100 than it is on the 5080 😭

I guess the next step would be to make a small sample that demonstrates the slowness, and that'll probably be very tricky...

[cool-RR]

Hi again @jakevdp . I reduced my code to a small runnable sample: bonchu.py The code looks long only because I put lots of useful output in it, but the JAX action is only a few lines.

Time on RTX 5080: 0.8422s
Time on A100: 1.3011s

Does this code give you an idea for what I'm doing wrong? In other words, why it's slower on the A100?

[jakevdp]

It might be that RTX vs A100 choose different matmul precisions when it's unspecified. You might try running the benchmarks after setting jax_default_matmul_precision to bfloat16.