| """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)) |
| 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) |
|
|
| 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 |
|
|
| |
| 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) |
|
|
