81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LlamaFactoryCommand:
|
|
command: list[str]
|
|
work_dir: str
|
|
env: dict[str, str]
|
|
|
|
|
|
def validate_config(config: dict[str, Any]) -> list[str]:
|
|
errors: list[str] = []
|
|
if not config.get("base_model") and not config.get("model_name_or_path"):
|
|
errors.append("base_model or model_name_or_path is required")
|
|
if not config.get("dataset") and not config.get("dataset_dir"):
|
|
errors.append("dataset or dataset_dir is required")
|
|
learning_rate = float(config.get("learning_rate", 0.0002))
|
|
if learning_rate <= 0:
|
|
errors.append("learning_rate must be greater than zero")
|
|
epochs = int(config.get("n_epochs", config.get("num_train_epochs", 1)))
|
|
if epochs <= 0:
|
|
errors.append("n_epochs must be greater than zero")
|
|
return errors
|
|
|
|
|
|
def build_command(config: dict[str, Any], llama_factory_home: str = "/app/LLaMA-Factory") -> LlamaFactoryCommand:
|
|
errors = validate_config(config)
|
|
if errors:
|
|
raise ValueError("; ".join(errors))
|
|
|
|
model_path = config.get("base_model") or config.get("model_name_or_path")
|
|
dataset = config.get("dataset") or config.get("dataset_dir")
|
|
output_dir = config.get("output_dir") or f"/data/yg-ft/outputs/{config.get('name', 'training-job')}"
|
|
command = [
|
|
"llamafactory-cli",
|
|
"train",
|
|
"--stage",
|
|
str(config.get("stage", "sft")).lower(),
|
|
"--do_train",
|
|
"true",
|
|
"--model_name_or_path",
|
|
str(model_path),
|
|
"--dataset",
|
|
str(dataset),
|
|
"--template",
|
|
str(config.get("template", "qwen")),
|
|
"--finetuning_type",
|
|
str(config.get("train_method", config.get("finetuning_type", "lora"))),
|
|
"--output_dir",
|
|
str(output_dir),
|
|
"--per_device_train_batch_size",
|
|
str(config.get("batch_size", 2)),
|
|
"--learning_rate",
|
|
str(config.get("learning_rate", 0.0002)),
|
|
"--num_train_epochs",
|
|
str(config.get("n_epochs", 3)),
|
|
"--save_steps",
|
|
str(config.get("save_steps", 50)),
|
|
]
|
|
quantization_bit = int(config.get("quantization_bit", 0) or 0)
|
|
if quantization_bit in {4, 8}:
|
|
command.extend(["--quantization_bit", str(quantization_bit)])
|
|
return LlamaFactoryCommand(command=command, work_dir=str(Path(llama_factory_home)), env={})
|
|
|
|
|
|
def parse_log_line(line: str) -> dict[str, float] | None:
|
|
if "loss" not in line or "learning_rate" not in line:
|
|
return None
|
|
result: dict[str, float] = {}
|
|
for key in ["loss", "grad_norm", "learning_rate", "epoch"]:
|
|
match = re.search(rf"['\"]?{key}['\"]?\s*:\s*([-+]?\d+(?:\.\d+)?(?:[eE][-+]?\d+)?)", line)
|
|
if match:
|
|
result[key] = float(match.group(1))
|
|
return result or None
|
|
|