mirror of https://github.com/razor-ai/soup.git
Add Unsloth backend for 2-5x faster training (Phase 5) — v0.4.3
- Add `backend: unsloth` config option (default: transformers) - Create utils/unsloth.py with FastLanguageModel integration - Update SFT/DPO/GRPO trainers with _setup_unsloth() path - Auto-detect unsloth and suggest enabling in `soup train` - Add `fast` extra: pip install 'soup-cli[fast]' - Add unsloth to `soup doctor` dependency check - Add `backend` sweep shortcut - 33 new tests (404 total), ruff clean Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Alpamys
parent
dd73d91356
commit
6de66fb28c
|
|
@ -43,7 +43,9 @@ soup train --config soup.yaml
|
|||
|
||||
**Data pipeline:** `data/loader.py` handles local files (JSONL/JSON/CSV/Parquet) and HuggingFace datasets. `data/formats.py` auto-detects and normalizes alpaca/sharegpt/chatml formats into a unified `{"messages": [...]}` structure. Also supports reverse conversion via `messages_to_format()`.
|
||||
|
||||
**Trainer:** `trainer/sft.py` (`SFTTrainerWrapper`), `trainer/dpo.py` (`DPOTrainerWrapper`), and `trainer/grpo.py` (`GRPOTrainerWrapper`) wrap HuggingFace's SFTTrainer/DPOTrainer/GRPOTrainer with auto quantization (BitsAndBytes), LoRA (PEFT), and batch size estimation. Heavy ML imports are lazy (inside methods) so CLI stays fast for non-training commands. All trainers enable Rich progress bars for HuggingFace Hub model downloads via `_enable_hf_transfer_progress()`.
|
||||
**Trainer:** `trainer/sft.py` (`SFTTrainerWrapper`), `trainer/dpo.py` (`DPOTrainerWrapper`), and `trainer/grpo.py` (`GRPOTrainerWrapper`) wrap HuggingFace's SFTTrainer/DPOTrainer/GRPOTrainer with auto quantization (BitsAndBytes), LoRA (PEFT), and batch size estimation. Heavy ML imports are lazy (inside methods) so CLI stays fast for non-training commands. All trainers enable Rich progress bars for HuggingFace Hub model downloads via `_enable_hf_transfer_progress()`. All trainers support `backend: unsloth` — when enabled, model loading is delegated to `utils/unsloth.py` which uses `unsloth.FastLanguageModel` for 2-5x faster training.
|
||||
|
||||
**Unsloth backend:** `utils/unsloth.py` provides `is_unsloth_available()`, `get_unsloth_version()`, and `load_model_and_tokenizer()`. The latter uses `FastLanguageModel.from_pretrained()` + `FastLanguageModel.get_peft_model()` which handles quantization, LoRA patching, and kernel optimization internally. Config: `backend: unsloth` (default: `transformers`). Install: `pip install 'soup-cli[fast]'`. Auto-detection: if unsloth is installed but not enabled, `commands/train.py` shows a hint.
|
||||
|
||||
**GRPO (Group Relative Policy Optimization):** `trainer/grpo.py` implements reasoning model training (DeepSeek-R1 style). Generates multiple completions per prompt, scores them with reward functions, and optimizes using group-relative advantages. `trainer/rewards.py` provides built-in reward functions (`accuracy` — checks final answer, `format` — checks `<think>` blocks) and supports custom rewards via Python files. Config: `task: grpo`, `grpo_beta`, `num_generations`, `reward_fn`.
|
||||
|
||||
|
|
@ -164,3 +166,4 @@ Test suite lives in `tests/`:
|
|||
| `test_quickstart.py` | `soup quickstart` demo, data/config creation, --dry-run |
|
||||
| `test_grpo.py` | GRPO config, rewards, data prep, template, sweep shortcuts |
|
||||
| `test_progress.py` | Rich download progress bar, `_enable_hf_transfer_progress` |
|
||||
| `test_unsloth.py` | Unsloth backend config, detection, trainer integration, templates |
|
||||
|
|
|
|||
35
README.md
35
README.md
|
|
@ -104,6 +104,7 @@ soup export --model ./output --format gguf --quant q4_k_m
|
|||
```yaml
|
||||
base: meta-llama/Llama-3.1-8B-Instruct
|
||||
task: sft
|
||||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||||
|
||||
data:
|
||||
train: ./data/train.jsonl
|
||||
|
|
@ -122,6 +123,39 @@ training:
|
|||
output: ./output
|
||||
```
|
||||
|
||||
## Unsloth Backend (2-5x Faster Training)
|
||||
|
||||
Use the [Unsloth](https://github.com/unslothai/unsloth) backend for significantly faster training and up to 80% less VRAM:
|
||||
|
||||
```bash
|
||||
# Install unsloth support
|
||||
pip install 'soup-cli[fast]'
|
||||
```
|
||||
|
||||
Then add one line to your config:
|
||||
|
||||
```yaml
|
||||
base: meta-llama/Llama-3.1-8B-Instruct
|
||||
task: sft
|
||||
backend: unsloth # 2-5x faster, -80% VRAM
|
||||
|
||||
data:
|
||||
train: ./data/train.jsonl
|
||||
format: alpaca
|
||||
|
||||
training:
|
||||
epochs: 3
|
||||
lr: 2e-5
|
||||
quantization: 4bit
|
||||
lora:
|
||||
r: 64
|
||||
alpha: 16
|
||||
```
|
||||
|
||||
Works with all training tasks: SFT, DPO, and GRPO. If unsloth is installed but not enabled, Soup will suggest it automatically.
|
||||
|
||||
> **Tip:** Soup auto-detects unsloth. When installed, you'll see a hint during `soup train` if you haven't enabled it yet.
|
||||
|
||||
## DPO Training
|
||||
|
||||
Train with preference data using Direct Preference Optimization:
|
||||
|
|
@ -545,6 +579,7 @@ soup --verbose <command> Full traceback on errors
|
|||
|
||||
| Extra | Install | What it adds |
|
||||
|---|---|---|
|
||||
| `fast` | `pip install 'soup-cli[fast]'` | Unsloth backend (2-5x faster, -80% VRAM) |
|
||||
| `serve` | `pip install 'soup-cli[serve]'` | Inference server (FastAPI + uvicorn) |
|
||||
| `data` | `pip install 'soup-cli[data]'` | Deduplication (MinHash via datasketch) |
|
||||
| `eval` | `pip install 'soup-cli[eval]'` | Benchmark evaluation (lm-evaluation-harness) |
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|||
|
||||
[project]
|
||||
name = "soup-cli"
|
||||
version = "0.4.2"
|
||||
version = "0.4.3"
|
||||
description = "Fine-tune LLMs in one command. No SSH, no config hell."
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
|
|
@ -46,6 +46,7 @@ ui = ["fastapi>=0.104.0", "uvicorn>=0.24.0"]
|
|||
serve = ["fastapi>=0.104.0", "uvicorn>=0.24.0"]
|
||||
generate = ["httpx>=0.24.0"]
|
||||
deepspeed = ["deepspeed>=0.12.0"]
|
||||
fast = ["unsloth>=2024.8"]
|
||||
|
||||
[project.scripts]
|
||||
soup = "soup_cli.cli:run"
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
"""Soup CLI — Fine-tune LLMs in one command."""
|
||||
|
||||
__version__ = "0.4.2"
|
||||
__version__ = "0.4.3"
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ DEPS = [
|
|||
("wandb", "wandb", "0.15.0", False),
|
||||
("deepspeed", "deepspeed", "0.12.0", False),
|
||||
("httpx", "httpx", "0.24.0", False),
|
||||
("unsloth", "unsloth", "2024.8", False),
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -290,6 +290,7 @@ def _set_nested_param(config_dict: dict, key: str, value) -> dict:
|
|||
"grpo_beta": "training.grpo_beta",
|
||||
"num_generations": "training.num_generations",
|
||||
"reward_fn": "training.reward_fn",
|
||||
"backend": "backend",
|
||||
}
|
||||
|
||||
full_key = shortcuts.get(key, key)
|
||||
|
|
|
|||
|
|
@ -100,18 +100,34 @@ def train(
|
|||
# Detect hardware
|
||||
device, device_name = detect_device()
|
||||
gpu_info = get_gpu_info()
|
||||
|
||||
backend_label = cfg.backend
|
||||
if cfg.backend == "unsloth":
|
||||
backend_label = "unsloth [green](fast mode)[/]"
|
||||
|
||||
console.print(
|
||||
Panel(
|
||||
f"Device: [bold]{device_name}[/]\n"
|
||||
f"Memory: [bold]{gpu_info['memory_total']}[/]\n"
|
||||
f"Model: [bold]{cfg.base}[/]\n"
|
||||
f"Task: [bold]{cfg.task}[/]\n"
|
||||
f"LoRA: [bold]r={cfg.training.lora.r}, alpha={cfg.training.lora.alpha}[/]\n"
|
||||
f"Quant: [bold]{cfg.training.quantization}[/]",
|
||||
f"Device: [bold]{device_name}[/]\n"
|
||||
f"Memory: [bold]{gpu_info['memory_total']}[/]\n"
|
||||
f"Model: [bold]{cfg.base}[/]\n"
|
||||
f"Task: [bold]{cfg.task}[/]\n"
|
||||
f"Backend: [bold]{backend_label}[/]\n"
|
||||
f"LoRA: [bold]r={cfg.training.lora.r}, alpha={cfg.training.lora.alpha}[/]\n"
|
||||
f"Quant: [bold]{cfg.training.quantization}[/]",
|
||||
title="Training Setup",
|
||||
)
|
||||
)
|
||||
|
||||
# Suggest unsloth if available but not being used
|
||||
if cfg.backend == "transformers":
|
||||
from soup_cli.utils.unsloth import is_unsloth_available
|
||||
|
||||
if is_unsloth_available():
|
||||
console.print(
|
||||
"[dim]Tip: unsloth is installed. Add [bold]backend: unsloth[/dim]"
|
||||
"[dim] to soup.yaml for 2-5x faster training.[/]"
|
||||
)
|
||||
|
||||
if not dry_run and not yes:
|
||||
if not typer.confirm("Start training?", default=True):
|
||||
console.print("[yellow]Cancelled.[/]")
|
||||
|
|
|
|||
|
|
@ -67,6 +67,10 @@ class SoupConfig(BaseModel):
|
|||
|
||||
base: str = Field(..., description="Base model name or path (HF model ID)")
|
||||
task: Literal["sft", "dpo", "grpo"] = Field(default="sft", description="Training task type")
|
||||
backend: Literal["transformers", "unsloth"] = Field(
|
||||
default="transformers",
|
||||
description="Training backend: transformers (default) or unsloth (2-5x faster)",
|
||||
)
|
||||
data: DataConfig
|
||||
training: TrainingConfig = Field(default_factory=TrainingConfig)
|
||||
output: str = Field(default="./output", description="Output directory for trained model")
|
||||
|
|
@ -81,6 +85,7 @@ TEMPLATES: dict[str, str] = {
|
|||
|
||||
base: meta-llama/Llama-3.1-8B-Instruct
|
||||
task: sft
|
||||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||||
|
||||
data:
|
||||
train: ./data/train.jsonl
|
||||
|
|
@ -105,6 +110,7 @@ output: ./output
|
|||
|
||||
base: codellama/CodeLlama-7b-Instruct-hf
|
||||
task: sft
|
||||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||||
|
||||
data:
|
||||
train: ./data/code_train.jsonl
|
||||
|
|
@ -129,6 +135,7 @@ output: ./output
|
|||
|
||||
base: meta-llama/Llama-3.1-8B-Instruct
|
||||
task: grpo
|
||||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||||
|
||||
data:
|
||||
train: ./data/reasoning_train.jsonl
|
||||
|
|
@ -157,6 +164,7 @@ output: ./output
|
|||
|
||||
base: meta-llama/Llama-3.1-8B-Instruct
|
||||
task: sft
|
||||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||||
|
||||
data:
|
||||
train: ./data/medical_train.jsonl
|
||||
|
|
|
|||
|
|
@ -40,12 +40,6 @@ class DPOTrainerWrapper:
|
|||
def setup(self, dataset: dict):
|
||||
"""Load model, tokenizer, apply LoRA, create DPO trainer."""
|
||||
from datasets import Dataset
|
||||
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
|
||||
from transformers import (
|
||||
AutoModelForCausalLM,
|
||||
AutoTokenizer,
|
||||
BitsAndBytesConfig,
|
||||
)
|
||||
from trl import DPOConfig, DPOTrainer
|
||||
|
||||
# Enable Rich progress bar for HuggingFace downloads
|
||||
|
|
@ -55,52 +49,13 @@ class DPOTrainerWrapper:
|
|||
|
||||
cfg = self.config
|
||||
tcfg = cfg.training
|
||||
use_unsloth = cfg.backend == "unsloth"
|
||||
|
||||
# --- Tokenizer ---
|
||||
console.print(f"[dim]Loading tokenizer: {cfg.base}[/]")
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True)
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
if use_unsloth:
|
||||
self._setup_unsloth(cfg, tcfg)
|
||||
else:
|
||||
self._setup_transformers(cfg, tcfg)
|
||||
|
||||
# --- Quantization ---
|
||||
bnb_config = None
|
||||
if tcfg.quantization == "4bit":
|
||||
import torch
|
||||
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
)
|
||||
elif tcfg.quantization == "8bit":
|
||||
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
|
||||
|
||||
# --- Model ---
|
||||
console.print(f"[dim]Loading model: {cfg.base}[/]")
|
||||
model_kwargs = {"trust_remote_code": True, "device_map": "auto"}
|
||||
if bnb_config:
|
||||
model_kwargs["quantization_config"] = bnb_config
|
||||
|
||||
self.model = AutoModelForCausalLM.from_pretrained(cfg.base, **model_kwargs)
|
||||
|
||||
if tcfg.quantization in ("4bit", "8bit"):
|
||||
self.model = prepare_model_for_kbit_training(self.model)
|
||||
|
||||
# --- LoRA ---
|
||||
target_modules = tcfg.lora.target_modules
|
||||
if target_modules == "auto":
|
||||
target_modules = None # peft will auto-detect
|
||||
|
||||
lora_config = LoraConfig(
|
||||
r=tcfg.lora.r,
|
||||
lora_alpha=tcfg.lora.alpha,
|
||||
lora_dropout=tcfg.lora.dropout,
|
||||
target_modules=target_modules,
|
||||
task_type=TaskType.CAUSAL_LM,
|
||||
bias="none",
|
||||
)
|
||||
self.model = get_peft_model(self.model, lora_config)
|
||||
trainable, total = self.model.get_nb_trainable_parameters()
|
||||
pct = 100 * trainable / total
|
||||
console.print(
|
||||
|
|
@ -183,6 +138,70 @@ class DPOTrainerWrapper:
|
|||
|
||||
self._output_dir = str(output_dir)
|
||||
|
||||
def _setup_transformers(self, cfg, tcfg):
|
||||
"""Load model via standard transformers + peft pipeline."""
|
||||
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
||||
|
||||
console.print(f"[dim]Loading tokenizer: {cfg.base}[/]")
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True)
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
|
||||
bnb_config = None
|
||||
if tcfg.quantization == "4bit":
|
||||
import torch
|
||||
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
)
|
||||
elif tcfg.quantization == "8bit":
|
||||
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
|
||||
|
||||
console.print(f"[dim]Loading model: {cfg.base}[/]")
|
||||
model_kwargs = {"trust_remote_code": True, "device_map": "auto"}
|
||||
if bnb_config:
|
||||
model_kwargs["quantization_config"] = bnb_config
|
||||
|
||||
self.model = AutoModelForCausalLM.from_pretrained(cfg.base, **model_kwargs)
|
||||
|
||||
if tcfg.quantization in ("4bit", "8bit"):
|
||||
self.model = prepare_model_for_kbit_training(self.model)
|
||||
|
||||
target_modules = tcfg.lora.target_modules
|
||||
if target_modules == "auto":
|
||||
target_modules = None
|
||||
|
||||
lora_config = LoraConfig(
|
||||
r=tcfg.lora.r,
|
||||
lora_alpha=tcfg.lora.alpha,
|
||||
lora_dropout=tcfg.lora.dropout,
|
||||
target_modules=target_modules,
|
||||
task_type=TaskType.CAUSAL_LM,
|
||||
bias="none",
|
||||
)
|
||||
self.model = get_peft_model(self.model, lora_config)
|
||||
|
||||
def _setup_unsloth(self, cfg, tcfg):
|
||||
"""Load model via unsloth FastLanguageModel (2-5x faster)."""
|
||||
from soup_cli.utils.unsloth import load_model_and_tokenizer
|
||||
|
||||
console.print(f"[dim]Loading model via [bold]unsloth[/]: {cfg.base}[/]")
|
||||
self.model, self.tokenizer = load_model_and_tokenizer(
|
||||
model_name=cfg.base,
|
||||
max_seq_length=cfg.data.max_length,
|
||||
quantization=tcfg.quantization,
|
||||
lora_r=tcfg.lora.r,
|
||||
lora_alpha=tcfg.lora.alpha,
|
||||
lora_dropout=tcfg.lora.dropout,
|
||||
target_modules=tcfg.lora.target_modules,
|
||||
)
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
|
||||
def train(
|
||||
self,
|
||||
display: Optional[object] = None,
|
||||
|
|
|
|||
|
|
@ -41,12 +41,6 @@ class GRPOTrainerWrapper:
|
|||
def setup(self, dataset: dict):
|
||||
"""Load model, tokenizer, apply LoRA, create GRPO trainer."""
|
||||
from datasets import Dataset
|
||||
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
|
||||
from transformers import (
|
||||
AutoModelForCausalLM,
|
||||
AutoTokenizer,
|
||||
BitsAndBytesConfig,
|
||||
)
|
||||
from trl import GRPOConfig, GRPOTrainer
|
||||
|
||||
# Enable Rich progress bar for HuggingFace downloads
|
||||
|
|
@ -56,57 +50,18 @@ class GRPOTrainerWrapper:
|
|||
|
||||
cfg = self.config
|
||||
tcfg = cfg.training
|
||||
use_unsloth = cfg.backend == "unsloth"
|
||||
|
||||
# --- Load reward function ---
|
||||
from soup_cli.trainer.rewards import load_reward_fn
|
||||
|
||||
reward_fn = load_reward_fn(tcfg.reward_fn)
|
||||
|
||||
# --- Tokenizer ---
|
||||
console.print(f"[dim]Loading tokenizer: {cfg.base}[/]")
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True)
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
if use_unsloth:
|
||||
self._setup_unsloth(cfg, tcfg)
|
||||
else:
|
||||
self._setup_transformers(cfg, tcfg)
|
||||
|
||||
# --- Quantization ---
|
||||
bnb_config = None
|
||||
if tcfg.quantization == "4bit":
|
||||
import torch
|
||||
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
)
|
||||
elif tcfg.quantization == "8bit":
|
||||
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
|
||||
|
||||
# --- Model ---
|
||||
console.print(f"[dim]Loading model: {cfg.base}[/]")
|
||||
model_kwargs = {"trust_remote_code": True, "device_map": "auto"}
|
||||
if bnb_config:
|
||||
model_kwargs["quantization_config"] = bnb_config
|
||||
|
||||
self.model = AutoModelForCausalLM.from_pretrained(cfg.base, **model_kwargs)
|
||||
|
||||
if tcfg.quantization in ("4bit", "8bit"):
|
||||
self.model = prepare_model_for_kbit_training(self.model)
|
||||
|
||||
# --- LoRA ---
|
||||
target_modules = tcfg.lora.target_modules
|
||||
if target_modules == "auto":
|
||||
target_modules = None # peft will auto-detect
|
||||
|
||||
lora_config = LoraConfig(
|
||||
r=tcfg.lora.r,
|
||||
lora_alpha=tcfg.lora.alpha,
|
||||
lora_dropout=tcfg.lora.dropout,
|
||||
target_modules=target_modules,
|
||||
task_type=TaskType.CAUSAL_LM,
|
||||
bias="none",
|
||||
)
|
||||
self.model = get_peft_model(self.model, lora_config)
|
||||
trainable, total = self.model.get_nb_trainable_parameters()
|
||||
pct = 100 * trainable / total
|
||||
console.print(
|
||||
|
|
@ -192,6 +147,70 @@ class GRPOTrainerWrapper:
|
|||
|
||||
self._output_dir = str(output_dir)
|
||||
|
||||
def _setup_transformers(self, cfg, tcfg):
|
||||
"""Load model via standard transformers + peft pipeline."""
|
||||
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
||||
|
||||
console.print(f"[dim]Loading tokenizer: {cfg.base}[/]")
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True)
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
|
||||
bnb_config = None
|
||||
if tcfg.quantization == "4bit":
|
||||
import torch
|
||||
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
)
|
||||
elif tcfg.quantization == "8bit":
|
||||
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
|
||||
|
||||
console.print(f"[dim]Loading model: {cfg.base}[/]")
|
||||
model_kwargs = {"trust_remote_code": True, "device_map": "auto"}
|
||||
if bnb_config:
|
||||
model_kwargs["quantization_config"] = bnb_config
|
||||
|
||||
self.model = AutoModelForCausalLM.from_pretrained(cfg.base, **model_kwargs)
|
||||
|
||||
if tcfg.quantization in ("4bit", "8bit"):
|
||||
self.model = prepare_model_for_kbit_training(self.model)
|
||||
|
||||
target_modules = tcfg.lora.target_modules
|
||||
if target_modules == "auto":
|
||||
target_modules = None
|
||||
|
||||
lora_config = LoraConfig(
|
||||
r=tcfg.lora.r,
|
||||
lora_alpha=tcfg.lora.alpha,
|
||||
lora_dropout=tcfg.lora.dropout,
|
||||
target_modules=target_modules,
|
||||
task_type=TaskType.CAUSAL_LM,
|
||||
bias="none",
|
||||
)
|
||||
self.model = get_peft_model(self.model, lora_config)
|
||||
|
||||
def _setup_unsloth(self, cfg, tcfg):
|
||||
"""Load model via unsloth FastLanguageModel (2-5x faster)."""
|
||||
from soup_cli.utils.unsloth import load_model_and_tokenizer
|
||||
|
||||
console.print(f"[dim]Loading model via [bold]unsloth[/]: {cfg.base}[/]")
|
||||
self.model, self.tokenizer = load_model_and_tokenizer(
|
||||
model_name=cfg.base,
|
||||
max_seq_length=cfg.data.max_length,
|
||||
quantization=tcfg.quantization,
|
||||
lora_r=tcfg.lora.r,
|
||||
lora_alpha=tcfg.lora.alpha,
|
||||
lora_dropout=tcfg.lora.dropout,
|
||||
target_modules=tcfg.lora.target_modules,
|
||||
)
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
|
||||
def train(
|
||||
self,
|
||||
display: Optional[object] = None,
|
||||
|
|
|
|||
|
|
@ -33,13 +33,7 @@ class SFTTrainerWrapper:
|
|||
def setup(self, dataset: dict):
|
||||
"""Load model, tokenizer, apply LoRA, create trainer."""
|
||||
from datasets import Dataset
|
||||
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
|
||||
from transformers import (
|
||||
AutoModelForCausalLM,
|
||||
AutoTokenizer,
|
||||
BitsAndBytesConfig,
|
||||
TrainingArguments,
|
||||
)
|
||||
from transformers import TrainingArguments
|
||||
from trl import SFTTrainer
|
||||
|
||||
# Enable Rich progress bar for HuggingFace downloads
|
||||
|
|
@ -47,52 +41,13 @@ class SFTTrainerWrapper:
|
|||
|
||||
cfg = self.config
|
||||
tcfg = cfg.training
|
||||
use_unsloth = cfg.backend == "unsloth"
|
||||
|
||||
# --- Tokenizer ---
|
||||
console.print(f"[dim]Loading tokenizer: {cfg.base}[/]")
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True)
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
if use_unsloth:
|
||||
self._setup_unsloth(cfg, tcfg)
|
||||
else:
|
||||
self._setup_transformers(cfg, tcfg)
|
||||
|
||||
# --- Quantization ---
|
||||
bnb_config = None
|
||||
if tcfg.quantization == "4bit":
|
||||
import torch
|
||||
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
)
|
||||
elif tcfg.quantization == "8bit":
|
||||
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
|
||||
|
||||
# --- Model ---
|
||||
console.print(f"[dim]Loading model: {cfg.base}[/]")
|
||||
model_kwargs = {"trust_remote_code": True, "device_map": "auto"}
|
||||
if bnb_config:
|
||||
model_kwargs["quantization_config"] = bnb_config
|
||||
|
||||
self.model = AutoModelForCausalLM.from_pretrained(cfg.base, **model_kwargs)
|
||||
|
||||
if tcfg.quantization in ("4bit", "8bit"):
|
||||
self.model = prepare_model_for_kbit_training(self.model)
|
||||
|
||||
# --- LoRA ---
|
||||
target_modules = tcfg.lora.target_modules
|
||||
if target_modules == "auto":
|
||||
target_modules = None # peft will auto-detect
|
||||
|
||||
lora_config = LoraConfig(
|
||||
r=tcfg.lora.r,
|
||||
lora_alpha=tcfg.lora.alpha,
|
||||
lora_dropout=tcfg.lora.dropout,
|
||||
target_modules=target_modules,
|
||||
task_type=TaskType.CAUSAL_LM,
|
||||
bias="none",
|
||||
)
|
||||
self.model = get_peft_model(self.model, lora_config)
|
||||
trainable, total = self.model.get_nb_trainable_parameters()
|
||||
pct = 100 * trainable / total
|
||||
console.print(
|
||||
|
|
@ -188,6 +143,72 @@ class SFTTrainerWrapper:
|
|||
|
||||
self._output_dir = str(output_dir)
|
||||
|
||||
def _setup_transformers(self, cfg, tcfg):
|
||||
"""Load model via standard transformers + peft pipeline."""
|
||||
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
||||
|
||||
console.print(f"[dim]Loading tokenizer: {cfg.base}[/]")
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True)
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
|
||||
# Quantization
|
||||
bnb_config = None
|
||||
if tcfg.quantization == "4bit":
|
||||
import torch
|
||||
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
)
|
||||
elif tcfg.quantization == "8bit":
|
||||
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
|
||||
|
||||
console.print(f"[dim]Loading model: {cfg.base}[/]")
|
||||
model_kwargs = {"trust_remote_code": True, "device_map": "auto"}
|
||||
if bnb_config:
|
||||
model_kwargs["quantization_config"] = bnb_config
|
||||
|
||||
self.model = AutoModelForCausalLM.from_pretrained(cfg.base, **model_kwargs)
|
||||
|
||||
if tcfg.quantization in ("4bit", "8bit"):
|
||||
self.model = prepare_model_for_kbit_training(self.model)
|
||||
|
||||
# LoRA
|
||||
target_modules = tcfg.lora.target_modules
|
||||
if target_modules == "auto":
|
||||
target_modules = None
|
||||
|
||||
lora_config = LoraConfig(
|
||||
r=tcfg.lora.r,
|
||||
lora_alpha=tcfg.lora.alpha,
|
||||
lora_dropout=tcfg.lora.dropout,
|
||||
target_modules=target_modules,
|
||||
task_type=TaskType.CAUSAL_LM,
|
||||
bias="none",
|
||||
)
|
||||
self.model = get_peft_model(self.model, lora_config)
|
||||
|
||||
def _setup_unsloth(self, cfg, tcfg):
|
||||
"""Load model via unsloth FastLanguageModel (2-5x faster)."""
|
||||
from soup_cli.utils.unsloth import load_model_and_tokenizer
|
||||
|
||||
console.print(f"[dim]Loading model via [bold]unsloth[/]: {cfg.base}[/]")
|
||||
self.model, self.tokenizer = load_model_and_tokenizer(
|
||||
model_name=cfg.base,
|
||||
max_seq_length=cfg.data.max_length,
|
||||
quantization=tcfg.quantization,
|
||||
lora_r=tcfg.lora.r,
|
||||
lora_alpha=tcfg.lora.alpha,
|
||||
lora_dropout=tcfg.lora.dropout,
|
||||
target_modules=tcfg.lora.target_modules,
|
||||
)
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
|
||||
def train(
|
||||
self,
|
||||
display: Optional[object] = None,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,74 @@
|
|||
"""Unsloth backend utilities — detection, model loading, LoRA patching."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
def is_unsloth_available() -> bool:
|
||||
"""Check if unsloth is installed and importable."""
|
||||
try:
|
||||
import unsloth # noqa: F401
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
def get_unsloth_version() -> str | None:
|
||||
"""Return unsloth version string, or None if not installed."""
|
||||
try:
|
||||
import unsloth
|
||||
|
||||
return getattr(unsloth, "__version__", "unknown")
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
|
||||
def load_model_and_tokenizer(
|
||||
model_name: str,
|
||||
max_seq_length: int,
|
||||
quantization: str = "4bit",
|
||||
lora_r: int = 64,
|
||||
lora_alpha: int = 16,
|
||||
lora_dropout: float = 0.05,
|
||||
target_modules: str | list[str] | None = "auto",
|
||||
):
|
||||
"""Load model + tokenizer via unsloth FastLanguageModel with LoRA already applied.
|
||||
|
||||
Returns (model, tokenizer) — model already has LoRA adapters attached.
|
||||
Unsloth handles quantization, LoRA patching, and kernel optimization internally.
|
||||
"""
|
||||
from unsloth import FastLanguageModel
|
||||
|
||||
load_in_4bit = quantization == "4bit"
|
||||
|
||||
# Unsloth's FastLanguageModel.from_pretrained handles quantization internally
|
||||
dtype = None # auto-detect
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name=model_name,
|
||||
max_seq_length=max_seq_length,
|
||||
dtype=dtype,
|
||||
load_in_4bit=load_in_4bit,
|
||||
)
|
||||
|
||||
# Resolve target_modules for LoRA
|
||||
if target_modules == "auto" or target_modules is None:
|
||||
# Unsloth default: all linear layers for maximum performance
|
||||
target_modules = [
|
||||
"q_proj", "k_proj", "v_proj", "o_proj",
|
||||
"gate_proj", "up_proj", "down_proj",
|
||||
]
|
||||
elif isinstance(target_modules, str):
|
||||
target_modules = [target_modules]
|
||||
|
||||
# Apply LoRA via unsloth's optimized path
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r=lora_r,
|
||||
lora_alpha=lora_alpha,
|
||||
lora_dropout=lora_dropout,
|
||||
target_modules=target_modules,
|
||||
bias="none",
|
||||
use_gradient_checkpointing="unsloth", # 2x longer context for free
|
||||
)
|
||||
|
||||
return model, tokenizer
|
||||
|
|
@ -0,0 +1,491 @@
|
|||
"""Tests for Unsloth backend — config, detection, trainer integration, templates."""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from soup_cli.config.schema import TEMPLATES, SoupConfig
|
||||
|
||||
# ─── Config Tests ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestUnslothConfig:
|
||||
"""Test backend config field validation."""
|
||||
|
||||
def test_backend_default_is_transformers(self):
|
||||
"""Default backend should be 'transformers'."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
assert cfg.backend == "transformers"
|
||||
|
||||
def test_backend_unsloth_accepted(self):
|
||||
"""backend: unsloth should be valid."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
backend="unsloth",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
assert cfg.backend == "unsloth"
|
||||
|
||||
def test_backend_transformers_accepted(self):
|
||||
"""backend: transformers should be valid."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
backend="transformers",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
assert cfg.backend == "transformers"
|
||||
|
||||
def test_backend_invalid_rejected(self):
|
||||
"""Invalid backend should raise validation error."""
|
||||
with pytest.raises(Exception):
|
||||
SoupConfig(
|
||||
base="some-model",
|
||||
backend="invalid",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
|
||||
def test_backend_with_sft(self):
|
||||
"""Unsloth backend should work with SFT task."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="sft",
|
||||
backend="unsloth",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
assert cfg.task == "sft"
|
||||
assert cfg.backend == "unsloth"
|
||||
|
||||
def test_backend_with_dpo(self):
|
||||
"""Unsloth backend should work with DPO task."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="dpo",
|
||||
backend="unsloth",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
assert cfg.task == "dpo"
|
||||
assert cfg.backend == "unsloth"
|
||||
|
||||
def test_backend_with_grpo(self):
|
||||
"""Unsloth backend should work with GRPO task."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="grpo",
|
||||
backend="unsloth",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
assert cfg.task == "grpo"
|
||||
assert cfg.backend == "unsloth"
|
||||
|
||||
def test_backend_in_model_dump(self):
|
||||
"""backend field should appear in model_dump output."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
backend="unsloth",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
dump = cfg.model_dump()
|
||||
assert dump["backend"] == "unsloth"
|
||||
|
||||
def test_full_unsloth_config(self):
|
||||
"""Full config with unsloth backend should validate."""
|
||||
cfg = SoupConfig(
|
||||
base="meta-llama/Llama-3.1-8B-Instruct",
|
||||
task="sft",
|
||||
backend="unsloth",
|
||||
data={"train": "./data.jsonl", "format": "alpaca", "max_length": 4096},
|
||||
training={
|
||||
"epochs": 3,
|
||||
"lr": 2e-5,
|
||||
"quantization": "4bit",
|
||||
"lora": {"r": 64, "alpha": 16},
|
||||
},
|
||||
)
|
||||
assert cfg.backend == "unsloth"
|
||||
assert cfg.training.quantization == "4bit"
|
||||
assert cfg.data.max_length == 4096
|
||||
|
||||
|
||||
# ─── Detection Tests ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestUnslothDetection:
|
||||
"""Test unsloth availability detection."""
|
||||
|
||||
def test_is_unsloth_available_when_installed(self):
|
||||
"""Should return True when unsloth is importable."""
|
||||
mock_module = MagicMock()
|
||||
with patch.dict("sys.modules", {"unsloth": mock_module}):
|
||||
|
||||
# Need to reimport to avoid cached result
|
||||
import importlib
|
||||
|
||||
import soup_cli.utils.unsloth
|
||||
|
||||
importlib.reload(soup_cli.utils.unsloth)
|
||||
assert soup_cli.utils.unsloth.is_unsloth_available() is True
|
||||
|
||||
def test_is_unsloth_available_when_not_installed(self):
|
||||
"""Should return False when unsloth is not importable."""
|
||||
from soup_cli.utils.unsloth import is_unsloth_available
|
||||
|
||||
# Default environment doesn't have unsloth
|
||||
# This test works because unsloth isn't installed in test env
|
||||
result = is_unsloth_available()
|
||||
assert isinstance(result, bool)
|
||||
|
||||
def test_get_unsloth_version_when_not_installed(self):
|
||||
"""Should return None when unsloth is not installed."""
|
||||
from soup_cli.utils.unsloth import get_unsloth_version
|
||||
|
||||
result = get_unsloth_version()
|
||||
# In test env, unsloth is not installed
|
||||
assert result is None or isinstance(result, str)
|
||||
|
||||
def test_get_unsloth_version_when_installed(self):
|
||||
"""Should return version string when unsloth is installed."""
|
||||
mock_module = MagicMock()
|
||||
mock_module.__version__ = "2024.11.0"
|
||||
with patch.dict("sys.modules", {"unsloth": mock_module}):
|
||||
import importlib
|
||||
|
||||
import soup_cli.utils.unsloth
|
||||
|
||||
importlib.reload(soup_cli.utils.unsloth)
|
||||
result = soup_cli.utils.unsloth.get_unsloth_version()
|
||||
assert result == "2024.11.0"
|
||||
|
||||
|
||||
# ─── Trainer Integration Tests ──────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestSFTUnslothIntegration:
|
||||
"""Test SFT trainer with unsloth backend."""
|
||||
|
||||
def test_sft_wrapper_init_with_unsloth(self):
|
||||
"""SFTTrainerWrapper should accept unsloth backend config."""
|
||||
from soup_cli.trainer.sft import SFTTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
backend="unsloth",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
wrapper = SFTTrainerWrapper(cfg, device="cuda")
|
||||
assert wrapper.config.backend == "unsloth"
|
||||
|
||||
def test_sft_wrapper_init_with_transformers(self):
|
||||
"""SFTTrainerWrapper should work with default transformers backend."""
|
||||
from soup_cli.trainer.sft import SFTTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
wrapper = SFTTrainerWrapper(cfg, device="cuda")
|
||||
assert wrapper.config.backend == "transformers"
|
||||
|
||||
def test_sft_setup_unsloth_calls_load(self):
|
||||
"""_setup_unsloth should call utils.unsloth.load_model_and_tokenizer."""
|
||||
from soup_cli.trainer.sft import SFTTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
backend="unsloth",
|
||||
data={"train": "./data.jsonl", "max_length": 2048},
|
||||
training={"lora": {"r": 64, "alpha": 16, "dropout": 0.05}},
|
||||
)
|
||||
wrapper = SFTTrainerWrapper(cfg, device="cuda")
|
||||
|
||||
mock_model = MagicMock()
|
||||
mock_model.get_nb_trainable_parameters.return_value = (1000, 100000)
|
||||
mock_tokenizer = MagicMock()
|
||||
mock_tokenizer.pad_token = "pad"
|
||||
|
||||
with patch(
|
||||
"soup_cli.utils.unsloth.load_model_and_tokenizer",
|
||||
return_value=(mock_model, mock_tokenizer),
|
||||
) as mock_load:
|
||||
wrapper._setup_unsloth(cfg, cfg.training)
|
||||
mock_load.assert_called_once_with(
|
||||
model_name="some-model",
|
||||
max_seq_length=2048,
|
||||
quantization="4bit",
|
||||
lora_r=64,
|
||||
lora_alpha=16,
|
||||
lora_dropout=0.05,
|
||||
target_modules="auto",
|
||||
)
|
||||
assert wrapper.model is mock_model
|
||||
assert wrapper.tokenizer is mock_tokenizer
|
||||
|
||||
|
||||
class TestDPOUnslothIntegration:
|
||||
"""Test DPO trainer with unsloth backend."""
|
||||
|
||||
def test_dpo_wrapper_init_with_unsloth(self):
|
||||
"""DPOTrainerWrapper should accept unsloth backend config."""
|
||||
from soup_cli.trainer.dpo import DPOTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="dpo",
|
||||
backend="unsloth",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
wrapper = DPOTrainerWrapper(cfg, device="cuda")
|
||||
assert wrapper.config.backend == "unsloth"
|
||||
|
||||
def test_dpo_setup_unsloth_calls_load(self):
|
||||
"""_setup_unsloth should call utils.unsloth.load_model_and_tokenizer."""
|
||||
from soup_cli.trainer.dpo import DPOTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="dpo",
|
||||
backend="unsloth",
|
||||
data={"train": "./data.jsonl", "max_length": 2048},
|
||||
)
|
||||
wrapper = DPOTrainerWrapper(cfg, device="cuda")
|
||||
|
||||
mock_model = MagicMock()
|
||||
mock_tokenizer = MagicMock()
|
||||
mock_tokenizer.pad_token = "pad"
|
||||
|
||||
with patch(
|
||||
"soup_cli.utils.unsloth.load_model_and_tokenizer",
|
||||
return_value=(mock_model, mock_tokenizer),
|
||||
) as mock_load:
|
||||
wrapper._setup_unsloth(cfg, cfg.training)
|
||||
mock_load.assert_called_once()
|
||||
assert wrapper.model is mock_model
|
||||
|
||||
|
||||
class TestGRPOUnslothIntegration:
|
||||
"""Test GRPO trainer with unsloth backend."""
|
||||
|
||||
def test_grpo_wrapper_init_with_unsloth(self):
|
||||
"""GRPOTrainerWrapper should accept unsloth backend config."""
|
||||
from soup_cli.trainer.grpo import GRPOTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="grpo",
|
||||
backend="unsloth",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
wrapper = GRPOTrainerWrapper(cfg, device="cuda")
|
||||
assert wrapper.config.backend == "unsloth"
|
||||
|
||||
def test_grpo_setup_unsloth_calls_load(self):
|
||||
"""_setup_unsloth should call utils.unsloth.load_model_and_tokenizer."""
|
||||
from soup_cli.trainer.grpo import GRPOTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="grpo",
|
||||
backend="unsloth",
|
||||
data={"train": "./data.jsonl", "max_length": 4096},
|
||||
)
|
||||
wrapper = GRPOTrainerWrapper(cfg, device="cuda")
|
||||
|
||||
mock_model = MagicMock()
|
||||
mock_tokenizer = MagicMock()
|
||||
mock_tokenizer.pad_token = "pad"
|
||||
|
||||
with patch(
|
||||
"soup_cli.utils.unsloth.load_model_and_tokenizer",
|
||||
return_value=(mock_model, mock_tokenizer),
|
||||
) as mock_load:
|
||||
wrapper._setup_unsloth(cfg, cfg.training)
|
||||
mock_load.assert_called_once()
|
||||
assert wrapper.model is mock_model
|
||||
|
||||
|
||||
# ─── Template Tests ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestTemplatesHaveUnslothHint:
|
||||
"""Test that templates mention unsloth backend as an option."""
|
||||
|
||||
def test_chat_template_mentions_unsloth(self):
|
||||
assert "unsloth" in TEMPLATES["chat"]
|
||||
|
||||
def test_code_template_mentions_unsloth(self):
|
||||
assert "unsloth" in TEMPLATES["code"]
|
||||
|
||||
def test_reasoning_template_mentions_unsloth(self):
|
||||
assert "unsloth" in TEMPLATES["reasoning"]
|
||||
|
||||
def test_medical_template_mentions_unsloth(self):
|
||||
assert "unsloth" in TEMPLATES["medical"]
|
||||
|
||||
def test_templates_default_backend_commented(self):
|
||||
"""Templates should have unsloth commented out (not active by default)."""
|
||||
for name, template in TEMPLATES.items():
|
||||
assert "# backend: unsloth" in template, f"{name} template missing unsloth hint"
|
||||
|
||||
|
||||
# ─── Sweep Shortcut Tests ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestBackendSweepParam:
|
||||
"""Test backend parameter in sweep shortcuts."""
|
||||
|
||||
def test_backend_shortcut(self):
|
||||
from soup_cli.commands.sweep import _set_nested_param
|
||||
|
||||
config = {"backend": "transformers"}
|
||||
_set_nested_param(config, "backend", "unsloth")
|
||||
assert config["backend"] == "unsloth"
|
||||
|
||||
|
||||
# ─── Doctor Tests ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestDoctorUnsloth:
|
||||
"""Test that doctor checks for unsloth."""
|
||||
|
||||
def test_unsloth_in_deps_list(self):
|
||||
from soup_cli.commands.doctor import DEPS
|
||||
|
||||
pkg_names = [pkg_name for _, pkg_name, _, _ in DEPS]
|
||||
assert "unsloth" in pkg_names
|
||||
|
||||
def test_unsloth_is_optional(self):
|
||||
from soup_cli.commands.doctor import DEPS
|
||||
|
||||
for import_name, pkg_name, _, required in DEPS:
|
||||
if pkg_name == "unsloth":
|
||||
assert required is False
|
||||
break
|
||||
else:
|
||||
pytest.fail("unsloth not found in DEPS")
|
||||
|
||||
|
||||
# ─── Load Function Tests ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestLoadModelAndTokenizer:
|
||||
"""Test the load_model_and_tokenizer function with mocked unsloth."""
|
||||
|
||||
def test_load_with_4bit(self):
|
||||
"""Should pass load_in_4bit=True for 4bit quantization."""
|
||||
mock_flm = MagicMock()
|
||||
mock_model = MagicMock()
|
||||
mock_tokenizer = MagicMock()
|
||||
mock_flm.from_pretrained.return_value = (mock_model, mock_tokenizer)
|
||||
mock_flm.get_peft_model.return_value = mock_model
|
||||
|
||||
with patch.dict("sys.modules", {"unsloth": MagicMock(FastLanguageModel=mock_flm)}):
|
||||
import importlib
|
||||
|
||||
import soup_cli.utils.unsloth
|
||||
|
||||
importlib.reload(soup_cli.utils.unsloth)
|
||||
model, tokenizer = soup_cli.utils.unsloth.load_model_and_tokenizer(
|
||||
model_name="test-model",
|
||||
max_seq_length=2048,
|
||||
quantization="4bit",
|
||||
)
|
||||
mock_flm.from_pretrained.assert_called_once()
|
||||
call_kwargs = mock_flm.from_pretrained.call_args
|
||||
assert call_kwargs[1]["load_in_4bit"] is True
|
||||
|
||||
def test_load_with_none_quantization(self):
|
||||
"""Should pass load_in_4bit=False for 'none' quantization."""
|
||||
mock_flm = MagicMock()
|
||||
mock_model = MagicMock()
|
||||
mock_tokenizer = MagicMock()
|
||||
mock_flm.from_pretrained.return_value = (mock_model, mock_tokenizer)
|
||||
mock_flm.get_peft_model.return_value = mock_model
|
||||
|
||||
with patch.dict("sys.modules", {"unsloth": MagicMock(FastLanguageModel=mock_flm)}):
|
||||
import importlib
|
||||
|
||||
import soup_cli.utils.unsloth
|
||||
|
||||
importlib.reload(soup_cli.utils.unsloth)
|
||||
soup_cli.utils.unsloth.load_model_and_tokenizer(
|
||||
model_name="test-model",
|
||||
max_seq_length=2048,
|
||||
quantization="none",
|
||||
)
|
||||
call_kwargs = mock_flm.from_pretrained.call_args
|
||||
assert call_kwargs[1]["load_in_4bit"] is False
|
||||
|
||||
def test_load_auto_target_modules(self):
|
||||
"""auto target_modules should expand to default linear layers."""
|
||||
mock_flm = MagicMock()
|
||||
mock_model = MagicMock()
|
||||
mock_tokenizer = MagicMock()
|
||||
mock_flm.from_pretrained.return_value = (mock_model, mock_tokenizer)
|
||||
mock_flm.get_peft_model.return_value = mock_model
|
||||
|
||||
with patch.dict("sys.modules", {"unsloth": MagicMock(FastLanguageModel=mock_flm)}):
|
||||
import importlib
|
||||
|
||||
import soup_cli.utils.unsloth
|
||||
|
||||
importlib.reload(soup_cli.utils.unsloth)
|
||||
soup_cli.utils.unsloth.load_model_and_tokenizer(
|
||||
model_name="test-model",
|
||||
max_seq_length=2048,
|
||||
target_modules="auto",
|
||||
)
|
||||
peft_call_kwargs = mock_flm.get_peft_model.call_args
|
||||
target = peft_call_kwargs[1]["target_modules"]
|
||||
assert "q_proj" in target
|
||||
assert "v_proj" in target
|
||||
|
||||
def test_load_custom_target_modules(self):
|
||||
"""Custom target_modules list should be passed through."""
|
||||
mock_flm = MagicMock()
|
||||
mock_model = MagicMock()
|
||||
mock_tokenizer = MagicMock()
|
||||
mock_flm.from_pretrained.return_value = (mock_model, mock_tokenizer)
|
||||
mock_flm.get_peft_model.return_value = mock_model
|
||||
|
||||
with patch.dict("sys.modules", {"unsloth": MagicMock(FastLanguageModel=mock_flm)}):
|
||||
import importlib
|
||||
|
||||
import soup_cli.utils.unsloth
|
||||
|
||||
importlib.reload(soup_cli.utils.unsloth)
|
||||
custom_modules = ["q_proj", "k_proj"]
|
||||
soup_cli.utils.unsloth.load_model_and_tokenizer(
|
||||
model_name="test-model",
|
||||
max_seq_length=2048,
|
||||
target_modules=custom_modules,
|
||||
)
|
||||
peft_call_kwargs = mock_flm.get_peft_model.call_args
|
||||
assert peft_call_kwargs[1]["target_modules"] == custom_modules
|
||||
|
||||
def test_load_lora_params_passed(self):
|
||||
"""LoRA params should be forwarded to get_peft_model."""
|
||||
mock_flm = MagicMock()
|
||||
mock_model = MagicMock()
|
||||
mock_tokenizer = MagicMock()
|
||||
mock_flm.from_pretrained.return_value = (mock_model, mock_tokenizer)
|
||||
mock_flm.get_peft_model.return_value = mock_model
|
||||
|
||||
with patch.dict("sys.modules", {"unsloth": MagicMock(FastLanguageModel=mock_flm)}):
|
||||
import importlib
|
||||
|
||||
import soup_cli.utils.unsloth
|
||||
|
||||
importlib.reload(soup_cli.utils.unsloth)
|
||||
soup_cli.utils.unsloth.load_model_and_tokenizer(
|
||||
model_name="test-model",
|
||||
max_seq_length=2048,
|
||||
lora_r=128,
|
||||
lora_alpha=32,
|
||||
lora_dropout=0.1,
|
||||
)
|
||||
peft_kwargs = mock_flm.get_peft_model.call_args[1]
|
||||
assert peft_kwargs["r"] == 128
|
||||
assert peft_kwargs["lora_alpha"] == 32
|
||||
assert peft_kwargs["lora_dropout"] == 0.1
|
||||
Loading…
Reference in New Issue