mirror of https://github.com/razor-ai/soup.git
Initial project setup: CLI skeleton + config + trainer + data pipeline
- Typer CLI: soup init, soup train, soup data inspect/validate - Pydantic config schema with YAML loader and validation - Data pipeline: JSONL/JSON/CSV/Parquet + HuggingFace datasets - Format detection: Alpaca, ShareGPT, ChatML (auto-detect) - SFT trainer wrapper over transformers + peft + trl - QLoRA/LoRA support with auto batch size estimation - GPU detection (CUDA/MPS/CPU) and memory calculation - Rich live terminal dashboard for training monitoring - Config templates: chat, code, medical - Tests (pytest) + GitHub Actions CI - MIT license Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Alpamys
commit
d6e932a1d3
|
|
@ -0,0 +1,31 @@
|
|||
name: CI
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
pull_request:
|
||||
branches: [main]
|
||||
|
||||
jobs:
|
||||
lint:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.11"
|
||||
- run: pip install ruff
|
||||
- run: ruff check soup_cli/ tests/
|
||||
|
||||
test:
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: ["3.9", "3.11", "3.12"]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- run: pip install -e ".[dev]"
|
||||
- run: pytest tests/ -v --tb=short
|
||||
|
|
@ -0,0 +1,36 @@
|
|||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.egg-info/
|
||||
dist/
|
||||
build/
|
||||
*.egg
|
||||
.eggs/
|
||||
|
||||
# Virtual env
|
||||
.venv/
|
||||
venv/
|
||||
env/
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# Training outputs
|
||||
output/
|
||||
checkpoints/
|
||||
wandb/
|
||||
|
||||
# Soup local data
|
||||
.soup/
|
||||
*.db
|
||||
|
||||
# Secrets
|
||||
.env
|
||||
*.key
|
||||
|
|
@ -0,0 +1,21 @@
|
|||
MIT License
|
||||
|
||||
Copyright (c) 2026 Soup Team
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
|
|
@ -0,0 +1,147 @@
|
|||
# 🍜 Soup
|
||||
|
||||
**Fine-tune LLMs in one command. No SSH, no config hell.**
|
||||
|
||||
Soup turns the pain of LLM fine-tuning into a simple workflow. One config, one command, done.
|
||||
|
||||
```bash
|
||||
pip install soup-cli
|
||||
soup init --template chat
|
||||
soup train
|
||||
```
|
||||
|
||||
## Why Soup?
|
||||
|
||||
Training LLMs is still painful. Even experienced teams spend 30-50% of their time fighting infrastructure instead of improving models. Soup fixes that.
|
||||
|
||||
- **Zero SSH.** Never SSH into a broken GPU box again.
|
||||
- **One config.** A simple YAML file is all you need.
|
||||
- **Auto everything.** Batch size, GPU detection, quantization — handled.
|
||||
- **Works locally.** Train on your own GPU with QLoRA. No cloud required.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### 1. Install
|
||||
|
||||
```bash
|
||||
pip install soup-cli
|
||||
```
|
||||
|
||||
### 2. Create config
|
||||
|
||||
```bash
|
||||
# Interactive wizard
|
||||
soup init
|
||||
|
||||
# Or use a template
|
||||
soup init --template chat # conversational fine-tune
|
||||
soup init --template code # code generation
|
||||
soup init --template medical # domain expert
|
||||
```
|
||||
|
||||
### 3. Train
|
||||
|
||||
```bash
|
||||
soup train --config soup.yaml
|
||||
```
|
||||
|
||||
That's it. Soup handles LoRA setup, quantization, batch size, monitoring, and checkpoints.
|
||||
|
||||
### 4. Test your model
|
||||
|
||||
```bash
|
||||
soup chat --model ./output
|
||||
```
|
||||
|
||||
### 5. Push to HuggingFace
|
||||
|
||||
```bash
|
||||
soup push --model ./output --repo your-username/my-model
|
||||
```
|
||||
|
||||
## Config Example
|
||||
|
||||
```yaml
|
||||
base: meta-llama/Llama-3.1-8B-Instruct
|
||||
task: sft
|
||||
|
||||
data:
|
||||
train: ./data/train.jsonl
|
||||
format: alpaca
|
||||
val_split: 0.1
|
||||
|
||||
training:
|
||||
epochs: 3
|
||||
lr: 2e-5
|
||||
batch_size: auto
|
||||
lora:
|
||||
r: 64
|
||||
alpha: 16
|
||||
quantization: 4bit
|
||||
|
||||
output: ./output
|
||||
```
|
||||
|
||||
## Data Formats
|
||||
|
||||
Soup supports these formats (auto-detected):
|
||||
|
||||
**Alpaca:**
|
||||
```json
|
||||
{"instruction": "Explain gravity", "input": "", "output": "Gravity is..."}
|
||||
```
|
||||
|
||||
**ShareGPT:**
|
||||
```json
|
||||
{"conversations": [{"from": "human", "value": "Hi"}, {"from": "gpt", "value": "Hello!"}]}
|
||||
```
|
||||
|
||||
**ChatML:**
|
||||
```json
|
||||
{"messages": [{"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Hello!"}]}
|
||||
```
|
||||
|
||||
## Data Tools
|
||||
|
||||
```bash
|
||||
# Inspect a dataset
|
||||
soup data inspect ./data/train.jsonl
|
||||
|
||||
# Validate format
|
||||
soup data validate ./data/train.jsonl --format alpaca
|
||||
```
|
||||
|
||||
## Features
|
||||
|
||||
| Feature | Status |
|
||||
|---|---|
|
||||
| LoRA / QLoRA fine-tuning | ✅ |
|
||||
| SFT (Supervised Fine-Tune) | ✅ |
|
||||
| DPO (Direct Preference Optimization) | 🔜 |
|
||||
| Auto batch size | ✅ |
|
||||
| Auto GPU detection (CUDA/MPS/CPU) | ✅ |
|
||||
| Live terminal dashboard | ✅ |
|
||||
| Alpaca / ShareGPT / ChatML formats | ✅ |
|
||||
| HuggingFace datasets support | ✅ |
|
||||
| Experiment tracking | 🔜 |
|
||||
| Web dashboard | 🔜 |
|
||||
| Cloud mode (BYOG) | 🔜 |
|
||||
|
||||
## Requirements
|
||||
|
||||
- Python 3.9+
|
||||
- GPU with CUDA (recommended) or Apple Silicon (MPS) or CPU (slow)
|
||||
- 8 GB+ VRAM for 7B models with QLoRA
|
||||
|
||||
## Development
|
||||
|
||||
```bash
|
||||
git clone https://github.com/MakazhanAlpamys/Soup.git
|
||||
cd Soup
|
||||
pip install -e ".[dev]"
|
||||
pytest tests/ -v
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
|
|
@ -0,0 +1,59 @@
|
|||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "soup-cli"
|
||||
version = "0.1.0"
|
||||
description = "Fine-tune LLMs in one command. No SSH, no config hell."
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
requires-python = ">=3.9"
|
||||
authors = [
|
||||
{ name = "Soup Team" },
|
||||
]
|
||||
keywords = ["llm", "fine-tuning", "lora", "qlora", "machine-learning"]
|
||||
classifiers = [
|
||||
"Development Status :: 3 - Alpha",
|
||||
"Intended Audience :: Developers",
|
||||
"Intended Audience :: Science/Research",
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||
]
|
||||
dependencies = [
|
||||
"typer>=0.9.0",
|
||||
"rich>=13.0.0",
|
||||
"pydantic>=2.0.0",
|
||||
"pyyaml>=6.0",
|
||||
"torch>=2.0.0",
|
||||
"transformers>=4.36.0",
|
||||
"peft>=0.7.0",
|
||||
"trl>=0.7.0",
|
||||
"datasets>=2.14.0",
|
||||
"bitsandbytes>=0.41.0",
|
||||
"accelerate>=0.25.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
eval = ["lm-eval>=0.4.0"]
|
||||
dev = ["pytest>=7.0", "ruff>=0.1.0", "pytest-cov>=4.0"]
|
||||
ui = ["fastapi>=0.104.0", "uvicorn>=0.24.0"]
|
||||
|
||||
[project.scripts]
|
||||
soup = "soup_cli.cli:app"
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://github.com/MakazhanAlpamys/Soup"
|
||||
Repository = "https://github.com/MakazhanAlpamys/Soup"
|
||||
Issues = "https://github.com/MakazhanAlpamys/Soup/issues"
|
||||
|
||||
[tool.ruff]
|
||||
target-version = "py39"
|
||||
line-length = 100
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["E", "F", "I", "N", "W"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = ["tests"]
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
"""Soup CLI — Fine-tune LLMs in one command."""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
"""Allow running as `python -m soup_cli`."""
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
app()
|
||||
|
|
@ -0,0 +1,33 @@
|
|||
"""Main CLI entry point — all commands registered here."""
|
||||
|
||||
import typer
|
||||
from rich.console import Console
|
||||
|
||||
from soup_cli import __version__
|
||||
from soup_cli.commands import data, init, train
|
||||
|
||||
console = Console()
|
||||
|
||||
app = typer.Typer(
|
||||
name="soup",
|
||||
help="Fine-tune LLMs in one command. No SSH, no config hell.",
|
||||
no_args_is_help=True,
|
||||
rich_markup_mode="rich",
|
||||
)
|
||||
|
||||
# Register sub-commands
|
||||
app.command()(init.init)
|
||||
app.command()(train.train)
|
||||
app.add_typer(data.app, name="data", help="Dataset tools: inspect, convert, validate.")
|
||||
|
||||
|
||||
@app.command()
|
||||
def version():
|
||||
"""Show Soup CLI version."""
|
||||
console.print(f"[bold green]soup[/] v{__version__}")
|
||||
|
||||
|
||||
@app.callback(invoke_without_command=True)
|
||||
def main(ctx: typer.Context):
|
||||
"""Soup — fine-tune LLMs in one command."""
|
||||
pass
|
||||
|
|
@ -0,0 +1,78 @@
|
|||
"""soup data — dataset inspection and tools."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import typer
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
from soup_cli.data.loader import load_raw_data
|
||||
from soup_cli.data.validator import validate_and_stats
|
||||
|
||||
console = Console()
|
||||
|
||||
app = typer.Typer(no_args_is_help=True)
|
||||
|
||||
|
||||
@app.command()
|
||||
def inspect(
|
||||
path: str = typer.Argument(..., help="Path to dataset file (jsonl, csv, parquet)"),
|
||||
rows: int = typer.Option(5, "--rows", "-r", help="Number of sample rows to show"),
|
||||
):
|
||||
"""Inspect a dataset: show stats and sample rows."""
|
||||
file_path = Path(path)
|
||||
if not file_path.exists():
|
||||
console.print(f"[red]File not found: {file_path}[/]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
console.print(f"[dim]Inspecting {file_path}...[/]\n")
|
||||
data = load_raw_data(file_path)
|
||||
stats = validate_and_stats(data)
|
||||
|
||||
# Print stats
|
||||
stats_table = Table(title="Dataset Stats")
|
||||
stats_table.add_column("Metric", style="bold")
|
||||
stats_table.add_column("Value")
|
||||
stats_table.add_row("Total samples", str(stats["total"]))
|
||||
stats_table.add_row("Columns", ", ".join(stats["columns"]))
|
||||
stats_table.add_row("Avg length (chars)", str(stats["avg_length"]))
|
||||
stats_table.add_row("Min length", str(stats["min_length"]))
|
||||
stats_table.add_row("Max length", str(stats["max_length"]))
|
||||
stats_table.add_row("Empty fields", str(stats["empty_fields"]))
|
||||
stats_table.add_row("Duplicates", str(stats["duplicates"]))
|
||||
console.print(stats_table)
|
||||
|
||||
# Print sample rows
|
||||
if rows > 0 and len(data) > 0:
|
||||
console.print(f"\n[bold]Sample rows ({min(rows, len(data))}):[/]")
|
||||
sample_table = Table(show_lines=True)
|
||||
for col in stats["columns"][:5]: # max 5 columns
|
||||
sample_table.add_column(col, max_width=60)
|
||||
for row in data[: min(rows, len(data))]:
|
||||
values = [str(row.get(col, ""))[:60] for col in stats["columns"][:5]]
|
||||
sample_table.add_row(*values)
|
||||
console.print(sample_table)
|
||||
|
||||
|
||||
@app.command()
|
||||
def validate(
|
||||
path: str = typer.Argument(..., help="Path to dataset file"),
|
||||
format: str = typer.Option("alpaca", "--format", "-f", help="Expected format: alpaca, sharegpt, chatml"),
|
||||
):
|
||||
"""Validate dataset format and report issues."""
|
||||
file_path = Path(path)
|
||||
if not file_path.exists():
|
||||
console.print(f"[red]File not found: {file_path}[/]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
data = load_raw_data(file_path)
|
||||
stats = validate_and_stats(data, expected_format=format)
|
||||
|
||||
if stats["issues"]:
|
||||
console.print("[yellow]Issues found:[/]")
|
||||
for issue in stats["issues"]:
|
||||
console.print(f" [yellow]![/] {issue}")
|
||||
else:
|
||||
console.print("[bold green]Dataset is valid![/]")
|
||||
|
||||
console.print(f"\n[green]{stats['valid_rows']}/{stats['total']} rows valid for {format} format[/]")
|
||||
|
|
@ -0,0 +1,95 @@
|
|||
"""soup init — interactive project setup wizard."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import typer
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.prompt import Prompt
|
||||
|
||||
from soup_cli.config.schema import TEMPLATES
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
def init(
|
||||
template: str = typer.Option(
|
||||
None,
|
||||
"--template",
|
||||
"-t",
|
||||
help="Use a template: chat, code, medical",
|
||||
),
|
||||
output: str = typer.Option(
|
||||
"soup.yaml",
|
||||
"--output",
|
||||
"-o",
|
||||
help="Output config file path",
|
||||
),
|
||||
):
|
||||
"""Create a new soup.yaml config interactively or from a template."""
|
||||
output_path = Path(output)
|
||||
|
||||
if output_path.exists():
|
||||
overwrite = typer.confirm(f"{output_path} already exists. Overwrite?")
|
||||
if not overwrite:
|
||||
raise typer.Exit()
|
||||
|
||||
if template:
|
||||
if template not in TEMPLATES:
|
||||
console.print(f"[red]Unknown template: {template}[/]")
|
||||
console.print(f"Available: {', '.join(TEMPLATES.keys())}")
|
||||
raise typer.Exit(1)
|
||||
config_text = TEMPLATES[template]
|
||||
console.print(f"[green]Using template:[/] {template}")
|
||||
else:
|
||||
config_text = _interactive_wizard()
|
||||
|
||||
output_path.write_text(config_text, encoding="utf-8")
|
||||
console.print(
|
||||
Panel(
|
||||
f"[bold green]Config saved to {output_path}[/]\n\n"
|
||||
f"Next step: [bold]soup train --config {output_path}[/]",
|
||||
title="Ready!",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _interactive_wizard() -> str:
|
||||
"""Walk user through config creation."""
|
||||
console.print(Panel("[bold]Soup Config Wizard[/]", subtitle="Let's set up your training"))
|
||||
|
||||
base_model = Prompt.ask(
|
||||
"Base model",
|
||||
default="meta-llama/Llama-3.1-8B-Instruct",
|
||||
)
|
||||
task = Prompt.ask("Task", choices=["sft", "dpo"], default="sft")
|
||||
data_path = Prompt.ask("Training data path", default="./data/train.jsonl")
|
||||
data_format = Prompt.ask("Data format", choices=["alpaca", "sharegpt", "chatml"], default="alpaca")
|
||||
epochs = Prompt.ask("Epochs", default="3")
|
||||
use_qlora = Prompt.ask("Use QLoRA (4-bit)?", choices=["yes", "no"], default="yes")
|
||||
|
||||
quantization = "4bit" if use_qlora == "yes" else "none"
|
||||
|
||||
return f"""# Soup training config
|
||||
# Docs: https://github.com/MakazhanAlpamys/Soup
|
||||
|
||||
base: {base_model}
|
||||
task: {task}
|
||||
|
||||
data:
|
||||
train: {data_path}
|
||||
format: {data_format}
|
||||
val_split: 0.1
|
||||
|
||||
training:
|
||||
epochs: {epochs}
|
||||
lr: 2e-5
|
||||
batch_size: auto
|
||||
lora:
|
||||
r: 64
|
||||
alpha: 16
|
||||
target_modules: auto
|
||||
quantization: {quantization}
|
||||
|
||||
output: ./output
|
||||
"""
|
||||
|
|
@ -0,0 +1,97 @@
|
|||
"""soup train — the main training command."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import typer
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
|
||||
from soup_cli.config.loader import load_config
|
||||
from soup_cli.data.loader import load_dataset
|
||||
from soup_cli.monitoring.display import TrainingDisplay
|
||||
from soup_cli.trainer.sft import SFTTrainerWrapper
|
||||
from soup_cli.utils.gpu import detect_device, get_gpu_info
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
def train(
|
||||
config: str = typer.Option(
|
||||
"soup.yaml",
|
||||
"--config",
|
||||
"-c",
|
||||
help="Path to soup.yaml config file",
|
||||
),
|
||||
name: str = typer.Option(
|
||||
None,
|
||||
"--name",
|
||||
"-n",
|
||||
help="Experiment name (auto-generated if not set)",
|
||||
),
|
||||
dry_run: bool = typer.Option(
|
||||
False,
|
||||
"--dry-run",
|
||||
help="Validate config and data without training",
|
||||
),
|
||||
):
|
||||
"""Start training from a soup.yaml config."""
|
||||
config_path = Path(config)
|
||||
if not config_path.exists():
|
||||
console.print(f"[red]Config not found: {config_path}[/]")
|
||||
console.print("Run [bold]soup init[/] to create one.")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Load & validate config
|
||||
console.print(f"[dim]Loading config from {config_path}...[/]")
|
||||
cfg = load_config(config_path)
|
||||
|
||||
# Detect hardware
|
||||
device, device_name = detect_device()
|
||||
gpu_info = get_gpu_info()
|
||||
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}[/]",
|
||||
title="Training Setup",
|
||||
)
|
||||
)
|
||||
|
||||
if dry_run:
|
||||
console.print("[yellow]Dry run — validating data...[/]")
|
||||
dataset = load_dataset(cfg.data)
|
||||
console.print(f"[green]Data OK:[/] {len(dataset['train'])} train samples")
|
||||
if "val" in dataset:
|
||||
console.print(f"[green]Val:[/] {len(dataset['val'])} samples")
|
||||
console.print("[green]Config valid. Ready to train![/]")
|
||||
raise typer.Exit()
|
||||
|
||||
# Load data
|
||||
console.print("[dim]Loading dataset...[/]")
|
||||
dataset = load_dataset(cfg.data)
|
||||
console.print(f"[green]Loaded:[/] {len(dataset['train'])} train samples")
|
||||
|
||||
# Build trainer
|
||||
console.print("[dim]Setting up model + trainer...[/]")
|
||||
trainer_wrapper = SFTTrainerWrapper(cfg, device=device)
|
||||
trainer_wrapper.setup(dataset)
|
||||
|
||||
# Train with live display
|
||||
display = TrainingDisplay(cfg, device_name=device_name)
|
||||
console.print("[bold green]Training started![/]\n")
|
||||
result = trainer_wrapper.train(display=display)
|
||||
|
||||
# Report
|
||||
console.print(
|
||||
Panel(
|
||||
f"Loss: [bold]{result['initial_loss']:.4f} → {result['final_loss']:.4f}[/]\n"
|
||||
f"Duration: [bold]{result['duration']}[/]\n"
|
||||
f"Output: [bold]{result['output_dir']}[/]\n\n"
|
||||
f"Quick test: [bold]soup chat --model {result['output_dir']}[/]\n"
|
||||
f"Push to HF: [bold]soup push --model {result['output_dir']}[/]",
|
||||
title="[bold green]Training Complete![/]",
|
||||
)
|
||||
)
|
||||
|
|
@ -0,0 +1,31 @@
|
|||
"""Load and validate soup.yaml configs."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from pydantic import ValidationError
|
||||
from rich.console import Console
|
||||
|
||||
from soup_cli.config.schema import SoupConfig
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
def load_config(path: Path) -> SoupConfig:
|
||||
"""Load a soup.yaml file and return validated SoupConfig."""
|
||||
raw = yaml.safe_load(path.read_text(encoding="utf-8"))
|
||||
|
||||
if raw is None:
|
||||
console.print("[red]Config file is empty[/]")
|
||||
raise SystemExit(1)
|
||||
|
||||
try:
|
||||
config = SoupConfig(**raw)
|
||||
except ValidationError as e:
|
||||
console.print("[red bold]Config validation error:[/]\n")
|
||||
for err in e.errors():
|
||||
loc = " → ".join(str(l) for l in err["loc"])
|
||||
console.print(f" [red]{loc}:[/] {err['msg']}")
|
||||
raise SystemExit(1)
|
||||
|
||||
return config
|
||||
|
|
@ -0,0 +1,137 @@
|
|||
"""Pydantic schemas for soup.yaml config — single source of truth."""
|
||||
|
||||
from typing import Literal, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class LoraConfig(BaseModel):
|
||||
r: int = Field(default=64, description="LoRA rank")
|
||||
alpha: int = Field(default=16, description="LoRA alpha")
|
||||
dropout: float = Field(default=0.05, description="LoRA dropout")
|
||||
target_modules: str | list[str] = Field(
|
||||
default="auto",
|
||||
description="Target modules for LoRA. 'auto' = let peft decide.",
|
||||
)
|
||||
|
||||
|
||||
class DataConfig(BaseModel):
|
||||
train: str = Field(..., description="Path to training data or HF dataset name")
|
||||
format: Literal["alpaca", "sharegpt", "chatml", "auto"] = Field(
|
||||
default="auto",
|
||||
description="Data format",
|
||||
)
|
||||
val_split: float = Field(default=0.1, ge=0.0, le=0.5, description="Validation split ratio")
|
||||
max_length: int = Field(default=2048, description="Max sequence length in tokens")
|
||||
|
||||
|
||||
class TrainingConfig(BaseModel):
|
||||
epochs: int = Field(default=3, ge=1, description="Number of training epochs")
|
||||
lr: float = Field(default=2e-5, gt=0, description="Learning rate")
|
||||
batch_size: int | Literal["auto"] = Field(
|
||||
default="auto",
|
||||
description="Batch size. 'auto' = find max that fits in memory.",
|
||||
)
|
||||
gradient_accumulation_steps: int = Field(default=4, ge=1)
|
||||
warmup_ratio: float = Field(default=0.03, ge=0.0, le=0.5)
|
||||
weight_decay: float = Field(default=0.01, ge=0.0)
|
||||
max_grad_norm: float = Field(default=1.0, gt=0)
|
||||
lora: LoraConfig = Field(default_factory=LoraConfig)
|
||||
quantization: Literal["4bit", "8bit", "none"] = Field(
|
||||
default="4bit",
|
||||
description="Quantization: 4bit (QLoRA), 8bit, or none (full precision)",
|
||||
)
|
||||
optimizer: str = Field(default="adamw_torch", description="Optimizer name")
|
||||
scheduler: str = Field(default="cosine", description="LR scheduler type")
|
||||
save_steps: int = Field(default=100, description="Save checkpoint every N steps")
|
||||
logging_steps: int = Field(default=10, description="Log metrics every N steps")
|
||||
|
||||
|
||||
class SoupConfig(BaseModel):
|
||||
"""Root config for soup.yaml."""
|
||||
|
||||
base: str = Field(..., description="Base model name or path (HF model ID)")
|
||||
task: Literal["sft", "dpo"] = Field(default="sft", description="Training task type")
|
||||
data: DataConfig
|
||||
training: TrainingConfig = Field(default_factory=TrainingConfig)
|
||||
output: str = Field(default="./output", description="Output directory for trained model")
|
||||
experiment_name: Optional[str] = Field(default=None, description="Experiment name for tracking")
|
||||
|
||||
|
||||
# --- Built-in templates ---
|
||||
|
||||
TEMPLATES: dict[str, str] = {
|
||||
"chat": """# Soup template: Chat Assistant
|
||||
# Fine-tune a model for conversational chat
|
||||
|
||||
base: meta-llama/Llama-3.1-8B-Instruct
|
||||
task: sft
|
||||
|
||||
data:
|
||||
train: ./data/train.jsonl
|
||||
format: alpaca
|
||||
val_split: 0.1
|
||||
max_length: 2048
|
||||
|
||||
training:
|
||||
epochs: 3
|
||||
lr: 2e-5
|
||||
batch_size: auto
|
||||
lora:
|
||||
r: 64
|
||||
alpha: 16
|
||||
target_modules: auto
|
||||
quantization: 4bit
|
||||
|
||||
output: ./output
|
||||
""",
|
||||
"code": """# Soup template: Code Model
|
||||
# Fine-tune a model for code generation / completion
|
||||
|
||||
base: codellama/CodeLlama-7b-Instruct-hf
|
||||
task: sft
|
||||
|
||||
data:
|
||||
train: ./data/code_train.jsonl
|
||||
format: alpaca
|
||||
val_split: 0.1
|
||||
max_length: 4096
|
||||
|
||||
training:
|
||||
epochs: 2
|
||||
lr: 1e-5
|
||||
batch_size: auto
|
||||
lora:
|
||||
r: 128
|
||||
alpha: 32
|
||||
target_modules: auto
|
||||
quantization: 4bit
|
||||
|
||||
output: ./output
|
||||
""",
|
||||
"medical": """# Soup template: Medical / Domain Expert
|
||||
# Fine-tune a model with domain-specific knowledge
|
||||
|
||||
base: meta-llama/Llama-3.1-8B-Instruct
|
||||
task: sft
|
||||
|
||||
data:
|
||||
train: ./data/medical_train.jsonl
|
||||
format: alpaca
|
||||
val_split: 0.15
|
||||
max_length: 2048
|
||||
|
||||
training:
|
||||
epochs: 5
|
||||
lr: 1e-5
|
||||
batch_size: auto
|
||||
gradient_accumulation_steps: 8
|
||||
lora:
|
||||
r: 128
|
||||
alpha: 32
|
||||
target_modules: auto
|
||||
quantization: 4bit
|
||||
|
||||
output: ./output
|
||||
""",
|
||||
}
|
||||
|
|
@ -0,0 +1,92 @@
|
|||
"""Dataset format detection and conversion.
|
||||
|
||||
Supported formats:
|
||||
- alpaca: {"instruction": ..., "input": ..., "output": ...}
|
||||
- sharegpt: {"conversations": [{"from": "human", "value": ...}, ...]}
|
||||
- chatml: {"messages": [{"role": "user", "content": ...}, ...]}
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
# Required keys per format
|
||||
FORMAT_SIGNATURES = {
|
||||
"alpaca": {"instruction", "output"},
|
||||
"sharegpt": {"conversations"},
|
||||
"chatml": {"messages"},
|
||||
}
|
||||
|
||||
|
||||
def detect_format(data: list[dict]) -> str:
|
||||
"""Auto-detect dataset format from first few rows."""
|
||||
if not data:
|
||||
raise ValueError("Empty dataset — cannot detect format")
|
||||
|
||||
sample = data[0]
|
||||
keys = set(sample.keys())
|
||||
|
||||
for fmt, required_keys in FORMAT_SIGNATURES.items():
|
||||
if required_keys.issubset(keys):
|
||||
return fmt
|
||||
|
||||
raise ValueError(
|
||||
f"Cannot detect format. Keys found: {keys}. "
|
||||
f"Expected one of: alpaca (instruction, output), "
|
||||
f"sharegpt (conversations), chatml (messages)"
|
||||
)
|
||||
|
||||
|
||||
def format_to_messages(row: dict, fmt: str) -> Optional[dict]:
|
||||
"""Convert any format to unified messages format for training.
|
||||
|
||||
Returns: {"messages": [{"role": ..., "content": ...}, ...]}
|
||||
"""
|
||||
try:
|
||||
if fmt == "chatml":
|
||||
return _convert_chatml(row)
|
||||
elif fmt == "alpaca":
|
||||
return _convert_alpaca(row)
|
||||
elif fmt == "sharegpt":
|
||||
return _convert_sharegpt(row)
|
||||
else:
|
||||
raise ValueError(f"Unknown format: {fmt}")
|
||||
except (KeyError, TypeError, IndexError):
|
||||
return None
|
||||
|
||||
|
||||
def _convert_alpaca(row: dict) -> dict:
|
||||
instruction = row["instruction"]
|
||||
input_text = row.get("input", "")
|
||||
output = row["output"]
|
||||
|
||||
user_content = f"{instruction}\n{input_text}".strip() if input_text else instruction
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": user_content},
|
||||
{"role": "assistant", "content": output},
|
||||
]
|
||||
|
||||
if row.get("system"):
|
||||
messages.insert(0, {"role": "system", "content": row["system"]})
|
||||
|
||||
return {"messages": messages}
|
||||
|
||||
|
||||
def _convert_sharegpt(row: dict) -> dict:
|
||||
conversations = row["conversations"]
|
||||
role_map = {"human": "user", "gpt": "assistant", "system": "system"}
|
||||
|
||||
messages = []
|
||||
for turn in conversations:
|
||||
role = role_map.get(turn["from"], turn["from"])
|
||||
messages.append({"role": role, "content": turn["value"]})
|
||||
|
||||
return {"messages": messages}
|
||||
|
||||
|
||||
def _convert_chatml(row: dict) -> dict:
|
||||
# Already in the right format
|
||||
return {"messages": row["messages"]}
|
||||
|
|
@ -0,0 +1,146 @@
|
|||
"""Data loading from local files and HuggingFace."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
from soup_cli.config.schema import DataConfig
|
||||
from soup_cli.data.formats import detect_format, format_to_messages
|
||||
|
||||
console = Console()
|
||||
|
||||
# File extensions we support
|
||||
SUPPORTED_EXTENSIONS = {".jsonl", ".json", ".csv", ".parquet"}
|
||||
|
||||
|
||||
def load_raw_data(path: Path) -> list[dict]:
|
||||
"""Load raw data from a file into list of dicts."""
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Data file not found: {path}")
|
||||
|
||||
ext = path.suffix.lower()
|
||||
if ext not in SUPPORTED_EXTENSIONS:
|
||||
raise ValueError(f"Unsupported file format: {ext}. Supported: {SUPPORTED_EXTENSIONS}")
|
||||
|
||||
if ext == ".jsonl":
|
||||
return _load_jsonl(path)
|
||||
elif ext == ".json":
|
||||
return _load_json(path)
|
||||
elif ext == ".csv":
|
||||
return _load_csv(path)
|
||||
elif ext == ".parquet":
|
||||
return _load_parquet(path)
|
||||
|
||||
raise ValueError(f"Unsupported format: {ext}")
|
||||
|
||||
|
||||
def _load_jsonl(path: Path) -> list[dict]:
|
||||
data = []
|
||||
with open(path, encoding="utf-8") as f:
|
||||
for i, line in enumerate(f):
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
data.append(json.loads(line))
|
||||
except json.JSONDecodeError as e:
|
||||
console.print(f"[yellow]Warning: invalid JSON on line {i + 1}: {e}[/]")
|
||||
return data
|
||||
|
||||
|
||||
def _load_json(path: Path) -> list[dict]:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
if isinstance(raw, list):
|
||||
return raw
|
||||
raise ValueError("JSON file must contain a list of objects")
|
||||
|
||||
|
||||
def _load_csv(path: Path) -> list[dict]:
|
||||
import csv
|
||||
|
||||
with open(path, encoding="utf-8") as f:
|
||||
reader = csv.DictReader(f)
|
||||
return list(reader)
|
||||
|
||||
|
||||
def _load_parquet(path: Path) -> list[dict]:
|
||||
try:
|
||||
import pandas as pd
|
||||
except ImportError:
|
||||
raise ImportError("Install pandas to read parquet files: pip install pandas pyarrow")
|
||||
df = pd.read_parquet(path)
|
||||
return df.to_dict(orient="records")
|
||||
|
||||
|
||||
def load_dataset(data_config: DataConfig) -> dict:
|
||||
"""Load dataset for training. Returns dict with 'train' and optionally 'val' keys.
|
||||
|
||||
Supports:
|
||||
- Local files (.jsonl, .json, .csv, .parquet)
|
||||
- HuggingFace dataset names (auto-detected if no file extension)
|
||||
"""
|
||||
train_path = data_config.train
|
||||
|
||||
# Check if it's a HuggingFace dataset
|
||||
if not Path(train_path).suffix:
|
||||
return _load_hf_dataset(train_path, data_config)
|
||||
|
||||
# Local file
|
||||
path = Path(train_path)
|
||||
raw_data = load_raw_data(path)
|
||||
|
||||
# Detect or use specified format
|
||||
fmt = data_config.format
|
||||
if fmt == "auto":
|
||||
fmt = detect_format(raw_data)
|
||||
console.print(f"[dim]Auto-detected format: {fmt}[/]")
|
||||
|
||||
# Convert to standard message format
|
||||
formatted = [format_to_messages(row, fmt) for row in raw_data]
|
||||
formatted = [r for r in formatted if r is not None] # filter failed rows
|
||||
|
||||
# Split into train/val
|
||||
if data_config.val_split > 0:
|
||||
split_idx = int(len(formatted) * (1 - data_config.val_split))
|
||||
return {
|
||||
"train": formatted[:split_idx],
|
||||
"val": formatted[split_idx:],
|
||||
}
|
||||
|
||||
return {"train": formatted}
|
||||
|
||||
|
||||
def _load_hf_dataset(name: str, data_config: DataConfig) -> dict:
|
||||
"""Load a dataset from HuggingFace Hub."""
|
||||
try:
|
||||
from datasets import load_dataset as hf_load
|
||||
except ImportError:
|
||||
raise ImportError("Install datasets: pip install datasets")
|
||||
|
||||
console.print(f"[dim]Loading from HuggingFace: {name}[/]")
|
||||
ds = hf_load(name)
|
||||
|
||||
if "train" not in ds:
|
||||
raise ValueError(f"Dataset {name} has no 'train' split")
|
||||
|
||||
raw_data = [dict(row) for row in ds["train"]]
|
||||
fmt = data_config.format
|
||||
if fmt == "auto":
|
||||
fmt = detect_format(raw_data)
|
||||
|
||||
formatted = [format_to_messages(row, fmt) for row in raw_data]
|
||||
formatted = [r for r in formatted if r is not None]
|
||||
|
||||
if data_config.val_split > 0 and "validation" not in ds:
|
||||
split_idx = int(len(formatted) * (1 - data_config.val_split))
|
||||
return {"train": formatted[:split_idx], "val": formatted[split_idx:]}
|
||||
|
||||
result = {"train": formatted}
|
||||
if "validation" in ds:
|
||||
val_data = [dict(row) for row in ds["validation"]]
|
||||
val_formatted = [format_to_messages(row, fmt) for row in val_data]
|
||||
result["val"] = [r for r in val_formatted if r is not None]
|
||||
|
||||
return result
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
"""Dataset validation and statistics."""
|
||||
|
||||
from collections import Counter
|
||||
|
||||
from soup_cli.data.formats import FORMAT_SIGNATURES
|
||||
|
||||
|
||||
def validate_and_stats(data: list[dict], expected_format: str | None = None) -> dict:
|
||||
"""Compute stats and validate dataset."""
|
||||
if not data:
|
||||
return {
|
||||
"total": 0,
|
||||
"columns": [],
|
||||
"avg_length": 0,
|
||||
"min_length": 0,
|
||||
"max_length": 0,
|
||||
"empty_fields": 0,
|
||||
"duplicates": 0,
|
||||
"issues": ["Dataset is empty"],
|
||||
"valid_rows": 0,
|
||||
}
|
||||
|
||||
columns = list(data[0].keys())
|
||||
|
||||
# Compute text lengths (join all string values)
|
||||
lengths = []
|
||||
empty_count = 0
|
||||
for row in data:
|
||||
text = " ".join(str(v) for v in row.values() if v)
|
||||
lengths.append(len(text))
|
||||
for v in row.values():
|
||||
if v is None or (isinstance(v, str) and not v.strip()):
|
||||
empty_count += 1
|
||||
|
||||
# Detect duplicates by stringifying rows
|
||||
row_strs = [str(sorted(row.items())) for row in data]
|
||||
dup_count = len(row_strs) - len(set(row_strs))
|
||||
|
||||
# Validate format
|
||||
issues = []
|
||||
valid_rows = len(data)
|
||||
if expected_format and expected_format in FORMAT_SIGNATURES:
|
||||
required = FORMAT_SIGNATURES[expected_format]
|
||||
invalid = 0
|
||||
for row in data:
|
||||
if not required.issubset(row.keys()):
|
||||
invalid += 1
|
||||
valid_rows = len(data) - invalid
|
||||
if invalid > 0:
|
||||
issues.append(
|
||||
f"{invalid} rows missing required keys for '{expected_format}' format: {required}"
|
||||
)
|
||||
|
||||
if dup_count > 0:
|
||||
issues.append(f"{dup_count} duplicate rows found")
|
||||
if empty_count > 0:
|
||||
issues.append(f"{empty_count} empty fields found")
|
||||
|
||||
# Check for very short samples
|
||||
short = sum(1 for l in lengths if l < 10)
|
||||
if short > 0:
|
||||
issues.append(f"{short} samples are very short (<10 chars)")
|
||||
|
||||
return {
|
||||
"total": len(data),
|
||||
"columns": columns,
|
||||
"avg_length": round(sum(lengths) / len(lengths)),
|
||||
"min_length": min(lengths),
|
||||
"max_length": max(lengths),
|
||||
"empty_fields": empty_count,
|
||||
"duplicates": dup_count,
|
||||
"issues": issues,
|
||||
"valid_rows": valid_rows,
|
||||
}
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
"""Experiment tracking — stores runs in local SQLite. (Phase 2)"""
|
||||
|
||||
# Placeholder for Phase 2 implementation.
|
||||
# Will store: config, metrics, hardware info, checkpoints per run.
|
||||
# Storage: ~/.soup/experiments.db (SQLite)
|
||||
|
|
@ -0,0 +1,44 @@
|
|||
"""HuggingFace Trainer callback that feeds metrics to our display."""
|
||||
|
||||
from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments
|
||||
|
||||
from soup_cli.monitoring.display import TrainingDisplay
|
||||
|
||||
|
||||
class SoupTrainerCallback(TrainerCallback):
|
||||
"""Bridges HF Trainer events to Soup's Rich live display."""
|
||||
|
||||
def __init__(self, display: TrainingDisplay):
|
||||
self.display = display
|
||||
|
||||
def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
|
||||
self.display.start(total_steps=state.max_steps)
|
||||
|
||||
def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, logs=None, **kwargs):
|
||||
if logs is None:
|
||||
return
|
||||
|
||||
# Try to get GPU memory
|
||||
gpu_mem = ""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
used = torch.cuda.memory_allocated() / (1024**3)
|
||||
total = torch.cuda.get_device_properties(0).total_mem / (1024**3)
|
||||
gpu_mem = f"{used:.1f}/{total:.1f} GB"
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self.display.update(
|
||||
step=state.global_step,
|
||||
epoch=state.epoch or 0,
|
||||
loss=logs.get("loss", 0.0),
|
||||
lr=logs.get("learning_rate", 0.0),
|
||||
grad_norm=logs.get("grad_norm", 0.0),
|
||||
speed=logs.get("train_steps_per_second", 0.0),
|
||||
gpu_mem=gpu_mem,
|
||||
)
|
||||
|
||||
def on_train_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
|
||||
self.display.stop()
|
||||
|
|
@ -0,0 +1,80 @@
|
|||
"""Rich live training dashboard in the terminal."""
|
||||
|
||||
from rich.console import Console
|
||||
from rich.live import Live
|
||||
from rich.panel import Panel
|
||||
from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
|
||||
from rich.table import Table
|
||||
from rich.text import Text
|
||||
|
||||
from soup_cli.config.schema import SoupConfig
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
class TrainingDisplay:
|
||||
"""Live-updating terminal dashboard for training progress."""
|
||||
|
||||
def __init__(self, config: SoupConfig, device_name: str = ""):
|
||||
self.config = config
|
||||
self.device_name = device_name
|
||||
self.current_step = 0
|
||||
self.total_steps = 0
|
||||
self.current_epoch = 0
|
||||
self.loss = 0.0
|
||||
self.lr = 0.0
|
||||
self.grad_norm = 0.0
|
||||
self.gpu_mem = ""
|
||||
self.speed = 0.0
|
||||
self._live: Live | None = None
|
||||
|
||||
def start(self, total_steps: int):
|
||||
"""Start the live display."""
|
||||
self.total_steps = total_steps
|
||||
self._live = Live(self._render(), console=console, refresh_per_second=2)
|
||||
self._live.start()
|
||||
|
||||
def update(self, step: int, epoch: float, loss: float, lr: float, **kwargs):
|
||||
"""Update display with new metrics."""
|
||||
self.current_step = step
|
||||
self.current_epoch = epoch
|
||||
self.loss = loss
|
||||
self.lr = lr
|
||||
self.grad_norm = kwargs.get("grad_norm", 0.0)
|
||||
self.speed = kwargs.get("speed", 0.0)
|
||||
self.gpu_mem = kwargs.get("gpu_mem", "")
|
||||
|
||||
if self._live:
|
||||
self._live.update(self._render())
|
||||
|
||||
def stop(self):
|
||||
"""Stop the live display."""
|
||||
if self._live:
|
||||
self._live.stop()
|
||||
|
||||
def _render(self) -> Panel:
|
||||
"""Render the dashboard panel."""
|
||||
progress_pct = (self.current_step / self.total_steps * 100) if self.total_steps > 0 else 0
|
||||
bar_width = 30
|
||||
filled = int(bar_width * progress_pct / 100)
|
||||
bar = "█" * filled + "░" * (bar_width - filled)
|
||||
|
||||
lines = []
|
||||
lines.append(f"Epoch {self.current_epoch:.1f}/{self.config.training.epochs} [{bar}] {progress_pct:.0f}%")
|
||||
lines.append(f"Step: {self.current_step}/{self.total_steps}")
|
||||
lines.append(f"Loss: {self.loss:.4f} LR: {self.lr:.2e}")
|
||||
|
||||
if self.speed > 0:
|
||||
lines.append(f"Speed: {self.speed:.2f} it/s")
|
||||
if self.gpu_mem:
|
||||
lines.append(f"GPU: {self.gpu_mem}")
|
||||
if self.grad_norm > 0:
|
||||
lines.append(f"Grad: {self.grad_norm:.4f}")
|
||||
|
||||
content = "\n".join(lines)
|
||||
return Panel(
|
||||
content,
|
||||
title=f"[bold green]Soup Training: {self.config.experiment_name or self.config.base}[/]",
|
||||
subtitle=f"[dim]{self.device_name}[/]",
|
||||
border_style="green",
|
||||
)
|
||||
|
|
@ -0,0 +1,185 @@
|
|||
"""SFT (Supervised Fine-Tuning) trainer — wraps HuggingFace transformers + peft + trl."""
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
from soup_cli.config.schema import SoupConfig
|
||||
from soup_cli.utils.gpu import estimate_batch_size, model_size_from_name
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
class SFTTrainerWrapper:
|
||||
"""High-level wrapper that sets up model + tokenizer + trainer from SoupConfig."""
|
||||
|
||||
def __init__(self, config: SoupConfig, device: str = "cuda"):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.trainer = None
|
||||
|
||||
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 trl import SFTTrainer
|
||||
|
||||
cfg = self.config
|
||||
tcfg = cfg.training
|
||||
|
||||
# --- 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
|
||||
|
||||
# --- 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(f"[green]LoRA applied:[/] {trainable:,} trainable / {total:,} total ({pct:.2f}%)")
|
||||
|
||||
# --- Batch size ---
|
||||
batch_size = tcfg.batch_size
|
||||
if batch_size == "auto":
|
||||
from soup_cli.utils.gpu import get_gpu_info
|
||||
|
||||
gpu_info = get_gpu_info()
|
||||
model_size = model_size_from_name(cfg.base)
|
||||
batch_size = estimate_batch_size(
|
||||
model_params_b=model_size,
|
||||
seq_length=cfg.data.max_length,
|
||||
gpu_memory_bytes=gpu_info["memory_total_bytes"],
|
||||
quantization=tcfg.quantization,
|
||||
lora_r=tcfg.lora.r,
|
||||
)
|
||||
console.print(f"[green]Auto batch size:[/] {batch_size}")
|
||||
|
||||
# --- Dataset ---
|
||||
def format_row(example):
|
||||
text = self.tokenizer.apply_chat_template(
|
||||
example["messages"], tokenize=False, add_generation_prompt=False
|
||||
)
|
||||
return {"text": text}
|
||||
|
||||
train_ds = Dataset.from_list(dataset["train"]).map(format_row)
|
||||
eval_ds = None
|
||||
if "val" in dataset and dataset["val"]:
|
||||
eval_ds = Dataset.from_list(dataset["val"]).map(format_row)
|
||||
|
||||
# --- Output dir ---
|
||||
output_dir = Path(cfg.output)
|
||||
if cfg.experiment_name:
|
||||
output_dir = output_dir / cfg.experiment_name
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# --- Training args ---
|
||||
training_args = TrainingArguments(
|
||||
output_dir=str(output_dir),
|
||||
num_train_epochs=tcfg.epochs,
|
||||
per_device_train_batch_size=batch_size,
|
||||
gradient_accumulation_steps=tcfg.gradient_accumulation_steps,
|
||||
learning_rate=tcfg.lr,
|
||||
warmup_ratio=tcfg.warmup_ratio,
|
||||
weight_decay=tcfg.weight_decay,
|
||||
max_grad_norm=tcfg.max_grad_norm,
|
||||
optim=tcfg.optimizer,
|
||||
lr_scheduler_type=tcfg.scheduler,
|
||||
logging_steps=tcfg.logging_steps,
|
||||
save_steps=tcfg.save_steps,
|
||||
save_total_limit=3,
|
||||
bf16=self.device == "cuda",
|
||||
report_to="none",
|
||||
remove_unused_columns=False,
|
||||
)
|
||||
|
||||
# --- Trainer ---
|
||||
self.trainer = SFTTrainer(
|
||||
model=self.model,
|
||||
args=training_args,
|
||||
train_dataset=train_ds,
|
||||
eval_dataset=eval_ds,
|
||||
processing_class=self.tokenizer,
|
||||
)
|
||||
|
||||
self._output_dir = str(output_dir)
|
||||
|
||||
def train(self, display: Optional[object] = None) -> dict:
|
||||
"""Run training and return results summary."""
|
||||
start = time.time()
|
||||
|
||||
# Add callback for live display
|
||||
if display:
|
||||
from soup_cli.monitoring.callback import SoupTrainerCallback
|
||||
|
||||
self.trainer.add_callback(SoupTrainerCallback(display))
|
||||
|
||||
result = self.trainer.train()
|
||||
duration = time.time() - start
|
||||
|
||||
# Save final model (LoRA adapter)
|
||||
self.trainer.save_model(self._output_dir)
|
||||
self.tokenizer.save_pretrained(self._output_dir)
|
||||
|
||||
# Extract metrics
|
||||
logs = self.trainer.state.log_history
|
||||
train_losses = [l["loss"] for l in logs if "loss" in l]
|
||||
|
||||
hours = int(duration // 3600)
|
||||
minutes = int((duration % 3600) // 60)
|
||||
duration_str = f"{hours}h {minutes}m" if hours > 0 else f"{minutes}m"
|
||||
|
||||
return {
|
||||
"initial_loss": train_losses[0] if train_losses else 0,
|
||||
"final_loss": train_losses[-1] if train_losses else 0,
|
||||
"duration": duration_str,
|
||||
"output_dir": self._output_dir,
|
||||
"total_steps": self.trainer.state.global_step,
|
||||
}
|
||||
|
|
@ -0,0 +1,8 @@
|
|||
"""Global constants and defaults."""
|
||||
|
||||
APP_NAME = "soup"
|
||||
CONFIG_FILE = "soup.yaml"
|
||||
SOUP_DIR = ".soup"
|
||||
EXPERIMENTS_DB = "experiments.db"
|
||||
|
||||
DEFAULT_CHAT_TEMPLATE = """{% for message in messages %}{% if message['role'] == 'system' %}{{ message['content'] + '\\n' }}{% elif message['role'] == 'user' %}{{ 'User: ' + message['content'] + '\\n' }}{% elif message['role'] == 'assistant' %}{{ 'Assistant: ' + message['content'] + '\\n' }}{% endif %}{% endfor %}"""
|
||||
|
|
@ -0,0 +1,115 @@
|
|||
"""GPU detection, memory calculation, and auto batch size."""
|
||||
|
||||
import math
|
||||
|
||||
|
||||
def detect_device() -> tuple[str, str]:
|
||||
"""Detect available device. Returns (device_string, human_name)."""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
name = torch.cuda.get_device_name(0)
|
||||
return "cuda", name
|
||||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
return "mps", "Apple Silicon (MPS)"
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
return "cpu", "CPU (no GPU detected)"
|
||||
|
||||
|
||||
def get_gpu_info() -> dict:
|
||||
"""Get GPU memory info."""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
total = torch.cuda.get_device_properties(0).total_mem
|
||||
total_gb = total / (1024**3)
|
||||
return {
|
||||
"memory_total": f"{total_gb:.1f} GB",
|
||||
"memory_total_bytes": total,
|
||||
"gpu_count": torch.cuda.device_count(),
|
||||
}
|
||||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
# MPS doesn't expose memory easily, estimate from system
|
||||
return {
|
||||
"memory_total": "shared (Apple Silicon)",
|
||||
"memory_total_bytes": 0,
|
||||
"gpu_count": 1,
|
||||
}
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
return {
|
||||
"memory_total": "N/A (CPU mode)",
|
||||
"memory_total_bytes": 0,
|
||||
"gpu_count": 0,
|
||||
}
|
||||
|
||||
|
||||
def estimate_batch_size(
|
||||
model_params_b: float,
|
||||
seq_length: int,
|
||||
gpu_memory_bytes: int,
|
||||
quantization: str = "4bit",
|
||||
lora_r: int = 64,
|
||||
) -> int:
|
||||
"""Estimate max batch size that fits in GPU memory.
|
||||
|
||||
Conservative estimate — better to start smaller and gradient accumulate.
|
||||
"""
|
||||
if gpu_memory_bytes == 0:
|
||||
return 1 # CPU fallback
|
||||
|
||||
gpu_gb = gpu_memory_bytes / (1024**3)
|
||||
|
||||
# Rough memory per param based on quantization
|
||||
bytes_per_param = {"4bit": 0.5, "8bit": 1.0, "none": 2.0} # FP16
|
||||
bpp = bytes_per_param.get(quantization, 2.0)
|
||||
|
||||
# Model memory (static)
|
||||
model_mem_gb = model_params_b * bpp
|
||||
|
||||
# LoRA trainable params (usually ~1-3% of total)
|
||||
lora_ratio = min(lora_r * 2 / 4096, 0.05) # rough estimate
|
||||
trainable_mem_gb = model_params_b * 2 * lora_ratio # FP16 for trainable
|
||||
|
||||
# Optimizer states (Adam: 2x params)
|
||||
optimizer_mem_gb = trainable_mem_gb * 2
|
||||
|
||||
# Available for activations
|
||||
overhead_gb = 1.5 # CUDA overhead, fragmentation
|
||||
available_gb = gpu_gb - model_mem_gb - trainable_mem_gb - optimizer_mem_gb - overhead_gb
|
||||
|
||||
if available_gb <= 0:
|
||||
return 1
|
||||
|
||||
# Rough activation memory per sample per token
|
||||
# ~2 bytes per hidden dim per layer per token for a transformer
|
||||
activation_per_sample_gb = (seq_length * model_params_b * 0.001) # very rough
|
||||
activation_per_sample_gb = max(activation_per_sample_gb, 0.5) # minimum 0.5 GB
|
||||
|
||||
batch_size = max(1, int(available_gb / activation_per_sample_gb))
|
||||
# Clamp to power of 2 (common practice)
|
||||
batch_size = 2 ** int(math.log2(batch_size)) if batch_size > 1 else 1
|
||||
|
||||
return min(batch_size, 32) # cap at 32
|
||||
|
||||
|
||||
def model_size_from_name(model_name: str) -> float:
|
||||
"""Guess model size in billions from model name."""
|
||||
name_lower = model_name.lower()
|
||||
|
||||
size_markers = [
|
||||
("70b", 70), ("65b", 65), ("34b", 34), ("33b", 33),
|
||||
("13b", 13), ("8b", 8), ("7b", 7), ("3b", 3),
|
||||
("1.5b", 1.5), ("1b", 1), ("0.5b", 0.5),
|
||||
]
|
||||
|
||||
for marker, size in size_markers:
|
||||
if marker in name_lower:
|
||||
return size
|
||||
|
||||
return 7.0 # default guess
|
||||
|
|
@ -0,0 +1,24 @@
|
|||
# Soup template: Chat Assistant
|
||||
# Fine-tune a model for conversational chat
|
||||
# Usage: soup init --template chat
|
||||
|
||||
base: meta-llama/Llama-3.1-8B-Instruct
|
||||
task: sft
|
||||
|
||||
data:
|
||||
train: ./data/train.jsonl
|
||||
format: alpaca
|
||||
val_split: 0.1
|
||||
max_length: 2048
|
||||
|
||||
training:
|
||||
epochs: 3
|
||||
lr: 2e-5
|
||||
batch_size: auto
|
||||
lora:
|
||||
r: 64
|
||||
alpha: 16
|
||||
target_modules: auto
|
||||
quantization: 4bit
|
||||
|
||||
output: ./output
|
||||
|
|
@ -0,0 +1,24 @@
|
|||
# Soup template: Code Model
|
||||
# Fine-tune a model for code generation
|
||||
# Usage: soup init --template code
|
||||
|
||||
base: codellama/CodeLlama-7b-Instruct-hf
|
||||
task: sft
|
||||
|
||||
data:
|
||||
train: ./data/code_train.jsonl
|
||||
format: alpaca
|
||||
val_split: 0.1
|
||||
max_length: 4096
|
||||
|
||||
training:
|
||||
epochs: 2
|
||||
lr: 1e-5
|
||||
batch_size: auto
|
||||
lora:
|
||||
r: 128
|
||||
alpha: 32
|
||||
target_modules: auto
|
||||
quantization: 4bit
|
||||
|
||||
output: ./output
|
||||
|
|
@ -0,0 +1,25 @@
|
|||
# Soup template: Medical / Domain Expert
|
||||
# Fine-tune a model with domain-specific knowledge
|
||||
# Usage: soup init --template medical
|
||||
|
||||
base: meta-llama/Llama-3.1-8B-Instruct
|
||||
task: sft
|
||||
|
||||
data:
|
||||
train: ./data/medical_train.jsonl
|
||||
format: alpaca
|
||||
val_split: 0.15
|
||||
max_length: 2048
|
||||
|
||||
training:
|
||||
epochs: 5
|
||||
lr: 1e-5
|
||||
batch_size: auto
|
||||
gradient_accumulation_steps: 8
|
||||
lora:
|
||||
r: 128
|
||||
alpha: 32
|
||||
target_modules: auto
|
||||
quantization: 4bit
|
||||
|
||||
output: ./output
|
||||
|
|
@ -0,0 +1,55 @@
|
|||
"""Shared test fixtures."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tmp_data_dir(tmp_path: Path) -> Path:
|
||||
"""Create a temp directory with sample training data."""
|
||||
data_dir = tmp_path / "data"
|
||||
data_dir.mkdir()
|
||||
return data_dir
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_alpaca_data(tmp_data_dir: Path) -> Path:
|
||||
"""Create a sample alpaca-format JSONL file."""
|
||||
path = tmp_data_dir / "train.jsonl"
|
||||
samples = [
|
||||
{"instruction": "What is Python?", "input": "", "output": "Python is a programming language."},
|
||||
{"instruction": "Explain gravity", "input": "", "output": "Gravity is a fundamental force."},
|
||||
{"instruction": "Translate hello to Spanish", "input": "hello", "output": "hola"},
|
||||
]
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
for s in samples:
|
||||
f.write(json.dumps(s) + "\n")
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_config(tmp_path: Path, sample_alpaca_data: Path) -> Path:
|
||||
"""Create a sample soup.yaml config."""
|
||||
config_path = tmp_path / "soup.yaml"
|
||||
config_path.write_text(
|
||||
f"""base: meta-llama/Llama-3.1-8B-Instruct
|
||||
task: sft
|
||||
data:
|
||||
train: {sample_alpaca_data}
|
||||
format: alpaca
|
||||
val_split: 0.1
|
||||
training:
|
||||
epochs: 1
|
||||
lr: 2e-5
|
||||
batch_size: 1
|
||||
lora:
|
||||
r: 8
|
||||
alpha: 16
|
||||
quantization: 4bit
|
||||
output: {tmp_path / 'output'}
|
||||
""",
|
||||
encoding="utf-8",
|
||||
)
|
||||
return config_path
|
||||
|
|
@ -0,0 +1,29 @@
|
|||
"""Tests for CLI commands."""
|
||||
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
|
||||
def test_version():
|
||||
result = runner.invoke(app, ["version"])
|
||||
assert result.exit_code == 0
|
||||
assert "0.1.0" in result.output
|
||||
|
||||
|
||||
def test_help():
|
||||
result = runner.invoke(app, ["--help"])
|
||||
assert result.exit_code == 0
|
||||
assert "Fine-tune" in result.output
|
||||
|
||||
|
||||
def test_init_unknown_template():
|
||||
result = runner.invoke(app, ["init", "--template", "nonexistent"])
|
||||
assert result.exit_code == 1
|
||||
|
||||
|
||||
def test_train_missing_config():
|
||||
result = runner.invoke(app, ["train", "--config", "nonexistent.yaml"])
|
||||
assert result.exit_code == 1
|
||||
|
|
@ -0,0 +1,53 @@
|
|||
"""Tests for config loading and validation."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from soup_cli.config.loader import load_config
|
||||
from soup_cli.config.schema import SoupConfig
|
||||
|
||||
|
||||
def test_load_valid_config(sample_config: Path):
|
||||
"""Valid config should parse without errors."""
|
||||
cfg = load_config(sample_config)
|
||||
assert isinstance(cfg, SoupConfig)
|
||||
assert cfg.base == "meta-llama/Llama-3.1-8B-Instruct"
|
||||
assert cfg.task == "sft"
|
||||
assert cfg.training.epochs == 1
|
||||
assert cfg.training.lora.r == 8
|
||||
assert cfg.training.quantization == "4bit"
|
||||
|
||||
|
||||
def test_config_defaults():
|
||||
"""Config should fill in defaults for optional fields."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
assert cfg.task == "sft"
|
||||
assert cfg.training.epochs == 3
|
||||
assert cfg.training.lr == 2e-5
|
||||
assert cfg.training.batch_size == "auto"
|
||||
assert cfg.training.lora.r == 64
|
||||
assert cfg.training.quantization == "4bit"
|
||||
assert cfg.output == "./output"
|
||||
|
||||
|
||||
def test_config_invalid_task():
|
||||
"""Invalid task should raise validation error."""
|
||||
with pytest.raises(Exception):
|
||||
SoupConfig(
|
||||
base="some-model",
|
||||
task="invalid_task",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
|
||||
|
||||
def test_config_val_split_bounds():
|
||||
"""val_split must be between 0 and 0.5."""
|
||||
with pytest.raises(Exception):
|
||||
SoupConfig(
|
||||
base="some-model",
|
||||
data={"train": "./data.jsonl", "val_split": 0.9},
|
||||
)
|
||||
|
|
@ -0,0 +1,71 @@
|
|||
"""Tests for data loading, format detection, and validation."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from soup_cli.data.formats import detect_format, format_to_messages
|
||||
from soup_cli.data.loader import load_raw_data
|
||||
from soup_cli.data.validator import validate_and_stats
|
||||
|
||||
|
||||
def test_load_jsonl(sample_alpaca_data: Path):
|
||||
data = load_raw_data(sample_alpaca_data)
|
||||
assert len(data) == 3
|
||||
assert "instruction" in data[0]
|
||||
assert "output" in data[0]
|
||||
|
||||
|
||||
def test_detect_alpaca_format():
|
||||
data = [{"instruction": "test", "input": "", "output": "result"}]
|
||||
assert detect_format(data) == "alpaca"
|
||||
|
||||
|
||||
def test_detect_sharegpt_format():
|
||||
data = [{"conversations": [{"from": "human", "value": "hi"}]}]
|
||||
assert detect_format(data) == "sharegpt"
|
||||
|
||||
|
||||
def test_detect_chatml_format():
|
||||
data = [{"messages": [{"role": "user", "content": "hi"}]}]
|
||||
assert detect_format(data) == "chatml"
|
||||
|
||||
|
||||
def test_convert_alpaca():
|
||||
row = {"instruction": "Explain AI", "input": "", "output": "AI is..."}
|
||||
result = format_to_messages(row, "alpaca")
|
||||
assert result is not None
|
||||
assert len(result["messages"]) == 2
|
||||
assert result["messages"][0]["role"] == "user"
|
||||
assert result["messages"][1]["role"] == "assistant"
|
||||
|
||||
|
||||
def test_convert_alpaca_with_input():
|
||||
row = {"instruction": "Translate", "input": "hello", "output": "привет"}
|
||||
result = format_to_messages(row, "alpaca")
|
||||
assert "hello" in result["messages"][0]["content"]
|
||||
|
||||
|
||||
def test_convert_sharegpt():
|
||||
row = {
|
||||
"conversations": [
|
||||
{"from": "human", "value": "What is 2+2?"},
|
||||
{"from": "gpt", "value": "4"},
|
||||
]
|
||||
}
|
||||
result = format_to_messages(row, "sharegpt")
|
||||
assert result["messages"][0]["role"] == "user"
|
||||
assert result["messages"][1]["role"] == "assistant"
|
||||
|
||||
|
||||
def test_validate_stats(sample_alpaca_data: Path):
|
||||
data = load_raw_data(sample_alpaca_data)
|
||||
stats = validate_and_stats(data)
|
||||
assert stats["total"] == 3
|
||||
assert "instruction" in stats["columns"]
|
||||
assert stats["avg_length"] > 0
|
||||
|
||||
|
||||
def test_validate_with_format(sample_alpaca_data: Path):
|
||||
data = load_raw_data(sample_alpaca_data)
|
||||
stats = validate_and_stats(data, expected_format="alpaca")
|
||||
assert stats["valid_rows"] == 3
|
||||
assert len(stats["issues"]) == 0 # no issues for valid data
|
||||
|
|
@ -0,0 +1,24 @@
|
|||
"""Tests for GPU utils."""
|
||||
|
||||
from soup_cli.utils.gpu import estimate_batch_size, model_size_from_name
|
||||
|
||||
|
||||
def test_model_size_detection():
|
||||
assert model_size_from_name("meta-llama/Llama-3.1-8B-Instruct") == 8
|
||||
assert model_size_from_name("meta-llama/Llama-3.1-70B") == 70
|
||||
assert model_size_from_name("codellama/CodeLlama-7b-hf") == 7
|
||||
assert model_size_from_name("some-unknown-model") == 7.0 # default
|
||||
|
||||
|
||||
def test_batch_size_cpu():
|
||||
"""CPU (0 memory) should return batch_size=1."""
|
||||
bs = estimate_batch_size(7.0, 2048, 0, "4bit", 64)
|
||||
assert bs == 1
|
||||
|
||||
|
||||
def test_batch_size_24gb():
|
||||
"""24 GB GPU with 8B model QLoRA should fit batch > 1."""
|
||||
mem = 24 * (1024**3)
|
||||
bs = estimate_batch_size(8.0, 2048, mem, "4bit", 64)
|
||||
assert bs >= 1
|
||||
assert bs <= 32
|
||||
Loading…
Reference in New Issue