Skip to content

feat(target_utils): auto-detect embed/lm_head keys for LLM and VLM checkpoints #666

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
curnane-lab wants to merge 3 commits into
base: main
Choose a base branch
from
+77 −14
Open

feat(target_utils): auto-detect embed/lm_head keys for LLM and VLM checkpoints #666

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d6b561f
fix(target_utils): auto-detect embed/lm_head weight keys
mingliangfu Jul 7, 2026
c062f3c
fix(target_utils): add VLM embed key variants for A3B-style models
mingliangfu Jul 7, 2026
6d9cf6e
fix(target_utils): use framework=pt and extract candidate keys as con…
mingliangfu Jul 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
125 changes: 116 additions & 9 deletions specforge/modeling/target/target_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,20 +98,61 @@ def _load_weights(
index = json.load(f)
weight_map = index.get("weight_map", {})

if embed_key in weight_map:
files_to_load[embed_key] = weight_map[embed_key]
else:
# Auto-detect the embedding key if the supplied/default key is missing.
# This handles plain LLMs (model.embed_tokens.weight), MLLMs
# (model.language_model.model.embed_tokens.weight), and renamed checkpoints.
candidate_embed_keys = [
embed_key,
"model.embed_tokens.weight",
"embed_tokens.weight",
"model.language_model.model.embed_tokens.weight",
"model.language_model.embed_tokens.weight",
"language_model.model.embed_tokens.weight",
"language_model.embed_tokens.weight",
"model.model.embed_tokens.weight",
]

@gemini-code-assist gemini-code-assist Bot Jul 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The list of candidate embedding keys is duplicated here and in the single-file checkpoint loading block (lines 174-183). To improve maintainability and avoid potential desynchronization in the future, consider defining these candidate keys as module-level constants (e.g., CANDIDATE_EMBED_KEYS and CANDIDATE_HEAD_KEYS) at the top of the file.

resolved_embed_key = None
for key in candidate_embed_keys:
if key in weight_map:
resolved_embed_key = key
break
if resolved_embed_key is None:
raise ValueError(
f"Embedding key '{embed_key}' not found in weight map."
f"Embedding key '{embed_key}' not found in weight map and no "
f"candidate embed key matched. Available keys (first 20): "
f"{list(weight_map.keys())[:20]}"
)
if resolved_embed_key != embed_key:
print(
f"Resolved embedding key '{embed_key}' -> '{resolved_embed_key}'"
)
files_to_load[resolved_embed_key] = weight_map[resolved_embed_key]
embed_key = resolved_embed_key

if not tie_weights:
if lm_head_key in weight_map:
files_to_load[lm_head_key] = weight_map[lm_head_key]
else:
candidate_head_keys = [
lm_head_key,
"lm_head.weight",
"model.lm_head.weight",
"model.language_model.lm_head.weight",
"language_model.lm_head.weight",
]
resolved_head_key = None
for key in candidate_head_keys:
if key in weight_map:
resolved_head_key = key
break
if resolved_head_key is None:
print(
f"Warning: {lm_head_key} not found. Ensure model doesn't use tied weights manually."
)
else:
if resolved_head_key != lm_head_key:
print(
f"Resolved lm_head key '{lm_head_key}' -> '{resolved_head_key}'"
)
files_to_load[resolved_head_key] = weight_map[resolved_head_key]
lm_head_key = resolved_head_key
else:
safetensors = glob.glob(os.path.join(model_path, "*.safetensors"))
bins = glob.glob(os.path.join(model_path, "*.bin"))
Expand All @@ -120,9 +161,75 @@ def _load_weights(
if not target_file:
raise FileNotFoundError("No checkpoint found.")

files_to_load[embed_key] = os.path.basename(target_file)
# Read the available keys so we can auto-detect embed/lm_head names
# in single-file checkpoints too.
if target_file.endswith(".safetensors"):
with safe_open(target_file, framework="np") as f:

@gemini-code-assist gemini-code-assist Bot Jul 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The safe_open function from the safetensors library does not support "np" as a valid framework. Passing framework="np" will raise a ValueError: Invalid framework np at runtime. Please use framework="pt" (or framework="numpy") instead.

Suggested change
with safe_open(target_file, framework="np") as f:
with safe_open(target_file, framework="pt") as f:

available_keys = set(f.keys())
else:
# For .bin files we fall back to the provided keys; auto-detection
# would require loading the whole state dict up front.
available_keys = None

candidate_embed_keys = [
embed_key,
"model.embed_tokens.weight",
"embed_tokens.weight",
"model.language_model.model.embed_tokens.weight",
"model.language_model.embed_tokens.weight",
"language_model.model.embed_tokens.weight",
"language_model.embed_tokens.weight",
"model.model.embed_tokens.weight",
]
resolved_embed_key = None
if available_keys is not None:
for key in candidate_embed_keys:
if key in available_keys:
resolved_embed_key = key
break
else:
resolved_embed_key = embed_key

if resolved_embed_key is None:
raise ValueError(
f"Embedding key '{embed_key}' not found in checkpoint and no "
f"candidate embed key matched."
)
if resolved_embed_key != embed_key:
print(
f"Resolved embedding key '{embed_key}' -> '{resolved_embed_key}'"
)
files_to_load[resolved_embed_key] = os.path.basename(target_file)
embed_key = resolved_embed_key

if not tie_weights:
files_to_load[lm_head_key] = os.path.basename(target_file)
candidate_head_keys = [
lm_head_key,
"lm_head.weight",
"model.lm_head.weight",
"model.language_model.lm_head.weight",
"language_model.lm_head.weight",
]
resolved_head_key = None
if available_keys is not None:
for key in candidate_head_keys:
if key in available_keys:
resolved_head_key = key
break
else:
resolved_head_key = lm_head_key

if resolved_head_key is None:
print(
f"Warning: {lm_head_key} not found. Ensure model doesn't use tied weights manually."
)
else:
if resolved_head_key != lm_head_key:
print(
f"Resolved lm_head key '{lm_head_key}' -> '{resolved_head_key}'"
)
files_to_load[resolved_head_key] = os.path.basename(target_file)
lm_head_key = resolved_head_key

loaded_keys = set()

Expand Down
Loading