Conversation
Summary of ChangesHello, 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 initiates a significant refactoring effort aimed at unifying the CausalLM classes within the project. By introducing a common base class and consolidating shared functionalities, the changes reduce code duplication, improve maintainability, and standardize the implementation of various language models. This is an initial step in a broader refactor, laying the groundwork for a more streamlined and consistent model architecture. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe 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 by creating a comment using either
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 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
|
![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 introduces a significant and valuable refactoring by creating a BaseForCausalLM class to unify the various Causal Language Model implementations. This change greatly reduces code duplication and improves maintainability. The common logic for forward passes, KV cache creation, and module specification is now centralized in the base class. The refactoring has been applied consistently across numerous models, with appropriate overrides for models that have unique naming conventions for their components. Overall, this is a solid architectural improvement. I have one suggestion to make the new base class even more robust.
Extract only the truly shared methods (embed, prefill, decode, batch_forward, etc.) into a base class. Model-specific logic stays in each model file.
babusid
force-pushed
the
causal_lm_refactor
branch
from
f865f31 to
8ae622a
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request is a great refactoring effort to unify the CausalLM classes by introducing a CausalLMABC base class. This significantly reduces code duplication across various model implementations and improves maintainability. The changes are consistent and well-structured, and I appreciate the cleanups like removing duplicated attribute assignments.
I've found one critical issue regarding the initialization of the dtype attribute which would cause a runtime error. Please see the detailed comment.
| dtype: str | ||
| hidden_size: int | ||
| tensor_parallel_shards: int | ||
|
|
There was a problem hiding this comment.
The dtype attribute is used in create_paged_kv_cache and get_default_spec in subclasses, but it is not initialized in this base class. The original model implementations set self.dtype = "float32" in their __init__ methods, but this line was removed during refactoring and not added to the CausalLMABC. This will lead to an AttributeError if methods like get_default_spec are called before to(dtype=...) has been invoked.
Please initialize self.dtype in the base class __init__ to provide a default value.
| def __init__(self): | |
| super().__init__() | |
| self.dtype = "float32" | |
1 participant
This refactor might be a bit tougher to do fully correctly, but the idea here is that we have a good amount of duplication across the CausalLM classes that we can reduce. I've implemented the more basic / easier ones so far, but some of the models will be trickier I think.
Not yet ready for merge, but I've marked it non-draft for gemini review