@tuned-tensor/local 0.2.3 → 0.2.5

package/training/local-runner/src/train_dpo.py

@@ -0,0 +1,270 @@
+from __future__ import annotations
+
+import inspect
+import json
+import os
+import tarfile
+import time
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any
+
+import torch
+from datasets import Dataset as HFDataset
+from peft import LoraConfig, PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from trl import DPOConfig, DPOTrainer
+
+
+TRAINING_DIR = Path(os.environ.get("SM_CHANNEL_TRAINING", "/opt/ml/input/data/training"))
+BASE_MODEL_DIR = Path(os.environ.get("SM_CHANNEL_BASE_MODEL", "/opt/ml/input/data/base_model"))
+HYPERPARAMETERS_PATH = Path(
+    os.environ.get("TT_HYPERPARAMETERS_PATH", "/opt/ml/input/config/hyperparameters.json")
+)
+MODEL_DIR = Path(os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
+OUTPUT_DIR = Path(os.environ.get("SM_OUTPUT_DIR", "/opt/ml/output"))
+
+
+def load_hyperparameters() -> dict[str, str]:
+    if not HYPERPARAMETERS_PATH.is_file():
+        return {}
+    raw = json.loads(HYPERPARAMETERS_PATH.read_text())
+    return {str(key): str(value) for key, value in raw.items()}
+
+
+HP = load_hyperparameters()
+
+
+def hp(name: str, default: str | None = None) -> str | None:
+    value = os.getenv(f"SM_HP_{name.upper()}", HP.get(name, default))
+    if value is None:
+        return None
+    value = str(value).strip()
+    if len(value) >= 2 and value[0] == value[-1] == '"':
+        return value[1:-1]
+    return value
+
+
+def hp_int(name: str, default: int) -> int:
+    return int(hp(name, str(default)) or default)
+
+
+def hp_float(name: str, default: float) -> float:
+    return float(hp(name, str(default)) or default)
+
+
+def hp_bool(name: str, default: bool) -> bool:
+    return (hp(name, str(default)) or "").lower() in {"1", "true", "yes", "y"}
+
+
+def supported_kwargs(callable_obj: Any, values: dict[str, Any]) -> dict[str, Any]:
+    accepted = set(inspect.signature(callable_obj).parameters)
+    return {key: value for key, value in values.items() if key in accepted}
+
+
+def load_preference_rows() -> list[dict[str, str]]:
+    rows: list[dict[str, str]] = []
+    for path in sorted(TRAINING_DIR.rglob("*.jsonl")):
+        for line_number, line in enumerate(path.read_text().splitlines(), start=1):
+            if not line.strip():
+                continue
+            row = json.loads(line)
+            cleaned: dict[str, str] = {}
+            for key in ("prompt", "chosen", "rejected"):
+                value = row.get(key)
+                if not isinstance(value, str) or not value.strip():
+                    raise ValueError(f"{path}:{line_number} missing non-empty string field {key}")
+                cleaned[key] = value
+            rows.append(cleaned)
+    if not rows:
+        raise ValueError(f"No preference JSONL rows found under {TRAINING_DIR}")
+    return rows
+
+
+def resolve_model_source() -> str:
+    archive = BASE_MODEL_DIR / "model.tar.gz"
+    if archive.is_file():
+        extracted = Path("/tmp/base-model")
+        extracted.mkdir(parents=True, exist_ok=True)
+        with tarfile.open(archive, "r:gz") as tar:
+            extracted_root = extracted.resolve()
+            for member in tar.getmembers():
+                destination = (extracted / member.name).resolve()
+                try:
+                    destination.relative_to(extracted_root)
+                except ValueError:
+                    raise ValueError(f"Unsafe archive member: {member.name}")
+            tar.extractall(extracted)
+        candidates = [path for path in extracted.iterdir() if path.is_dir()]
+        return str(candidates[0] if len(candidates) == 1 else extracted)
+    return hp("base_model") or "Qwen/Qwen3.5-2B"
+
+
+def strip_file_uri(value: str | None) -> str | None:
+    if not value:
+        return value
+    if value.startswith("file://"):
+        return value[7:]
+    return value
+
+
+def resolve_adapter_path(value: str | None, tmp: Path) -> str | None:
+    path_value = strip_file_uri(value)
+    if not path_value:
+        return None
+    path = Path(path_value)
+    if path.is_file() and path.name.endswith(".tar.gz"):
+        extracted = tmp / "parent-adapter"
+        extracted.mkdir(parents=True, exist_ok=True)
+        with tarfile.open(path, "r:gz") as tar:
+            extracted_root = extracted.resolve()
+            for member in tar.getmembers():
+                destination = (extracted / member.name).resolve()
+                try:
+                    destination.relative_to(extracted_root)
+                except ValueError:
+                    raise ValueError(f"Unsafe archive member: {member.name}")
+            tar.extractall(extracted)
+        candidates = [candidate for candidate in extracted.rglob("adapter_config.json")]
+        if candidates:
+            return str(candidates[0].parent)
+        directories = [candidate for candidate in extracted.iterdir() if candidate.is_dir()]
+        return str(directories[0] if len(directories) == 1 else extracted)
+    return str(path)
+
+
+def create_model_and_tokenizer(model_source: str):
+    trust_remote_code = hp_bool("trust_remote_code", True)
+    token = os.getenv("HF_TOKEN")
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_source,
+        trust_remote_code=trust_remote_code,
+        token=token,
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
+    model = AutoModelForCausalLM.from_pretrained(
+        model_source,
+        trust_remote_code=trust_remote_code,
+        token=token,
+        torch_dtype=dtype if torch.cuda.is_available() else None,
+        device_map="auto" if torch.cuda.is_available() else None,
+    )
+    model.config.use_cache = False
+    return model, tokenizer
+
+
+def lora_target_modules(default: str) -> list[str] | str:
+    raw = hp("lora_target_modules", default) or default
+    if raw == "all-linear":
+        return raw
+    return [item.strip() for item in raw.split(",") if item.strip()]
+
+
+def lora_config() -> LoraConfig:
+    return LoraConfig(
+        r=hp_int("lora_rank", 16),
+        lora_alpha=hp_int("lora_alpha", 32),
+        lora_dropout=hp_float("lora_dropout", 0.05),
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules=lora_target_modules("all-linear"),
+    )
+
+
+def create_model_archive() -> Path:
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    archive_path = OUTPUT_DIR / "model.tar.gz"
+    with tarfile.open(archive_path, "w:gz") as tar:
+        tar.add(MODEL_DIR, arcname="model")
+    return archive_path
+
+
+def dpo_config(output_dir: str) -> DPOConfig:
+    max_prompt_length = hp("max_prompt_length")
+    max_completion_length = hp("max_completion_length")
+    config_values: dict[str, Any] = {
+        "output_dir": output_dir,
+        "num_train_epochs": hp_int("n_epochs", 3),
+        "learning_rate": hp_float("learning_rate", 0.00001),
+        "per_device_train_batch_size": hp_int("per_device_train_batch_size", 1),
+        "gradient_accumulation_steps": hp_int("gradient_accumulation_steps", 8),
+        "logging_steps": 1,
+        "save_strategy": "no",
+        "report_to": "none",
+        "remove_unused_columns": False,
+        "bf16": torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
+        "fp16": torch.cuda.is_available() and not torch.cuda.is_bf16_supported(),
+        "gradient_checkpointing": True,
+        "gradient_checkpointing_kwargs": {"use_reentrant": False},
+        "beta": hp_float("dpo_beta", 0.1),
+        "loss_type": hp("dpo_loss_type", "sigmoid"),
+        "label_smoothing": hp_float("dpo_label_smoothing", 0.0),
+        "reference_free": hp_bool("dpo_reference_free", False),
+        "max_length": hp_int("max_seq_length", 2048),
+        "max_prompt_length": int(max_prompt_length) if max_prompt_length else None,
+        "max_completion_length": int(max_completion_length) if max_completion_length else None,
+    }
+    return DPOConfig(**supported_kwargs(
+        DPOConfig,
+        {key: value for key, value in config_values.items() if value is not None},
+    ))
+
+
+def run_training(rows: list[dict[str, str]], model_source: str) -> dict[str, Any]:
+    with TemporaryDirectory() as tmp:
+        parent_adapter = resolve_adapter_path(hp("parent_model_artifact"), Path(tmp))
+        model, tokenizer = create_model_and_tokenizer(model_source)
+        if parent_adapter:
+            model = PeftModel.from_pretrained(model, parent_adapter, is_trainable=True)
+        dataset = HFDataset.from_list(rows)
+
+        args = dpo_config(tmp)
+        trainer_values: dict[str, Any] = {
+            "model": model,
+            "ref_model": None,
+            "args": args,
+            "train_dataset": dataset,
+            "processing_class": tokenizer,
+            "tokenizer": tokenizer,
+        }
+        if not parent_adapter:
+            trainer_values["peft_config"] = lora_config()
+        trainer = DPOTrainer(**supported_kwargs(DPOTrainer, trainer_values))
+        result = trainer.train()
+        trainer.save_model(MODEL_DIR)
+
+    tokenizer.save_pretrained(MODEL_DIR)
+    return result.metrics
+
+
+def main() -> None:
+    started = time.time()
+    MODEL_DIR.mkdir(parents=True, exist_ok=True)
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+    rows = load_preference_rows()
+    model_source = resolve_model_source()
+    metrics = run_training(rows, model_source)
+
+    archive_path = create_model_archive()
+    output_metrics = {
+        "training_method": "dpo",
+        "preference_rows": len(rows),
+        "model_source": model_source,
+        "parent_model_artifact": hp("parent_model_artifact"),
+        "dpo_beta": hp_float("dpo_beta", 0.1),
+        "dpo_loss_type": hp("dpo_loss_type", "sigmoid"),
+        "train_runtime": round(time.time() - started, 3),
+        **{key: float(value) for key, value in metrics.items() if isinstance(value, (int, float))},
+        "model_archive": str(archive_path),
+    }
+    (MODEL_DIR / "training-metrics.json").write_text(json.dumps(output_metrics, indent=2))
+    (OUTPUT_DIR / "training-metrics.json").write_text(json.dumps(output_metrics, indent=2))
+    print(json.dumps(output_metrics, indent=2))
+
+
+if __name__ == "__main__":
+    main()