183 lines
10 KiB
Python
183 lines
10 KiB
Python
# /// script
|
|
# requires-python = ">=3.10"
|
|
# dependencies = ["unsloth", "datasets", "transformers", "huggingface_hub"]
|
|
# ///
|
|
"""Stage 1 + stage 2 mixed bf16 LoRA run for Qwen 3.8 27B on Hugging Face Jobs (Opus item E, 2026-10-06). NOT started: no job without Kral's go.
|
|
|
|
Memory test first (a few dollars, finds the largest sequence length that fits, no data needed):
|
|
hf jobs uv run --flavor h200 --timeout 40m --secrets HF_TOKEN train/hf_train_bf16.py -- --memory-test --sweep 16000,32000,48000,64000
|
|
Real run (after the memory test and Kral's decision on the ratio):
|
|
hf jobs uv run --flavor <flavor> --timeout 10h --secrets HF_TOKEN train/hf_train_bf16.py -- \\
|
|
--stage1 erhankeseli/abap-stage1-data --stage2 erhankeseli/abap-stage2-data --s2-epochs 3 --s2-loss-share 0.6 --out erhankeseli/abap-mixed-adapter
|
|
|
|
Data: stage 1 rows have `text` (loss on every token); stage 2 rows have `text` and `assistant_spans` (loss only inside the spans: assistant turns;
|
|
none on the system turn with the tool schemas, the user turn or the tool results). A sample longer than --max-seq is skipped, never cut.
|
|
Base model in bf16 (no nf4): the adapter then fits the MLX 4-bit base better (stage 1 finding, train/STATE.md).
|
|
"""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import random
|
|
import time
|
|
|
|
|
|
def tokenize_masked(tok, text, spans, max_len):
|
|
"""input_ids and labels; labels are -100 outside the spans (spans=None: loss on every token). None when longer than max_len."""
|
|
enc = tok(text, add_special_tokens=False, return_offsets_mapping=True)
|
|
ids = enc["input_ids"]
|
|
if len(ids) > max_len:
|
|
return None
|
|
if spans is None:
|
|
return ids, list(ids)
|
|
labels = []
|
|
for tid, (a, b) in zip(ids, enc["offset_mapping"]):
|
|
labels.append(tid if any(a >= s and b <= e for s, e in spans) else -100)
|
|
return ids, labels
|
|
|
|
|
|
def s1_epochs_for(share, s2_loss_tokens, s2_epochs, s1_tokens):
|
|
"""Epochs of stage 1 so that stage 2 carries `share` of all loss tokens (Kral + Opus 2026-10-06: share = 0.6)."""
|
|
s2_total = s2_loss_tokens * s2_epochs
|
|
return (1 - share) / share * s2_total / max(s1_tokens, 1)
|
|
|
|
|
|
def copies(weight, rnd):
|
|
"""Weight 1 = one copy per epoch, 2 = two, 0.5 = a copy in half of the epochs (a down-weight, not a drop)."""
|
|
w = max(float(weight), 0.0)
|
|
return int(w) + (1 if rnd.random() < w - int(w) else 0)
|
|
|
|
|
|
def build_examples(tok, stage1, stage2, s1_epochs, s2_epochs, max_len, seed):
|
|
rnd = random.Random(seed)
|
|
rows = []
|
|
for ep in range(int(s1_epochs)):
|
|
rows += [("s1", r["text"], None, 1.0) for r in stage1]
|
|
frac = s1_epochs - int(s1_epochs)
|
|
if frac:
|
|
rows += [("s1", r["text"], None, 1.0) for r in rnd.sample(stage1, int(len(stage1) * frac))]
|
|
for ep in range(int(s2_epochs)):
|
|
for r in stage2:
|
|
rows += [("s2", r["text"], r["assistant_spans"], 1.0)] * copies(r.get("weight", 1.0), rnd)
|
|
rnd.shuffle(rows)
|
|
out, skipped = [], {"s1": 0, "s2": 0}
|
|
for src, text, spans, _ in rows:
|
|
t = tokenize_masked(tok, text, spans, max_len)
|
|
if t is None:
|
|
skipped[src] += 1
|
|
continue
|
|
out.append({"input_ids": t[0], "labels": t[1], "src": src})
|
|
return out, skipped
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--model", default="Qwen/Qwen3.8-27B")
|
|
ap.add_argument("--stage1", default="erhankeseli/abap-stage1-data")
|
|
ap.add_argument("--stage2", default="erhankeseli/abap-stage2-data")
|
|
ap.add_argument("--s1-file", default="train.jsonl")
|
|
ap.add_argument("--s2-file", default="stage2_train.jsonl")
|
|
ap.add_argument("--s2-valid", default="stage2_valid.jsonl")
|
|
ap.add_argument("--s1-epochs", type=float, default=None, help="epochs of stage 1; default: computed from --s2-loss-share")
|
|
ap.add_argument("--s2-epochs", type=float, default=3.0)
|
|
ap.add_argument("--s2-loss-share", type=float, default=0.6, help="share of the loss tokens that stage 2 carries (decision 2026-10-06: 0.6)")
|
|
ap.add_argument("--max-seq", type=int, default=48000)
|
|
ap.add_argument("--rank", type=int, default=16)
|
|
ap.add_argument("--alpha", type=int, default=32)
|
|
ap.add_argument("--lr", type=float, default=5e-5)
|
|
ap.add_argument("--grad-accum", type=int, default=4)
|
|
ap.add_argument("--warmup-steps", type=int, default=20)
|
|
ap.add_argument("--no-offload", action="store_true", help="gradient checkpointing on the GPU instead of the Unsloth offload")
|
|
ap.add_argument("--out", default="erhankeseli/abap-mixed-adapter")
|
|
ap.add_argument("--save-every", type=int, default=0)
|
|
ap.add_argument("--memory-test", action="store_true")
|
|
ap.add_argument("--sweep", default="16000,32000,48000,64000", help="memory test: sequence lengths, tried in this order, stops at the first OOM")
|
|
ap.add_argument("--steps", type=int, default=3, help="memory test: optimizer steps per length")
|
|
a = ap.parse_args()
|
|
|
|
import torch
|
|
from unsloth import FastLanguageModel
|
|
from huggingface_hub import HfApi, hf_hub_download
|
|
token = os.environ["HF_TOKEN"]
|
|
seq_for_model = max([a.max_seq] + ([int(x) for x in a.sweep.split(",")] if a.memory_test else []))
|
|
model, tok = FastLanguageModel.from_pretrained(a.model, max_seq_length=seq_for_model, load_in_4bit=False, dtype=torch.bfloat16, token=token)
|
|
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=True if a.no_offload else "unsloth", random_state=20261006)
|
|
print("GPU", torch.cuda.get_device_name(0), "x", torch.cuda.device_count(), "| weights allocated GB", round(torch.cuda.memory_allocated() / 2**30, 1), flush=True)
|
|
|
|
if a.memory_test:
|
|
text = "CLASS zcl_demo DEFINITION PUBLIC. METHODS run IMPORTING iv TYPE i. ENDCLASS.\n" * 4000
|
|
ids_all = tok(text, add_special_tokens=False)["input_ids"]
|
|
opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-5)
|
|
model.train()
|
|
result = {"gpu": torch.cuda.get_device_name(0), "gpus": torch.cuda.device_count(), "offload": not a.no_offload, "sweep": []}
|
|
for n in [int(x) for x in a.sweep.split(",")]:
|
|
ids = torch.tensor([(ids_all * (n // len(ids_all) + 1))[:n]], device="cuda")
|
|
labels = ids.clone()
|
|
labels[:, : int(n * 0.68)] = -100 # about 32 % loss tokens, as in the stage 2 data
|
|
torch.cuda.reset_peak_memory_stats()
|
|
try:
|
|
times = []
|
|
for _ in range(a.steps):
|
|
t0 = time.time()
|
|
loss = model(input_ids=ids, labels=labels).loss
|
|
loss.backward()
|
|
opt.step()
|
|
opt.zero_grad(set_to_none=True)
|
|
torch.cuda.synchronize()
|
|
times.append(round(time.time() - t0, 1))
|
|
row = {"seq_len": n, "ok": True, "peak_gb": round(torch.cuda.max_memory_allocated() / 2**30, 1),
|
|
"reserved_gb": round(torch.cuda.max_memory_reserved() / 2**30, 1), "step_seconds": times, "loss": float(loss)}
|
|
except torch.cuda.OutOfMemoryError as e:
|
|
row = {"seq_len": n, "ok": False, "error": "OOM", "peak_gb": round(torch.cuda.max_memory_allocated() / 2**30, 1)}
|
|
result["sweep"].append(row)
|
|
print("MEMTEST", json.dumps(row), flush=True)
|
|
break
|
|
result["sweep"].append(row)
|
|
print("MEMTEST", json.dumps(row), flush=True)
|
|
torch.cuda.empty_cache()
|
|
print("RESULT", json.dumps(result), flush=True)
|
|
HfApi(token=token).upload_file(path_or_fileobj=json.dumps(result, indent=1).encode(), path_in_repo="memtest_%s_%dx.json" % (
|
|
result["gpu"].replace(" ", "_"), result["gpus"]), repo_id=a.out, repo_type="model", create_pr=False) if a.out else None
|
|
return
|
|
|
|
def rows(repo, fname):
|
|
p = hf_hub_download(repo, fname, repo_type="dataset", token=token)
|
|
return [json.loads(l) for l in open(p)]
|
|
s1, s2 = rows(a.stage1, a.s1_file), rows(a.stage2, a.s2_file)
|
|
if a.s1_epochs is None:
|
|
s1_tok = sum(len(tok(r["text"], add_special_tokens=False)["input_ids"]) for r in s1)
|
|
s2_loss = sum(r["assistant_tokens"] * float(r.get("weight", 1.0)) for r in s2)
|
|
a.s1_epochs = s1_epochs_for(a.s2_loss_share, s2_loss, a.s2_epochs, s1_tok)
|
|
print("STAGE1 EPOCHS", round(a.s1_epochs, 3), "(stage 1 tokens", s1_tok, ", stage 2 loss tokens per epoch", round(s2_loss), ", share", a.s2_loss_share, ")", flush=True)
|
|
ex, skipped = build_examples(tok, s1, s2, a.s1_epochs, a.s2_epochs, a.max_seq, 20261006)
|
|
loss_tok = sum(sum(1 for x in e["labels"] if x != -100) for e in ex)
|
|
print("EXAMPLES", len(ex), "skipped", skipped, "loss tokens", loss_tok, "stage 2 share of loss tokens",
|
|
round(sum(sum(1 for x in e["labels"] if x != -100) for e in ex if e["src"] == "s2") / max(loss_tok, 1), 2), flush=True)
|
|
from datasets import Dataset
|
|
from transformers import Trainer, TrainingArguments
|
|
|
|
class Collate:
|
|
def __call__(self, batch): # batch size 1; no padding needed
|
|
b = batch[0]
|
|
return {"input_ids": torch.tensor([b["input_ids"]]), "labels": torch.tensor([b["labels"]]),
|
|
"attention_mask": torch.ones(1, len(b["input_ids"]), dtype=torch.long)}
|
|
args = TrainingArguments(output_dir="out", per_device_train_batch_size=1, gradient_accumulation_steps=a.grad_accum, num_train_epochs=1,
|
|
learning_rate=a.lr, lr_scheduler_type="cosine", warmup_steps=a.warmup_steps, optim="adamw_8bit", bf16=True,
|
|
logging_steps=1, save_strategy="no", report_to="none", seed=20261006, remove_unused_columns=False)
|
|
tr = Trainer(model=model, args=args, train_dataset=Dataset.from_list(ex), data_collator=Collate())
|
|
t0 = time.time()
|
|
tr.train()
|
|
info = {"examples": len(ex), "skipped": skipped, "train_seconds": round(time.time() - t0, 1), "rank": a.rank, "alpha": a.alpha, "lr": a.lr,
|
|
"peak_gpu_gb": round(torch.cuda.max_memory_allocated() / 2**30, 1), "s1_epochs": a.s1_epochs, "s2_epochs": a.s2_epochs}
|
|
print("RESULT", json.dumps(info), flush=True)
|
|
model.save_pretrained("adapter")
|
|
json.dump(info, open("adapter/job_result.json", "w"))
|
|
HfApi(token=token).upload_folder(folder_path="adapter", repo_id=a.out, repo_type="model")
|
|
print("PUSHED", a.out)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|