riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-08-01 10:41:07 +06:00
Code Issues Actions 3 Packages Projects Releases Wiki Activity

Cast bfloat16 to float32 for Numpy conversions (#29755)

Browse Source 
* Cast bfloat16 to float32 for Numpy conversions

* Add test
This commit is contained in:
Matt 2024-03-21 14:04:11 +00:00 committed by GitHub
parent 73a73b415e
commit de627f5a14
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 15 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/transformers/modeling_tf_pytorch_utils.py
View File

@ -249,7 +249,10 @@ def load_pytorch_weights_in_tf2_model(
) )
raise raise
pt_state_dict = {k: v.numpy() for k, v in pt_state_dict.items()} # Numpy doesn't understand bfloat16, so upcast to a dtype that doesn't lose precision
pt_state_dict = {
k: v.numpy() if v.dtype != torch.bfloat16 else v.float().numpy() for k, v in pt_state_dict.items()
}
return load_pytorch_state_dict_in_tf2_model( return load_pytorch_state_dict_in_tf2_model(
tf_model, tf_model,
pt_state_dict, pt_state_dict,

11
tests/test_modeling_tf_utils.py
View File

@ -63,6 +63,7 @@ if is_tf_available():
PreTrainedModel, PreTrainedModel,
PushToHubCallback, PushToHubCallback,
RagRetriever, RagRetriever,
TFAutoModel,
TFBertForMaskedLM, TFBertForMaskedLM,
TFBertForSequenceClassification, TFBertForSequenceClassification,
TFBertModel, TFBertModel,
@ -435,6 +436,16 @@ class TFModelUtilsTest(unittest.TestCase):
for p1, p2 in zip(model.weights, new_model.weights): for p1, p2 in zip(model.weights, new_model.weights):
self.assertTrue(np.allclose(p1.numpy(), p2.numpy())) self.assertTrue(np.allclose(p1.numpy(), p2.numpy()))
@is_pt_tf_cross_test
@require_safetensors
def test_bfloat16_torch_loading(self):
# Assert that neither of these raise an error - both repos contain bfloat16 tensors
model1 = TFAutoModel.from_pretrained("Rocketknight1/tiny-random-gpt2-bfloat16-pt", from_pt=True)
model2 = TFAutoModel.from_pretrained("Rocketknight1/tiny-random-gpt2-bfloat16") # PT-format safetensors
# Check that PT and safetensors loading paths end up with the same values
for weight1, weight2 in zip(model1.weights, model2.weights):
self.assertTrue(tf.reduce_all(weight1 == weight2))
@slow @slow
def test_save_pretrained_signatures(self): def test_save_pretrained_signatures(self):
model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert") model = TFBertModel.from_pretrained("hf-internal-testing/tiny-random-bert")