File size: 4,923 Bytes
5814195 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """SFT LoRA for text->SVG. Runs on a Runpod pod. Wall-clock bounded; pushes adapter to HF; caller kills pod."""
import os, sys, time, json, random, math
import torch
from datasets import load_dataset
from transformers import (AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments,
TrainerCallback)
from peft import LoraConfig, get_peft_model
BASE = os.environ.get("BASE", "Qwen/Qwen3.5-4B")
DATA = os.environ.get("DATA", "aixonlab/svg-sft-v1")
OUT_REPO = os.environ.get("OUT_REPO", "aixonlab/svg-slm-4b-lora")
MAX_LEN = int(os.environ.get("MAX_LEN", 6400))
MAX_SECONDS = float(os.environ.get("TRAIN_SECONDS", 3 * 3600)) # hard wall-clock for training loop
MAX_STEPS = int(os.environ.get("MAX_STEPS", -1))
LR = float(os.environ.get("LR", 1e-4))
TOKEN = os.environ["HF_TOKEN"]
t0 = time.time()
tok = AutoTokenizer.from_pretrained(BASE, token=TOKEN)
model = AutoModelForCausalLM.from_pretrained(BASE, torch_dtype=torch.bfloat16, token=TOKEN,
attn_implementation="sdpa", device_map={"": 0})
model.gradient_checkpointing_enable()
model.enable_input_require_grads()
cfg = LoraConfig(r=32, lora_alpha=64, lora_dropout=0.0, task_type="CAUSAL_LM",
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"])
model = get_peft_model(model, cfg)
model.print_trainable_parameters()
ds = load_dataset(DATA, token=TOKEN)
train = ds["train"].shuffle(seed=0)
def encode(ex):
msgs = ex["messages"]
try:
prompt = tok.apply_chat_template(msgs[:1], tokenize=False, add_generation_prompt=True, enable_thinking=False)
except TypeError:
prompt = tok.apply_chat_template(msgs[:1], tokenize=False, add_generation_prompt=True)
p = tok(prompt, add_special_tokens=False)["input_ids"]
a = tok(msgs[1]["content"] + tok.eos_token if False else msgs[1]["content"] + "<|im_end|>\n",
add_special_tokens=False)["input_ids"]
ids = (p + a)[:MAX_LEN]
labels = ([-100] * len(p) + a)[:MAX_LEN]
return {"input_ids": ids, "labels": labels}
train = train.map(encode, remove_columns=train.column_names, num_proc=4)
train = train.filter(lambda e: len(e["input_ids"]) < MAX_LEN) # drop truncated
def collate(batch):
m = max(len(b["input_ids"]) for b in batch)
pad = tok.pad_token_id or 0
ids = torch.tensor([b["input_ids"] + [pad] * (m - len(b["input_ids"])) for b in batch])
lab = torch.tensor([b["labels"] + [-100] * (m - len(b["labels"])) for b in batch])
att = torch.tensor([[1] * len(b["input_ids"]) + [0] * (m - len(b["input_ids"])) for b in batch])
return {"input_ids": ids, "labels": lab, "attention_mask": att}
class TimeLimit(TrainerCallback):
def on_step_end(self, args, state, control, **kw):
if time.time() - t0 > MAX_SECONDS:
print("TIME LIMIT HIT -> stopping", flush=True)
control.should_training_stop = True
return control
# Estimate steps so cosine schedule completes inside the time budget: use max_steps after a probe.
bs, ga = int(os.environ.get("BS",1)), int(os.environ.get("GA",16))
n = len(train)
steps_full = n // (bs * ga)
args = TrainingArguments(
output_dir="/workspace/out", per_device_train_batch_size=bs, gradient_accumulation_steps=ga,
learning_rate=LR, lr_scheduler_type="cosine", warmup_steps=20, bf16=True, logging_steps=5,
save_strategy="steps", save_steps=100, save_total_limit=1, max_steps=MAX_STEPS if MAX_STEPS > 0 else steps_full,
report_to="none", dataloader_num_workers=2, remove_unused_columns=False,
gradient_checkpointing=True, optim="adamw_torch_fused", weight_decay=0.0, max_grad_norm=1.0)
trainer = Trainer(model=model, args=args, train_dataset=train, data_collator=collate, callbacks=[TimeLimit()])
def push(tag):
try:
from huggingface_hub import HfApi
api = HfApi(token=TOKEN)
api.create_repo(OUT_REPO, exist_ok=True, private=False)
model.save_pretrained("/workspace/adapter")
tok.save_pretrained("/workspace/adapter")
api.upload_folder(folder_path="/workspace/adapter", repo_id=OUT_REPO, commit_message=f"adapter {tag}")
print("PUSHED", tag, flush=True)
except Exception as e:
print("PUSH FAILED", e, flush=True)
class PushCkpt(TrainerCallback):
def on_save(self, args, state, control, **kw):
if state.global_step % 300 == 0:
push(f"step{state.global_step}")
trainer.add_callback(PushCkpt())
print(f"examples={n} steps_planned={args.max_steps} setup_s={time.time()-t0:.0f}", flush=True)
if os.environ.get("PROBE"):
t1 = time.time(); trainer.train(); sps = (time.time() - t1) / max(1, args.max_steps)
json.dump({"sec_per_step": sps, "full_steps": steps_full}, open("/workspace/probe.json", "w"))
print("PROBE", sps, steps_full, flush=True); sys.exit(0)
trainer.train()
push("final")
print("TRAINING DONE", flush=True)
|