C+E: stage 2 builder (mask, 48k, CLAS cap, family split, hook), private HF dataset, bf16 mixed training script with memory test, memory table, ratio proposal
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
Kral
163
train/hf_train_bf16.py
Normal file
163
train/hf_train_bf16.py
Normal file
@@ -0,0 +1,163 @@
|
||||
# /// 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
|
||||
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 --s1-epochs 1 --s2-epochs 3 --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 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)] * max(1, round(float(r.get("weight", 1.0))))
|
||||
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=1.0)
|
||||
ap.add_argument("--s2-epochs", type=float, default=3.0)
|
||||
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", 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"]
|
||||
model, tok = FastLanguageModel.from_pretrained(a.model, max_seq_length=a.max_seq, 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)
|
||||
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()
|
||||
Reference in New Issue
Block a user