@tuned-tensor/local 0.1.0

package/training/sft-local/src/evaluate.py

@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import tarfile
+import time
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any
+
+import torch
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+def load_json(path: Path) -> dict[str, Any]:
+    return json.loads(path.read_text())
+
+
+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 / "adapter"
+        extracted.mkdir(parents=True, exist_ok=True)
+        with tarfile.open(path, "r:gz") as tar:
+            for member in tar.getmembers():
+                destination = (extracted / member.name).resolve()
+                if not str(destination).startswith(str(extracted.resolve())):
+                    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 fallback_prompt(system: str, prompt: str) -> str:
+    return f"system: {system}\nuser: {prompt}\nassistant:"
+
+
+def format_prompt(tokenizer: Any, system: str, prompt: str) -> str:
+    messages = [
+        {"role": "system", "content": system},
+        {"role": "user", "content": prompt},
+    ]
+    try:
+        return tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+    except Exception:
+        return fallback_prompt(system, prompt)
+
+
+def resolve_device(device: str) -> str:
+    if device == "auto":
+        if torch.cuda.is_available():
+            return "cuda"
+        if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+            return "mps"
+        return "cpu"
+    return device
+
+
+def load_model(payload: dict[str, Any], adapter_path: str | None):
+    base_model = payload["base_model"]
+    trust_remote_code = bool(payload.get("trust_remote_code", True))
+    device = resolve_device(str(payload.get("device", "auto")))
+    token = os.getenv("HF_TOKEN")
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        base_model,
+        trust_remote_code=trust_remote_code,
+        token=token,
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    dtype = None
+    if device == "cuda":
+        dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+
+    kwargs: dict[str, Any] = {
+        "trust_remote_code": trust_remote_code,
+        "token": token,
+    }
+    if dtype is not None:
+        kwargs["torch_dtype"] = dtype
+    if device == "cuda":
+        kwargs["device_map"] = "auto"
+
+    model = AutoModelForCausalLM.from_pretrained(base_model, **kwargs)
+    if adapter_path:
+        model = PeftModel.from_pretrained(model, adapter_path)
+    if device != "cuda":
+        model = model.to(device)
+    model.eval()
+    return model, tokenizer, device
+
+
+def generate_one(model: Any, tokenizer: Any, device: str, system: str, example: dict[str, Any], generation: dict[str, Any]) -> dict[str, Any]:
+    prompt = str(example["input"])
+    expected = str(example["output"])
+    formatted = format_prompt(tokenizer, system, prompt)
+    inputs = tokenizer(formatted, return_tensors="pt")
+    target_device = next(model.parameters()).device
+    inputs = {key: value.to(target_device) for key, value in inputs.items()}
+    started = time.perf_counter()
+    with torch.no_grad():
+        generated = model.generate(
+            **inputs,
+            max_new_tokens=int(generation.get("max_new_tokens", 256)),
+            do_sample=float(generation.get("temperature", 0)) > 0,
+            temperature=max(float(generation.get("temperature", 0)), 1e-5),
+            top_p=float(generation.get("top_p", 1)),
+            pad_token_id=tokenizer.eos_token_id,
+        )
+    latency_ms = max(0, round((time.perf_counter() - started) * 1000))
+    input_length = int(inputs["input_ids"].shape[-1])
+    output_tokens = generated[0][input_length:]
+    actual = tokenizer.decode(output_tokens, skip_special_tokens=True).strip()
+    return {
+        "prompt": prompt,
+        "expected": expected,
+        "actual": actual,
+        "latency_ms": latency_ms,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", required=True)
+    parser.add_argument("--output", required=True)
+    args = parser.parse_args()
+
+    payload = load_json(Path(args.input))
+    if payload.get("model_cache"):
+        os.environ["HF_HOME"] = str(payload["model_cache"])
+
+    with TemporaryDirectory() as tmp_dir:
+        adapter_path = resolve_adapter_path(payload.get("adapter_path"), Path(tmp_dir))
+        model, tokenizer, device = load_model(payload, adapter_path)
+        generation = payload.get("generation", {})
+        results = [
+            generate_one(model, tokenizer, device, str(payload.get("system", "")), example, generation)
+            for example in payload.get("examples", [])
+        ]
+
+    output = {
+        "provider": "transformers",
+        "model_id": payload.get("model_id"),
+        "base_model": payload.get("base_model"),
+        "adapter_path": payload.get("adapter_path"),
+        "generation_config": generation,
+        "results": results,
+    }
+    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+    Path(args.output).write_text(json.dumps(output, indent=2))
+
+
+if __name__ == "__main__":
+    main()

package/training/sft-local/src/train.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+import json
+import os
+import tarfile
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any
+
+import torch
+from peft import LoraConfig, get_peft_model
+from torch.utils.data import Dataset
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    DataCollatorForLanguageModeling,
+    Trainer,
+    TrainingArguments,
+)
+
+
+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 load_rows() -> list[dict[str, Any]]:
+    rows: list[dict[str, Any]] = []
+    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)
+            if "messages" not in row:
+                raise ValueError(f"{path}:{line_number} missing messages")
+            rows.append(row)
+    if not rows:
+        raise ValueError(f"No chat 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:
+            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 row_to_text(tokenizer: Any, row: dict[str, Any]) -> str:
+    messages = row["messages"]
+    try:
+        return tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+    except Exception:
+        parts: list[str] = []
+        for message in messages:
+            role = message.get("role", "user")
+            content = message.get("content", "")
+            if not isinstance(content, str):
+                content = json.dumps(content, ensure_ascii=False)
+            parts.append(f"{role}: {content}")
+        return "\n".join(parts)
+
+
+class TextDataset(Dataset[dict[str, torch.Tensor]]):
+    def __init__(self, texts: list[str], tokenizer: Any, max_length: int):
+        self.examples = [
+            tokenizer(
+                text,
+                truncation=True,
+                max_length=max_length,
+                padding=False,
+                return_tensors=None,
+            )
+            for text in texts
+        ]
+
+    def __len__(self) -> int:
+        return len(self.examples)
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
+        item = self.examples[index]
+        return {
+            "input_ids": torch.tensor(item["input_ids"], dtype=torch.long),
+            "attention_mask": torch.tensor(item["attention_mask"], dtype=torch.long),
+        }
+
+
+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 maybe_apply_lora(model: Any):
+    rank = hp_int("lora_rank", 16)
+    alpha = hp_int("lora_alpha", 32)
+    dropout = hp_float("lora_dropout", 0.05)
+    config = LoraConfig(
+        r=rank,
+        lora_alpha=alpha,
+        lora_dropout=dropout,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules="all-linear",
+    )
+    return get_peft_model(model, config)
+
+
+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 main() -> None:
+    started = time.time()
+    MODEL_DIR.mkdir(parents=True, exist_ok=True)
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+    rows = load_rows()
+    model_source = resolve_model_source()
+    model, tokenizer = create_model_and_tokenizer(model_source)
+    model = maybe_apply_lora(model)
+
+    texts = [row_to_text(tokenizer, row) for row in rows]
+    dataset = TextDataset(texts, tokenizer, max_length=hp_int("max_seq_length", 2048))
+    collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+    with TemporaryDirectory() as tmp:
+        args = TrainingArguments(
+            output_dir=tmp,
+            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=[],
+            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(),
+        )
+        trainer = Trainer(
+            model=model,
+            args=args,
+            train_dataset=dataset,
+            data_collator=collator,
+        )
+        result = trainer.train()
+
+    model.save_pretrained(MODEL_DIR)
+    tokenizer.save_pretrained(MODEL_DIR)
+    archive_path = create_model_archive()
+    metrics = {
+        "training_rows": len(rows),
+        "model_source": model_source,
+        "train_runtime": round(time.time() - started, 3),
+        **{key: float(value) for key, value in result.metrics.items() if isinstance(value, (int, float))},
+        "model_archive": str(archive_path),
+    }
+    (MODEL_DIR / "training-metrics.json").write_text(json.dumps(metrics, indent=2))
+    (OUTPUT_DIR / "training-metrics.json").write_text(json.dumps(metrics, indent=2))
+    print(json.dumps(metrics, indent=2))
+
+
+if __name__ == "__main__":
+    main()