[WIP] Add tinker backend. #448
Conversation
Summary of ChangesHello @chenyushuo, 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 a comprehensive integration of a new 'tinker' backend into the system. It extends the framework's capabilities to support Tinker models for both inference and training, complete with dedicated configuration, model handling, and training wrappers. The changes also involve significant refactoring to enable asynchronous operations and ensure proper synchronization and checkpointing mechanisms are in place for the new backend. 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. 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 two major changes. First, it refactors the module registration mechanism, moving from a decorator-based approach to a centralized default_mapping in __init__.py files. This is a significant improvement for code clarity and maintainability. Second, it adds a new "tinker" backend, which is a substantial piece of work but appears to be in a work-in-progress state. My review focuses on the new tinker backend implementation and some general observations. I've identified a critical issue in the new tinker trainer that will cause a crash, along with a few other areas for improvement.
| batch, kl_metrics = self.kl_fn.apply_kl_penalty_to_reward(batch) | ||
| metrics.update(prefix_metrics(kl_metrics, prefix="critic")) | ||
| # compute advantages, executed on the driver process | ||
| batch, _ = self.advantage_fn(batch) |
There was a problem hiding this comment.
There's a critical issue here. The batch variable is a list of tinker.types.Datum, but self.kl_fn.apply_kl_penalty_to_reward and self.advantage_fn expect a verl.DataProto object. This type mismatch will cause a runtime crash. The TODO comment also indicates this is known. This part of the implementation needs to be adapted to handle the tinker data format correctly.
| self.logger.warning( | ||
| f"Prompt was truncated to {self.config.max_prompt_tokens} tokens" | ||
| ) | ||
| token_ids = token_ids[: self.config.max_prompt_tokens + 1] # leave one for response |
There was a problem hiding this comment.
The prompt truncation logic seems a bit unusual. It truncates the token_ids to max_prompt_tokens + 1. If the intention is to limit the prompt to max_prompt_tokens, the slice should be [:self.config.max_prompt_tokens]. The current implementation takes one extra token, which might be unintentional and could lead to off-by-one errors or unexpected behavior in prompt processing. Could you clarify the intent here or correct it if it's a bug?
| print(supported_models) | ||
| raise ValueError(f"{model.tinker.base_model} is not supported by tinker!") |
There was a problem hiding this comment.
Using print() for error messages can make logs harder to parse and manage, especially in a distributed environment. It's better to use the configured logger, for instance logger.error(), to provide more context and allow for centralized log management. This also ensures the message is routed to the correct output stream and includes timestamps and other metadata.
| print(supported_models) | |
| raise ValueError(f"{model.tinker.base_model} is not supported by tinker!") | |
| logger.error(f"Supported models: {supported_models}") | |
| raise ValueError(f"{model.tinker.base_model} is not supported by tinker!") | |
1 participant
Description
As the title says.
Checklist
Please check the following items before code is ready to be reviewed.