diff --git a/CLAUDE.md b/CLAUDE.md index 884abca..769c174 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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 `` 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 | diff --git a/README.md b/README.md index 2124442..df7c12e 100644 --- a/README.md +++ b/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 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) | diff --git a/pyproject.toml b/pyproject.toml index be0b3e0..48ffd3c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/soup_cli/__init__.py b/soup_cli/__init__.py index b1398d3..1311874 100644 --- a/soup_cli/__init__.py +++ b/soup_cli/__init__.py @@ -1,3 +1,3 @@ """Soup CLI — Fine-tune LLMs in one command.""" -__version__ = "0.4.2" +__version__ = "0.4.3" diff --git a/soup_cli/commands/doctor.py b/soup_cli/commands/doctor.py index a6d3e06..6ab08e1 100644 --- a/soup_cli/commands/doctor.py +++ b/soup_cli/commands/doctor.py @@ -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), ] diff --git a/soup_cli/commands/sweep.py b/soup_cli/commands/sweep.py index bb393b7..e3da1aa 100644 --- a/soup_cli/commands/sweep.py +++ b/soup_cli/commands/sweep.py @@ -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) diff --git a/soup_cli/commands/train.py b/soup_cli/commands/train.py index fa00b81..8e667b8 100644 --- a/soup_cli/commands/train.py +++ b/soup_cli/commands/train.py @@ -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.[/]") diff --git a/soup_cli/config/schema.py b/soup_cli/config/schema.py index 4c1e788..abc0788 100644 --- a/soup_cli/config/schema.py +++ b/soup_cli/config/schema.py @@ -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 diff --git a/soup_cli/trainer/dpo.py b/soup_cli/trainer/dpo.py index de96d96..1c086f3 100644 --- a/soup_cli/trainer/dpo.py +++ b/soup_cli/trainer/dpo.py @@ -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, diff --git a/soup_cli/trainer/grpo.py b/soup_cli/trainer/grpo.py index 22d18e9..617bfd4 100644 --- a/soup_cli/trainer/grpo.py +++ b/soup_cli/trainer/grpo.py @@ -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, diff --git a/soup_cli/trainer/sft.py b/soup_cli/trainer/sft.py index e618703..c5179e3 100644 --- a/soup_cli/trainer/sft.py +++ b/soup_cli/trainer/sft.py @@ -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, diff --git a/soup_cli/utils/unsloth.py b/soup_cli/utils/unsloth.py new file mode 100644 index 0000000..89befb1 --- /dev/null +++ b/soup_cli/utils/unsloth.py @@ -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 diff --git a/tests/test_unsloth.py b/tests/test_unsloth.py new file mode 100644 index 0000000..b203ee1 --- /dev/null +++ b/tests/test_unsloth.py @@ -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