Support bfloat16 torch fromto flax conversions

[yhavinga]

I noticed when working with a bfloat16 model, converstion between torch and flax would raise errors like:

TypeError: can't convert np.ndarray of type "bfloat16." The only supported types are: float64, float32, float16, complex64, complex128, int64, int32, int16, int8, uint8, and bool.

The following commit tries to implement this for Mistral, both directions.
main...yhavinga:EasyDeL:mistral_bfloat16_conversions

The test python_test/mistral_torch_bf16_conversions.py runs these conversions:

• torch -> flax (with mistral_convert_hf_to_flax)
• flax->torch (with easystate_to_huggingface_model)

then compares output of the two torch and the flax model.
In both float32 and bfloat16 mode there are differences, the bf16 differences are larger than the fp32.

Since the fp32 numbers are really close, I'm inclined to believe the model implementations are correct, and differences in output are due to differences in algorithm implementations causing slightly different values for the highest precisions and compounding these.

I am not sure if this commit (converstions of bfloat16 parameters) is worth pursueing. Maybe in the future numpy will support bfloat16 conversions, rendering this code useless. And it's not a general solution for all models. Maybe a more general approach for model conversions would be: always upcast to fp32, convert between torch and flax, then downcast to datatype the model was trained in again.

================== <class 'jax.numpy.float32'> conversions ==================
Convert HF Torch model to params
Initialize EasyState with config and params
Saving State to File: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 47/47 [00:00<00:00, 147.70it/s, gather_functions_mismatch=0]
Convert EasyState to HF Torch model
Loading Checkpoints From File: 47it [00:00, 280.36it/s, shard_functions_mismatch=0]
Converting EasyDelState to torch state_dict: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 39/39 [00:00<00:00, 4520.35it/s]
Generate random input_ids and compare predictions
Mean Error : -7.294333470575509e-10
Flax jax_dtype=<class 'jax.numpy.float32'> Predictions :
 [[[-0.24755938  0.15139946  0.20477434 ... -0.17459351 -0.67721415  0.22053908]
  [-0.27330583 -0.08644339  0.28088355 ... -0.15419051 -0.4318745   0.19456264]
  [-0.24764328 -0.12680095  0.28403017 ...  0.05390351 -0.4325688   0.12856238]
  ...
  [-0.37154594 -0.08576845  0.23800655 ...  0.36986047 -0.34367213 -0.14205466]
  [-0.39135382 -0.11681088  0.0957402  ...  0.33429816 -0.32594    -0.09676895]
  [-0.23300466 -0.29766357  0.14068954 ...  0.31457457 -0.44966078  0.03667825]]]
Torch torch_dtype=torch.float32 Predictions :
 [[[-0.2475598   0.15139951  0.20477438 ... -0.17459336 -0.6772145   0.22053893]
  [-0.27330598 -0.08644354  0.2808835  ... -0.15419045 -0.43187448  0.19456246]
  [-0.24764341 -0.12680094  0.28403038 ...  0.0539038  -0.43256867  0.12856245]
  ...
  [-0.3715456  -0.0857687   0.23800664 ...  0.36986023 -0.3436723  -0.14205465]
  [-0.391354   -0.1168113   0.09574009 ...  0.33429793 -0.32594    -0.09676883]
  [-0.23300485 -0.29766363  0.14068961 ...  0.31457466 -0.44966125  0.03667808]]]
================== <class 'jax.numpy.bfloat16'> conversions ==================
Convert HF Torch model to params
Initialize EasyState with config and params
Saving State to File: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 47/47 [00:00<00:00, 489.22it/s, gather_functions_mismatch=0]
Convert EasyState to HF Torch model
Loading Checkpoints From File: 47it [00:00, 808.91it/s, shard_functions_mismatch=0]
Converting EasyDelState to torch state_dict: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 39/39 [00:00<00:00, 4096.31it/s]
Generate random input_ids and compare predictions
Mean Error : -8.60682575876126e-06
Flax jax_dtype=<class 'jax.numpy.bfloat16'> Predictions :
 [[[-0.244141 0.15332 0.204102 ... -0.172852 -0.679688 0.21875]
  [-0.273438 -0.0849609 0.28125 ... -0.15625 -0.433594 0.191406]
  [-0.245117 -0.125977 0.283203 ... 0.0532227 -0.435547 0.12793]
  ...
  [-0.371094 -0.0834961 0.238281 ... 0.373047 -0.351562 -0.136719]
  [-0.392578 -0.112305 0.09375 ... 0.337891 -0.332031 -0.0917969]
  [-0.232422 -0.294922 0.137695 ... 0.318359 -0.451172 0.0400391]]]
Torch torch_dtype=torch.bfloat16 Predictions :
 [[[-0.24511719  0.15332031  0.20507812 ... -0.17480469 -0.68359375  0.21972656]
  [-0.2734375  -0.08642578  0.28125    ... -0.15527344 -0.43554688  0.19140625]
  [-0.24707031 -0.125       0.28125    ...  0.05200195 -0.43554688  0.12988281]
  ...
  [-0.37304688 -0.08203125  0.23535156 ...  0.37109375 -0.34960938 -0.13574219]
  [-0.39257812 -0.11132812  0.09326172 ...  0.3359375  -0.33007812 -0.09179688]
  [-0.234375   -0.29492188  0.13867188 ...  0.31445312 -0.453125    0.04296875]]]

[erfanzar]

hello and thanks for the idea
I guess that would be something nice too and you can make a pull request for those changes to be merged but in the next version I'm about to remove all the static model converts like mistral_convert_hf_to_flax and ...

[yhavinga]

Ok I'll wait for the static model convert changes and also want to try if I can iron out a generic way for model upcast / convert / downcast. As a sidenote, one thing that puzzled a bit is the conversion of layernorms: after HF model.bfloat16() they're also bfloat16, but I thought layernorms should be kept in float32 -- thats something I also wanted to look into. # parameters is hardly increased with having it in fp32, so any avoidable bfloat16-related precision errors would be a plus I guess.