| |
| """LoRA/QLoRA SFT runner for the Qwen CyberGym project. |
| |
| This script is intentionally framework-light: it uses Transformers Trainer plus |
| PEFT, and consumes the YAML contracts in training/configs. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import inspect |
| import json |
| import os |
| import subprocess |
| import sys |
| import time |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| import yaml |
|
|
|
|
| ASSISTANT_START = "<|im_start|>assistant\n" |
| IM_END = "<|im_end|>" |
|
|
|
|
| @dataclass |
| class TokenizedExample: |
| input_ids: list[int] |
| attention_mask: list[int] |
| labels: list[int] |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", required=True, help="YAML config under training/configs.") |
| parser.add_argument( |
| "--allow-missing-cybergym-baseline", |
| action="store_true", |
| help="Dry-run escape hatch. Real training should not use this.", |
| ) |
| parser.add_argument("--dry-run", action="store_true", help="Validate config/data/model class imports without training.") |
| return parser.parse_args() |
|
|
|
|
| def read_yaml(path: str | Path) -> dict[str, Any]: |
| with Path(path).open("r", encoding="utf-8") as fh: |
| payload = yaml.safe_load(fh) or {} |
| if not isinstance(payload, dict): |
| raise TypeError(f"Expected a YAML mapping in {path}") |
| return payload |
|
|
|
|
| def run_gate_check(config_path: str, allow_missing: bool) -> None: |
| cmd = [sys.executable, "training/scripts/check_training_gates.py", "--config", config_path] |
| if allow_missing: |
| cmd.append("--allow-missing-cybergym-baseline") |
| subprocess.run(cmd, check=True) |
|
|
|
|
| def import_training_deps(): |
| try: |
| import transformers |
| from datasets import Dataset |
| from peft import LoraConfig, TaskType, get_peft_model |
| from transformers import ( |
| AutoTokenizer, |
| BitsAndBytesConfig, |
| Trainer, |
| TrainingArguments, |
| ) |
| except Exception as exc: |
| raise RuntimeError(f"Missing training dependency: {exc!r}") from exc |
|
|
| return { |
| "transformers": transformers, |
| "Dataset": Dataset, |
| "AutoTokenizer": AutoTokenizer, |
| "BitsAndBytesConfig": BitsAndBytesConfig, |
| "Trainer": Trainer, |
| "TrainingArguments": TrainingArguments, |
| "LoraConfig": LoraConfig, |
| "TaskType": TaskType, |
| "get_peft_model": get_peft_model, |
| } |
|
|
|
|
| def torch_dtype(name: str): |
| if name == "auto": |
| return "auto" |
| return { |
| "bfloat16": torch.bfloat16, |
| "float16": torch.float16, |
| "float32": torch.float32, |
| }[name] |
|
|
|
|
| def load_model(transformers_module, model_cfg: dict[str, Any], quantization_cfg: dict[str, Any] | None): |
| model_name = model_cfg["name_or_path"] |
| dtype = torch_dtype(str(model_cfg.get("dtype", "bfloat16"))) |
| kwargs: dict[str, Any] = { |
| "trust_remote_code": bool(model_cfg.get("trust_remote_code", True)), |
| "torch_dtype": dtype, |
| } |
| if quantization_cfg: |
| deps = import_training_deps() |
| kwargs["quantization_config"] = deps["BitsAndBytesConfig"](**quantization_cfg) |
| kwargs["device_map"] = model_cfg.get("device_map", "auto") |
|
|
| attn = model_cfg.get("attn_implementation") |
| if attn: |
| kwargs["attn_implementation"] = attn |
|
|
| candidate_class_names = [ |
| "AutoModelForMultimodalLM", |
| "AutoModelForImageTextToText", |
| "AutoModelForVision2Seq", |
| "AutoModelForCausalLM", |
| ] |
| errors: list[str] = [] |
| for class_name in candidate_class_names: |
| model_cls = getattr(transformers_module, class_name, None) |
| if model_cls is None: |
| errors.append(f"{class_name}: not available") |
| continue |
| try: |
| return model_cls.from_pretrained(model_name, **kwargs) |
| except Exception as exc: |
| errors.append(f"{class_name}: {exc!r}") |
|
|
| fallback_attn = model_cfg.get("fallback_attn_implementation") |
| if fallback_attn and attn and fallback_attn != attn: |
| kwargs["attn_implementation"] = fallback_attn |
| for class_name in candidate_class_names: |
| model_cls = getattr(transformers_module, class_name, None) |
| if model_cls is None: |
| continue |
| try: |
| return model_cls.from_pretrained(model_name, **kwargs) |
| except Exception as exc: |
| errors.append(f"{class_name} with fallback attn: {exc!r}") |
|
|
| raise RuntimeError("Could not load model:\n" + "\n".join(errors)) |
|
|
|
|
| def freeze_by_name(model, patterns: list[str]) -> int: |
| frozen = 0 |
| lowered = [pattern.lower() for pattern in patterns] |
| for name, param in model.named_parameters(): |
| if any(pattern in name.lower() for pattern in lowered): |
| param.requires_grad = False |
| frozen += param.numel() |
| return frozen |
|
|
|
|
| def build_lora_config(lora_cls, task_type, lora_cfg: dict[str, Any]): |
| payload: dict[str, Any] = { |
| "task_type": task_type.CAUSAL_LM, |
| "r": int(lora_cfg["r"]), |
| "lora_alpha": int(lora_cfg["alpha"]), |
| "lora_dropout": float(lora_cfg.get("dropout", 0.0)), |
| "bias": "none", |
| } |
| target_modules = lora_cfg.get("target_modules", "all-linear") |
| payload["target_modules"] = target_modules |
|
|
| if lora_cfg.get("use_rslora") is not None: |
| payload["use_rslora"] = bool(lora_cfg["use_rslora"]) |
|
|
| signature = inspect.signature(lora_cls) |
| if "exclude_modules" in signature.parameters and lora_cfg.get("exclude_modules"): |
| payload["exclude_modules"] = lora_cfg["exclude_modules"] |
|
|
| accepted = {key: value for key, value in payload.items() if key in signature.parameters} |
| return lora_cls(**accepted) |
|
|
|
|
| def read_jsonl(path: str | Path) -> list[dict[str, Any]]: |
| rows: list[dict[str, Any]] = [] |
| with Path(path).open("r", encoding="utf-8") as fh: |
| for line_no, line in enumerate(fh, start=1): |
| line = line.strip() |
| if not line: |
| continue |
| try: |
| payload = json.loads(line) |
| except json.JSONDecodeError as exc: |
| raise ValueError(f"Invalid JSON in {path}:{line_no}: {exc}") from exc |
| rows.append(payload) |
| return rows |
|
|
|
|
| def validate_think_blocks(row: dict[str, Any], require: bool) -> None: |
| if not require or "messages" not in row: |
| return |
| for message in row["messages"]: |
| if message.get("role") == "assistant": |
| content = str(message.get("content", "")) |
| if "<think>" not in content or "</think>" not in content: |
| row_id = row.get("id", "<unknown>") |
| raise ValueError(f"Assistant message missing <think> block in row {row_id}") |
|
|
|
|
| def assistant_char_mask(rendered: str) -> list[bool]: |
| mask = [False] * len(rendered) |
| cursor = 0 |
| while True: |
| start = rendered.find(ASSISTANT_START, cursor) |
| if start == -1: |
| break |
| content_start = start + len(ASSISTANT_START) |
| end = rendered.find(IM_END, content_start) |
| if end == -1: |
| end = len(rendered) |
| for idx in range(content_start, end): |
| mask[idx] = True |
| cursor = end + len(IM_END) |
| return mask |
|
|
|
|
| def tokenize_row(tokenizer, row: dict[str, Any], max_seq_length: int, require_think: bool) -> TokenizedExample: |
| validate_think_blocks(row, require_think) |
|
|
| if "messages" in row: |
| try: |
| rendered = tokenizer.apply_chat_template( |
| row["messages"], |
| tokenize=False, |
| add_generation_prompt=False, |
| preserve_thinking=True, |
| ) |
| except TypeError: |
| rendered = tokenizer.apply_chat_template( |
| row["messages"], |
| tokenize=False, |
| add_generation_prompt=False, |
| ) |
| mask = assistant_char_mask(rendered) |
| elif "text" in row: |
| rendered = str(row["text"]) |
| mask = [True] * len(rendered) |
| else: |
| raise ValueError("Each row needs either messages or text") |
|
|
| encoded = tokenizer( |
| rendered, |
| add_special_tokens=False, |
| truncation=True, |
| max_length=max_seq_length, |
| return_offsets_mapping=True, |
| ) |
|
|
| labels: list[int] = [] |
| for token_id, (start, end) in zip(encoded["input_ids"], encoded["offset_mapping"], strict=True): |
| if end <= start: |
| labels.append(-100) |
| continue |
| supervised = any(mask[idx] for idx in range(start, min(end, len(mask)))) |
| labels.append(token_id if supervised else -100) |
|
|
| return TokenizedExample( |
| input_ids=list(encoded["input_ids"]), |
| attention_mask=[1] * len(encoded["input_ids"]), |
| labels=labels, |
| ) |
|
|
|
|
| def pack_examples(examples: list[TokenizedExample], max_seq_length: int) -> list[TokenizedExample]: |
| packed: list[TokenizedExample] = [] |
| cur_ids: list[int] = [] |
| cur_labels: list[int] = [] |
|
|
| def flush() -> None: |
| nonlocal cur_ids, cur_labels |
| if cur_ids: |
| packed.append(TokenizedExample(cur_ids, [1] * len(cur_ids), cur_labels)) |
| cur_ids = [] |
| cur_labels = [] |
|
|
| for example in examples: |
| ids = example.input_ids |
| labels = example.labels |
| if len(ids) > max_seq_length: |
| ids = ids[:max_seq_length] |
| labels = labels[:max_seq_length] |
| if cur_ids and len(cur_ids) + len(ids) > max_seq_length: |
| flush() |
| if len(ids) == max_seq_length: |
| packed.append(TokenizedExample(ids, [1] * len(ids), labels)) |
| else: |
| cur_ids.extend(ids) |
| cur_labels.extend(labels) |
| flush() |
| return packed |
|
|
|
|
| class CausalCollator: |
| def __init__(self, pad_token_id: int, label_pad_token_id: int = -100): |
| self.pad_token_id = pad_token_id |
| self.label_pad_token_id = label_pad_token_id |
|
|
| def __call__(self, features: list[dict[str, list[int]]]) -> dict[str, torch.Tensor]: |
| max_len = max(len(feature["input_ids"]) for feature in features) |
| input_ids = [] |
| attention_mask = [] |
| labels = [] |
| for feature in features: |
| pad = max_len - len(feature["input_ids"]) |
| input_ids.append(feature["input_ids"] + [self.pad_token_id] * pad) |
| attention_mask.append(feature["attention_mask"] + [0] * pad) |
| labels.append(feature["labels"] + [self.label_pad_token_id] * pad) |
| return { |
| "input_ids": torch.tensor(input_ids, dtype=torch.long), |
| "attention_mask": torch.tensor(attention_mask, dtype=torch.long), |
| "labels": torch.tensor(labels, dtype=torch.long), |
| } |
|
|
|
|
| def make_dataset(dataset_cls, rows: list[dict[str, Any]], tokenizer, data_cfg: dict[str, Any]): |
| max_seq_length = int(data_cfg["max_seq_length"]) |
| require_think = bool(data_cfg.get("require_think_blocks", False)) |
| tokenized = [tokenize_row(tokenizer, row, max_seq_length, require_think) for row in rows] |
| if data_cfg.get("packing", False): |
| tokenized = pack_examples(tokenized, max_seq_length) |
| payload = [ |
| {"input_ids": item.input_ids, "attention_mask": item.attention_mask, "labels": item.labels} |
| for item in tokenized |
| ] |
| return dataset_cls.from_list(payload) |
|
|
|
|
| def training_args_kwargs(training_arguments_cls, run_cfg: dict[str, Any], training_cfg: dict[str, Any]) -> dict[str, Any]: |
| output_dir = run_cfg["output_dir"] |
| payload: dict[str, Any] = { |
| "output_dir": output_dir, |
| "overwrite_output_dir": False, |
| "learning_rate": float(training_cfg["learning_rate"]), |
| "lr_scheduler_type": training_cfg.get("lr_scheduler_type", "cosine"), |
| "warmup_ratio": float(training_cfg.get("warmup_ratio", 0.03)), |
| "num_train_epochs": float(training_cfg["num_train_epochs"]), |
| "per_device_train_batch_size": int(training_cfg["per_device_train_batch_size"]), |
| "per_device_eval_batch_size": int(training_cfg.get("per_device_eval_batch_size", 1)), |
| "gradient_accumulation_steps": int(training_cfg.get("gradient_accumulation_steps", 1)), |
| "gradient_checkpointing": bool(training_cfg.get("gradient_checkpointing", True)), |
| "max_grad_norm": float(training_cfg.get("max_grad_norm", 1.0)), |
| "logging_steps": int(training_cfg.get("logging_steps", 10)), |
| "save_strategy": training_cfg.get("save_strategy", "steps"), |
| "bf16": bool(training_cfg.get("bf16", True)), |
| "tf32": bool(training_cfg.get("tf32", True)), |
| "report_to": ["wandb"] if os.getenv("WANDB_API_KEY") else [], |
| "remove_unused_columns": False, |
| "seed": int(run_cfg.get("seed", 1337)), |
| } |
|
|
| if "max_steps" in training_cfg: |
| payload["max_steps"] = int(training_cfg["max_steps"]) |
| if "save_steps" in training_cfg: |
| payload["save_steps"] = int(training_cfg["save_steps"]) |
| if "eval_steps" in training_cfg: |
| payload["eval_steps"] = int(training_cfg["eval_steps"]) |
| if "eval_strategy" in training_cfg: |
| payload["eval_strategy"] = training_cfg["eval_strategy"] |
| elif "evaluation_strategy" in training_cfg: |
| payload["evaluation_strategy"] = training_cfg["evaluation_strategy"] |
| elif "eval_steps" in training_cfg: |
| payload["eval_strategy"] = "steps" |
| else: |
| payload["eval_strategy"] = "epoch" |
|
|
| signature = inspect.signature(training_arguments_cls) |
| return {key: value for key, value in payload.items() if key in signature.parameters} |
|
|
|
|
| def build_callbacks(run_cfg: dict[str, Any], config: dict[str, Any]): |
| """Metrics logger (always) + optional in-training held-out benchmark.""" |
| from transformers import TrainerCallback |
|
|
| output_dir = Path(run_cfg["output_dir"]) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| metrics_path = output_dir / "metrics.jsonl" |
| progress_path = output_dir / "eval_progress.jsonl" |
|
|
| class JsonlMetricsCallback(TrainerCallback): |
| """Append every Trainer log line to metrics.jsonl for the watcher.""" |
|
|
| def on_log(self, args, state, control, logs=None, **kwargs): |
| if not logs or not state.is_world_process_zero: |
| return |
| row = {k: v for k, v in logs.items() if isinstance(v, (int, float))} |
| row.update({"step": state.global_step, "epoch": state.epoch, "ts": time.time()}) |
| with metrics_path.open("a", encoding="utf-8") as fh: |
| fh.write(json.dumps(row) + "\n") |
|
|
| callbacks = [JsonlMetricsCallback()] |
|
|
| eval_cfg = config.get("in_training_eval") or {} |
| if not eval_cfg.get("enabled"): |
| return callbacks |
|
|
| eval_files = eval_cfg.get("eval_files", []) |
| sample = int(eval_cfg.get("sample_per_set", 60)) |
| max_new = int(eval_cfg.get("max_new_tokens", 256)) |
| enable_thinking = bool(eval_cfg.get("enable_thinking", False)) |
| eval_every_steps = int(eval_cfg.get("eval_every_steps", 0)) |
| base_acc: dict[str, float] = {} |
| base_json = eval_cfg.get("base_eval_json") |
| if base_json and Path(base_json).is_file(): |
| try: |
| payload = json.loads(Path(base_json).read_text(encoding="utf-8")) |
| base_acc = {r["kind"]: r["accuracy"] for r in payload.get("results", []) if "kind" in r} |
| except Exception: |
| base_acc = {} |
|
|
| class PeriodicEvalCallback(TrainerCallback): |
| """Run the held-out benchmark on the training model at a step interval + each epoch.""" |
|
|
| def _run(self, state, kwargs): |
| if not state.is_world_process_zero: |
| return |
| model = kwargs.get("model") |
| tokenizer = kwargs.get("processing_class") or kwargs.get("tokenizer") |
| if model is None or tokenizer is None: |
| return |
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| try: |
| from intraining_eval import run_eval_sets |
|
|
| sets = run_eval_sets(model, tokenizer, eval_files, sample, max_new, enable_thinking) |
| deltas = {} |
| for metrics in sets.values(): |
| kind = metrics.get("kind") |
| if kind in base_acc and "accuracy" in metrics: |
| deltas[kind] = round(metrics["accuracy"] - base_acc[kind], 4) |
| row = { |
| "epoch": state.epoch, "step": state.global_step, "ts": time.time(), |
| "sets": sets, "base": base_acc, "deltas_vs_base": deltas, |
| } |
| except Exception as exc: |
| row = {"epoch": state.epoch, "step": state.global_step, "ts": time.time(), |
| "error": repr(exc)} |
| with progress_path.open("a", encoding="utf-8") as fh: |
| fh.write(json.dumps(row) + "\n") |
| print(f"[in-training-eval] step={state.global_step} epoch={state.epoch}: " |
| f"{row.get('deltas_vs_base', row.get('error'))}") |
|
|
| def on_step_end(self, args, state, control, **kwargs): |
| if eval_every_steps and state.global_step > 0 and state.global_step % eval_every_steps == 0: |
| self._run(state, kwargs) |
|
|
| def on_epoch_end(self, args, state, control, **kwargs): |
| self._run(state, kwargs) |
|
|
| callbacks.append(PeriodicEvalCallback()) |
| return callbacks |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| run_gate_check(args.config, args.allow_missing_cybergym_baseline) |
|
|
| config = read_yaml(args.config) |
| deps = import_training_deps() |
|
|
| run_cfg = config["run"] |
| model_cfg = config["model"] |
| data_cfg = config["data"] |
| training_cfg = config["training"] |
|
|
| tokenizer = deps["AutoTokenizer"].from_pretrained( |
| model_cfg["name_or_path"], |
| trust_remote_code=bool(model_cfg.get("trust_remote_code", True)), |
| ) |
| if tokenizer.pad_token_id is None: |
| tokenizer.pad_token = tokenizer.eos_token |
|
|
| train_rows = read_jsonl(data_cfg["train_jsonl"]) |
| val_rows = read_jsonl(data_cfg["validation_jsonl"]) |
| train_ds = make_dataset(deps["Dataset"], train_rows, tokenizer, data_cfg) |
| eval_ds = make_dataset(deps["Dataset"], val_rows, tokenizer, data_cfg) |
|
|
| if args.dry_run: |
| print(f"Dry run ok: train={len(train_ds)} eval={len(eval_ds)}") |
| return 0 |
|
|
| quantization_cfg = training_cfg.get("quantization") if training_cfg.get("method") == "qlora" else None |
| model = load_model(deps["transformers"], model_cfg, quantization_cfg) |
|
|
| freeze_patterns: list[str] = [] |
| if model_cfg.get("freeze_vision_tower", True): |
| freeze_patterns.extend(["visual", "vision_tower", "multi_modal_projector"]) |
| if model_cfg.get("freeze_mtp_head", True): |
| freeze_patterns.extend(["mtp"]) |
| frozen_params = freeze_by_name(model, freeze_patterns) |
| print(f"Frozen parameter elements by name pattern: {frozen_params}") |
|
|
| lora_config = build_lora_config(deps["LoraConfig"], deps["TaskType"], training_cfg["lora"]) |
| model = deps["get_peft_model"](model, lora_config) |
| model.print_trainable_parameters() |
|
|
| if training_cfg.get("gradient_checkpointing", True): |
| model.config.use_cache = False |
|
|
| training_args = deps["TrainingArguments"]( |
| **training_args_kwargs(deps["TrainingArguments"], run_cfg, training_cfg) |
| ) |
| trainer_kwargs = { |
| "model": model, |
| "args": training_args, |
| "train_dataset": train_ds, |
| "eval_dataset": eval_ds, |
| "data_collator": CausalCollator(tokenizer.pad_token_id), |
| } |
| trainer_signature = inspect.signature(deps["Trainer"]) |
| if "processing_class" in trainer_signature.parameters: |
| trainer_kwargs["processing_class"] = tokenizer |
| elif "tokenizer" in trainer_signature.parameters: |
| trainer_kwargs["tokenizer"] = tokenizer |
| trainer_kwargs["callbacks"] = build_callbacks(run_cfg, config) |
| trainer = deps["Trainer"](**trainer_kwargs) |
| trainer.train() |
| trainer.save_model(run_cfg["output_dir"]) |
| tokenizer.save_pretrained(run_cfg["output_dir"]) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
