diff --git a/tests/python/test_fast_model_config_passthrough.py b/tests/python/test_fast_model_config_passthrough.py new file mode 100644 index 0000000000..b2ba3d2eef --- /dev/null +++ b/tests/python/test_fast_model_config_passthrough.py @@ -0,0 +1,215 @@ +"""FastModel config passthrough and nested task config handling.""" + +import ast +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[2] +LOADER_PATH = REPO_ROOT / "unsloth" / "models" / "loader.py" +VISION_PATH = REPO_ROOT / "unsloth" / "models" / "vision.py" +UTILS_PATH = REPO_ROOT / "unsloth" / "models" / "_utils.py" +LLAMA_PATH = REPO_ROOT / "unsloth" / "models" / "llama.py" + + +def _source(path): + return path.read_text() + + +def _class_method(tree, class_name, method_name): + for node in tree.body: + if isinstance(node, ast.ClassDef) and node.name == class_name: + for item in node.body: + if isinstance(item, ast.FunctionDef) and item.name == method_name: + return item + raise AssertionError(f"{class_name}.{method_name} not found") + + +def _assigns_from_kwargs_pop(method, target_name, key_name): + for node in ast.walk(method): + if not isinstance(node, ast.Assign): + continue + if not any( + isinstance(target, ast.Name) and target.id == target_name for target in node.targets + ): + continue + value = node.value + if not ( + isinstance(value, ast.Call) + and isinstance(value.func, ast.Attribute) + and value.func.attr == "pop" + and isinstance(value.func.value, ast.Name) + and value.func.value.id == "kwargs" + and value.args + and isinstance(value.args[0], ast.Constant) + and value.args[0].value == key_name + ): + continue + return True + return False + + +def _calls_name(method, name): + return any( + isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == name + for node in ast.walk(method) + ) + + +def _load_task_attr_helper(): + source = _source(UTILS_PATH) + funcs = { + node.name: ast.get_source_segment(source, node) + for node in ast.parse(source).body + if isinstance(node, ast.FunctionDef) + } + ns = {} + for name in ("_config_set", "set_task_config_attr"): + exec(funcs[name], ns) + return ns["set_task_config_attr"] + + +def _load_loader_task_helpers(): + source = _source(LOADER_PATH) + funcs = { + node.name: ast.get_source_segment(source, node) + for node in ast.parse(source).body + if isinstance(node, ast.FunctionDef) + } + ns = {} + for name in ( + "_config_get", + "_config_diff", + "_has_sequence_classification_architecture", + "_get_user_task_config_attrs", + ): + exec(funcs[name], ns) + return ns["_get_user_task_config_attrs"] + + +def test_fast_model_consumes_user_config_kwarg(): + tree = ast.parse(_source(LOADER_PATH)) + method = _class_method(tree, "FastModel", "from_pretrained") + + assert _assigns_from_kwargs_pop(method, "user_config", "config") + + +def test_fast_base_model_consumes_user_config_kwarg(): + tree = ast.parse(_source(VISION_PATH)) + method = _class_method(tree, "FastBaseModel", "from_pretrained") + + assert _assigns_from_kwargs_pop(method, "user_config", "config") + + +def test_fast_llama_model_consumes_user_config_kwarg(): + tree = ast.parse(_source(LLAMA_PATH)) + method = _class_method(tree, "FastLlamaModel", "from_pretrained") + + assert _assigns_from_kwargs_pop(method, "user_config", "config") + + +def test_fast_base_model_sets_task_attrs_on_nested_text_config(): + tree = ast.parse(_source(VISION_PATH)) + method = _class_method(tree, "FastBaseModel", "from_pretrained") + + assert _calls_name(method, "set_task_config_attr") + + +def test_fast_base_model_pops_problem_type_as_config_attr(): + source = _source(VISION_PATH) + + assert '("id2label", "label2id", "problem_type")' in source + + +def test_fast_model_uses_user_config_num_labels_for_task_model_selection(): + tree = ast.parse(_source(LOADER_PATH)) + method = _class_method(tree, "FastModel", "from_pretrained") + + assert _calls_name(method, "_get_user_task_config_attrs") + + +def test_fast_model_captures_user_config_num_labels_before_text_only_switch(): + source = _source(LOADER_PATH) + + fallback = source.index("task_config_attrs = _get_user_task_config_attrs(user_config)") + text_only_switch = source.index("model_config = text_config") + + assert fallback < text_only_switch + + +def test_user_task_config_attrs_ignore_default_num_labels(): + get_user_task_config_attrs = _load_loader_task_helpers() + + class Config: + num_labels = 2 + id2label = {0: "LABEL_0", 1: "LABEL_1"} + label2id = {"LABEL_0": 0, "LABEL_1": 1} + + def to_diff_dict(self): + return {} + + assert get_user_task_config_attrs(Config()) == {} + + +def test_user_task_config_attrs_preserve_custom_label_maps(): + get_user_task_config_attrs = _load_loader_task_helpers() + + class Config: + num_labels = 2 + id2label = {0: "negative", 1: "positive"} + label2id = {"negative": 0, "positive": 1} + + def to_diff_dict(self): + return {"id2label": self.id2label, "label2id": self.label2id} + + attrs = get_user_task_config_attrs(Config()) + + assert attrs["num_labels"] == 2 + assert attrs["id2label"] == {0: "negative", 1: "positive"} + assert attrs["label2id"] == {"negative": 0, "positive": 1} + + +def test_user_task_config_attrs_preserve_explicit_dict_num_labels(): + get_user_task_config_attrs = _load_loader_task_helpers() + + assert get_user_task_config_attrs({"num_labels": 2}) == {"num_labels": 2} + + +def test_task_config_attr_updates_parent_and_text_config_objects(): + set_task_config_attr = _load_task_attr_helper() + + class TextConfig: + pass + + class ParentConfig: + def __init__(self): + self.text_config = TextConfig() + + def get_text_config(self): + return self.text_config + + config = ParentConfig() + + set_task_config_attr(config, "num_labels", 3) + + assert config.num_labels == 3 + assert config.text_config.num_labels == 3 + + +def test_task_config_attr_updates_parent_and_text_config_dicts(): + set_task_config_attr = _load_task_attr_helper() + config = {"text_config": {}} + + set_task_config_attr(config, "label2id", {"negative": 0, "positive": 1}) + + assert config["label2id"] == {"negative": 0, "positive": 1} + assert config["text_config"]["label2id"] == {"negative": 0, "positive": 1} + + +def test_task_config_attr_ignores_primitive_text_config(): + set_task_config_attr = _load_task_attr_helper() + config = {"text_config": "not-a-config"} + + set_task_config_attr(config, "num_labels", 2) + + assert config["num_labels"] == 2 + assert config["text_config"] == "not-a-config" diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 6baa9a1398..2f4e3a069e 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -69,6 +69,7 @@ __all__ = [ "resolve_attention_implementation", "resolve_encoder_attention_implementation", "_set_attn_impl", + "set_task_config_attr", "patch_fast_lora", "validate_loftq_config", "RaiseUninitialized", @@ -306,6 +307,28 @@ def _config_set(config, field_name, value): setattr(config, field_name, value) +def set_task_config_attr(config, field_name, value): + _config_set(config, field_name, value) + text_config = None + if isinstance(config, dict): + text_config = config.get("text_config", None) + elif config is not None: + get_text_config = getattr(config, "get_text_config", None) + if callable(get_text_config): + try: + text_config = get_text_config() + except Exception: + text_config = None + if text_config is None: + text_config = getattr(config, "text_config", None) + if ( + text_config is not None + and text_config is not config + and (isinstance(text_config, dict) or hasattr(text_config, "__dict__")) + ): + _config_set(text_config, field_name, value) + + def _iter_attention_configs(config, seen = None): if config is None or (not isinstance(config, dict) and not hasattr(config, "__dict__")): return diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index a60edf3fbe..08802f030e 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2383,11 +2383,30 @@ class FastLlamaModel: assert dtype == torch.float16 or dtype == torch.bfloat16 or dtype == torch.float32 # RoPE Scaling - model_config = AutoConfig.from_pretrained( - model_name, - token = token, - attn_implementation = "sdpa", - ) + # Respect a user-provided config so it is the single config object used + # everywhere below; otherwise HF would receive it again through **kwargs + # alongside our own config= and fail with a duplicate-kwarg TypeError. + user_config = kwargs.pop("config", None) + if user_config is not None: + model_config = user_config + # model_name may have been remapped to a prequantized repo whose + # checkpoint needs its quantization_config; graft it onto the user + # config or the 4bit weights load without their quant state. + if getattr(model_config, "quantization_config", None) is None: + _checkpoint_config = AutoConfig.from_pretrained( + model_name, + token = token, + attn_implementation = "sdpa", + ) + _checkpoint_quant = getattr(_checkpoint_config, "quantization_config", None) + if _checkpoint_quant is not None: + model_config.quantization_config = _checkpoint_quant + else: + model_config = AutoConfig.from_pretrained( + model_name, + token = token, + attn_implementation = "sdpa", + ) model_config.model_name = model_name model_max_seq_length = model_config.max_position_embeddings @@ -2504,14 +2523,17 @@ class FastLlamaModel: # Transformers 5.x @strict config classes reject unexpected kwargs # like num_labels and max_position_embeddings. Set on the config # object directly and pass config= instead. - model_config.num_labels = num_labels + set_task_config_attr(model_config, "num_labels", num_labels) if max_position_embeddings is not None: model_config.max_position_embeddings = max_position_embeddings # Pop config-level attrs that would be rejected by @strict model init for _cfg_key in ("id2label", "label2id", "rope_scaling"): _cfg_val = kwargs.pop(_cfg_key, None) if _cfg_val is not None: - setattr(model_config, _cfg_key, _cfg_val) + if _cfg_key in ("id2label", "label2id"): + set_task_config_attr(model_config, _cfg_key, _cfg_val) + else: + setattr(model_config, _cfg_key, _cfg_val) model = AutoModelForSequenceClassification.from_pretrained( model_name, config = model_config, @@ -2544,17 +2566,33 @@ class FastLlamaModel: fast_inference = fast_inference, ) elif not fast_inference: - model = AutoModelForCausalLM.from_pretrained( - model_name, - device_map = device_map, - # torch_dtype = dtype, # transformers changed torch_dtype to dtype - # quantization_config = bnb_config, - token = token, - max_position_embeddings = max_position_embeddings, - trust_remote_code = trust_remote_code, - attn_implementation = preferred_attn_impl, - **kwargs, - ) + if user_config is not None: + # Transformers 5.x @strict model init rejects extra kwargs next + # to config=; set the override on the config and pass the single + # config object through so user overrides reach the actual load. + if max_position_embeddings is not None: + model_config.max_position_embeddings = max_position_embeddings + model = AutoModelForCausalLM.from_pretrained( + model_name, + config = model_config, + device_map = device_map, + token = token, + trust_remote_code = trust_remote_code, + attn_implementation = preferred_attn_impl, + **kwargs, + ) + else: + model = AutoModelForCausalLM.from_pretrained( + model_name, + device_map = device_map, + # torch_dtype = dtype, # transformers changed torch_dtype to dtype + # quantization_config = bnb_config, + token = token, + max_position_embeddings = max_position_embeddings, + trust_remote_code = trust_remote_code, + attn_implementation = preferred_attn_impl, + **kwargs, + ) # Attach dispatch hooks for bnb multi-device loads. from unsloth.models.vision import _attach_bnb_multidevice_hooks diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 4dc3046928..cfcfcae505 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -101,6 +101,7 @@ from ._utils import ( resolve_model_class, _is_family_text_decoder, _apply_text_only_key_mapping, + set_task_config_attr, ) # Single source of truth is unsloth_zoo.model_lists. Re-exported so callers @@ -133,6 +134,57 @@ def _strip_unsloth_bnb_4bit_suffix(model_name: str) -> str: return s +def _config_get( + config, + field_name, + default = None, +): + if isinstance(config, dict): + return config.get(field_name, default) + return getattr(config, field_name, default) + + +def _config_diff(config): + if isinstance(config, dict): + return config + to_diff_dict = getattr(config, "to_diff_dict", None) + if callable(to_diff_dict): + try: + diff = to_diff_dict() + if isinstance(diff, dict): + return diff + except Exception: + pass + return {} + + +def _has_sequence_classification_architecture(config): + architectures = _config_get(config, "architectures", None) or [] + return any(str(arch).endswith("ForSequenceClassification") for arch in architectures) + + +def _get_user_task_config_attrs(user_config): + if user_config is None: + return {} + diff = _config_diff(user_config) + attrs = {} + for key in ("id2label", "label2id", "problem_type"): + if key in diff: + attrs[key] = _config_get(user_config, key, diff.get(key)) + if isinstance(user_config, dict) and "num_labels" in user_config: + attrs["num_labels"] = user_config["num_labels"] + elif _has_sequence_classification_architecture(user_config): + num_labels = _config_get(user_config, "num_labels", None) + if num_labels is not None: + attrs["num_labels"] = num_labels + elif "id2label" in attrs: + try: + attrs["num_labels"] = len(attrs["id2label"]) + except TypeError: + pass + return attrs + + DISABLE_COMPILE_MODEL_NAMES = [ "aya_vision", "modernbert", @@ -907,6 +959,7 @@ class FastModel(FastBaseModel): *args, **kwargs, ): + user_config = kwargs.pop("config", None) # Respect user-provided quantization_config (e.g. BitsAndBytesConfig) quantization_config = kwargs.get("quantization_config", None) if quantization_config is not None: @@ -1104,13 +1157,15 @@ class FastModel(FastBaseModel): ) try: - model_config = AutoConfig.from_pretrained( - model_name, - token = token, - revision = revision, - trust_remote_code = trust_remote_code, - local_files_only = local_files_only, - ) + model_config = user_config + if model_config is None: + model_config = AutoConfig.from_pretrained( + model_name, + token = token, + revision = revision, + trust_remote_code = trust_remote_code, + local_files_only = local_files_only, + ) is_model = True except ImportError: raise @@ -1384,12 +1439,15 @@ class FastModel(FastBaseModel): load_in_fp8 = False load_in_16bit = True - model_config = AutoConfig.from_pretrained( - model_name, - token = token, - trust_remote_code = trust_remote_code, - local_files_only = local_files_only, - ) + if user_config is not None: + model_config = user_config + else: + model_config = AutoConfig.from_pretrained( + model_name, + token = token, + trust_remote_code = trust_remote_code, + local_files_only = local_files_only, + ) if not was_disabled: enable_progress_bars() @@ -1469,6 +1527,17 @@ class FastModel(FastBaseModel): else: tokenizer_name = kwargs.pop("tokenizer_name", None) + # Capture task intent before text_only can replace a parent VLM config + # with its nested text config. + task_config_attrs = _get_user_task_config_attrs(user_config) + for _cfg_key in ("num_labels", "id2label", "label2id", "problem_type"): + _cfg_val = kwargs.get(_cfg_key, None) + if _cfg_val is not None: + task_config_attrs[_cfg_key] = _cfg_val + _num_labels = task_config_attrs.get("num_labels", None) + for _cfg_key, _cfg_val in task_config_attrs.items(): + set_task_config_attr(model_config, _cfg_key, _cfg_val) + # Check if VLM architectures = getattr(model_config, "architectures", None) if architectures is None: @@ -1499,7 +1568,8 @@ class FastModel(FastBaseModel): else: is_vlm = False # If num_labels is set, use AutoModelForSequenceClassification - _num_labels = kwargs.get("num_labels", None) + for _cfg_key, _cfg_val in task_config_attrs.items(): + set_task_config_attr(model_config, _cfg_key, _cfg_val) if auto_model is None: if _num_labels is not None: from transformers import AutoModelForSequenceClassification diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index a52ab359bd..e8161427d5 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -37,6 +37,7 @@ from ._utils import ( _get_text_only_config, _is_family_text_decoder, _apply_text_only_key_mapping, + set_task_config_attr, ) from ._utils import * from .loader_utils import _get_fp8_mode_and_check_settings @@ -592,6 +593,10 @@ class FastBaseModel: text_only = False, **kwargs, ): + user_config = kwargs.pop("config", None) + if auto_config is None and user_config is not None: + auto_config = user_config + if unsloth_vllm_standby and os.environ.get("UNSLOTH_VLLM_STANDBY", "0") != "1": raise RuntimeError( "Unsloth: UNSLOTH_VLLM_STANDBY is True, but UNSLOTH_VLLM_STANDBY is not set to 1!" @@ -950,11 +955,14 @@ class FastBaseModel: # Move config-level attributes onto the config object directly. _num_labels = kwargs.pop("num_labels", None) if _num_labels is not None: - model_config.num_labels = _num_labels - for _cfg_key in ("id2label", "label2id", "max_position_embeddings"): + set_task_config_attr(model_config, "num_labels", _num_labels) + for _cfg_key in ("id2label", "label2id", "problem_type"): _cfg_val = kwargs.pop(_cfg_key, None) if _cfg_val is not None: - setattr(model_config, _cfg_key, _cfg_val) + set_task_config_attr(model_config, _cfg_key, _cfg_val) + _cfg_val = kwargs.pop("max_position_embeddings", None) + if _cfg_val is not None: + setattr(model_config, "max_position_embeddings", _cfg_val) model = auto_model.from_pretrained( model_name, config = model_config,