Skip to content

fix: MambaLMHeadModel and create_block robustness - #1008

Open
Chessing234 wants to merge 3 commits into
state-spaces:mainfrom
Chessing234:fix/mixer-seq-simple-robustness
Open

fix: MambaLMHeadModel and create_block robustness#1008
Chessing234 wants to merge 3 commits into
state-spaces:mainfrom
Chessing234:fix/mixer-seq-simple-robustness

Conversation

@Chessing234

Copy link
Copy Markdown
Contributor

Summary

  • Accept legacy ssm_cfg.layer=\"Mamba\" alias for Mamba-1 blocks and list supported layer names on error.
  • Expose strict= in from_pretrained for partial checkpoint loads.
  • Write save_pretrained artifacts atomically via temp files + os.replace.

Test plan

  • Manual review of create_block / from_pretrained / save_pretrained paths
  • GPU generation tests (unchanged behavior expected)

Made with Cursor

Configs using ssm_cfg.layer=\"Mamba\" now map to Mamba1, and invalid layer
names list all supported values.
Allow partial checkpoint loads for fine-tuned heads while keeping strict=True
as the default for full pretrained restores.
Use temporary files plus os.replace so interrupted saves do not leave
partial pytorch_model.bin or config.json behind.
@Chessing234
Chessing234 force-pushed the fix/mixer-seq-simple-robustness branch from 9350225 to 6489c23 Compare August 5, 2026 10:23
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.

1 participant