Add Qwen Moe #2163
Add Qwen Moe #2163
Conversation
heyyanshuman
requested review from
mattdangerw,
abheesht17 and
divyashreepathihalli
| @@ -79,7 +79,7 @@ def build(self, decoder_sequence_shape): | |||
| self.hidden_dim = decoder_sequence_shape[-1] | |||
|
|
|||
| # Self attention layer. | |||
| self._self_attention_layer = QwenAttention( | |||
| self._self_attention_layer = QwenMoeAttention( | |||
There was a problem hiding this comment.
I see you have forked this file in qwen_moe folder - why is this being edited?
There was a problem hiding this comment.
looks like this slipped in in find & replace, fixed it.
|
|
||
| for expert_idx in range(self.num_experts): | ||
| expert_layer = self.experts[expert_idx] | ||
| idx, top_x = ops.where(expert_mask[expert_idx]) |
There was a problem hiding this comment.
will the ops.where return different output shape here on different forward passes? - if so that would not work with JAX XLA
There was a problem hiding this comment.
I don't think that this should be the case. Why do you think so?
There was a problem hiding this comment.
ops.where won't work on JAX at all, if you don't provide x and y.
ops.where calls jnp.where. Check the note here:
Because the size of the output of nonzero is data-dependent, the function is not
compatible with JIT and other transformations. The JAX version adds the optional
size argument which must be specified statically for jnp.nonzero to be used within
JAX’s transformations.
You can, however, consider passing the size argument. That might make it work.
| ) | ||
| expert_mask = ops.transpose(expert_mask, axes=[2, 1, 0]) | ||
|
|
||
| for expert_idx in range(self.num_experts): |
There was a problem hiding this comment.
The for loop here to go over the experts are inefficient for XLA compilation. This implementation would need to updated - I had tried out a dummy MOE implementation in JAX here - https://colab.sandbox.google.com/drive/1r0rscZK_2bNpDmFLC1POEEQoKcqWQYlQ
in order to bring this to KearsHub we are missing ragged_dot op.
There was a problem hiding this comment.
I have prototyped the implementation - will add this op soon
There was a problem hiding this comment.
When can we expect this to be available as a part of keras op?
|
@divyashreepathihalli How should we accomodate aux_loss for CausalLM task here model here? We are specifying Sparse Categorical CrossEntropy Loss here: keras-hub/keras_hub/src/models/causal_lm.py Lines 109 to 119 in b997444 |
There was a problem hiding this comment.
Thanks for the updates @heyyanshuman!
I left some comments on the PR regarding tf ops
please add tests for the layers, backbones and tasks
I am curious to know if model.fit works, do you have a demo colab for inference and FT? - looking for the aux loss implementation
| ) | ||
| self._query_dense.build(inputs_shape) | ||
|
|
||
| self._key_dense = keras.layers.EinsumDense( |
There was a problem hiding this comment.
you might want to rename this to to match other KH models here - value_dense and query_dense
this will allow enabling LoRA on this Model - https://docs.google.com/document/d/1BSjDMuSP9N0e2sw83E5ujSbiwiDpoKzMBLCRN9xdtLY/edit?resourcekey=0-OMMpYhmIDBMxQtVuV5saSg&tab=t.0#heading=h.fttfr1z4ln91
There was a problem hiding this comment.
I don't have access to this document :(
| @keras_hub_export( | ||
| "keras_hub.models.QwenMoeCausalLM", | ||
| ) | ||
| class QwenMoeCausalLM(CausalLM): |
There was a problem hiding this comment.
add docstring and example to show model.fit and generate
This PR adds Qwen Mixture-of-expert model to Keras Hub.
Huggingface Reference : link