Why not use low precision matmul for reverse embedding in gemma model

[lingzhi98]

Why not use low precision matmul for reverse embedding in gemma model

[mattdangerw]

@lingzhi98 thanks for filing and timely. We actually just hit this as it was causing some numerical differences at half precision for llama 2. I do wonder if we could flip this safely, the best default behavior would be to follow the default dtype strategy whatever that is. cc @tirthasheshpatel who was working on this for llama.

I'll try to check out if this is a safe change to make (particularly with Gemma). As long as we make sure to cast to float32 before any output softmax during sampling, seems like we should be ok.

[SuryanarayanaY]

Hi @lingzhi98 ,

As per documentation that states the reason for this as below.

keras-nlp/keras_nlp/layers/modeling/reversible_embedding.py

Lines 52 to 54 in 1889369

	reverse_dtype: The dtype for the reverse projection computation.
	For stability, it is usually best to use full precision even when
	working with half or mixed precision training.

[lingzhi98]

Thanks. My understanding is that use float32 here to ensure the stability of training, is it necessary for inference also?

[mattdangerw]

Here's the fix, will push this out next week #1548