From 1bdc7e644597ede7e6043e98593b3e8b0078e9c2 Mon Sep 17 00:00:00 2001 From: Taksh Date: Fri, 31 Jul 2026 13:17:14 +0300 Subject: [PATCH 1/3] fix: accept legacy Mamba alias and clarify ssm_layer errors Configs using ssm_cfg.layer=\"Mamba\" now map to Mamba1, and invalid layer names list all supported values. --- mamba_ssm/models/mixer_seq_simple.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mamba_ssm/models/mixer_seq_simple.py b/mamba_ssm/models/mixer_seq_simple.py index b178f48f9..4e2b1bd2f 100644 --- a/mamba_ssm/models/mixer_seq_simple.py +++ b/mamba_ssm/models/mixer_seq_simple.py @@ -54,11 +54,15 @@ def create_block( ssm_layer = ssm_cfg.pop("layer", "Mamba1") ssm_layer_map = { "Mamba1": Mamba, + "Mamba": Mamba, "Mamba2": Mamba2, "Mamba3": Mamba3, } if ssm_layer not in ssm_layer_map: - raise ValueError(f"Invalid ssm_layer: {ssm_layer}, only support Mamba1, Mamba2, and Mamba3") + supported = ", ".join(sorted(ssm_layer_map)) + raise ValueError( + f"Invalid ssm_layer: {ssm_layer!r}. Supported values: {supported}." + ) mixer_cls = partial( ssm_layer_map[ssm_layer], layer_idx=layer_idx, From 5d1b107c0a69abc66b4d56436f2bdb7295fd7646 Mon Sep 17 00:00:00 2001 From: Taksh Date: Fri, 31 Jul 2026 13:17:16 +0300 Subject: [PATCH 2/3] fix: expose strict flag in MambaLMHeadModel.from_pretrained Allow partial checkpoint loads for fine-tuned heads while keeping strict=True as the default for full pretrained restores. --- mamba_ssm/models/mixer_seq_simple.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mamba_ssm/models/mixer_seq_simple.py b/mamba_ssm/models/mixer_seq_simple.py index 4e2b1bd2f..a1af0e3d0 100644 --- a/mamba_ssm/models/mixer_seq_simple.py +++ b/mamba_ssm/models/mixer_seq_simple.py @@ -294,11 +294,12 @@ def forward(self, input_ids, position_ids=None, inference_params=None, num_last_ return CausalLMOutput(logits=lm_logits) @classmethod - def from_pretrained(cls, pretrained_model_name, device=None, dtype=None, **kwargs): + def from_pretrained(cls, pretrained_model_name, device=None, dtype=None, strict=True, **kwargs): config_data = load_config_hf(pretrained_model_name) config = MambaConfig(**config_data) model = cls(config, device=device, dtype=dtype, **kwargs) - model.load_state_dict(load_state_dict_hf(pretrained_model_name, device=device, dtype=dtype)) + state_dict = load_state_dict_hf(pretrained_model_name, device=device, dtype=dtype) + model.load_state_dict(state_dict, strict=strict) return model def save_pretrained(self, save_directory): From 6489c23eaac9e58e4020cbe5db00dbd647c0dfa7 Mon Sep 17 00:00:00 2001 From: Taksh Date: Fri, 31 Jul 2026 13:17:22 +0300 Subject: [PATCH 3/3] fix: write save_pretrained artifacts atomically Use temporary files plus os.replace so interrupted saves do not leave partial pytorch_model.bin or config.json behind. --- mamba_ssm/models/mixer_seq_simple.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/mamba_ssm/models/mixer_seq_simple.py b/mamba_ssm/models/mixer_seq_simple.py index a1af0e3d0..3e50fa64f 100644 --- a/mamba_ssm/models/mixer_seq_simple.py +++ b/mamba_ssm/models/mixer_seq_simple.py @@ -312,9 +312,13 @@ def save_pretrained(self, save_directory): # Save the model's state_dict model_path = os.path.join(save_directory, 'pytorch_model.bin') - torch.save(self.state_dict(), model_path) + tmp_model_path = model_path + ".tmp" + torch.save(self.state_dict(), tmp_model_path) + os.replace(tmp_model_path, model_path) # Save the configuration of the model config_path = os.path.join(save_directory, 'config.json') - with open(config_path, 'w') as f: + tmp_config_path = config_path + ".tmp" + with open(tmp_config_path, 'w', encoding='utf-8') as f: json.dump(self.config.__dict__, f, indent=4) + os.replace(tmp_config_path, config_path)