Add a mixed precision test and fix mixed precision errors for layers #1242
Conversation
|
/gcbrun |
|
Ah I think this needs a keras-core release to pass. Will sync with Francois. |
mattdangerw
force-pushed
the
mixed-precision-layer-fixes
branch
from
44b1005
to
fc4f0f5
Compare
| ) | ||
| intermediate_shape = list(decoder_sequence_shape) | ||
| intermediate_shape[-1] = self.intermediate_dim | ||
| self._feedforward_output_dense.build(tuple(intermediate_shape)) | ||
| self._feedforward_layer_norm = keras.layers.LayerNormalization( | ||
| epsilon=self.layer_norm_epsilon, | ||
| name="output_layer_norm", | ||
| dtype=self.dtype_policy, | ||
| name="feedforward_layer_norm", |
There was a problem hiding this comment.
Just a random thought while I saw these
With the new distribution API, we're going to need to be careful about changing layer names moving forward!
There was a problem hiding this comment.
Yep! Trying to get us in a nice consistent state before we do.
keras_nlp/samplers/sampler.py
Outdated
| This will always be done in full precision, regardless of dtype, and | ||
| scale by `temperature`. | ||
| """ | ||
| dtype = logits.dtype |
There was a problem hiding this comment.
Maybe logits_dtype? (For consistency of style with inputs_dtype a few files up)
keras_nlp/tests/test_case.py
Outdated
| output_data = layer(input_data) | ||
| for tensor in tree.flatten(output_data): | ||
| dtype = standardize_dtype(tensor.dtype) | ||
| if "float" in dtype: |
There was a problem hiding this comment.
Seems like assertDType should be good here no?
mattdangerw
force-pushed
the
mixed-precision-layer-fixes
branch
from
03c1c85
to
79c0219
Compare
|
Is mixed precision working in Keras Core now? |
mattdangerw
force-pushed
the
mixed-precision-layer-fixes
branch
from
aefe4ed
to
926a737
Compare
Yes, or at least it should be largely with the latest release. Landed the loss scaling optimizer, which was the main piece we were missing, as well as a few other fixes. |
|
/gcbrun |
5 similar comments
|
/gcbrun |
|
/gcbrun |
|
/gcbrun |
|
/gcbrun |
|
/gcbrun |
|
Merging, last breakage unrelated #1251 |
No description provided.