mirror of https://github.com/razor-ai/soup.git
72 lines
2.1 KiB
Python
72 lines
2.1 KiB
Python
"""PEFT config builder — unifies LoRA / DoRA / VeRA / OLoRA construction.
|
|
|
|
Returns an intermediate spec dict so trainers can import the right class and
|
|
instantiate it without every trainer needing to replicate the branching logic.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from soup_cli.config.schema import LoraConfig as SchemaLoraConfig
|
|
|
|
|
|
def build_peft_config(
|
|
lora_cfg: SchemaLoraConfig,
|
|
target_modules: "str | list[str]",
|
|
task_type: str,
|
|
) -> dict[str, Any]:
|
|
"""Build a peft config spec from Soup's schema LoraConfig.
|
|
|
|
Returns:
|
|
Dict with keys:
|
|
- ``peft_cls``: str — class name to import from ``peft`` (``LoraConfig``
|
|
or ``VeraConfig``).
|
|
- ``init_kwargs``: dict — kwargs to pass to the constructor.
|
|
|
|
Trainers can use this spec to instantiate the right peft config without
|
|
duplicating the branching logic for DoRA / VeRA / OLoRA.
|
|
"""
|
|
if lora_cfg.use_vera:
|
|
return {
|
|
"peft_cls": "VeraConfig",
|
|
"init_kwargs": {
|
|
"r": lora_cfg.r,
|
|
"target_modules": target_modules,
|
|
"task_type": task_type,
|
|
"vera_dropout": lora_cfg.dropout,
|
|
"bias": "none",
|
|
},
|
|
}
|
|
|
|
init_kwargs: dict = {
|
|
"r": lora_cfg.r,
|
|
"lora_alpha": lora_cfg.alpha,
|
|
"lora_dropout": lora_cfg.dropout,
|
|
"target_modules": target_modules,
|
|
"task_type": task_type,
|
|
"bias": "none",
|
|
"use_dora": lora_cfg.use_dora,
|
|
"use_rslora": lora_cfg.use_rslora,
|
|
}
|
|
|
|
if lora_cfg.use_olora:
|
|
init_kwargs["init_lora_weights"] = "olora"
|
|
|
|
return {
|
|
"peft_cls": "LoraConfig",
|
|
"init_kwargs": init_kwargs,
|
|
}
|
|
|
|
|
|
def instantiate_peft_config(spec: dict[str, Any]) -> Any:
|
|
"""Instantiate the peft config from a spec dict (lazy import).
|
|
|
|
Returns ``peft.PeftConfig`` (return type is ``Any`` because ``peft`` is a
|
|
lazy import and cannot be referenced at module scope).
|
|
"""
|
|
import peft # lazy
|
|
|
|
cls = getattr(peft, spec["peft_cls"])
|
|
return cls(**spec["init_kwargs"])
|