diff --git a/cli/commands/train.py b/cli/commands/train.py index dd357c55b9..efc034d6ea 100644 --- a/cli/commands/train.py +++ b/cli/commands/train.py @@ -19,6 +19,9 @@ def train( hf_token: Optional[str] = typer.Option( None, "--hf-token", envvar="HF_TOKEN", help="Hugging Face token if needed." ), + wandb_token: Optional[str] = typer.Option( + None, "--wandb-token", envvar="WANDB_API_KEY", help="Weights & Biases API key." + ), dry_run: bool = typer.Option( False, "--dry-run", @@ -35,6 +38,10 @@ def train( cfg.apply_overrides(**config_overrides) + # CLI/env tokens take precedence over config + hf_token = hf_token or cfg.logging.hf_token + wandb_token = wandb_token or cfg.logging.wandb_token + if dry_run: import yaml data = cfg.model_dump() @@ -52,13 +59,18 @@ def train( ) raise typer.Exit(code=2) + from pathlib import Path as PathlibPath from backend.trainer import UnslothTrainer from backend.model_config import ModelConfig trainer = UnslothTrainer() + # Check if the model path is a LoRA adapter (has adapter_config.json) + model_path = PathlibPath(cfg.model) if cfg.model else None + model_is_lora = model_path and model_path.is_dir() and (model_path / "adapter_config.json").exists() + model_config = ModelConfig.from_ui_selection( - dropdown_value=cfg.model, search_value=None, hf_token=hf_token, is_lora=False + dropdown_value=cfg.model, search_value=None, hf_token=hf_token, is_lora=model_is_lora ) if not model_config: typer.echo("Could not resolve model config", err=True) @@ -67,6 +79,14 @@ def train( is_vision = model_config.is_vision use_lora = cfg.training.training_type.lower() == "lora" + if model_is_lora and not use_lora: + typer.echo( + "Error: Cannot do full finetuning on a LoRA adapter. " + "Use --training-type lora or provide a base model.", + err=True, + ) + raise typer.Exit(code=2) + if not trainer.load_model( model_name=model_config.identifier, max_seq_length=cfg.training.max_seq_length, @@ -89,7 +109,9 @@ def train( typer.echo("Dataset load failed", err=True) raise typer.Exit(code=1) - started = trainer.start_training(dataset=ds, **cfg.training_kwargs()) + training_kwargs = cfg.training_kwargs() + training_kwargs["wandb_token"] = wandb_token # CLI/env takes precedence + started = trainer.start_training(dataset=ds, **training_kwargs) if not started: typer.echo("Training failed to start", err=True) diff --git a/cli/config.py b/cli/config.py index d80f3a2f63..aa402f32ef 100644 --- a/cli/config.py +++ b/cli/config.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Optional, List +from typing import Literal, Optional, List import yaml from pydantic import BaseModel, Field @@ -8,11 +8,11 @@ from pydantic import BaseModel, Field class DataConfig(BaseModel): dataset: Optional[str] = None local_dataset: Optional[List[str]] = None - format_type: str = "auto" + format_type: Literal["auto", "alpaca", "chatml", "sharegpt"] = "auto" class TrainingConfig(BaseModel): - training_type: str = "lora" + training_type: Literal["lora", "full"] = "lora" max_seq_length: int = 2048 load_in_4bit: bool = True output_dir: Path = Path("./outputs") @@ -27,7 +27,7 @@ class TrainingConfig(BaseModel): random_seed: int = 3407 packing: bool = False train_on_completions: bool = False - gradient_checkpointing: bool = True + gradient_checkpointing: Literal["unsloth", "true", "none"] = "unsloth" class LoraConfig(BaseModel): @@ -53,6 +53,7 @@ class LoggingConfig(BaseModel): wandb_token: Optional[str] = None enable_tensorboard: bool = False tensorboard_dir: str = "runs" + hf_token: Optional[str] = None class Config(BaseModel): diff --git a/cli/options.py b/cli/options.py index e9269a6cd1..c8a08e59f2 100644 --- a/cli/options.py +++ b/cli/options.py @@ -82,11 +82,15 @@ def add_options_from_config(config_class: type[BaseModel]) -> Callable: def decorator(func: Callable) -> Callable: sig = inspect.signature(func) original_params = list(sig.parameters.values()) + original_param_names = {p.name for p in original_params} # Build new parameters: config fields first, then original params new_params = [] for field_name, field_info in fields: + # Skip fields already defined in function signature (e.g., with envvar) + if field_name in original_param_names: + continue annotation = field_info.annotation if _is_list_type(annotation): continue