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

[curnane-lab]

Motivation

TargetEmbeddingsAndHead.from_pretrained() currently hardcodes model.embed_tokens.weight and lm_head.weight as the keys to load from the target checkpoint. This breaks on multimodal models (e.g. Qwen3.5-4B / Qwen3.5-A3B VLM checkpoints) where the text decoder is nested under model.language_model or model.language_model.model, so the embedding key becomes model.language_model.model.embed_tokens.weight or model.language_model.embed_tokens.weight.

Without auto-detection, users must remember to pass --embedding-key and --lm-head-key for every VLM target, which is error-prone and makes example scripts model-specific.

Modifications

• specforge/modeling/target/target_utils.py:
  • Add candidate key lists for both embed_tokens and lm_head, covering:
    • Plain causal LMs: model.embed_tokens.weight, embed_tokens.weight
    • Qwen3.5-4B-style VLMs: model.language_model.model.embed_tokens.weight
    • Qwen3.5-A3B-style VLMs: model.language_model.embed_tokens.weight, language_model.*
  • Try candidates in order for both sharded checkpoints (model.safetensors.index.json) and single-file checkpoints (*.safetensors).
  • Print a message when a non-default key is resolved.
  • Improve error diagnostics by listing available keys when no candidate matches.

Related Issues

N/A

Accuracy Test

N/A (data-loading utility, no model math changes)

Benchmark & Profiling

N/A

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

[gemini-code-assist]

Code Review

This pull request introduces auto-detection for embedding and language model head keys in _load_weights to support various LLM/MLLM checkpoint structures. The feedback highlights a critical runtime error where an invalid framework argument ("np") is passed to safe_open, and suggests refactoring duplicated candidate key lists into module-level constants to improve maintainability.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

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:

[gemini-code-assist]

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.