Skip to content

Conversation

@davidkoski
Copy link
Collaborator

No description provided.

newKey = newKey.replacingOccurrences(of: "lm_head", with: "language_model.lm_head")
guard !config.textConfiguration.tieWordEmbeddings else { continue }
adjusted[newKey] = value
continue
Copy link
Collaborator Author

@davidkoski davidkoski Oct 28, 2025

Choose a reason for hiding this comment

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

Make the sanitize match the python version, see #431. It was converting

(lldb) po weights.keys.sorted()
▿ 1258 elements
  - 0 : "language_model.lm_head.biases"
  - 1 : "language_model.lm_head.scales"
  - 2 : "language_model.lm_head.weight"

to

(lldb) po weights.keys.sorted()
▿ 1258 elements
  - 0 : "language_language_model.lm_head.biases"
  - 1 : "language_language_model.lm_head.scales"
  - 2 : "language_language_model.lm_head.weight"

This omits the logic of guard !config.textConfiguration.tieWordEmbeddings else { continue } (these keys are unused if tieWordEmbeddings is present). If we need that I think it should be re-inserted right before line 1725 (adjusted[newKey] = value)

@DePasqualeOrg @rudrankriyam

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

I think the mlx_vlm code doesn't end up caring if the lm_head keys are consumed so we do need that logic.

@davidkoski davidkoski requested a review from awni October 28, 2025 16:07
@rudrankriyam
Copy link
Contributor

Just went through this; LGTM! Apologies for the error in the first place

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants