Fix FastModel config passthrough for sequence classification (#6203)
* add FastModel config passthrough * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix fastmodel config passthrough for task configs * fix config-driven FastModel task model selection * fix text only fastmodel task config selection * fix fastmodel task config inference from user configs * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix fastmodel problem_type config passthrough * fix fastlanguagemodel config passthrough: FastLlamaModel owns user config * fix fastlanguagemodel config passthrough: forward user config to causal loads and keep checkpoint quantization_config --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
This commit is contained in:
alkinun
GitHub
parent
a24c9987ca
commit
14ed91e39a
5 changed files with 389 additions and 35 deletions
215
tests/python/test_fast_model_config_passthrough.py
Normal file
215
tests/python/test_fast_model_config_passthrough.py
Normal file
|
|
@ -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"
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue