diff --git a/tests/saving/test_qwen3_5_vlm_full_finetune_key_remap.py b/tests/saving/test_qwen3_5_vlm_full_finetune_key_remap.py new file mode 100644 index 0000000000..ba3b6874a4 --- /dev/null +++ b/tests/saving/test_qwen3_5_vlm_full_finetune_key_remap.py @@ -0,0 +1,71 @@ +import ast +import types +from pathlib import Path + +import torch + + +def _load_qwen3_5_vlm_save_helpers(): + source = Path(__file__).parents[2] / "unsloth" / "save.py" + tree = ast.parse(source.read_text(encoding = "utf-8")) + helpers = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) + and node.name + in { + "_is_qwen3_5_vlm", + "_qwen3_5_vlm_state_dict_for_save", + } + ] + module = ast.Module(body = helpers, type_ignores = []) + ast.fix_missing_locations(module) + namespace = {} + exec(compile(module, str(source), "exec"), namespace) + return namespace + + +def _qwen3_5_vlm_model(): + return types.SimpleNamespace( + config = types.SimpleNamespace( + architectures = ["Qwen3_5ForConditionalGeneration"], + model_type = "qwen3_5", + vision_config = types.SimpleNamespace(), + ) + ) + + +def test_qwen3_5_vlm_state_dict_uses_hf_checkpoint_namespace(): + helpers = _load_qwen3_5_vlm_save_helpers() + state_dict = { + "language_model.model.embed_tokens.weight": torch.ones(2, 2), + "language_model.model.layers.0.input_layernorm.weight": torch.ones(2), + "language_model.lm_head.weight": torch.ones(2, 2), + "visual.blocks.0.norm1.weight": torch.ones(2), + "other.weight": torch.ones(2), + } + + remapped = helpers["_qwen3_5_vlm_state_dict_for_save"](state_dict) + + assert "model.language_model.embed_tokens.weight" in remapped + assert "model.language_model.layers.0.input_layernorm.weight" in remapped + assert "lm_head.weight" in remapped + assert "model.visual.blocks.0.norm1.weight" in remapped + assert "other.weight" in remapped + assert "language_model.model.embed_tokens.weight" not in remapped + assert "language_model.lm_head.weight" not in remapped + assert "visual.blocks.0.norm1.weight" not in remapped + + +def test_qwen3_5_vlm_detection_requires_vision_config(): + helpers = _load_qwen3_5_vlm_save_helpers() + assert helpers["_is_qwen3_5_vlm"](_qwen3_5_vlm_model()) + + model = types.SimpleNamespace( + config = types.SimpleNamespace( + architectures = ["Qwen3_5ForCausalLM"], + model_type = "qwen3_5_text", + ) + ) + + assert not helpers["_is_qwen3_5_vlm"](model) diff --git a/tests/version_compat/test_trl_grpo_pinned_symbols.py b/tests/version_compat/test_trl_grpo_pinned_symbols.py index e49f9292e0..ce5e2644fd 100644 --- a/tests/version_compat/test_trl_grpo_pinned_symbols.py +++ b/tests/version_compat/test_trl_grpo_pinned_symbols.py @@ -560,6 +560,7 @@ def test_trl_kto_get_batch_logps_signature(tag: str): if src is None: continue checked_sources.append((path, src)) + # Legacy: explicit get_batch_logps method. if has_def(src, "get_batch_logps", "func"): return # TRL 1.x: refactored into _compute_logps + selective_log_softmax. diff --git a/unsloth/save.py b/unsloth/save.py index d6f69e6a6c..a44f186e40 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -457,6 +457,36 @@ def _preserve_tokenizer_eos_token( ) +def _is_qwen3_5_vlm(model): + config = getattr(model, "config", None) + if config is None or not hasattr(config, "vision_config"): + return False + architectures = getattr(config, "architectures", None) or () + return any( + architecture + in ( + "Qwen3_5ForConditionalGeneration", + "Qwen3_5MoeForConditionalGeneration", + ) + for architecture in architectures + ) or getattr(config, "model_type", None) in ("qwen3_5", "qwen3_5_moe") + + +def _qwen3_5_vlm_state_dict_for_save(state_dict): + remapped_state_dict = {} + for key, value in state_dict.items(): + if key.startswith("language_model.model."): + new_key = "model.language_model." + key[len("language_model.model.") :] + elif key.startswith("visual."): + new_key = "model.visual." + key[len("visual.") :] + elif key.startswith("language_model.lm_head."): + new_key = "lm_head." + key[len("language_model.lm_head.") :] + else: + new_key = key + remapped_state_dict[new_key] = value + return remapped_state_dict + + @torch.inference_mode def unsloth_save_model( model, @@ -2983,18 +3013,24 @@ def unsloth_generic_save( if not is_main_process: return - # Honor merged_16bit by casting to the target dtype if needed _save_kwargs = dict( safe_serialization = safe_serialization, max_shard_size = max_shard_size, variant = variant, ) + is_qwen3_5_vlm = _is_qwen3_5_vlm(model) + if ("16bit" in save_method or is_qwen3_5_vlm) and state_dict is None: + state_dict = model.state_dict() if "16bit" in save_method: _target_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 - _save_kwargs["state_dict"] = { + state_dict = { k: v.to(dtype = _target_dtype) if v.is_floating_point() else v - for k, v in model.state_dict().items() + for k, v in state_dict.items() } + if is_qwen3_5_vlm: + state_dict = _qwen3_5_vlm_state_dict_for_save(state_dict) + if state_dict is not None: + _save_kwargs["state_dict"] = state_dict if push_to_hub: print(f"Unsloth: Pushing full fine-tuned model to '{save_directory}' ...")