#!/usr/bin/env python3 # Drop-in replacement: # - NO freezing hacks # - NO PEFT/LoRA # - Full-parameter regression fine-tuning from a Qwen LM checkpoint # - Tiny dataset -> regularization sweep (dropout/weight_decay/etc.) # - Labels are used AS-IS (no normalization) # # Saves: # Sweep logs: ./qwen_regression_ckpt//sweep_results.jsonl # TensorBoard: tensorboard//... # Final model: ./qwen_regression_ckpt//final/model # # Run: # python train_regression_sweep.py import os import json import random import gc import inspect from pathlib import Path from typing import Dict, Any, List, Optional import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from datasets import Dataset from sklearn.metrics import mean_squared_error, r2_score from scipy.stats import pearsonr from transformers import ( AutoTokenizer, AutoConfig, AutoModel, AutoModelForCausalLM, TrainingArguments, Trainer, TrainerCallback, set_seed, EvalPrediction, ) from transformers.data.data_collator import DataCollatorWithPadding from transformers.modeling_outputs import SequenceClassifierOutput from transformers import PreTrainedModel # ========================= # 0) PATHS (edit if needed) # ========================= BASE_MODEL_PATH = "./checkpoint-388560" TOKENIZER_PATH = "/opt/platform/regression_efficiency/checkpoint-5956" TRAIN_JSON = "evenBetterDataFolded-tr.json" VALID_JSON = "evenBetterDataFolded-vl.json" # ========================= # 1) EXPERIMENT SETTINGS # ========================= EXPERIMENT_NAME = "fullft_regression_sweep_stanardunnorm" ROOT_OUT = Path(f"./qwen_regression_ckpt/{EXPERIMENT_NAME}") TB_ROOT = Path(f"tensorboard/{EXPERIMENT_NAME}") ROOT_OUT.mkdir(parents=True, exist_ok=True) TB_ROOT.mkdir(parents=True, exist_ok=True) PER_DEVICE_TRAIN_BATCH = 24 PER_DEVICE_EVAL_BATCH = 2 GRAD_ACCUM_STEPS = 1 NUM_EPOCHS_CAP = 200 # early stop will cut it EARLY_STOP_PATIENCE = 50 EARLY_STOP_MIN_DELTA = 0.0 # What we optimize (recommended for tiny regression) BEST_METRIC_KEY = "eval_pearson_r" # must match compute_metrics output with eval_ prefix GREATER_IS_BETTER = True # Sweep control SWEEP_MODE = "random" # "random" or "grid" NUM_TRIALS = 12 # if random SEEDS = [42] # add more seeds if you want robustness # Precision / perf #torch.backends.cuda.matmul.allow_tf32 = True USE_BF16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported() DEFAULT_OPTIM = "adamw_torch_fused" # ========================= # 2) HF ARG COMPATIBILITY # ========================= _TA_SIG = inspect.signature(TrainingArguments.__init__) _HAS_EVAL_STRATEGY = "eval_strategy" in _TA_SIG.parameters _HAS_EVALUATION_STRATEGY = "evaluation_strategy" in _TA_SIG.parameters def make_training_args(**kwargs): """ Make TrainingArguments across HF versions: - some versions use eval_strategy, others evaluation_strategy """ # Caller should pass eval_strategy="epoch" (preferred); we map if needed. if "eval_strategy" in kwargs and not _HAS_EVAL_STRATEGY and _HAS_EVALUATION_STRATEGY: kwargs["evaluation_strategy"] = kwargs.pop("eval_strategy") if "evaluation_strategy" in kwargs and _HAS_EVAL_STRATEGY and not _HAS_EVALUATION_STRATEGY: kwargs["eval_strategy"] = kwargs.pop("evaluation_strategy") return TrainingArguments(**kwargs) # ========================= # 3) DATA LOADING (labels unchanged) # ========================= def load_json_mapping(path: str) -> Dict[str, float]: with open(path, "r", encoding="utf-8") as f: d = json.load(f) return {k: float(v) for k, v in d.items()} train_map = load_json_mapping(TRAIN_JSON) valid_map = load_json_mapping(VALID_JSON) train_texts = list(train_map.keys()) train_labels = [train_map[k] for k in train_texts] valid_texts = list(valid_map.keys()) valid_labels = [valid_map[k] for k in valid_texts] # Save raw maps for inspection (like you were doing) with open(TB_ROOT / "train.json", "w", encoding="utf-8") as f: json.dump(train_map, f, indent=2) with open(TB_ROOT / "valid.json", "w", encoding="utf-8") as f: json.dump(valid_map, f, indent=2) print("Train size:", len(train_texts), " Valid size:", len(valid_texts)) print("Train labels stats:", f"mean={np.mean(train_labels):.6f}", f"std={np.std(train_labels):.6f}", f"min={np.min(train_labels):.6f}", f"max={np.max(train_labels):.6f}") baseline_mse = float(np.mean((np.array(train_labels) - np.mean(train_labels))**2)) print("Baseline MSE (predict train mean):", baseline_mse) tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH, use_fast=True) # Ensure pad token exists for dynamic padding if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token train_raw = Dataset.from_dict({"text": train_texts, "labels": train_labels}) valid_raw = Dataset.from_dict({"text": valid_texts, "labels": valid_labels}) def tok(batch): return tokenizer(batch["text"], truncation=True) train_ds = train_raw.map(tok, batched=True, remove_columns=["text"]) valid_ds = valid_raw.map(tok, batched=True, remove_columns=["text"]) train_ds.set_format(type="torch") valid_ds.set_format(type="torch") data_collator = DataCollatorWithPadding(tokenizer=tokenizer, pad_to_multiple_of=8, return_tensors="pt") # ========================= # 4) MODEL (PreTrainedModel for proper saving) # ========================= class LastTokenPooling(nn.Module): def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None): if attention_mask is None: return hidden_states[:, -1, :] B, T, H = hidden_states.size() # left-padding: last token always real if attention_mask[:, -1].sum().item() == B: return hidden_states[:, -1, :] # right-padding: gather last non-pad idx = (attention_mask.sum(dim=1).long() - 1).clamp(min=0) idx = idx.view(B, 1, 1).expand(-1, 1, H) return hidden_states.gather(1, idx).squeeze(1) class RegressionHead(nn.Module): def __init__(self, hidden_size: int, head_dropout: float): super().__init__() self.ln = nn.LayerNorm(hidden_size) self.drop = nn.Dropout(head_dropout) self.out = nn.Linear(hidden_size, 1) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.ln(x) x = self.drop(x) return self.out(x).squeeze(-1) def robust_set_dropout(config, p_hidden: float, p_attn: float, layerdrop: float): # set whichever exist; ignore missing fields hidden_fields = [ "hidden_dropout_prob", "hidden_dropout", "dropout", "emb_dropout", "resid_pdrop", "classifier_dropout", ] attn_fields = [ "attention_probs_dropout_prob", "attention_dropout", "attn_dropout", "attn_pdrop", ] for f in hidden_fields: if hasattr(config, f): setattr(config, f, float(p_hidden)) for f in attn_fields: if hasattr(config, f): setattr(config, f, float(p_attn)) if hasattr(config, "layerdrop"): setattr(config, "layerdrop", float(layerdrop)) def load_backbone(base_model_path: str, config: AutoConfig): dtype = torch.bfloat16 if USE_BF16 else (torch.float16 if USE_FP16 else None) try: return AutoModel.from_pretrained( base_model_path, config=config, torch_dtype=dtype, device_map=None, ) except Exception as e: print("[warn] AutoModel load failed, falling back to AutoModelForCausalLM().model") lm = AutoModelForCausalLM.from_pretrained( base_model_path, config=config, torch_dtype=dtype, device_map=None, ) if hasattr(lm, "model"): return lm.model if hasattr(lm, "transformer"): return lm.transformer raise RuntimeError("Could not locate backbone module on LM model.") from e class QwenForRegression(PreTrainedModel): config_class = AutoConfig base_model_prefix = "backbone" def __init__(self, config: AutoConfig, base_model_path: str, head_dropout: float): super().__init__(config) self.backbone = load_backbone(base_model_path, config) self.pool = LastTokenPooling() self.regression_head = RegressionHead(config.hidden_size, head_dropout=head_dropout) # Full-param training ON for p in self.parameters(): p.requires_grad = True def gradient_checkpointing_enable(self, **kwargs): if hasattr(self.backbone, "gradient_checkpointing_enable"): self.backbone.gradient_checkpointing_enable(**kwargs) def gradient_checkpointing_disable(self, **kwargs): if hasattr(self.backbone, "gradient_checkpointing_disable"): self.backbone.gradient_checkpointing_disable(**kwargs) def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs): out = self.backbone(input_ids=input_ids, attention_mask=attention_mask, return_dict=True) pooled = self.pool(out.last_hidden_state, attention_mask) logits = self.regression_head(pooled) # [B] loss = None if labels is not None: loss = F.mse_loss(logits.float(), labels.float()) return SequenceClassifierOutput(loss=loss, logits=logits) def save_pretrained(self, save_directory: str, state_dict=None, **kwargs): """ Save config + backbone (HF-style) + regression head separately. Works with Trainer saving. """ os.makedirs(save_directory, exist_ok=True) self.config.save_pretrained(save_directory) # If state_dict provided (e.g., FSDP), split it; else use live weights. if state_dict is None: state_dict = self.state_dict() backbone_sd = {k[len("backbone."):]: v for k, v in state_dict.items() if k.startswith("backbone.")} head_sd = {k[len("regression_head."):]: v for k, v in state_dict.items() if k.startswith("regression_head.")} # Save backbone in HF format self.backbone.save_pretrained(save_directory, state_dict=backbone_sd, **kwargs) # Save head as a torch file torch.save(head_sd, os.path.join(save_directory, "regression_head.pt")) # ========================= # 5) METRICS # ========================= def safe_pearson(x: np.ndarray, y: np.ndarray) -> float: x = np.asarray(x).reshape(-1) y = np.asarray(y).reshape(-1) if x.size < 2: return 0.0 if np.std(x) < 1e-12 or np.std(y) < 1e-12: return 0.0 r, _ = pearsonr(x, y) return float(r) def compute_metrics(eval_pred: EvalPrediction) -> Dict[str, float]: preds = np.asarray(eval_pred.predictions).reshape(-1) labels = np.asarray(eval_pred.label_ids).reshape(-1) mse = mean_squared_error(labels, preds) r2 = r2_score(labels, preds) pr = safe_pearson(preds, labels) return {"mse": float(mse), "r2": float(r2), "pearson_r": float(pr)} # ========================= # 6) NO-SAVE EARLY STOP CALLBACK # ========================= class EarlyStopNoSaveCallback(TrainerCallback): """ Early stopping WITHOUT requiring: - metric_for_best_model - load_best_model_at_end - checkpoint saving It just stops training when eval metric stops improving. """ def __init__(self, metric_key: str, patience: int, greater_is_better: bool, min_delta: float = 0.0): self.metric_key = metric_key self.patience = int(patience) self.greater_is_better = bool(greater_is_better) self.min_delta = float(min_delta) self.best = None self.bad_count = 0 def on_evaluate(self, args, state, control, metrics=None, **kwargs): if metrics is None: return control if self.metric_key not in metrics: # If metric missing, do nothing return control val = float(metrics[self.metric_key]) if self.best is None: self.best = val self.bad_count = 0 return control improved = (val - self.best) > self.min_delta if self.greater_is_better else (self.best - val) > self.min_delta if improved: self.best = val self.bad_count = 0 else: self.bad_count += 1 if self.bad_count >= self.patience: control.should_training_stop = True return control # ========================= # 7) SWEEP PARAMS # ========================= def sample_hparams(rng: random.Random) -> Dict[str, Any]: # Reasonable priors for huge model + tiny data return { "learning_rate": rng.choice([1e-5, 2e-5, 3e-5, 5e-5, 8e-5]), "weight_decay": rng.choice([0.0, 0.01, 0.03, 0.05, 0.1]), "hidden_dropout": rng.choice([0.0, 0.1, 0.2, 0.3]), "attn_dropout": rng.choice([0.0, 0.1, 0.2, 0.3]), "layerdrop": rng.choice([0.0, 0.05, 0.1, 0.2]), "head_dropout": rng.choice([0.0, 0.1, 0.2, 0.3, 0.5]), "max_grad_norm": rng.choice([0.5, 1.0, 2.0]), "warmup_ratio": rng.choice([0.0, 0.03, 0.05, 0.1]), "lr_scheduler_type": rng.choice(["cosine", "cosine_with_restarts"]), "num_cycles": rng.choice([1, 2, 4]), # only for cosine_with_restarts } def grid_hparams() -> List[Dict[str, Any]]: grid = [] for lr in [2e-5, 3e-5, 5e-5]: for wd in [0.0, 0.03, 0.1]: for dp in [0.1, 0.2, 0.3]: for head_dp in [0.1, 0.3]: grid.append({ "learning_rate": lr, "weight_decay": wd, "hidden_dropout": dp, "attn_dropout": dp, "layerdrop": 0.0, "head_dropout": head_dp, "max_grad_norm": 1.0, "warmup_ratio": 0.05, "lr_scheduler_type": "cosine", "num_cycles": 1, }) return grid HP_LIST = grid_hparams() if SWEEP_MODE == "grid" else [sample_hparams(random.Random(1234 + i)) for i in range(NUM_TRIALS)] # ========================= # 8) HELPERS # ========================= def cleanup(): gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # ========================= # 9) RUN ONE TRIAL (no checkpoint saves) # ========================= def run_trial(trial_dir: Path, tb_dir: Path, h: Dict[str, Any], seed: int) -> Dict[str, Any]: set_seed(seed) config = AutoConfig.from_pretrained(BASE_MODEL_PATH) robust_set_dropout(config, h["hidden_dropout"], h["attn_dropout"], h["layerdrop"]) model = QwenForRegression(config=config, base_model_path=BASE_MODEL_PATH, head_dropout=h["head_dropout"]) args = make_training_args( output_dir=str(trial_dir), per_device_train_batch_size=PER_DEVICE_TRAIN_BATCH, per_device_eval_batch_size=PER_DEVICE_EVAL_BATCH, gradient_accumulation_steps=GRAD_ACCUM_STEPS, num_train_epochs=NUM_EPOCHS_CAP, eval_strategy="epoch", save_strategy="no", # IMPORTANT: no huge checkpoints during sweep logging_dir=str(tb_dir), logging_steps=1, report_to=["tensorboard"], learning_rate=h["learning_rate"], weight_decay=h["weight_decay"], warmup_ratio=h["warmup_ratio"], lr_scheduler_type=h["lr_scheduler_type"], lr_scheduler_kwargs={"num_cycles": h["num_cycles"]} if h["lr_scheduler_type"] == "cosine_with_restarts" else {}, bf16=USE_BF16, fp16=USE_FP16, gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False}, max_grad_norm=h["max_grad_norm"], dataloader_pin_memory=True, remove_unused_columns=False, optim=DEFAULT_OPTIM, seed=seed, data_seed=seed, # NOTE: we do NOT set load_best_model_at_end in sweep (no saves). ) trainer = Trainer( model=model, args=args, train_dataset=train_ds, eval_dataset=valid_ds, tokenizer=tokenizer, data_collator=data_collator, compute_metrics=compute_metrics, callbacks=[EarlyStopNoSaveCallback( metric_key=BEST_METRIC_KEY, patience=EARLY_STOP_PATIENCE, greater_is_better=GREATER_IS_BETTER, min_delta=EARLY_STOP_MIN_DELTA, )], ) train_result = trainer.train() eval_metrics = trainer.evaluate() out = { "seed": int(seed), "hparams": dict(h), "train_metrics": {k: float(v) for k, v in train_result.metrics.items()}, "eval_metrics": {k: float(v) for k, v in eval_metrics.items()}, } del trainer, model cleanup() return out # ========================= # 10) SWEEP LOOP # ========================= results_file = ROOT_OUT / "sweep_results.jsonl" best = None # {"score":..., "trial_id":..., "record":...} trial_id = 0 for h in HP_LIST: for seed in SEEDS: trial_dir = ROOT_OUT / "trials" / f"trial_{trial_id:03d}_seed{seed}" tb_dir = TB_ROOT / "trials" / f"trial_{trial_id:03d}_seed{seed}" trial_dir.mkdir(parents=True, exist_ok=True) tb_dir.mkdir(parents=True, exist_ok=True) print(f"\n=== TRIAL {trial_id:03d} seed={seed} ===") print(json.dumps(h, indent=2)) record = run_trial(trial_dir, tb_dir, h, seed) with open(results_file, "a", encoding="utf-8") as f: f.write(json.dumps({"trial_id": trial_id, **record}) + "\n") # scoring by BEST_METRIC_KEY score = record["eval_metrics"].get(BEST_METRIC_KEY, None) if score is None: # fallback to -eval_loss if needed score = -record["eval_metrics"].get("eval_loss", 1e30) score = float(score) is_better = (best is None) or ((score > best["score"]) if GREATER_IS_BETTER else (score < best["score"])) if is_better: best = {"score": score, "trial_id": trial_id, "record": record} print(f"--> NEW BEST: {BEST_METRIC_KEY} = {score:.6f}") trial_id += 1 if best is None: raise RuntimeError("No trials ran.") print("\n====================") print("BEST TRIAL SUMMARY") print("====================") print(json.dumps( { "trial_id": best["trial_id"], "score": best["score"], "seed": best["record"]["seed"], "hparams": best["record"]["hparams"], "eval_metrics": best["record"]["eval_metrics"], }, indent=2, )) # ========================= # 11) FINAL TRAIN (save best model properly) # ========================= final_dir = ROOT_OUT / "final" / "model" final_tb = TB_ROOT / "final" final_dir.mkdir(parents=True, exist_ok=True) final_tb.mkdir(parents=True, exist_ok=True) best_h = best["record"]["hparams"] best_seed = int(best["record"]["seed"]) set_seed(best_seed) final_config = AutoConfig.from_pretrained(BASE_MODEL_PATH) robust_set_dropout(final_config, best_h["hidden_dropout"], best_h["attn_dropout"], best_h["layerdrop"]) final_model = QwenForRegression(config=final_config, base_model_path=BASE_MODEL_PATH, head_dropout=best_h["head_dropout"]) # For final training we CAN save + load best model. # To satisfy “best model” logic, we set metric_for_best_model. # (Some HF versions want "pearson_r" and will look for "eval_pearson_r"; both are ok. We'll set the non-prefixed name.) metric_for_best = "pearson_r" final_args = make_training_args( output_dir=str(final_dir), per_device_train_batch_size=PER_DEVICE_TRAIN_BATCH, per_device_eval_batch_size=PER_DEVICE_EVAL_BATCH, gradient_accumulation_steps=GRAD_ACCUM_STEPS, num_train_epochs=NUM_EPOCHS_CAP, eval_strategy="epoch", save_strategy="epoch", save_total_limit=2, load_best_model_at_end=True, metric_for_best_model=metric_for_best, greater_is_better=GREATER_IS_BETTER, logging_dir=str(final_tb), logging_steps=1, report_to=["tensorboard"], learning_rate=best_h["learning_rate"], weight_decay=best_h["weight_decay"], warmup_ratio=best_h["warmup_ratio"], lr_scheduler_type=best_h["lr_scheduler_type"], lr_scheduler_kwargs={"num_cycles": best_h["num_cycles"]} if best_h["lr_scheduler_type"] == "cosine_with_restarts" else {}, bf16=USE_BF16, fp16=USE_FP16, gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False}, max_grad_norm=best_h["max_grad_norm"], dataloader_pin_memory=True, remove_unused_columns=False, optim=DEFAULT_OPTIM, seed=best_seed, data_seed=best_seed, ) final_trainer = Trainer( model=final_model, args=final_args, train_dataset=train_ds, eval_dataset=valid_ds, tokenizer=tokenizer, data_collator=data_collator, compute_metrics=compute_metrics, # optional: also stop early in final run (and we DO save, so it's fine either way) callbacks=[EarlyStopNoSaveCallback( metric_key=BEST_METRIC_KEY, patience=EARLY_STOP_PATIENCE, greater_is_better=GREATER_IS_BETTER, min_delta=EARLY_STOP_MIN_DELTA, )], ) final_trainer.train() final_metrics = final_trainer.evaluate() print("\nFINAL EVAL METRICS:") print(json.dumps({k: float(v) for k, v in final_metrics.items()}, indent=2)) # Save tokenizer (handy) tokenizer.save_pretrained(str(final_dir)) # Save best sweep params with open(ROOT_OUT / "best_hparams.json", "w", encoding="utf-8") as f: json.dump(best_h, f, indent=2) print("\nSaved final model to:", str(final_dir)) print("Sweep results jsonl:", str(results_file)) print("Tensorboard root:", str(TB_ROOT))