253 lines
9.8 KiB
Python
253 lines
9.8 KiB
Python
# ---------------------------------------------------------------------
|
||
# helper: ensure Gemma-3 checkpoints use language_model.model.… keys
|
||
# ---------------------------------------------------------------------
|
||
import os, json, logging
|
||
from pathlib import Path
|
||
from safetensors.torch import safe_open, save_file
|
||
|
||
log = logging.getLogger(__name__)
|
||
|
||
|
||
def unwrap_clippable_linears(model) -> int:
|
||
"""
|
||
Replace remote-code clippable linear wrappers with their underlying
|
||
torch.nn.Linear modules so vanilla PEFT can inject adapters.
|
||
|
||
Some Gemma loaders expose projection modules as Gemma*ClippableLinear
|
||
wrappers with the real Linear stored on `.linear`. PEFT's LoRA injection
|
||
only accepts the inner Linear type.
|
||
"""
|
||
count = 0
|
||
|
||
for child_name, child in list(model.named_children()):
|
||
count += unwrap_clippable_linears(child)
|
||
|
||
inner = getattr(child, "linear", None)
|
||
if inner is None:
|
||
continue
|
||
if child.__class__.__name__.endswith("ClippableLinear"):
|
||
setattr(model, child_name, inner)
|
||
count += 1
|
||
|
||
return count
|
||
|
||
def fix_gemma3_checkpoint(ckpt_dir: str | Path) -> None:
|
||
"""
|
||
If `ckpt_dir` is a Gemma-3 checkpoint whose tensor keys look like
|
||
language_model.embed_tokens.weight
|
||
instead of
|
||
language_model.model.embed_tokens.weight
|
||
rewrite the shards and index file in-place.
|
||
|
||
No-op when:
|
||
• model_type ≠ 'gemma3'
|
||
• keys are already correct
|
||
• required files are missing
|
||
"""
|
||
ckpt_dir = Path(ckpt_dir)
|
||
index_file = ckpt_dir / "model.safetensors.index.json"
|
||
config_file = ckpt_dir / "config.json"
|
||
if not index_file.is_file() or not config_file.is_file():
|
||
return # nothing to do
|
||
|
||
# ── guard: only patch Gemma-3 checkpoints ───────────────────────────
|
||
try:
|
||
with open(config_file) as f:
|
||
cfg = json.load(f)
|
||
if (cfg.get("model_type") or "").lower() != "gemma3":
|
||
return
|
||
except Exception as e:
|
||
log.warning("Could not read %s (%s); skipping fix.", config_file, e)
|
||
return
|
||
|
||
# ── scan weight map ─────────────────────────────────────────────────
|
||
with open(index_file) as f:
|
||
idx = json.load(f)
|
||
|
||
wm = idx["weight_map"]
|
||
broken = [
|
||
k for k in wm
|
||
if k.startswith("language_model.") and not k.startswith("language_model.model.")
|
||
]
|
||
if not broken:
|
||
return # already fine
|
||
|
||
log.info("Repairing Gemma-3 key prefixes in %s", ckpt_dir)
|
||
|
||
def _fixed(k: str) -> str:
|
||
if k.startswith("language_model.") and not k.startswith("language_model.model."):
|
||
return "language_model.model." + k[len("language_model."):]
|
||
return k
|
||
|
||
# ── rewrite every shard exactly once ────────────────────────────────
|
||
repaired_shards = set()
|
||
for old_key, shard_name in wm.items():
|
||
wm[_fixed(old_key)] = wm.pop(old_key) # update key in dict
|
||
if shard_name in repaired_shards:
|
||
continue
|
||
repaired_shards.add(shard_name)
|
||
|
||
src = ckpt_dir / shard_name
|
||
tmp = ckpt_dir / (shard_name + ".tmp")
|
||
|
||
fixed_tensors = {}
|
||
with safe_open(src, framework="pt", device="cpu") as f:
|
||
for k in f.keys():
|
||
fixed_tensors[_fixed(k)] = f.get_tensor(k)
|
||
|
||
save_file(fixed_tensors, tmp, metadata={"format": "pt"})
|
||
tmp.replace(src) # atomic overwrite
|
||
|
||
# ── write new index ────────────────────────────────────────────────
|
||
with open(index_file, "w") as f:
|
||
json.dump(idx, f, indent=2)
|
||
|
||
log.info("✓ Gemma-3 checkpoint repaired.")
|
||
|
||
|
||
# fully detie lm_head from embeddings so safetensors can flatten
|
||
def detie_lm_head(model):
|
||
"""
|
||
Untie the logits projection from the input embeddings and register it
|
||
*exactly* where the model (and loaders like vLLM) expect it.
|
||
|
||
Works with HF models whose output head is either `lm_head` or some
|
||
nested attribute (e.g. `language_model.output_projection` in Gemma-3).
|
||
"""
|
||
import torch
|
||
from types import SimpleNamespace
|
||
|
||
emb = model.get_input_embeddings() # nn.Embedding
|
||
old_head = model.get_output_embeddings() # whatever Linear HF exposes
|
||
|
||
# nothing to do if they are already separate tensors
|
||
if old_head.weight.data_ptr() != emb.weight.data_ptr():
|
||
return
|
||
|
||
vocab_size, hidden_size = emb.weight.shape
|
||
new_head = torch.nn.Linear(hidden_size, vocab_size, bias=False)
|
||
new_head.weight = torch.nn.Parameter(emb.weight.detach().clone())
|
||
new_head.to(next(model.parameters()).dtype)
|
||
|
||
# ------------------------------------------------------------------
|
||
# find the *attribute path* of the existing output head
|
||
# ------------------------------------------------------------------
|
||
path = None
|
||
for name, module in model.named_modules():
|
||
if module is old_head:
|
||
path = name # e.g. "lm_head" or "language_model.output_projection"
|
||
break
|
||
if path is None: # very unusual, but fall back to "lm_head"
|
||
path = "lm_head"
|
||
|
||
print('!!', name)
|
||
|
||
# ------------------------------------------------------------------
|
||
# install the new head at that path
|
||
# ------------------------------------------------------------------
|
||
def set_by_path(root, dotted_name, value):
|
||
parts = dotted_name.split(".")
|
||
parent = root
|
||
for p in parts[:-1]:
|
||
parent = getattr(parent, p)
|
||
setattr(parent, parts[-1], value)
|
||
|
||
set_by_path(model, path, new_head)
|
||
|
||
# HF convenience: if the public attribute `lm_head` *is not* the main path,
|
||
# mirror it so code expecting `model.lm_head` still works. This does *not*
|
||
# duplicate weights – both names reference the same nn.Linear instance.
|
||
#if path != "lm_head":
|
||
# model.lm_head = new_head
|
||
|
||
model.config.tie_word_embeddings = False
|
||
|
||
|
||
# --------------------------------------------------------------------
|
||
# restore Gemma-3’s weight-tying and ensure only the embed_tokens
|
||
# key lands in the safetensors index (no lm_head, no duplication)
|
||
# --------------------------------------------------------------------
|
||
# ---------------------------------------------------------------
|
||
# Gemma-3: keep weight-tying *and* give vLLM the path it wants
|
||
# ---------------------------------------------------------------
|
||
def retie_gemma3_and_prune_alias(model):
|
||
"""
|
||
Re-establish tying between embeddings and logits projection and ensure
|
||
the projection is reachable at `language_model.output_projection`.
|
||
Removes the top-level `lm_head` alias so the serializer never emits
|
||
an `lm_head.*` key.
|
||
|
||
Call just before `save_pretrained(...)`.
|
||
"""
|
||
import torch.nn as nn
|
||
|
||
if (getattr(model.config, "model_type", "") or "").lower() != "gemma3":
|
||
return # skip for anything that isn't Gemma-3
|
||
|
||
emb = model.get_input_embeddings() # nn.Embedding
|
||
proj = getattr(model, "lm_head", None) # HF always defines this
|
||
|
||
if proj is None or not isinstance(proj, nn.Linear):
|
||
raise RuntimeError("Could not find lm_head on Gemma-3 model")
|
||
|
||
# ── tie weights if they were detied earlier ─────────────────────────
|
||
if proj.weight.data_ptr() != emb.weight.data_ptr():
|
||
proj.weight = emb.weight # share storage again
|
||
model.config.tie_word_embeddings = True
|
||
|
||
# ── ensure wrapper + attribute for vLLM ────────────────────────────
|
||
# 1. make / fetch `model.language_model`
|
||
if not hasattr(model, "language_model"):
|
||
wrapper = nn.Module()
|
||
model.add_module("language_model", wrapper)
|
||
else:
|
||
wrapper = model.language_model
|
||
|
||
# 2. register projection inside wrapper
|
||
wrapper.add_module("output_projection", proj)
|
||
|
||
# ── drop the top-level alias so it won't be serialised ─────────────
|
||
if hasattr(model, "lm_head"):
|
||
delattr(model, "lm_head")
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# Gemma-3 helper: detie + relabel head for vLLM + safetensors
|
||
# ------------------------------------------------------------------
|
||
def prepare_gemma3_for_save(model):
|
||
"""
|
||
• Makes the output projection an independent tensor if it still shares
|
||
storage with the embeddings.
|
||
• Registers it at `language_model.lm_head` (the path vLLM uses).
|
||
• Deletes the top-level `lm_head` alias so no `lm_head.*` key is saved.
|
||
• Sets `tie_word_embeddings=False` so Transformers knows they’re untied.
|
||
"""
|
||
import torch.nn as nn, torch
|
||
|
||
if (getattr(model.config, "model_type", "") or "").lower() != "gemma3":
|
||
return
|
||
|
||
emb = model.get_input_embeddings()
|
||
head = model.get_output_embeddings() # this is model.lm_head
|
||
|
||
# 1. Detie if they still share storage
|
||
if head.weight.data_ptr() == emb.weight.data_ptr():
|
||
vocab, hidden = emb.weight.shape
|
||
new_head = nn.Linear(hidden, vocab, bias=False)
|
||
new_head.weight = nn.Parameter(emb.weight.detach().clone())
|
||
new_head.to(next(model.parameters()).dtype)
|
||
head = new_head
|
||
|
||
# 2. Ensure `language_model` wrapper exists
|
||
if not hasattr(model, "language_model"):
|
||
model.add_module("language_model", nn.Module())
|
||
|
||
# 3. Register under vLLM path
|
||
#model.language_model.add_module("lm_head", head)
|
||
|
||
# 4. Drop the alias so no `lm_head.*` key lands in the state-dict
|
||
if hasattr(model, "lm_head"):
|
||
delattr(model, "lm_head")
|
||
|
||
model.config.tie_word_embeddings = False
|