feat: gemma4 support initial commit (WIP)#191
Conversation
![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 Google Gemma4 models, including a training cookbook for FSDP2 and a new template class. The review feedback identifies several improvement opportunities: a hardcoded Windows path in the cookbook should be replaced to ensure portability across environments, a redundant import of LatexOCRProcessor should be removed, and unused imports and type aliases in the new template file should be cleaned up. Additionally, an incomplete comment in the training script needs to be addressed.
| # MODEL_PATH = '/nas/disk1/Qwen3-30B-A3B' | ||
| # MODEL_PATH = '/nas/disk1/gemma-4-E2B-it' | ||
| # MODEL_PATH = '/nas/disk1/gemma-4-26B-A4B' | ||
| MODEL_PATH = r"C:\Users\wliuu\.cache\modelscope\hub\models\google\gemma-4-E2B-it" |
There was a problem hiding this comment.
The MODEL_PATH is hardcoded to a local Windows absolute path. This makes the script non-portable and will cause it to fail on other environments (e.g., Linux servers). Consider using a model ID from the hub (e.g., google/gemma-4-E2B-it) or a relative path.
| from twinkle.dataloader import DataLoader | ||
| from twinkle.dataset import Dataset, DatasetMeta | ||
| from twinkle.model import TransformersModel | ||
| from twinkle.preprocessor import SelfCognitionProcessor, LatexOCRProcessor |
There was a problem hiding this comment.
LatexOCRProcessor is imported here but then redefined on line 61. This is redundant and can lead to confusion. If the local definition is intended for this cookbook, please remove the import of LatexOCRProcessor from line 26.
| from twinkle.preprocessor import SelfCognitionProcessor, LatexOCRProcessor | |
| from twinkle.preprocessor import SelfCognitionProcessor |
| import json | ||
| import torch | ||
| import torch.nn.functional as F | ||
| from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Type, Union | ||
| from dataclasses import dataclass, field, fields | ||
| from PIL import Image | ||
| from copy import deepcopy | ||
|
|
||
| from twinkle import remote_class, requires | ||
| from twinkle.template import Template | ||
|
|
||
| Tool = Dict[str, Union[str, Dict]] | ||
| History = List[Union[Tuple[str, str], List[str]]] | ||
| Message = Dict[str, Union[str, List[Dict[str, Any]], List[int], None]] | ||
| Messages = List[Message] | ||
| Prompt = List[Union[str, List[int], List[str]]] | ||
| Word = Union[str, List[int]] | ||
| Context = Word |
There was a problem hiding this comment.
These imports and type aliases are unused in the current implementation of Gemma4Template. They should be removed to maintain code cleanliness and reduce overhead.
|
|
||
| if step % 10 == 0: | ||
| # Print metric | ||
| metric = model.calculate_metric(is_training=True) # 若npu显存不足, |
There was a problem hiding this comment.
The comment # 若npu显存不足, (If NPU memory is insufficient...) is incomplete. Please complete the explanation or remove the trailing comment if it is no longer necessary.
1 participant
PR type
PR information
Write the detail information belongs to this PR.
Experiment results
Paste your experiment result here(if needed).