Closed grasskin closed 5 months ago
I think we also need to cast back to the compute_dtype before returning. Only the computation part needs to happen in float32.
Yeah looks like this is causing test failures, probably due to this issue?
Hi :) I'm assuming this came about from my Twitter thread https://twitter.com/danielhanchen/status/1765446273661075609 :)
I added a fix into transformers 4.38.2 here: https://github.com/huggingface/transformers/pull/29285. So using mixed_bfloat16
causes torch.autocast
to cast all ops to bfloat16. I don't normally use Keras, so unsure if torch.autocast
affects operations, since I know even explicitly forcing float32 causes autocast to override it. However, unsure on Keras.
Also another problematic line is https://github.com/keras-team/keras-nlp/blob/v0.8.2/keras_nlp/models/gemma/gemma_attention.py#L159
seq_len = ops.shape(x)[1]
start_index = cache_update_index
positions = ops.cast(
ops.arange(seq_len, dtype="float32"), >>>>> self.compute_dtype <<<<
)
positions = positions + ops.cast(start_index, self.compute_dtype)
Which is wrong - Assume if someone did RoPE Scaling with float16 - this will cause 65504 to be the maximum, which in turn causes overflow ie infinities to occur. bfloat16 loses precision, but can represent larger numbers.
Hi @danielhanchen, enjoyed reading the blogpost it was a great in depth dive!
Switched all of RoPE to happen in "float32" and added downcasting before returning. This likely works until we replace the call with normal Keras RoPE?
Switched all of RoPE to happen in "float32" and added downcasting before returning. This likely works until we replace the call with normal Keras RoPE?
Yeah, this should be good enough for now. We can merge this and I can rebase my PR on top of your changes.
@mattdangerw Gemma still downcasts the tensors to
compute_dtype
since it uses it's own implementation of RoPE. I can submit a follow-up PR to use this layer instead.