84 lines
4.2 KiB
Python
84 lines
4.2 KiB
Python
# /// script
|
|
# requires-python = ">=3.10"
|
|
# dependencies = ["unsloth", "datasets", "trl", "huggingface_hub"]
|
|
# ///
|
|
"""Stage 1 LoRA training with Unsloth on Hugging Face Jobs (1x A100 80GB).
|
|
|
|
Run (pipeline test, 10 steps):
|
|
hf jobs uv run --flavor a100-large --timeout 45m --secrets HF_TOKEN train/hf_train.py -- --max-steps 10
|
|
Full run: no --max-steps (2 epochs). Settings follow train/config.yaml; alpha is a proposal, see STATE.md.
|
|
"""
|
|
import argparse, json, os, time
|
|
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--model", default="unsloth/Qwen3.8-27B-unsloth-bnb-4bit")
|
|
ap.add_argument("--data", default="erhankeseli/abap-stage1-data")
|
|
ap.add_argument("--out", default="erhankeseli/abap-stage1-adapter-test")
|
|
ap.add_argument("--max-steps", type=int, default=-1)
|
|
ap.add_argument("--epochs", type=float, default=2.0)
|
|
ap.add_argument("--rank", type=int, default=16)
|
|
ap.add_argument("--alpha", type=int, default=16)
|
|
ap.add_argument("--lr", type=float, default=5e-5)
|
|
ap.add_argument("--train-file", default="train.jsonl")
|
|
ap.add_argument("--eval-file", default="valid.jsonl")
|
|
ap.add_argument("--warmup-steps", type=int, default=30)
|
|
ap.add_argument("--save-every", type=int, default=0, help="upload the adapter to <out>/step<N> every N steps")
|
|
ap.add_argument("--max-seq-length", type=int, default=16384)
|
|
a = ap.parse_args()
|
|
|
|
from unsloth import FastLanguageModel # noqa: E402 (import first)
|
|
from datasets import load_dataset # noqa: E402
|
|
from trl import SFTConfig, SFTTrainer # noqa: E402
|
|
|
|
tok_hf = os.environ["HF_TOKEN"]
|
|
model, tok = FastLanguageModel.from_pretrained(
|
|
a.model, max_seq_length=a.max_seq_length, load_in_4bit=True, token=tok_hf)
|
|
model = FastLanguageModel.get_peft_model(
|
|
model, r=a.rank, lora_alpha=a.alpha, lora_dropout=0.0, bias="none",
|
|
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj",
|
|
"in_proj_qkv", "in_proj_z", "out_proj"],
|
|
use_gradient_checkpointing="unsloth", random_state=20261003)
|
|
|
|
ds = load_dataset(a.data, data_files={"train": a.train_file, "valid": a.eval_file}, token=tok_hf)
|
|
cfg = SFTConfig(
|
|
output_dir="out", per_device_train_batch_size=1, per_device_eval_batch_size=1,
|
|
gradient_accumulation_steps=1, num_train_epochs=a.epochs, max_steps=a.max_steps,
|
|
learning_rate=a.lr, lr_scheduler_type="cosine", warmup_steps=a.warmup_steps,
|
|
optim="adamw_8bit", weight_decay=0.0, logging_steps=1, eval_strategy="no", save_strategy="no",
|
|
max_length=a.max_seq_length, dataset_text_field="text", packing=False, seed=20261003, report_to="none")
|
|
tr = SFTTrainer(model=model, processing_class=tok, train_dataset=ds["train"], eval_dataset=ds["valid"], args=cfg)
|
|
|
|
from transformers import TrainerCallback # noqa: E402
|
|
from huggingface_hub import HfApi # noqa: E402
|
|
api = HfApi(token=tok_hf)
|
|
|
|
|
|
class SaveCb(TrainerCallback):
|
|
def on_step_end(self, args, state, control, **kw):
|
|
n = state.global_step
|
|
if a.save_every and n % a.save_every == 0:
|
|
model.save_pretrained(f"ckpt{n}")
|
|
api.upload_folder(folder_path=f"ckpt{n}", repo_id=a.out, repo_type="model", path_in_repo=f"step{n}")
|
|
ev = tr.evaluate()
|
|
print("EVAL", json.dumps({"step": n, "eval_loss": ev.get("eval_loss")}), flush=True)
|
|
api.upload_file(path_or_fileobj=json.dumps({"step": n, "eval_loss": ev.get("eval_loss")}).encode(),
|
|
path_in_repo=f"step{n}/eval.json", repo_id=a.out, repo_type="model")
|
|
|
|
|
|
tr.add_callback(SaveCb())
|
|
ev0 = tr.evaluate()
|
|
print("EVAL", json.dumps({"step": 0, "eval_loss": ev0.get("eval_loss")}), flush=True)
|
|
t0 = time.time()
|
|
tr.train()
|
|
train_s = time.time() - t0
|
|
steps = tr.state.global_step
|
|
ev = tr.evaluate() # valid loss of the adapter, for the conversion check
|
|
info = {"eval_loss_step0": ev0.get("eval_loss"), "steps": steps, "train_seconds": round(train_s, 1), "sec_per_step": round(train_s / max(steps, 1), 2),
|
|
"eval_loss": ev.get("eval_loss"), "rank": a.rank, "alpha": a.alpha, "lr": a.lr,
|
|
"peak_gpu_gb": round(__import__("torch").cuda.max_memory_allocated() / 2**30, 1)}
|
|
print("RESULT", json.dumps(info))
|
|
model.save_pretrained("adapter")
|
|
json.dump(info, open("adapter/job_result.json", "w"))
|
|
api.upload_folder(folder_path="adapter", repo_id=a.out, repo_type="model")
|
|
print("PUSHED", a.out)
|