326 lines
16 KiB
Python
326 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""Train from one LLaMA-Factory YAML on four GPUs; import Unsloth in each worker first."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from datetime import datetime, timezone
|
|
import functools
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
from types import MethodType
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
GPUS = 4
|
|
|
|
|
|
def read_config(path: Path) -> dict:
|
|
import yaml
|
|
with path.open(encoding="utf-8-sig") as handle:
|
|
config = yaml.safe_load(handle)
|
|
if not isinstance(config, dict):
|
|
raise ValueError("The configuration must be a YAML mapping.")
|
|
if config.get("use_unsloth") is not True:
|
|
raise ValueError("This package requires use_unsloth: true.")
|
|
if config.get("bf16") is not True or config.get("fp16", False):
|
|
raise ValueError("The H200 configuration must use bf16, not fp16.")
|
|
if "model_max_length" in config:
|
|
raise ValueError("Remove model_max_length from YAML; LLaMA-Factory derives it from cutoff_len.")
|
|
if int(config.get("cutoff_len", 0)) < 1:
|
|
raise ValueError("cutoff_len must be positive.")
|
|
for key in ("dataset_dir", "output_dir"):
|
|
if key not in config:
|
|
raise ValueError(f"Missing configuration key: {key}")
|
|
value = Path(config[key]).expanduser()
|
|
config[key] = str((ROOT / value).resolve() if not value.is_absolute() else value.resolve())
|
|
return config
|
|
|
|
|
|
def write_config(path: Path, config: dict) -> None:
|
|
import yaml
|
|
path.write_text(yaml.safe_dump(config, allow_unicode=True, sort_keys=False), encoding="utf-8")
|
|
|
|
|
|
def digest(path: Path) -> str:
|
|
result = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
|
result.update(block)
|
|
return result.hexdigest()
|
|
|
|
|
|
def verify_data(data_dir: Path, splits: tuple[str, ...]) -> dict:
|
|
manifest_path = data_dir / "manifest.json"
|
|
if not manifest_path.is_file():
|
|
raise ValueError(f"Run prepare_data.py first; missing {manifest_path}")
|
|
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
|
|
for split in splits:
|
|
if manifest.get("splits", {}).get(split, {}).get("rows", 0) < 1:
|
|
raise ValueError(f"The prepared {split} split is empty.")
|
|
names = ["dataset_info.json"]
|
|
for split in splits:
|
|
names.extend((f"{split}.jsonl", f"{split}.metadata.jsonl"))
|
|
for name in names:
|
|
item = manifest.get("files", {}).get(name, {})
|
|
if not item.get("sha256") or digest(data_dir / name) != item["sha256"]:
|
|
raise ValueError(f"Prepared data changed or is missing: {name}")
|
|
return manifest
|
|
|
|
|
|
def verify_length_filter(manifest: dict, config: dict) -> None:
|
|
report = manifest.get("length_filter", {})
|
|
if report.get("enabled") is not True:
|
|
raise ValueError("Prepared data has no shared length filter. Rerun prepare_data.py into data/aya_slovak_4096.")
|
|
candidates = [profile for profile in report.get("models", {}).values()
|
|
if profile.get("model_name_or_path") == config["model_name_or_path"]
|
|
and profile.get("model_revision") == config["model_revision"]
|
|
and profile.get("template") == config["template"]]
|
|
if not candidates or not any(0 < int(profile["cutoff_len"]) <= int(config["cutoff_len"]) for profile in candidates):
|
|
raise ValueError("The data was not length-filtered for this model, revision, template and context limit.")
|
|
|
|
|
|
def run_logged(command: list[str], log: Path, env: dict | None = None) -> None:
|
|
print("Running:", " ".join(command), flush=True)
|
|
with log.open("a", encoding="utf-8") as handle:
|
|
process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
|
text=True, encoding="utf-8", errors="replace", env=env)
|
|
try:
|
|
for line in process.stdout:
|
|
print(line, end="", flush=True)
|
|
handle.write(line)
|
|
handle.flush()
|
|
code = process.wait()
|
|
except BaseException:
|
|
process.terminate()
|
|
process.wait()
|
|
raise
|
|
if code:
|
|
raise subprocess.CalledProcessError(code, command)
|
|
|
|
|
|
def context_settings(settings: dict, original_config) -> dict:
|
|
requested = int(settings["cutoff_len"])
|
|
native = int(original_config.max_position_embeddings)
|
|
if requested > native or settings.get("rope_scaling"):
|
|
raise ValueError("The shared 4096-token scenario uses the original model contexts without RoPE extension.")
|
|
return {"requested_length": requested, "base_positions": native}
|
|
|
|
|
|
def reject_truncation(source_len: int, target_len: int, cutoff_len: int) -> None:
|
|
if source_len + target_len > cutoff_len:
|
|
raise ValueError(f"A complete SFT example requires {source_len + target_len} tokens, "
|
|
f"but only {cutoff_len} are available. Stopping instead of truncating input or answer.")
|
|
|
|
|
|
def require_complete_examples() -> None:
|
|
from llamafactory.data.processor import supervised, unsupervised
|
|
for module in (supervised, unsupervised):
|
|
original = module.infer_seqlen
|
|
|
|
def checked(source_len, target_len, cutoff_len, _original=original):
|
|
reject_truncation(source_len, target_len, cutoff_len)
|
|
return _original(source_len, target_len, cutoff_len)
|
|
module.infer_seqlen = checked
|
|
|
|
|
|
def bound_generation_context(model, context_length: int, default_max_new_tokens: int) -> None:
|
|
original_generate = model.generate
|
|
|
|
@functools.wraps(original_generate)
|
|
def bounded_generate(self, *args, **kwargs):
|
|
inputs = kwargs.get("input_ids")
|
|
if inputs is None:
|
|
inputs = kwargs.get("inputs")
|
|
if inputs is None and args:
|
|
inputs = args[0]
|
|
if inputs is None:
|
|
raise ValueError("Cannot verify the text input length before generation.")
|
|
prompt_length = int(inputs.shape[-1])
|
|
available = context_length - prompt_length
|
|
if available < 1:
|
|
raise ValueError("The complete input leaves no room for an answer inside the context limit.")
|
|
requested = kwargs.get("max_new_tokens")
|
|
requested = int(default_max_new_tokens if requested is None else requested)
|
|
if requested < 1:
|
|
raise ValueError("max_new_tokens must be positive.")
|
|
kwargs["max_new_tokens"] = min(requested, available)
|
|
if requested > available:
|
|
print(f"Generation limit: {available} new tokens after {prompt_length} input tokens "
|
|
f"(context {context_length}); the complete input is preserved.", flush=True)
|
|
return original_generate(*args, **kwargs)
|
|
|
|
model.generate = MethodType(bounded_generate, model)
|
|
|
|
|
|
def require_unsloth_backend(settings: dict, context: dict) -> dict:
|
|
# LLaMA-Factory may otherwise silently fall back when Unsloth rejects a model.
|
|
# The wrappers affect this process only; no installed package is edited.
|
|
from llamafactory.model import loader
|
|
from llamafactory.model.model_utils import unsloth as integration
|
|
loaded = {}
|
|
|
|
if hasattr(integration, "_get_unsloth_kwargs"):
|
|
original_kwargs = integration._get_unsloth_kwargs
|
|
|
|
@functools.wraps(original_kwargs)
|
|
def pinned_kwargs(config, model_name_or_path, model_args, finetuning_args):
|
|
values = original_kwargs(config, model_name_or_path, model_args, finetuning_args)
|
|
values["use_exact_model_name"] = True
|
|
if model_args.model_revision:
|
|
values["revision"] = model_args.model_revision
|
|
return values
|
|
|
|
integration._get_unsloth_kwargs = pinned_kwargs
|
|
|
|
def strict(original):
|
|
@functools.wraps(original)
|
|
def wrapped(*args, **kwargs):
|
|
model = original(*args, **kwargs)
|
|
if model is None:
|
|
raise RuntimeError("Unsloth did not load this model. Training/evaluation stopped; "
|
|
"use a compatible Unsloth/LLaMA-Factory environment.")
|
|
if model.config.max_position_embeddings < context["requested_length"]:
|
|
raise RuntimeError("The loaded model's configured context is shorter than the requested SFT context.")
|
|
if settings.get("do_train"):
|
|
loaded["base_config_json"] = model.config.to_json_string()
|
|
return model
|
|
return wrapped
|
|
|
|
for name in ("load_unsloth_pretrained_model", "load_unsloth_peft_model"):
|
|
original = getattr(loader, name, None)
|
|
if original is None:
|
|
raise RuntimeError("The installed LLaMA-Factory has an unsupported Unsloth integration.")
|
|
setattr(loader, name, strict(original))
|
|
if settings.get("do_predict") and settings.get("predict_with_generate"):
|
|
# Keep PEFT/Unsloth delegation: GenerationMixin.generate bound directly
|
|
# to a PeftModel fails in Transformers 5 because its class has no _sample.
|
|
original_patch_model = loader.patch_model
|
|
|
|
@functools.wraps(original_patch_model)
|
|
def patch_model_preserving_generate(model, *args, **kwargs):
|
|
original_generate = model.generate
|
|
result = original_patch_model(model, *args, **kwargs)
|
|
model.generate = original_generate
|
|
return result
|
|
|
|
loader.patch_model = patch_model_preserving_generate
|
|
# Attach the guard to the final model, after LLaMA-Factory patches generate
|
|
# and finishes adapter initialization. Earlier attachment can be replaced.
|
|
from llamafactory.train.sft import workflow
|
|
original_load_model = workflow.load_model
|
|
|
|
@functools.wraps(original_load_model)
|
|
def load_bounded_model(*args, **kwargs):
|
|
model = original_load_model(*args, **kwargs)
|
|
bound_generation_context(model, context["requested_length"], int(settings["max_new_tokens"]))
|
|
return model
|
|
|
|
workflow.load_model = load_bounded_model
|
|
return loaded
|
|
|
|
|
|
def worker(config_path: Path) -> None:
|
|
# This order is important for Unsloth's patches.
|
|
import unsloth # noqa: F401
|
|
import torch
|
|
config = read_config(config_path)
|
|
world = int(os.environ.get("WORLD_SIZE", "1"))
|
|
if config.get("do_train") and world != GPUS:
|
|
raise RuntimeError("This training scenario requires four processes / four GPUs.")
|
|
if not torch.cuda.is_available() or torch.cuda.device_count() < world:
|
|
raise RuntimeError(f"This job requires {world} visible CUDA GPU(s).")
|
|
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
|
torch.cuda.set_device(local_rank)
|
|
if not torch.cuda.is_bf16_supported():
|
|
raise RuntimeError("The allocated GPU does not support the requested bf16 precision.")
|
|
print(f"rank={local_rank}, world_size={world}, GPU={torch.cuda.get_device_name(local_rank)}, "
|
|
"precision=bf16, backend=Unsloth", flush=True)
|
|
from transformers import AutoConfig
|
|
original_config = AutoConfig.from_pretrained(config["model_name_or_path"], revision=config["model_revision"], trust_remote_code=False)
|
|
context = context_settings(config, original_config)
|
|
print("Model context:", json.dumps(context), flush=True)
|
|
loaded = require_unsloth_backend(config, context)
|
|
require_complete_examples()
|
|
from llamafactory.train.tuner import run_exp
|
|
run_exp(config)
|
|
if config.get("do_train") and local_rank == 0:
|
|
(Path(config["output_dir"]) / "base_model_config.json").write_text(loaded["base_config_json"], encoding="utf-8")
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("config", type=Path, help="configs/mistral.yaml or configs/olmo.yaml")
|
|
parser.add_argument("--run-dir", type=Path, help="New run directory; adapter is saved inside adapter/.")
|
|
parser.add_argument("--resume", type=Path, help="Checkpoint inside the original run's adapter/ directory.")
|
|
parser.add_argument("--dry-run", action="store_true", help="Validate data/config; do not import GPU libraries.")
|
|
parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS)
|
|
args = parser.parse_args()
|
|
if args.worker:
|
|
worker(args.config.resolve())
|
|
return 0
|
|
config = read_config(args.config.resolve())
|
|
if config.get("stage") != "sft" or config.get("finetuning_type") != "lora":
|
|
raise ValueError("Expected an SFT / LoRA training configuration.")
|
|
if not config.get("do_train") or config.get("eval_dataset") != "aya_eval":
|
|
raise ValueError("Training must use the prepared train and eval splits.")
|
|
if config.get("dataset") != "aya_train" or "val_size" in config:
|
|
raise ValueError("Use aya_train and aya_eval without another automatic split.")
|
|
if any(key in config for key in ("max_steps", "max_samples")):
|
|
raise ValueError("The full-data scenario must not impose a smoke-test sample/step limit.")
|
|
batch = GPUS * config["per_device_train_batch_size"] * config["gradient_accumulation_steps"]
|
|
data_dir = Path(config["dataset_dir"])
|
|
manifest = verify_data(data_dir, ("train", "eval"))
|
|
verify_length_filter(manifest, config)
|
|
if args.run_dir:
|
|
config["output_dir"] = str(args.run_dir.resolve() / "adapter")
|
|
run_dir = Path(config["output_dir"]).parent
|
|
if args.resume:
|
|
checkpoint = args.resume.resolve()
|
|
checkpoint.relative_to(Path(config["output_dir"]))
|
|
for name in ("trainer_state.json", "adapter_config.json", "optimizer.pt", "scheduler.pt"):
|
|
if not (checkpoint / name).is_file():
|
|
raise ValueError(f"Incomplete checkpoint: {checkpoint / name}")
|
|
previous = json.loads((run_dir / "run.json").read_text(encoding="utf-8"))
|
|
if previous["config"] != config or previous["data_manifest_sha256"] != digest(data_dir / "manifest.json"):
|
|
raise ValueError("Resume requires the original configuration and prepared data.")
|
|
elif run_dir.exists() and any(run_dir.iterdir()):
|
|
raise ValueError(f"Run directory is not empty; choose --run-dir for a new run: {run_dir}")
|
|
print(json.dumps({"gpus_per_training_run": GPUS, "effective_batch": batch,
|
|
"cutoff_len": config["cutoff_len"], "bf16": True, "use_unsloth": True,
|
|
"train_rows": manifest["splits"]["train"]["rows"],
|
|
"eval_rows": manifest["splits"]["eval"]["rows"],
|
|
"adapter_output": config["output_dir"]}, ensure_ascii=False, indent=2))
|
|
if args.dry_run:
|
|
print("Configuration/data checks passed. No training started.")
|
|
return 0
|
|
run_dir.mkdir(parents=True, exist_ok=True)
|
|
if not args.resume:
|
|
record = {"config": config, "gpus": GPUS, "effective_batch": batch,
|
|
"data_manifest_sha256": digest(data_dir / "manifest.json"),
|
|
"created_at_utc": datetime.now(timezone.utc).isoformat()}
|
|
(run_dir / "run.json").write_text(json.dumps(record, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
execution = dict(config)
|
|
if args.resume:
|
|
execution["resume_from_checkpoint"] = str(args.resume.resolve())
|
|
execution_path = run_dir / "used_config.yaml"
|
|
write_config(execution_path, execution)
|
|
env = os.environ.copy()
|
|
env.update({"PYTHONUNBUFFERED": "1", "TOKENIZERS_PARALLELISM": "false"})
|
|
env.setdefault("OMP_NUM_THREADS", "4")
|
|
command = [sys.executable, "-m", "torch.distributed.run", "--standalone",
|
|
f"--nproc_per_node={GPUS}", str(Path(__file__).resolve()),
|
|
str(execution_path), "--worker"]
|
|
run_logged(command, run_dir / "train.log", env)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
raise SystemExit(main())
|
|
except (OSError, ValueError, RuntimeError, ImportError, subprocess.CalledProcessError) as error:
|
|
print(f"ERROR: {error}", file=sys.stderr)
|
|
raise SystemExit(1)
|