Add initial T5Gemma S/S support#947
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
jan-wassenberg
left a comment
jan-wassenberg
left a comment
There was a problem hiding this comment.
Remarkably nice work, very clean code, congrats :) A few suggestions:
| s_ffw_in(config.num_layers, max_workers), | ||
| s_ffw_hidden(config.num_layers, max_workers), | ||
| s_ffw_out(config.num_layers, max_workers), | ||
| s_ffw_in(config.is_encoder_decoder ? config.decoder_num_layers |
There was a problem hiding this comment.
Let's add a member variable initialized to config.is_encoder_decoder ? config.decoder_num_layers : config.num_layers, then we can use that everywhere.
| visitor(partial_rotary_factor); | ||
|
|
||
| visitor(is_encoder_decoder); | ||
| visitor(encoder_num_layers); |
There was a problem hiding this comment.
We can make encoder/decoder* fields conditional, i.e. only send them if is_encoder_decoder == true. (It's enough to wrap them in an if().)
| }); | ||
| } | ||
|
|
||
| static constexpr int kT5GemmaPadId = 0; |
There was a problem hiding this comment.
Let's move these to tokenizer.h? We can also reuse BOS_ID = 2.
| return distance <= window_size; | ||
| } | ||
|
|
||
| static HWY_INLINE float T5GemmaMaybeSoftCap(float score, float cap) { |
There was a problem hiding this comment.
Reuse MaybeLogitsSoftCap from ops-inl.h?
|
|
||
| const LayerConfig& layer_config = config.encoder_layer_configs[0]; | ||
| const size_t source_len = encoder_cache.source_len; | ||
| MatStorageT<BF16> pre_att_rms_out( |
There was a problem hiding this comment.
Allocating is pretty expensive, probably we want to hoist this into something like Activations?
|
|
||
| for (size_t qi = 0; qi < qbatch.Size(); ++qi) { | ||
| const T5GemmaEncoderCache* encoder_cache = qbatch.T5EncoderCache(qi); | ||
| HWY_ASSERT(encoder_cache != nullptr); |
There was a problem hiding this comment.
Note that our HWY_ASSERT is always on, HWY_DASSERT only in debug/sanitizer builds. Might want to switch to HWY_DASSERT here?
2 participants
Summary
Adds initial T5Gemma S/S support to gemma.cpp.
This includes:
python/convert_from_safetensors.pyScope
This PR targets T5Gemma v1 S/S only.
Out of scope:
Local Validation
bazel test -c opt //:configs_test //:tensor_info_test //:kv_cache_testt5gemma-s-s-itHelloI am a new user, ...Notes
Opened as draft to get early maintainer feedback on scope, naming, and whether this should be split further.
Additional Note:
//:gemma_lib_testcurrently fails locally for me in an existingGemmaAttentionFlash/EMU128golden mismatch on Apple Silicon; this reproduced outside the T5Gemma changes.Related to #661