#!/usr/bin/env python3 """Second-stage continuation training for an existing QLoRA adapter. Loads the base model in 4-bit, prepares it for k-bit training, loads an existing LoRA adapter with is_trainable=True, and continues SFT on a local weak-layer dataset. """ import argparse import math import os import re from pathlib import Path import torch from datasets import load_dataset from peft import PeftConfig, PeftModel, get_model_status, prepare_model_for_kbit_training from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainerCallback, set_seed from trl import SFTConfig, SFTTrainer from tmf921_train.utils import load_config, write_json try: import trackio except Exception: trackio = None class TrackioAlertCallback(TrainerCallback): def on_log(self, args, state, control, logs=None, **kwargs): if not state.is_world_process_zero or not logs or trackio is None: return loss = logs.get("loss") if loss is not None and (math.isnan(float(loss)) or math.isinf(float(loss))): trackio.alert(title="NaN/Inf stage2 loss", text=f"step={state.global_step} loss={loss} — lower LR", level="ERROR") def require_cuda(): print("=== CUDA CHECK ===") print(f"torch={torch.__version__} torch.version.cuda={torch.version.cuda} CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')}") if not torch.cuda.is_available(): raise RuntimeError("CUDA unavailable. Refusing CPU training.") print(f"cuda device_count={torch.cuda.device_count()} gpu0={torch.cuda.get_device_name(0)}") def valid_hf_repo_id(repo_id): if not repo_id or not isinstance(repo_id, str): return False if repo_id.endswith("/") or repo_id.startswith("/") or "//" in repo_id: return False return re.match(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,95}/[A-Za-z0-9][A-Za-z0-9._-]{0,95}$", repo_id) is not None def sanitize_trackio_config(cfg): env_space = os.environ.get("TRACKIO_SPACE_ID", "").strip() cfg_space = str(cfg.get("trackio_space_id") or "").strip() chosen = env_space or cfg_space if chosen and valid_hf_repo_id(chosen): cfg["trackio_space_id"] = chosen else: if chosen: print(f"WARNING: ignoring invalid Trackio Space ID: {chosen!r}") cfg["trackio_space_id"] = None os.environ.pop("TRACKIO_SPACE_ID", None) if os.environ.get("DISABLE_TRACKIO", "0") == "1": print("Trackio disabled via DISABLE_TRACKIO=1") cfg["project"] = None cfg["trackio_space_id"] = None return cfg def parse_args(): p = argparse.ArgumentParser() p.add_argument("--config", default="configs/stage2_weak_layer_qwen3_8b.yaml") p.add_argument("--adapter_path", help="Existing adapter path/repo to continue training") p.add_argument("--dataset_dir", help="Local dir containing train.parquet and validation.parquet") p.add_argument("--output_dir") p.add_argument("--hub_model_id") p.add_argument("--max_steps", type=int, default=None) p.add_argument("--no_push", action="store_true") p.add_argument("--seed", type=int, default=43) p.add_argument("--resume_from_checkpoint", default=None) return p.parse_args() def main(): args = parse_args() require_cuda() cfg = sanitize_trackio_config(load_config(args.config)) for k in ["adapter_path", "dataset_dir", "output_dir", "hub_model_id"]: v = getattr(args, k) if v is not None: cfg[k] = v if args.max_steps is not None: cfg["max_steps"] = args.max_steps if args.no_push: cfg["push_to_hub"] = False set_seed(args.seed) Path(cfg["output_dir"]).mkdir(parents=True, exist_ok=True) write_json(Path(cfg["output_dir"]) / "resolved_config.json", cfg) print("Loading local stage2 dataset", cfg["dataset_dir"]) data_files = { "train": str(Path(cfg["dataset_dir"]) / "train.parquet"), "validation": str(Path(cfg["dataset_dir"]) / "validation.parquet"), } ds = load_dataset("parquet", data_files=data_files) train_dataset = ds["train"].select_columns(["messages"]) eval_dataset = ds["validation"].select_columns(["messages"]) print(train_dataset) print(eval_dataset) peft_cfg = PeftConfig.from_pretrained(cfg["adapter_path"]) base_model_id = cfg.get("model_name_or_path") or peft_cfg.base_model_name_or_path or "Qwen/Qwen3-8B" print("Base model:", base_model_id) print("Adapter:", cfg["adapter_path"]) tokenizer = AutoTokenizer.from_pretrained(cfg["adapter_path"], trust_remote_code=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type=cfg.get("bnb_4bit_quant_type", "nf4"), bnb_4bit_use_double_quant=bool(cfg.get("bnb_4bit_use_double_quant", True)), bnb_4bit_compute_dtype=torch.bfloat16, ) base_model = AutoModelForCausalLM.from_pretrained( base_model_id, quantization_config=bnb_config, device_map={"": 0}, dtype=torch.bfloat16, trust_remote_code=True, ) base_model.config.use_cache = False base_model = prepare_model_for_kbit_training(base_model, use_gradient_checkpointing=bool(cfg.get("gradient_checkpointing", True))) model = PeftModel.from_pretrained(base_model, cfg["adapter_path"], is_trainable=True) model.print_trainable_parameters() status = get_model_status(model) print(status) if status.trainable_params <= 0: raise RuntimeError("No trainable adapter parameters found; refusing to run stage2.") report_to = "trackio" if cfg.get("project") else "none" sft_args = SFTConfig( output_dir=cfg["output_dir"], max_length=int(cfg.get("max_length", 2048)), packing=bool(cfg.get("packing", False)), assistant_only_loss=bool(cfg.get("assistant_only_loss", True)), dataset_num_proc=int(cfg.get("dataset_num_proc", 8)), learning_rate=float(cfg.get("learning_rate", 5e-5)), lr_scheduler_type=cfg.get("lr_scheduler_type", "constant"), warmup_steps=int(cfg.get("warmup_steps", 0)), weight_decay=float(cfg.get("weight_decay", 0.0)), max_grad_norm=float(cfg.get("max_grad_norm", 0.3)), num_train_epochs=float(cfg.get("epochs", 1)), max_steps=int(cfg["max_steps"]) if cfg.get("max_steps") is not None else -1, per_device_train_batch_size=int(cfg.get("per_device_train_batch_size", 1)), gradient_accumulation_steps=int(cfg.get("gradient_accumulation_steps", 16)), per_device_eval_batch_size=int(cfg.get("per_device_eval_batch_size", 1)), bf16=True, gradient_checkpointing=bool(cfg.get("gradient_checkpointing", True)), gradient_checkpointing_kwargs={"use_reentrant": False}, optim=cfg.get("optim", "paged_adamw_32bit"), eval_strategy="steps", eval_steps=int(cfg.get("eval_steps", 100)), save_strategy="steps", save_steps=int(cfg.get("save_steps", 100)), save_total_limit=int(cfg.get("save_total_limit", 3)), load_best_model_at_end=True, metric_for_best_model="eval_loss", greater_is_better=False, logging_strategy="steps", logging_steps=int(cfg.get("logging_steps", 10)), logging_first_step=True, disable_tqdm=True, report_to=report_to, run_name=cfg.get("run_name"), project=cfg.get("project"), trackio_space_id=cfg.get("trackio_space_id"), push_to_hub=bool(cfg.get("push_to_hub", True)), hub_model_id=cfg.get("hub_model_id"), ) trainer = SFTTrainer( model=model, args=sft_args, train_dataset=train_dataset, eval_dataset=eval_dataset, processing_class=tokenizer, callbacks=[TrackioAlertCallback()], ) resume_arg = args.resume_from_checkpoint if resume_arg is not None and str(resume_arg).lower() == "true": resume_arg = True trainer.train(resume_from_checkpoint=resume_arg) metrics = trainer.evaluate() write_json(Path(cfg["output_dir"]) / "final_eval_metrics.json", metrics) trainer.save_model(cfg["output_dir"]) tokenizer.save_pretrained(cfg["output_dir"]) if bool(cfg.get("push_to_hub", True)): # dataset_name must be a valid HF dataset id for model-card metadata validation. trainer.push_to_hub( commit_message="Stage2 weak-layer QLoRA continuation", dataset_name="nraptisss/TMF921-intent-to-config-research-sota", ) print(f"Pushed stage2 adapter to https://huggingface.co/{cfg.get('hub_model_id')}") if __name__ == "__main__": main()