fix: gemma3 patch vocab_size not found#2741
Conversation
WalkthroughA conditional check in the code was updated to compare the image token index against the vocabulary size from the text configuration, rather than the model's general vocabulary size. No changes were made to public interfaces or function signatures. Changes
Poem
✨ Finishing Touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/axolotl/monkeypatch/gemma3.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (8)
- GitHub Check: PyTest from Source Dist (3.11, 2.5.1)
- GitHub Check: PyTest from Source Dist (3.11, 2.7.0)
- GitHub Check: PyTest from Source Dist (3.11, 2.6.0)
- GitHub Check: pre-commit
- GitHub Check: PyTest (3.11, 2.5.1)
- GitHub Check: PyTest (3.11, 2.7.0)
- GitHub Check: PyTest (3.11, 2.6.0)
- GitHub Check: pre-commit
| if ( | ||
| input_ids is not None | ||
| and self.config.image_token_index >= self.config.text_config.vocab_size | ||
| ): |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Fix for AttributeError looks correct and consistent.
The change from self.vocab_size to self.config.text_config.vocab_size properly resolves the "vocab_size not found" issue mentioned in the PR. This change is consistent with:
- The function's docstring (lines 41-42) which references
config.text_config.vocab_size - Other usage in the same function (lines 203, 211) where
self.config.text_config.vocab_sizeis used for the loss function - Standard transformers library structure where vocabulary size is typically accessed through the text configuration
To ensure this fix works as expected, please verify that the patched model handles image tokens correctly:
🏁 Script executed:
#!/bin/bash
# Description: Search for other usages of vocab_size in gemma3 related files to ensure consistency
# Search for vocab_size usage patterns in gemma3 files
rg -A 3 -B 3 "vocab_size" --type py | grep -i gemma
# Search for text_config.vocab_size usage patterns
rg -A 2 -B 2 "text_config\.vocab_size" --type pyLength of output: 12135
Ensure consistent use of config.text_config.vocab_size across all Gemma3 patches.
The fix in src/axolotl/monkeypatch/gemma3.py is correct and aligns with other usages. However, the integration patch under src/axolotl/integrations/cut_cross_entropy/monkeypatch/gemma3.py still references self.vocab_size in the image‐token check. Update it as well:
Locations to update:
- src/axolotl/integrations/cut_cross_entropy/monkeypatch/gemma3.py
- if input_ids is not None and self.config.image_token_index >= self.vocab_size: + if input_ids is not None and self.config.image_token_index >= self.config.text_config.vocab_size:
After applying, verify end‐to‐end that image tokens are handled without index errors.
Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In src/axolotl/integrations/cut_cross_entropy/monkeypatch/gemma3.py, locate the
image-token index check where self.vocab_size is referenced and update it to use
self.config.text_config.vocab_size instead. This change ensures consistency with
the fix applied in src/axolotl/monkeypatch/gemma3.py and aligns with the
standard way vocabulary size is accessed. After making this update, run
end-to-end tests to confirm image tokens are handled correctly without index
errors.
Codecov ReportAttention: Patch coverage is
📢 Thoughts on this report? Let us know! |
|
looks like it may not be needed anymore. I've removed this patch in the #2735 PR and the loss is still stable for the test |
|
Thanks! |
2 participants
Description
Fixes: https://github.com/axolotl-ai-cloud/axolotl/actions/runs/15310334187/job/43075132418
Motivation and Context
How has this been tested?
Screenshots (if appropriate)
Types of changes
Social Handles (Optional)
Summary by CodeRabbit