Add neo chat #1161
Add neo chat #1161
Conversation
Summary of ChangesHello @SangChengC, 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 significantly extends LightLLM's capabilities by integrating the Neo Chat multimodal model. The changes encompass modifications to the KV cache management, the introduction of a specialized inference state, and the implementation of custom transformer layers and weight loading logic tailored for Neo Chat. A key feature is the advanced image processing pipeline, which includes 2D Rotary Position Embeddings and optimized Triton kernels, allowing the model to efficiently process and understand visual information alongside text. 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 support for the neo_chat multimodal model, adding new model definitions, inference logic, custom Triton kernels, and vision processing components. The changes are extensive and well-structured for integrating the new model.
My review highlights several areas for improvement to enhance code quality and robustness. I've pointed out a number of debug print statements that should be removed from the production code. There is a bare except block that could hide potential errors and should be made more specific. Additionally, a local import within a function should be moved to the top level of the file.
More critically, I've identified a comment in the tokenizer logic that indicates uncertainty about a calculation, which could point to a potential bug. I've also found that the visual model's data type is hardcoded, ignoring the configured data_type, which could lead to issues. Addressing these points will improve the correctness and maintainability of the new model's implementation.
| try: | ||
| ntk_alpha = float(os.environ.get("LIGHTLLM_NTK_ALPHA", 1)) | ||
| assert ntk_alpha >= 1 | ||
| if ntk_alpha > 1: | ||
| logger.info(f"Note: NTK enabled, alpha set to {ntk_alpha}") | ||
| max_seq_len *= ntk_alpha | ||
| base = base * (ntk_alpha ** (partial_head_dim / (partial_head_dim - 2))) # Base change formula | ||
| except: | ||
| pass |
There was a problem hiding this comment.
Using a bare except: is generally discouraged as it can catch and silence a wide range of unexpected errors, making debugging difficult. It's better to catch specific exceptions that you expect might occur, such as ValueError or AssertionError, and log them for better diagnostics.
| try: | |
| ntk_alpha = float(os.environ.get("LIGHTLLM_NTK_ALPHA", 1)) | |
| assert ntk_alpha >= 1 | |
| if ntk_alpha > 1: | |
| logger.info(f"Note: NTK enabled, alpha set to {ntk_alpha}") | |
| max_seq_len *= ntk_alpha | |
| base = base * (ntk_alpha ** (partial_head_dim / (partial_head_dim - 2))) # Base change formula | |
| except: | |
| pass | |
| try: | |
| ntk_alpha = float(os.environ.get("LIGHTLLM_NTK_ALPHA", 1)) | |
| assert ntk_alpha >= 1 | |
| if ntk_alpha > 1: | |
| logger.info(f"Note: NTK enabled, alpha set to {ntk_alpha}") | |
| max_seq_len *= ntk_alpha | |
| base = base * (ntk_alpha ** (partial_head_dim / (partial_head_dim - 2))) # Base change formula | |
| except (ValueError, AssertionError) as e: | |
| logger.warning(f"Could not apply NTK scaling: {e}") |
| ) | ||
| grid_h, grid_w = resized_height // self.patch_size, resized_width // self.patch_size | ||
| token_num = int((grid_h * grid_w) * (self.downsample_ratio ** 2)) | ||
| # 这里的grid_h和grid_w需要* self.downsample_ratio么？再仔细看下代码 |
There was a problem hiding this comment.
This comment indicates uncertainty about the correctness of the calculation for grid_h and grid_w. This should be investigated and clarified to prevent potential bugs. If the logic is correct, the comment should be removed. If it's incorrect, it needs to be fixed.
| elif self.model_type == "gemma3": | ||
| self.model = Gemma3VisionModel() | ||
| elif self.model_type == "neo_chat": | ||
| self.model = NeoVisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval().bfloat16() |
There was a problem hiding this comment.
The model's data type is hardcoded to bfloat16, which ignores the data_type provided in kvargs. This can lead to unexpected behavior or errors if a different data type like float16 is intended. The model should be cast to the data type from the configuration to ensure correctness.
| self.model = NeoVisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval().bfloat16() | |
| self.model = NeoVisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval() | |
| if self.data_type in ["bf16", "bfloat16"]: | |
| self.model.bfloat16() | |
| elif self.data_type in ["fp16", "float16"]: | |
| self.model.half() |
|
|
||
| base = self.config.get("rope_theta", float(default_base)) | ||
|
|
||
| print(f"base is {base}") |
There was a problem hiding this comment.
This print statement appears to be for debugging purposes and should be removed before merging.
| rope_scaling_factor = self.config.get("rope_scaling", {}).get("factor", 1.0) | ||
|
|
||
| base = self.config.get("rope_theta_hw", float(default_base)) | ||
| print(f"hw_base is {base}") |
There was a problem hiding this comment.
This print statement appears to be for debugging purposes and should be removed before merging.
| a = img["start_idx"] | ||
| print(f"img start_idx: {a}") |
There was a problem hiding this comment.
These lines, including a print statement, appear to be for debugging and should be removed.
| from lightllm.models.llama.triton_kernel.token_attention_softmax_and_reducev import ( | ||
| token_softmax_reducev_fwd, | ||
| ) |
There was a problem hiding this comment.
Imports should be at the top of the file for better readability, to follow standard Python conventions, and to avoid the overhead of re-importing on every function call. Please move this import to the top of the file.
|
|
||
| # must devide merge_length | ||
| cur_num = int(img_tensors[-1].shape[0] * (self.downsample_ratio ** 2)) | ||
| print(f"cur_num is {cur_num}") |
There was a problem hiding this comment.
This print statement appears to be for debugging and should be removed.
| ) | ||
| pixel_values, grid_hw = preprocess_pixel_values(transform(new_image).to(torch.float32), patch_size=patch_size) | ||
|
|
||
| print(f"Transfer image_size from ({image.height, image.width}) to ({new_image.height, new_image.width})") |
There was a problem hiding this comment.
This print statement seems to be for debugging purposes. It should be removed or replaced with a proper logger call if this information is valuable for diagnostics.
No description provided.