File size: 8,749 Bytes
60bd01c 0187cea 60bd01c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 | #!/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()
|