saymrwulf/transformers
1
0
Fork 0
mirror of https://github.com/saymrwulf/transformers.git synced 2026-05-14 20:58:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix copies

Browse source
This commit is contained in:
Arthur Zucker 2025-01-23 16:40:33 +01:00
parent ef0b5e279c
commit 2d480eccc7
2 changed files with 18 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
src/transformers/models/gemma/configuration_gemma.py
View file

@ -93,6 +93,15 @@ class GemmaConfig(PretrainedConfig):
model_type = "gemma"
keys_to_ignore_at_inference = ["past_key_values"]
base_model_tp_plan = {
"layers.*.self_attn.q_proj": "colwise",
"layers.*.self_attn.k_proj": "colwise",
"layers.*.self_attn.v_proj": "colwise",
"layers.*.self_attn.o_proj": "rowwise",
"layers.*.mlp.gate_proj": "colwise",
"layers.*.mlp.up_proj": "colwise",
"layers.*.mlp.down_proj": "rowwise",
}
def __init__(
self,

9
src/transformers/models/olmo2/configuration_olmo2.py
View file

@ -89,6 +89,15 @@ class Olmo2Config(PretrainedConfig):
model_type = "olmo2"
keys_to_ignore_at_inference = ["past_key_values"]
base_model_tp_plan = {
"layers.*.self_attn.q_proj": "colwise",
"layers.*.self_attn.k_proj": "colwise",
"layers.*.self_attn.v_proj": "colwise",
"layers.*.self_attn.o_proj": "rowwise",
"layers.*.mlp.gate_proj": "colwise",
"layers.*.mlp.up_proj": "colwise",
"layers.*.mlp.down_proj": "rowwise",
}
def __init__(
self,