Add StableLM-3B 4E1T to Keras Hub #2151
Conversation
|
@divyashreepathihalli Here is a comparison of numerics with Hugging Face in Colab. The results match with an absolute tolerance of 1e-3, but they do not match when using 1e-5. Could you please take a look and suggest some improvements or explanations for this discrepancy? |
sachinprasadhs
added
the
WIP
|
The numerics is good enough! |
|
@Bond099 let's sync this with the latest changes and make sure to run our format script. I'm not exactly sure why non of our CI is running, but I don't think it ran. |
|
Let's clean up the PR. Can we fix the following minor things?
|
|
Looks like there are conflicts. Please pull in master and resolve conflicts |
| @@ -3,7 +3,6 @@ | |||
| from keras import random | |||
There was a problem hiding this comment.
I'd actually want sampler.py to be untouched, because all models use this. We'd want to make this change only for Stable LM. The changes will be less intrusive that way.
Is there a way you can pass the correct mask right here:
, https://github.com/keras-team/keras-hub/blob/master/keras_hub/src/samplers/sampler.py#L79?If you do pass the correct mask, this should evaluate to the correct bool: https://github.com/keras-team/keras-hub/blob/master/keras_hub/src/samplers/sampler.py#L100.
Take a look if this works.
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for the StableLM-3B-4E1T model. I've identified a critical bug in the StableLMTransformerDecoder's caching logic that would cause failures during training, and a few medium-to-high severity issues in the StableLMBackbone's documentation and the checkpoint conversion script.
| x, self_attention_cache = self.self_attention_layer( | ||
| hidden_states=x, | ||
| attention_mask=self_attention_mask, | ||
| cache=self_attention_cache, | ||
| cache_update_index=self_attention_cache_update_index, | ||
| ) | ||
| x = self.self_attention_dropout(x, training=training) | ||
| x = x + residual | ||
|
|
||
| residual = x | ||
|
|
||
| # Feedforward block | ||
| x = self.feedforward_layernorm(x) | ||
| gate_output = self.feedforward_gate_dense(x) | ||
| gate_output = self.activation(gate_output) | ||
| intermediate_output = self.feedforward_intermediate_dense(x) | ||
| x = self.feedforward_output_dense( | ||
| ops.multiply(intermediate_output, gate_output) | ||
| ) | ||
| decoder_output = x + residual | ||
| if self_attention_cache is not None: | ||
| return decoder_output, self_attention_cache | ||
| return decoder_output |
There was a problem hiding this comment.
The call method's handling of the attention cache has a bug. When call is invoked without a cache (i.e., self_attention_cache is initially None), self.self_attention_layer returns a tuple of (attention_output, attention_output). This makes the local self_attention_cache a tensor, causing the if condition to be true and the method to incorrectly return a tuple, breaking the StableLMBackbone functional model which expects a single tensor output. Use a different variable name for the cache returned by the attention layer to fix this.
attention_output, next_attention_cache = self.self_attention_layer(
hidden_states=x,
attention_mask=self_attention_mask,
cache=self_attention_cache,
cache_update_index=self_attention_cache_update_index,
)
x = self.self_attention_dropout(attention_output, training=training)
x = x + residual
residual = x
# Feedforward block
x = self.feedforward_layernorm(x)
gate_output = self.feedforward_gate_dense(x)
gate_output = self.activation(gate_output)
intermediate_output = self.feedforward_intermediate_dense(x)
x = self.feedforward_output_dense(
ops.multiply(intermediate_output, gate_output)
)
decoder_output = x + residual
if self_attention_cache is not None:
return decoder_output, next_attention_cache
return decoder_output| cfg["max_sequence_length"] = pt_cfg["max_position_embeddings"] | ||
| cfg["layer_norm_epsilon"] = pt_cfg["layer_norm_eps"] | ||
| cfg["rope_max_wavelength"] = pt_cfg["rope_theta"] | ||
| cfg["partial_rotary_factor"] = pt_cfg["partial_rotary_factor"] |
There was a problem hiding this comment.
The Hugging Face config parameter partial_rotary_factor should be mapped to rotary_percentage to ensure the Keras model is configured correctly.
| cfg["partial_rotary_factor"] = pt_cfg["partial_rotary_factor"] | |
| cfg["rotary_percentage"] = pt_cfg["partial_rotary_factor"] |
| tie_word_embeddings: bool, optional. Whether to tie input and output | ||
| embeddings. Defaults to False. |
There was a problem hiding this comment.
The docstring for __init__ includes the tie_word_embeddings argument, but it is not present in the method's signature. Either add the argument to the __init__ method, or remove it from the docstring.
| keras_model.get_layer( | ||
| f"transformer_layer_{layer_index}" | ||
| )._self_attention_layer._query_dense.set_weights(weights) |
There was a problem hiding this comment.
The script accesses layer attributes using names with a leading underscore (e.g., _self_attention_layer, _query_dense). Use the public attribute names for setting weights instead.
| keras_model.get_layer( | |
| f"transformer_layer_{layer_index}" | |
| )._self_attention_layer._query_dense.set_weights(weights) | |
| keras_model.get_layer( | |
| f"transformer_layer_{layer_index}" | |
| ).self_attention_layer.query_dense.set_weights(weights) |
This PR adds the StableLM-3B 4E1T model to Keras Hub. However, numerical matching with the Hugging Face implementation is still in progress.