efficiency_figure2/testtrain2.py

#!/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/<EXPERIMENT_NAME>/sweep_results.jsonl
#   TensorBoard:   tensorboard/<EXPERIMENT_NAME>/...
#   Final model:   ./qwen_regression_ckpt/<EXPERIMENT_NAME>/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))