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
2026-10-06 06:01:49 +02:00
parent f5611cbf56
commit 7f03849a86
5 changed files with 551 additions and 0 deletions

260
train/build_stage2.py Normal file
View File

@@ -0,0 +1,260 @@
"""Stage 2 training set builder (Opus item C, 2026-10-06).
train/.venv/bin/python train/build_stage2.py [--out runs/stage2_data] [--max-tokens 48000] [--clas-cap 0.35]
[--min-per-kind 25] [--valid-frac 0.10] [--hook module:function]
Input: runs/traj (DeepSeek trajectories; the local Qwen series is never read). Steps: acceptance filter (score, end reason, no
harness text, loop repeats trimmed) -> eval overlap check -> Qwen 3.8 chat template (thinking off) with the tokenizer round trip
check -> samples over --max-tokens are dropped (never cut) -> CLAS capped at --clas-cap of the stage 2 samples (the rest goes to
`stage2_reserve.jsonl`) -> shortage report for kinds under --min-per-kind -> validation split by task family (K variants and all
trajectories of a task on the same side) -> optional hook (own-test score, weights) -> files + data card.
Loss mask: `assistant_spans` are character spans of the assistant turns (tool calls and the final report, with <|im_end|>);
nothing of the system turn (tool schemas), the user turn or the tool results is in a span. `assistant_tokens` counts the loss tokens.
"""
import argparse
import importlib
import json
import os
import random
import statistics
import sys
import time
from collections import Counter, defaultdict
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
sys.path.insert(0, os.path.join(ROOT, "train"))
import accept as acc # noqa: E402
import to_qwen as tq # noqa: E402
from harness import mix, overlap # noqa: E402
from transformers import AutoTokenizer # noqa: E402
POOL = os.path.join(ROOT, "tasks_gen", "train")
def pct(v, p):
v = sorted(v)
return v[min(len(v) - 1, int(p * len(v)))] if v else None
def family_of(task_id):
try:
t = json.load(open(os.path.join(POOL, task_id, "task.json")))
except OSError:
return task_id
return t.get("base_task") or task_id
def loss_tokens(tok, text, spans):
enc = tok(text, add_special_tokens=False, return_offsets_mapping=True)
n = 0
for (a, b) in enc["offset_mapping"]:
if any(a >= s and b <= e for s, e in spans):
n += 1
return len(enc["input_ids"]), n
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--traj", default=os.path.join(ROOT, "runs", "traj"))
ap.add_argument("--out", default=os.path.join(ROOT, "runs", "stage2_data"))
ap.add_argument("--max-tokens", type=int, default=48000)
ap.add_argument("--clas-cap", type=float, default=0.35)
ap.add_argument("--min-per-kind", type=int, default=25)
ap.add_argument("--min-score", type=float, default=80)
ap.add_argument("--valid-frac", type=float, default=0.10)
ap.add_argument("--seed", type=int, default=20261006)
ap.add_argument("--hook", help="module:function; function(row) -> None to drop, or a float weight (repeat factor), or a dict "
"{'keep': bool, 'weight': float, 'extra': {...}} (for example the own-test mutation score)")
a = ap.parse_args()
os.makedirs(a.out, exist_ok=True)
hook = None
if a.hook:
mod, fn = a.hook.split(":")
sys.path.insert(0, os.path.join(ROOT, "train"))
hook = getattr(importlib.import_module(mod), fn)
tok = AutoTokenizer.from_pretrained(os.path.expanduser("~/models/Qwen3.8-27B-4bit"))
evals = overlap.load_pool("eval")
report = {"built": time.strftime("%F %T"), "settings": vars(a), "dropped": defaultdict(Counter), "steps": {}}
# 1 accepted trajectories
rows = [json.loads(l) for l in open(os.path.join(a.traj, "summary.jsonl"))]
cand = []
for r in rows:
p = os.path.join(a.traj, r.get("run_dir") or "-", "record.json")
if not os.path.exists(p):
continue
rec = json.load(open(p))
ok, why = acc.judge(rec, r, a.min_score)
if not ok:
continue
msgs, removed = acc.trim_loops(rec["messages"])
cand.append({"rec": rec, "msgs": msgs, "removed": removed, "row": r, "repair": acc.has_repair(msgs)})
report["steps"]["accepted_trajectories"] = len(cand)
# 2 eval overlap at build time (the task against every eval task)
kept = []
for c in cand:
tid = c["row"]["task"]
kind = mix.kind_of_task_dir(tid)
try:
hits = overlap.check(overlap.load_task(os.path.join(POOL, tid)), evals)
except (OSError, ValueError):
hits = []
if hits:
report["dropped"][kind]["eval_overlap"] += 1
continue
c["kind"] = kind
kept.append(c)
report["steps"]["after_eval_overlap"] = len(kept)
# 3 Qwen template + tokens + loss tokens; drop over the limit (never cut)
samples = []
for c in kept:
msgs = tq.to_template_messages(c["msgs"])
tools = c["rec"]["tools"]
text = tok.apply_chat_template(msgs, tools=tools, tokenize=False, enable_thinking=False)
prefix = tok.apply_chat_template(msgs[:2], tools=tools, tokenize=False, add_generation_prompt=True, enable_thinking=False)
if not text.startswith(prefix) or not tq.check_calls(msgs, text):
report["dropped"][c["kind"]]["template_roundtrip"] += 1
continue
spans = tq.spans(text)
n, nl = loss_tokens(tok, text, spans)
if n > a.max_tokens:
report["dropped"][c["kind"]]["over_%dk" % (a.max_tokens // 1000)] += 1
continue
t = c["rec"]["task"]
samples.append({"id": f"{c['row']['task']}_r{c['row']['run']}", "task": c["row"]["task"], "family": family_of(c["row"]["task"]),
"kind": c["kind"], "category": t.get("category"), "object_type": t.get("object_type"),
"attempt": c["row"]["attempt"], "repair": c["repair"], "loop_pairs_removed": c["removed"],
"score": c["rec"]["metadata"].get("score"), "teacher": c["rec"]["teacher"], "weight": 1.0,
"n_tokens": n, "assistant_tokens": nl, "text": text, "assistant_spans": spans})
report["steps"]["after_token_limit"] = len(samples)
# 4 CLAS cap: keep the CLAS samples with a repair and the highest score first; the rest is the reserve
by_kind = defaultdict(list)
for s in samples:
by_kind[s["kind"]].append(s)
non_clas = sum(len(v) for k, v in by_kind.items() if k != "CLAS")
clas = sorted(by_kind.get("CLAS", []), key=lambda s: (not s["repair"], -(s["score"] or 0), s["id"]))
limit = int(a.clas_cap / (1 - a.clas_cap) * non_clas) if non_clas else len(clas)
reserve = clas[limit:]
by_kind["CLAS"] = clas[:limit]
report["steps"]["clas_cap"] = {"clas_before": len(clas), "clas_kept": len(by_kind["CLAS"]), "clas_reserve": len(reserve), "limit": limit}
pool = [s for v in by_kind.values() for s in v]
# 5 shortage report (never filled by copies)
report["kinds"] = {}
for k in mix.TYPE_SHARE:
n = len(by_kind.get(k, []))
report["kinds"][k] = {"samples": n, "min_required": a.min_per_kind, "short": max(0, a.min_per_kind - n),
"families": len({s["family"] for s in by_kind.get(k, [])})}
# 6 hook (own-test score, weights)
if hook:
out = []
for s in pool:
r = hook(s)
if r is None or r is False:
report["dropped"][s["kind"]]["hook"] += 1
continue
if isinstance(r, dict):
if not r.get("keep", True):
report["dropped"][s["kind"]]["hook"] += 1
continue
s["weight"] = float(r.get("weight", 1.0))
s.update(r.get("extra") or {})
elif isinstance(r, (int, float)) and not isinstance(r, bool):
s["weight"] = float(r)
out.append(s)
pool = out
# 7 validation split by family, stratified by kind
rnd = random.Random(a.seed)
fam_kind = {}
for s in pool:
fam_kind.setdefault(s["family"], s["kind"])
valid_fam = set()
for k in mix.TYPE_SHARE:
fams = sorted(f for f, kk in fam_kind.items() if kk == k)
rnd.shuffle(fams)
nv = round(len(fams) * a.valid_frac)
if len(fams) >= 4:
nv = max(nv, 1)
valid_fam |= set(fams[:nv])
train = [s for s in pool if s["family"] not in valid_fam]
valid = [s for s in pool if s["family"] in valid_fam]
for name, data in (("stage2_train", train), ("stage2_valid", valid), ("stage2_reserve", reserve)):
with open(os.path.join(a.out, name + ".jsonl"), "w") as f:
for s in data:
f.write(json.dumps(s) + "\n")
# 8 report and data card
def summary(data):
toks = [s["n_tokens"] for s in data]
return {"samples": len(data), "tokens": sum(toks), "loss_tokens": sum(s["assistant_tokens"] for s in data),
"p50": pct(toks, .5), "p90": pct(toks, .9), "p95": pct(toks, .95), "max": max(toks or [0]),
"by_kind": dict(Counter(s["kind"] for s in data)), "by_category": dict(Counter(s["category"] for s in data)),
"repair_share": round(sum(s["repair"] for s in data) / max(len(data), 1), 2)}
report["train"], report["valid"], report["reserve"] = summary(train), summary(valid), summary(reserve)
report["dropped"] = {k: dict(v) for k, v in report["dropped"].items()}
s1 = json.load(open(os.path.join(ROOT, "train", "data", "stats.json"))) if os.path.exists(os.path.join(ROOT, "train", "data", "stats.json")) else {}
report["stage1_stats"] = s1
json.dump(report, open(os.path.join(a.out, "build_report.json"), "w"), indent=1, default=str)
open(os.path.join(a.out, "README.md"), "w").write(data_card(report))
print(json.dumps({k: report[k] for k in ("steps", "train", "valid", "reserve", "dropped")}, indent=1, default=str)[:3500])
print("short kinds:", {k: v["short"] for k, v in report["kinds"].items() if v["short"]})
def data_card(r):
t, v = r["train"], r["valid"]
kinds = "\n".join(f"| {k} | {x['samples']} | {x['families']} | {x['short'] or ''} |" for k, x in r["kinds"].items())
drop = "\n".join(f"- {k}: {d}" for k, d in r["dropped"].items()) or "- nothing dropped"
return f"""---
license: mit
task_categories: [text-generation]
tags: [abap, sap, agentic, tool-use, sft]
private: true
---
# ABAP stage 2 agent trajectories (built {r['built']})
Tool-using ABAP development trajectories for supervised fine-tuning of Qwen 3.8 27B. A teacher model solved generated ABAP tasks on a real
SAP ABAP Platform system (A4H, SAP_BASIS 816) through ADT tools; only runs that passed the harness gates and scored at least {r['settings']['min_score']:.0f} of 100 are kept.
## Sources and licenses
- Teacher: DeepSeek V4.1 Flash (MIT license), via Ollama cloud. No output of Claude or other restricted models.
- Tasks: generated by the same teacher (spec, seed objects, hidden ABAP Unit tests, reference), validated on A4H (oracle 100, null 0, mutation check). Eval tasks are not in this data (overlap check at build time).
- The local Qwen runs (series A) are not in this data.
## Format
One JSON per line: `text` (the whole conversation in the Qwen 3.8 chat template, thinking off, Qwen XML tool call format, the 20 tool schemas in the system turn),
`assistant_spans` (character spans that carry the loss: assistant turns with tool calls and the final report; none of the system turn, user turn or tool results),
`n_tokens`, `assistant_tokens`, `kind`, `category`, `object_type`, `family` (task family: a K variant has the family of its base task), `repair` (an error followed by a fix), `score`, `weight`.
## Size
| split | samples | tokens | loss tokens | p50 | p90 | p95 | max | repair share |
|---|---|---|---|---|---|---|---|---|
| train | {t['samples']} | {t['tokens']} | {t['loss_tokens']} | {t['p50']} | {t['p90']} | {t['p95']} | {t['max']} | {t['repair_share']} |
| valid | {v['samples']} | {v['tokens']} | {v['loss_tokens']} | {v['p50']} | {v['p90']} | {v['p95']} | {v['max']} | {v['repair_share']} |
Samples over {r['settings']['max_tokens']} tokens are dropped, never cut. CLAS is capped at {int(100 * r['settings']['clas_cap'])} % of the samples (the rest is in `stage2_reserve.jsonl`).
## Kinds (target share: CLAS 28, INTF 7, CDS 25, FUNC 15, PROG 10, TABL 8, STRU 2, MSAG 2.5, exception 2.5 percent)
| kind | samples | families | short of the minimum {r['settings']['min_per_kind']} |
|---|---|---|---|
{kinds}
## Dropped
{drop}
## Known limits
- Small data (see the table); kinds marked short have fewer than the minimum samples.
- The proxy added local abaplint messages (`syntaxCheck`) to some failed writes; the EPOD server does not do this itself.
- The teacher sees only the tool results of its own runs; a run that repaired an error is kept with its error (60 to 65 % of the samples).
- Validation split is by task family; do not mix valid samples into the training set.
"""
if __name__ == "__main__":
main()

163
train/hf_train_bf16.py Normal file
View 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()

28
train/hooks_example.py Normal file
View File

@@ -0,0 +1,28 @@
"""Hooks for train/build_stage2.py (--hook hooks_example:own_test_weight). A hook gets one sample (a dict) and returns
None / False (drop), a number (repeat weight), or {"keep": bool, "weight": float, "extra": {...}}."""
import json
import os
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
def identity(row):
return 1.0
def own_test_weight(row):
"""Item D (own-test mutation score, metadata only for now): reads runs/traj/<run>/own_test_mutation.json when it exists,
stores it as extra data and does NOT drop or reweight (Kral + Opus 2026-10-06: do not change the acceptance yet)."""
run = row["id"].split("_r")[-1]
for d in os.listdir(os.path.join(ROOT, "runs", "traj")):
if d.startswith(run + "_"):
p = os.path.join(ROOT, "runs", "traj", d, "own_test_mutation.json")
if os.path.exists(p):
m = json.load(open(p))
return {"keep": True, "weight": 1.0, "extra": {"own_test_mutation": m.get("score"), "own_test_mutants": m.get("mutants")}}
return {"keep": True, "weight": 1.0}
def repair_up(row):
"""Example of a weight: a trajectory with a repair counts twice."""
return 2.0 if row.get("repair") else 1.0