Ilker ManapandClaude Opus 5.5
85025dfb03
An adapter now has an owner (org, team, project, user or key): only callers there see it, and to anyone else it is a model that does not exist, absent from /v1/models too. A training job's owner defaults to the one collection its pairs come from, must be able to read every collection they come from, and pairs from a knowledge base are never made public. ctl train publish keeps the GGUF or PEFT files on the mirror, which accepts uploads now, and sets the owned adapter after the usual architecture check. The knowledge bases' cache_salt follows the collections searched rather than the organization, so two teams of one organization no longer share cached passages. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
102 lines
4.9 KiB
Python
102 lines
4.9 KiB
Python
# llmgw's LoRA trainer (BACKLOG F4), run by the node agent with an installed
|
|
# engine's Python (its torch and transformers) and a small overlay (peft,
|
|
# accelerate). Reads a job, trains, writes the PEFT adapter, and says how it
|
|
# goes as one JSON object a line on standard output.
|
|
import json, math, os, random, sys, threading, subprocess, time
|
|
|
|
def say(**kw):
|
|
print(json.dumps(kw), flush=True)
|
|
|
|
job = json.load(open(sys.argv[1], encoding="utf-8"))
|
|
p = {"rank": 16, "alpha": 32, "dropout": 0.05, "epochs": 6, "lr": 2e-4, "batch": 4, "max_len": 512, "seed": 0}
|
|
p.update(job.get("params") or {})
|
|
|
|
say(state="preparing", message="loading libraries")
|
|
import torch
|
|
from huggingface_hub import snapshot_download
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
from peft import LoraConfig, get_peft_model
|
|
|
|
random.seed(p["seed"]); torch.manual_seed(p["seed"])
|
|
if not torch.cuda.is_available():
|
|
raise SystemExit("no CUDA device is visible to the trainer")
|
|
|
|
say(state="preparing", message="fetching " + job["base_repo"])
|
|
base_dir = snapshot_download(job["base_repo"], revision=job.get("base_revision") or None,
|
|
allow_patterns=["*.json", "*.safetensors", "*.txt", "*.model", "*.tiktoken"])
|
|
say(state="preparing", message="loading the base model", base_dir=base_dir)
|
|
|
|
# bf16 where the card has it; fp16 with a loss scaler where it does not
|
|
# (Turing and older).
|
|
bf16 = torch.cuda.is_bf16_supported()
|
|
dtype = torch.bfloat16 if bf16 else torch.float16
|
|
tok = AutoTokenizer.from_pretrained(base_dir)
|
|
if tok.pad_token_id is None:
|
|
tok.pad_token = tok.eos_token
|
|
model = AutoModelForCausalLM.from_pretrained(base_dir, dtype=dtype).to("cuda")
|
|
|
|
rows = [json.loads(l) for l in open(job["data"], encoding="utf-8") if l.strip()]
|
|
rows = [r for r in rows if r.get("q") and r.get("a")]
|
|
if not rows:
|
|
raise SystemExit("no {q, a} pairs in the data")
|
|
|
|
def encode(r):
|
|
prompt = tok.apply_chat_template([{"role": "user", "content": r["q"]}], add_generation_prompt=True, tokenize=False)
|
|
full = tok.apply_chat_template([{"role": "user", "content": r["q"]}, {"role": "assistant", "content": r["a"]}], tokenize=False)
|
|
pi = tok(prompt, add_special_tokens=False)["input_ids"]
|
|
fi = tok(full, add_special_tokens=False)["input_ids"][: p["max_len"]]
|
|
return fi, ([-100] * len(pi) + fi[len(pi):])[: len(fi)]
|
|
|
|
data = [encode(r) for r in rows]
|
|
|
|
model.gradient_checkpointing_enable()
|
|
model.enable_input_require_grads()
|
|
model = get_peft_model(model, LoraConfig(r=int(p["rank"]), lora_alpha=int(p["alpha"]), lora_dropout=float(p["dropout"]),
|
|
task_type="CAUSAL_LM", target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]))
|
|
trainable = [x for x in model.parameters() if x.requires_grad]
|
|
for x in trainable: # fp32 trainable weights, so a loss scaler can unscale them
|
|
x.data = x.data.float()
|
|
|
|
peak = {"mib": 0}
|
|
stop = threading.Event()
|
|
def watch():
|
|
while not stop.is_set():
|
|
try:
|
|
out = subprocess.run(["nvidia-smi", "--query-compute-apps=pid,used_memory", "--format=csv,noheader,nounits"],
|
|
capture_output=True, text=True, timeout=5).stdout
|
|
for line in out.splitlines():
|
|
pid, mib = [v.strip() for v in line.split(",")]
|
|
if int(pid) == os.getpid():
|
|
peak["mib"] = max(peak["mib"], int(mib))
|
|
except Exception:
|
|
pass
|
|
time.sleep(0.5)
|
|
threading.Thread(target=watch, daemon=True).start()
|
|
|
|
opt = torch.optim.AdamW(trainable, lr=float(p["lr"]), weight_decay=0.0)
|
|
scaler = torch.amp.GradScaler("cuda", enabled=not bf16)
|
|
bs, epochs = int(p["batch"]), int(p["epochs"])
|
|
steps = epochs * math.ceil(len(data) / bs)
|
|
warm = max(1, steps // 10)
|
|
sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / warm) * max(0.0, 1 - s / steps))
|
|
model.train(); step = 0; last = time.time()
|
|
for ep in range(epochs):
|
|
random.shuffle(data)
|
|
for i in range(0, len(data), bs):
|
|
batch = data[i:i + bs]; n = max(len(x[0]) for x in batch)
|
|
ids = torch.tensor([x[0] + [tok.pad_token_id] * (n - len(x[0])) for x in batch], device="cuda")
|
|
lab = torch.tensor([x[1] + [-100] * (n - len(x[1])) for x in batch], device="cuda")
|
|
att = torch.tensor([[1] * len(x[0]) + [0] * (n - len(x[0])) for x in batch], device="cuda")
|
|
with torch.autocast("cuda", dtype=dtype):
|
|
loss = model(input_ids=ids, attention_mask=att, labels=lab).loss
|
|
scaler.scale(loss).backward(); scaler.unscale_(opt)
|
|
torch.nn.utils.clip_grad_norm_(trainable, 1.0)
|
|
scaler.step(opt); scaler.update(); opt.zero_grad(set_to_none=True); sched.step(); step += 1
|
|
if time.time() - last > 2 or step == steps:
|
|
last = time.time()
|
|
say(state="running", message="training", step=step, steps=steps, loss=round(loss.item(), 4), peak_mib=peak["mib"])
|
|
|
|
model.save_pretrained(job["out"])
|
|
stop.set()
|
|
say(state="trained", step=step, steps=steps, peak_mib=peak["mib"], base_dir=base_dir)
|