Update MaskedLMHead to support dtype=bfloat16/float16/float64 #1197
Conversation
| @@ -153,9 +153,11 @@ def build(self, inputs_shape, masked_positions_shape=None): | |||
| activation=self.intermediate_activation, | |||
| kernel_initializer=self.kernel_initializer, | |||
| bias_initializer=self.bias_initializer, | |||
| dtype=self._dtype_policy, | |||
There was a problem hiding this comment.
It looks like self.dtype_policy should just work. Can we do that?
| @@ -36,6 +39,30 @@ def test_valid_call(self): | |||
| position_data = ops.random.randint(minval=0, maxval=10, shape=(4, 5)) | |||
| model((token_data, position_data)) | |||
|
|
|||
| @parameterized.named_parameters( | |||
| ("bfloat16", tf.bfloat16), | |||
There was a problem hiding this comment.
because we now run our testing suite with jax/torch/tf with keras-core, we are generally just referring to these by string name, e.g. "float16" instead of tf.float16.
Does anything break if we switch to that?
| @@ -119,6 +146,32 @@ def test_one_train_step(self): | |||
| loss = model.train_on_batch(x=(token_data, position_data), y=label_data) | |||
| self.assertGreater(loss, 0) | |||
|
|
|||
| @parameterized.named_parameters( | |||
There was a problem hiding this comment.
I would kill this test. Compiling a real loss function can make for slower tests, and with the parameterized testing this could slow down our suite.
| @@ -153,9 +153,11 @@ def build(self, inputs_shape, masked_positions_shape=None): | |||
| activation=self.intermediate_activation, | |||
There was a problem hiding this comment.
It seems like we should really have this for all our our "composite" layers in KerasNLP, right?
- token and position embedding
- transformer decoder
- transformer encoder
- cached multi head attention
- f net encoder
Are you interested in following up for other layers? (same PR or split PRs fine!)
| ) | ||
| encoded_tokens = keras.Input(shape=(10, 16)) | ||
| positions = keras.Input(shape=(5,), dtype="int32") | ||
| outputs = head(encoded_tokens, masked_positions=positions) |
There was a problem hiding this comment.
This might need a rebase over master. This should be mask_positions now. This is causing a lot of test failures.
Inspired by keras-team/keras@397ad57
i.e. using the idiom (?) of
dtype=self._dtype_policy.This is to fix #1195
I had a previous try at this where I accidentally included print statements, sorry.