| """SemanticWiki v3 - LoRA SFT of 14B-class students on citation-verified data. |
| |
| Uses TRL's conversational prompt-completion format: `prompt` = the system+user |
| messages, `completion` = the assistant wiki page. Loss is computed on the |
| completion only (TRL default for prompt-completion datasets), so the model is |
| not trained to regurgitate the code context. |
| |
| Real run: |
| python sw_train.py --base_model Qwen/Qwen2.5-Coder-14B-Instruct \ |
| --hub_model_id GhostScientist/semanticwiki-coder-14b-v3 |
| Smoke (CPU): add --smoke (tiny subset, 1 step, no push, no trackio). |
| """ |
|
|
| import argparse |
| import json |
| import os |
|
|
| from datasets import load_dataset |
| from peft import LoraConfig |
| from transformers import AutoTokenizer, TrainerCallback |
| from trl import SFTConfig, SFTTrainer |
|
|
| DATASET = "GhostScientist/semanticwiki-data-v3" |
| MAX_LEN = 24576 |
|
|
| TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj", |
| "gate_proj", "up_proj", "down_proj"] |
|
|
|
|
| class TrackioCallback(TrainerCallback): |
| def on_log(self, args, state, control, logs=None, **kwargs): |
| if logs: |
| import trackio |
| trackio.log({k: v for k, v in logs.items() |
| if isinstance(v, (int, float))}, |
| step=state.global_step) |
|
|
|
|
| def _norm_content(c): |
| """Arrow may infer message content as Json(decode=True), turning some string |
| contents into e.g. ["the string"]. Normalize everything back to str.""" |
| if isinstance(c, str): |
| return c |
| if isinstance(c, list): |
| return "".join(x if isinstance(x, str) else x.get("text", "") for x in c) |
| if c is None: |
| return "" |
| return json.dumps(c) |
|
|
|
|
| def to_prompt_completion(ex): |
| msgs = ex["messages"] |
| if isinstance(msgs, str): |
| msgs = json.loads(msgs) |
| msgs = [{"role": m.get("role", "user"), "content": _norm_content(m.get("content"))} |
| for m in msgs] |
| return {"prompt": msgs[:-1], "completion": [msgs[-1]]} |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--base_model", required=True) |
| ap.add_argument("--hub_model_id", default=None) |
| ap.add_argument("--dataset", default=DATASET) |
| ap.add_argument("--epochs", type=float, default=2.0) |
| ap.add_argument("--lr", type=float, default=1e-4) |
| ap.add_argument("--max_steps", type=int, default=-1) |
| ap.add_argument("--max_len", type=int, default=MAX_LEN, |
| help="override sequence length (GPU smoke uses the real 24576)") |
| ap.add_argument("--smoke", action="store_true", |
| help="tiny subset, 1 step, no push, no trackio, short seqs") |
| args = ap.parse_args() |
|
|
| tok = AutoTokenizer.from_pretrained(args.base_model) |
| ds = load_dataset(args.dataset, split="train") |
| assert "messages" in ds.column_names, f"unexpected columns: {ds.column_names}" |
|
|
| max_len = 2048 if args.smoke else args.max_len |
| pc = ds.map(to_prompt_completion, remove_columns=ds.column_names, num_proc=4) |
| |
| for ex in pc.select(range(min(2, len(pc)))): |
| tok.apply_chat_template(ex["prompt"] + ex["completion"]) |
| print("[train] chat-template sanity check passed on 2 rows") |
| bad = [i for i, ex in enumerate(pc) |
| if not (isinstance(ex["prompt"], list) and ex["prompt"] |
| and isinstance(ex["completion"], list) and ex["completion"])] |
| print(f"[train] rows with malformed prompt/completion: {len(bad)}") |
|
|
| def keep(ex): |
| ids = tok.apply_chat_template(ex["prompt"] + ex["completion"], tokenize=True) |
| if isinstance(ids, dict): |
| ids = ids["input_ids"] |
| return len(ids) <= max_len |
|
|
| n0 = len(pc) |
| pc = pc.filter(keep, num_proc=4) |
| print(f"[train] dataset: {n0} -> {len(pc)} rows (length <= {max_len})") |
| assert len(pc) > 0, "length filter removed everything" |
|
|
| if args.smoke: |
| pc = pc.select(range(min(4, len(pc)))) |
|
|
| space_id = None if args.smoke else os.environ.get("TRACKIO_SPACE_ID") |
| if space_id: |
| import trackio |
| trackio.init(project="semanticwiki-sft-14b", space_id=space_id) |
|
|
| |
| dtype = "float32" if args.smoke else "bfloat16" |
| cfg = SFTConfig( |
| output_dir="/work/out", |
| per_device_train_batch_size=1, |
| gradient_accumulation_steps=16, |
| num_train_epochs=args.epochs, |
| max_steps=args.max_steps, |
| learning_rate=args.lr, |
| lr_scheduler_type="cosine", |
| warmup_steps=12, |
| logging_steps=5, |
| save_strategy="no", |
| max_length=max_len, |
| packing=False, |
| report_to="none", |
| bf16=not args.smoke, |
| fp16=False, |
| model_init_kwargs={"dtype": dtype, "attn_implementation": "sdpa"}, |
| push_to_hub=bool(args.hub_model_id) and not args.smoke, |
| hub_model_id=args.hub_model_id, |
| seed=42, |
| ) |
| lora = LoraConfig( |
| r=64, lora_alpha=128, lora_dropout=0.05, |
| target_modules=TARGET_MODULES, task_type="CAUSAL_LM", |
| ) |
|
|
| trainer = SFTTrainer( |
| model=args.base_model, |
| args=cfg, |
| train_dataset=pc, |
| peft_config=lora, |
| callbacks=[TrackioCallback()] if space_id else [], |
| ) |
| trainer.train() |
| trainer.save_model("/work/out/final") |
| print("[train] done; final loss history:", |
| [round(log.get("loss", -1), 4) for log in trainer.state.log_history |
| if "loss" in log][-5:]) |
|
|
| if args.hub_model_id and not args.smoke: |
| trainer.push_to_hub() |
| print("[train] pushed adapter to " + args.hub_model_id) |
| if space_id: |
| import trackio |
| trackio.finish() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
