Add dora support #2371
Add dora support #2371
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. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @Ajinkya-25, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request introduces DoRA (Weight-Decomposed Low-Rank Adaptation) support to Keras Hub, providing a novel approach for efficient fine-tuning of large language models. It adds new Keras layers, DoRADense, DoRAEmbedding, and DoRAPositionEmbedding, which implement the DoRA mechanism by decomposing weight matrices into magnitude and direction components. The PR also includes convenient utility functions to convert existing standard Dense and Embedding layers to their DoRA-enabled versions. To showcase its practical application, DoRA has been integrated into the BertBackbone model, allowing users to easily leverage this technique for more efficient model adaptation.
Highlights
- New DoRADense Layer: Introduced
DoRADense, a new Keras layer that implements Weight-Decomposed Low-Rank Adaptation for dense connections, allowing for efficient fine-tuning by decomposing weights into frozen pretrained weights, learnable low-rank matrices, and a learnable magnitude vector. - New DoRA Embedding Layers: Added
DoRAEmbeddingandDoRAPositionEmbeddinglayers, which extend the DoRA concept to embedding and positional embedding matrices, enabling efficient adaptation of these crucial components in neural networks. - Conversion Utilities for Existing Layers: Provided utility functions,
convert_dense_to_doraandconvert_embedding_to_dora, to seamlessly transform existing standard KerasDenseandEmbeddinglayers into their DoRA-enabled counterparts, preserving pretrained weights. - BertBackbone Integration: Integrated DoRA support into the
BertBackbonemodel, allowing users to enable DoRA for its token embeddings, positional embeddings, and pooled dense layer via a newenable_doraflag, demonstrating a practical application of the new layers. - Extensive Testing: Included comprehensive unit and integration tests for all new DoRA layers and their integration within the
BertBackbone, ensuring correctness, functional equivalence with original layers (when applicable), and proper model saving/loading.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces DoRA (Weight-Decomposed Low-Rank Adaptation) support to Keras Hub by adding new reusable layers: DoRADense, DoRAEmbedding, and DoRAPositionEmbedding. The changes also include conversion utilities and an integration example in the BertBackbone. The implementation is well-structured and includes comprehensive tests. My main feedback is to address a style guide violation regarding the use of type hints in function signatures in the new layer files. Addressing this will make the contribution align perfectly with the project's standards.
|
This PR is complete with all implementations, and the test cases are passing. Please let me know if anything else is required from my side. |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces comprehensive support for DoRA (Weight-Decomposed Low-Rank Adaptation) into Keras Hub. The changes include new reusable layers (DoRADense, DoRAEmbedding, DoRAPositionEmbedding), conversion utilities, and an example integration into the BertBackbone. The implementation is well-structured and the test coverage is excellent, ensuring the new components are robust and backend-agnostic. My review focuses on improving documentation to align with the repository's style guide, enhancing code clarity by removing unused code, and ensuring consistency in the public API. Overall, this is a high-quality contribution that significantly enhances the fine-tuning capabilities within Keras Hub.
| class DoRAPositionEmbedding(layers.Layer): | ||
| """DoRA-enabled position embedding layer. | ||
| This layer creates learnable positional embeddings | ||
| that are added to token embeddings, | ||
| using DoRA weight decomposition for efficient adaptation. | ||
| """ |
There was a problem hiding this comment.
The docstring for DoRAPositionEmbedding is incomplete. As per the repository style guide, it should include Args, Example, and Reference sections to be comprehensive.1
Style Guide References
Footnotes
-
The style guide requires thorough documentation for public classes, including arguments, usage examples, and references. ↩
| from keras_hub.src.layers.modeling.dora_embeddings import ( | ||
| convert_embedding_to_dora as embedding_to_dora, | ||
| ) |
There was a problem hiding this comment.
For consistency with convert_dense_to_dora, it would be better to alias convert_embedding_to_dora to convert_embedding_to_dora instead of embedding_to_dora. This makes the public API more predictable and consistent.
| from keras_hub.src.layers.modeling.dora_embeddings import ( | |
| convert_embedding_to_dora as embedding_to_dora, | |
| ) | |
| from keras_hub.src.layers.modeling.dora_embeddings import ( | |
| convert_embedding_to_dora as convert_embedding_to_dora, | |
| ) |
| """DoRA (Weight-Decomposed Low-Rank Adaptation) Dense layer. | ||
| DoRA decomposes the weight matrix W into magnitude and direction components | ||
| W = m * (W_0 + B @ A) / ||W_0 + B @ A||_c | ||
| Where: | ||
| - m: magnitude vector (learnable) | ||
| - W_0: frozen pretrained weights | ||
| - A, B: low-rank adaptation matrices (learnable) | ||
| - ||.||_c: column-wise L2 norm | ||
| Args: | ||
| units: Positive integer, dimensionality of the output space. | ||
| rank: Rank of the adaptation. Positive integer. | ||
| alpha: LoRA scaling parameter. Float. | ||
| use_bias: Boolean, whether the layer uses a bias vector. | ||
| dropout: Float between 0 and 1. Fraction of input units to drop. | ||
| activation: Activation function to use. | ||
| kernel_initializer: Initializer for the kernel weights matrix. | ||
| bias_initializer: Initializer for the bias vector. | ||
| lora_a_initializer: Initializer for the A matrix. | ||
| Defaults to 'he_uniform'. | ||
| lora_b_initializer: Initializer for the B matrix. | ||
| Defaults to 'zeros'. | ||
| magnitude_initializer: Initializer for magnitude vector. | ||
| Defaults to 'ones'. | ||
| kernel_regularizer: Regularizer function applied to kernel weights. | ||
| bias_regularizer: Regularizer function applied to bias. | ||
| activity_regularizer: Regularizer function applied to output. | ||
| kernel_constraint: Constraint function applied to kernel weights. | ||
| bias_constraint: Constraint function applied to bias. | ||
| **kwargs: Additional keyword arguments. | ||
| """ |
There was a problem hiding this comment.
|
|
||
| # Scaling factor | ||
| self.scaling = self.alpha / self.rank | ||
| self.input_spec = None |
There was a problem hiding this comment.
The self.input_spec = None assignment appears to be unnecessary as input_spec is not used elsewhere in the class. Removing it would improve code clarity.
| class DoRAEmbedding(layers.Layer): | ||
| """DoRA (Weight-Decomposed Low-Rank Adaptation) | ||
| Embedding layer. | ||
| DoRA decomposes the embedding weight | ||
| matrix W into magnitude and direction components: | ||
| W = m * (W_0 + B @ A) / ||W_0 + B @ A||_c | ||
| Where: | ||
| - m: magnitude vector (learnable) | ||
| - W_0: frozen pretrained embedding weights | ||
| - A, B: low-rank adaptation matrices (learnable) | ||
| - ||.||_c: column-wise L2 norm | ||
| Args: | ||
| input_dim: Size of the vocabulary (number of tokens). | ||
| output_dim: Dimension of the dense embedding vectors. | ||
| rank: Rank of the adaptation. Positive integer. | ||
| alpha: LoRA scaling parameter. Float. | ||
| embeddings_initializer: | ||
| Initializer for the embeddings matrix. | ||
| lora_a_initializer: | ||
| Initializer for the A matrix. Defaults to 'he_uniform'. | ||
| lora_b_initializer: | ||
| Initializer for the B matrix. Defaults to 'zeros'. | ||
| magnitude_initializer: | ||
| Initializer for magnitude vector. Defaults to 'ones'. | ||
| embeddings_regularizer: | ||
| Regularizer function applied to embeddings. | ||
| activity_regularizer: | ||
| Regularizer function applied to output. | ||
| embeddings_constraint: | ||
| Constraint function applied to embeddings. | ||
| mask_zero: | ||
| Whether input value 0 is a special "padding" value. | ||
| input_length: | ||
| Length of input sequences (for compatibility). | ||
| sparse: | ||
| Whether to use sparse embedding lookup (experimental). | ||
| **kwargs: Additional keyword arguments. | ||
| """ |
There was a problem hiding this comment.
| """def call(self, inputs, start_index=0): | ||
| Forward pass of DoRA position embedding. | ||
| Args: | ||
| inputs: Input tensor (token embeddings) | ||
| of shape [batch_size, seq_len, hidden_dim]. | ||
| start_index: Starting position index | ||
| (for compatibility with KerasHub). | ||
| Returns: | ||
| Position embeddings of shape [batch_size, seq_len, hidden_dim]. | ||
| input_shape = ops.shape(inputs) | ||
| seq_len = input_shape[-2] | ||
| # Get effective position embeddings using DoRA | ||
| effective_pos_embeddings = self._get_effective_position_embeddings() | ||
| # Create position indices using backend-agnostic operations | ||
| start_tensor = ops.convert_to_tensor(start_index, dtype="int32") | ||
| seq_len_tensor = ops.convert_to_tensor(seq_len, dtype="int32") | ||
| end_index = ops.add(start_tensor, seq_len_tensor) | ||
| positions = ops.arange(start_index, end_index, dtype="int32") | ||
| # Clip positions to valid range | ||
| max_pos = ops.convert_to_tensor( | ||
| self.sequence_length - 1, dtype="int32" | ||
| ) | ||
| min_pos = ops.convert_to_tensor(0, dtype="int32") | ||
| positions = ops.clip(positions, min_pos, max_pos) | ||
| # Lookup position embeddings | ||
| position_embeddings = ops.take( | ||
| effective_pos_embeddings, positions, axis=0 | ||
| ) | ||
| # Expand dimensions to match input batch size | ||
| position_embeddings = ops.expand_dims(position_embeddings, axis=0) | ||
| # Create target shape for broadcasting | ||
| batch_size = input_shape[0] | ||
| target_shape = [batch_size, seq_len, self.output_dim] | ||
| position_embeddings = ops.broadcast_to( | ||
| position_embeddings, target_shape | ||
| ) | ||
| return position_embeddings""" |
There was a problem hiding this comment.
This commented-out block of code appears to be a remnant of a previous implementation. It should be removed to improve code clarity and reduce clutter.
2 participants
Description of the change
This PR adds DoRA (Weight-Decomposed Low-Rank Adaptation) support to Keras Hub by introducing new reusable layers:
DoRADense: Drop-in replacement for Dense with DoRA decomposition.
DoRAEmbedding: Drop-in replacement for Embedding with DoRA adaptation.
DoRAPositionEmbedding: DoRA-based positional embeddings.
Conversion utilities (convert_dense_to_dora, convert_embedding_to_dora) for upgrading existing layers.
To demonstrate usage, we integrated DoRA into the BertBackbone (enable_dora=True) as an example of how existing models can adopt DoRA with minimal changes.
Reference
Paper: https://arxiv.org/abs/2402.09353
Checklist