Allows for Bumblebee.load_model({:hf, "..."}, type: :bf16) to set mixed precision policy on the model and cast params on load. The user can also pass the policy struct itself, type is just a shorthand.
Note: currently loading a f16 checkpoint returns params in f16, but that's somewhat a bug, specifically if all params are available we skip Axon init, which would otherwise cast the params to f32. Now it's going to always be f32 by default and the user can override it using :type.
hf/transformers look at "torch_dtype" in the config file to determine the type, but I'm not sure if we want to configure the policy automatically based on that. @seanmor5 thoughts?
Allows for
Bumblebee.load_model({:hf, "..."}, type: :bf16)to set mixed precision policy on the model and cast params on load. The user can also pass the policy struct itself, type is just a shorthand.Note: currently loading a f16 checkpoint returns params in f16, but that's somewhat a bug, specifically if all params are available we skip Axon init, which would otherwise cast the params to f32. Now it's going to always be f32 by default and the user can override it using
:type.hf/transformers look at
"torch_dtype"in the config file to determine the type, but I'm not sure if we want to configure the policy automatically based on that. @seanmor5 thoughts?