Fix target head weight load and refactor checkpoint loading

[cih9088]

Motivation

When loading the weights for TargetHead, the logic looks for model.safetensors.index.json and raises an error if it is not found. However model.safetensors.index.json exits only if the model weights are shards. (index creation, index saving in transformers).

Some of checkpoint loading logics implemented this feature and they are scattered around. I factored them out to a single utility function

Modifications

• added unified checkpoint load function
• replaced the checkpoint load logic with the unified load function

Related Issues

Closes #927

Accuracy Test

Benchmark & Profiling

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.

[lru49]

@cih9088 thank you for this PR. I filed issue #927 because even with your fix, Qwen2.5-0.5B still fails because from_pretrained() is looking for lm_head weights but the model is using tied embeddings. It's such a small fix to have from_pretrained() reference the embedding weights instead of lm_head for tied embedding models, that I thought you might want to integrate it into this PR. I implemented the fix & confirmed that the test suite passes on my branch (except for 4 tests that I cannot run on my single GPU setup), you can see the fix here:

cih9088/SpecForge@feature/robust-weight-load...lru49:SpecForge:fix/target-head-tied-embeddings

Tell me if you think this makes more sense as a separate PR.

edit: my changes have been merged into this PR branch, so the diff link above no longer shows any changes

[cih9088]

@lru49
Oh, I have missed that. I don't mind that you stack up additional commits on top of mine. I gave you a write permission on my fork temporarily. Please feel free to add commits of yours or open a separate PR.

[lru49]

Thanks @cih9088 ! I just merged my changes

[lru49]

@cih9088 Have you posted this PR in sglang slack for review? That could help get eyes on it

[cih9088]

@cih9088 Have you posted this PR in sglang slack for review? That could help get eyes on it

@lru49 no, I didn't join the slack. If you happened to be in the slack, be my guest.

I have no idea who to ping. Could you take a look at this PR? @curnane-lab @jiapingW

[curnane-lab]

For a tied checkpoint that still stores lm_head.weight, this prefers the stored lm_head while TargetEmbeddingsAndHead (and transformers, which re-ties after load) ignores it — worth picking one answer for both.