# lmr/training/bert_finetune_trainer.py """ BertFineTuneTrainer Features: - Supports classification (single-sentence), sentence-pair (GLUE-style), and next-sentence-prediction (synthesizes pairs). - Loads HF datasets via `datasets.load_dataset`. - Handles tokenization, dataloader creation, optimizer, scheduler, fp16, DDP (via torch.distributed if use_ddp True). - Evaluates accuracy/precision/recall/f1 and reports. - Saves best checkpoint by eval metric (F1 for binary/multi-class; accuracy fallback). """ from pathlib import Path import random import os import time import math import json import shutil from typing import Optional, Dict, Any import numpy as np import torch from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler from torch.optim import AdamW from torch.cuda.amp import GradScaler, autocast from transformers import ( AutoTokenizer, AutoConfig, AutoModelForSequenceClassification, BertForNextSentencePrediction, get_linear_schedule_with_warmup, ) from datasets import load_dataset, Dataset, DatasetDict from tqdm import tqdm from sklearn.metrics import accuracy_score, precision_recall_fscore_support # Basic logger def log(*args, **kwargs): print(time.strftime("%Y-%m-%d %H:%M:%S"), "-", *args, **kwargs) class BertFineTuneTrainer: def __init__(self, cfg: Dict[str, Any], device: Optional[torch.device] = None): """ cfg: dictionary with finetune settings (see DEFAULT_CONFIG in main.py) """ self.cfg = cfg self.device = device or (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")) self.output_dir = Path(cfg.get("output_dir", "./outputs/finetune")) self.output_dir.mkdir(parents=True, exist_ok=True) # random seed seed = cfg.get("seed", 42) random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) # DDP settings self.use_ddp = bool(cfg.get("use_ddp", False)) if self.use_ddp: # expect torchrun to have initialized the process group externally or we do it here if not torch.distributed.is_initialized(): init_method = os.environ.get("INIT_METHOD", "env://") torch.distributed.init_process_group(backend="nccl", init_method=init_method) self.rank = torch.distributed.get_rank() self.world_size = torch.distributed.get_world_size() torch.cuda.set_device(self.rank % torch.cuda.device_count()) self.device = torch.device(f"cuda:{torch.cuda.current_device()}") else: self.rank = 0 self.world_size = 1 # Load tokenizer & model model_name = cfg["model_name_or_path"] task = cfg.get("task", "sentence_pair") num_labels = cfg.get("num_labels", None) print('-----------------------------123123123123123123123') print(model_name,'----------------------------------------------------------') log(f"[rank {self.rank}] Loading tokenizer & model: {model_name}, task={task}") self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True) # ensure mask token exists for BERT-like tokenizers (for MLM/if needed) if getattr(self.tokenizer, "mask_token", None) is None: # add typical BERT mask token if missing self.tokenizer.add_special_tokens({"mask_token": "[MASK]"}) # create model if task == "next_sentence_prediction": self.model = BertForNextSentencePrediction.from_pretrained(model_name) else: # classification or sentence_pair => AutoModelForSequenceClassification # infer num_labels later from dataset if None cfg_model = AutoConfig.from_pretrained(model_name) if num_labels is None and hasattr(cfg_model, "num_labels"): num_labels = getattr(cfg_model, "num_labels", None) if num_labels is None: # default to 2 num_labels = 2 self.model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels) # Move model to device or DDP wrapping at later stage self.model.to(self.device) # fp16 scaler self.fp16 = bool(cfg.get("fp16", True)) self.scaler = GradScaler() if self.fp16 and torch.cuda.is_available() else None # training hyperparams self.batch_size = int(cfg.get("batch_size", 16)) self.eval_batch_size = int(cfg.get("eval_batch_size", max(32, self.batch_size))) self.num_epochs = int(cfg.get("num_epochs", 3)) self.learning_rate = float(cfg.get("lr", 2e-5)) self.weight_decay = float(cfg.get("weight_decay", 0.01)) self.gradient_accumulation_steps = int(cfg.get("gradient_accumulation_steps", 1)) self.max_grad_norm = float(cfg.get("max_grad_norm", 1.0)) self.max_length = int(cfg.get("max_length", 128)) self.num_workers = int(cfg.get("num_workers", 4)) self.logging_steps = int(cfg.get("logging_steps", 100)) self.eval_steps = int(cfg.get("eval_steps", 500)) self.save_steps = int(cfg.get("save_steps", 1000)) self.warmup_steps = int(cfg.get("warmup_steps", 0)) self.max_train_samples = cfg.get("max_train_samples", None) self.max_eval_samples = cfg.get("max_eval_samples", None) self.nsp_negatives_ratio = int(cfg.get("nsp_negatives_ratio", 1)) # dataset params self.dataset_name = cfg.get("dataset") self.dataset_config_name = cfg.get("dataset_config_name", None) # internal state self.best_metric = -1.0 self.global_step = 0 self.total_steps = 0 # ------------------------- # Dataset loading / preprocessing # ------------------------- def _load_hf_dataset(self): # Support forms: # - "glue/mrpc" -> load_dataset("glue", "mrpc") # - "glue", dataset_config_name="mrpc" # - "imdb" -> load_dataset("imdb") ds_name = self.dataset_name cfg_name = self.dataset_config_name if ds_name is None: raise ValueError("Please specify cfg['dataset'] (Hugging Face dataset id).") if "/" in ds_name and not ds_name.startswith("glue/"): # user provided dataset/config like "glue/mrpc" or "squad/v1" parts = ds_name.split("/", 1) ds = load_dataset(parts[0], parts[1]) elif ds_name.startswith("glue/"): # glue/mrpc form parts = ds_name.split("/", 1) ds = load_dataset(parts[0], parts[1]) else: # try load with dataset name and optional config if cfg_name: ds = load_dataset(ds_name, cfg_name) else: ds = load_dataset(ds_name) # Expect ds to be Dataset or DatasetDict if isinstance(ds, Dataset): ds = DatasetDict({"train": ds}) if isinstance(ds, dict) and not isinstance(ds, DatasetDict): ds = DatasetDict(ds) return ds def _prepare_sentence_pair(self, dataset: Dataset, text1_key: str, text2_key: str): # tokenization for sentence pair tokenizer = self.tokenizer max_length = self.max_length def fn_examples(examples): texts1 = examples[text1_key] texts2 = examples[text2_key] # handle lists / strings enc = tokenizer(texts1, texts2, truncation=True, padding="max_length", max_length=max_length) # ensure label present out = {"input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"]} if "label" in examples: out["labels"] = examples["label"] return out return dataset.map(fn_examples, batched=True, remove_columns=[c for c in dataset.column_names if c not in (text1_key, text2_key, "label")], num_proc=1) def _prepare_classification(self, dataset: Dataset, text_key: str): tokenizer = self.tokenizer max_length = self.max_length def fn_examples(examples): texts = examples[text_key] enc = tokenizer(texts, truncation=True, padding="max_length", max_length=max_length) out = {"input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"]} if "label" in examples: out["labels"] = examples["label"] return out return dataset.map(fn_examples, batched=True, remove_columns=[c for c in dataset.column_names if c not in (text_key, "label")], num_proc=1) def _synthesize_nsp_dataset(self, ds: Dataset): """ Build next-sentence pairs from a text dataset: - consecutive sentences -> label=1 (is_next) - random sentence pair -> label=0 This is a simple heuristic synthesizer. """ tokenizer = self.tokenizer max_length = self.max_length neg_ratio = self.nsp_negatives_ratio texts = [] # Collect text lines for ex in ds: # prefer fields 'text' or 'content' if "text" in ex: t = ex["text"] elif "content" in ex: t = ex["content"] else: # if dataset contains list-of-sentences or docstrings, try to convert # fallback to join all string fields t = " ".join(str(v) for k, v in ex.items() if isinstance(v, str)) if not t: continue # split into sentences roughly by punctuation sents = [s.strip() for s in t.replace("\n", " ").split(". ") if s.strip()] for i in range(len(sents)-1): texts.append((sents[i], sents[i+1], 1)) # add negatives by random pairing n_pos = len(texts) if n_pos == 0: raise RuntimeError("No sentence pairs extracted for NSP. Use a dataset with 'text' or 'content' fields.") n_neg = n_pos * neg_ratio rng = random.Random(42) all_sents = [s for pair in texts for s in pair[:2]] for _ in range(n_neg): a = rng.choice(all_sents) b = rng.choice(all_sents) texts.append((a, b, 0)) # build HF Dataset rows = {"sentence1": [], "sentence2": [], "label": []} for a,b,l in texts: rows["sentence1"].append(a) rows["sentence2"].append(b) rows["label"].append(int(l)) nrows = len(rows["label"]) ds_new = Dataset.from_dict(rows) # tokenize def fn(examples): enc = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, padding="max_length", max_length=max_length) return {"input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"], "labels": examples["label"]} ds_new = ds_new.map(fn, batched=True, remove_columns=["sentence1", "sentence2", "label"]) return ds_new def _build_datasets_and_loaders(self): ds = self._load_hf_dataset() # Determine which split keys exist; standard HF datasets use train/validation/test or train/validation train_key = "train" if "train" in ds else list(ds.keys())[0] valid_key = "validation" if "validation" in ds else ("validation_matched" if "validation_matched" in ds else None) test_key = "test" if "test" in ds else None # Limit samples if requested if self.max_train_samples: ds[train_key] = ds[train_key].select(range(min(len(ds[train_key]), int(self.max_train_samples)))) if valid_key and self.max_eval_samples: ds[valid_key] = ds[valid_key].select(range(min(len(ds[valid_key]), int(self.max_eval_samples)))) task = self.cfg.get("task", "sentence_pair") # Heuristics for fields # For sentence pair tasks we look for columns like 'sentence1' and 'sentence2' or 'premise'/'hypothesis' or 'text_a'/'text_b' train_ds = ds[train_key] valid_ds = ds[valid_key] if valid_key else None if task == "sentence_pair": # find keys candidates = [("sentence1","sentence2"), ("premise","hypothesis"), ("text_a","text_b"), ("sentence_a","sentence_b"), ("question","sentence")] found = None for a,b in candidates: if a in train_ds.column_names and b in train_ds.column_names: found = (a,b); break if found is None: # try GLUE style keys sentence1/2 a = "sentence1" if "sentence1" in train_ds.column_names else None b = "sentence2" if "sentence2" in train_ds.column_names else None if a is None or b is None: raise RuntimeError(f"Could not find sentence-pair fields in dataset columns: {train_ds.column_names}") found = (a,b) text1_key, text2_key = found log(f"[rank {self.rank}] Using fields {text1_key}/{text2_key} for sentence_pair task.") # Map and tokenize train_tok = self._prepare_sentence_pair(train_ds, text1_key, text2_key) valid_tok = self._prepare_sentence_pair(valid_ds, text1_key, text2_key) if valid_ds is not None else None elif task == "classification": # find a single text field possible_text = [k for k in train_ds.column_names if k in ("text", "sentence", "content", "review")] text_key = possible_text[0] if possible_text else train_ds.column_names[0] log(f"[rank {self.rank}] Using field {text_key} for classification task.") train_tok = self._prepare_classification(train_ds, text_key) valid_tok = self._prepare_classification(valid_ds, text_key) if valid_ds is not None else None elif task == "next_sentence_prediction": # synthesize NSP dataset from raw text # prefer 'train' dataset that contains long documents or paragraphs log(f"[rank {self.rank}] Synthesizing NSP dataset for next_sentence_prediction.") train_tok = self._synthesize_nsp_dataset(train_ds) valid_tok = None else: raise ValueError(f"Unsupported task: {task}") # final step: ensure columns -> tensors and proper names def collate_fn_train(batch): # default collator handled by DataLoader (we want tensors) return {k: torch.tensor([x[k] for x in batch]) for k in batch[0].keys()} # Wrap into DataLoader # Use DistributedSampler in DDP train_sampler = DistributedSampler(train_tok) if self.use_ddp else None train_loader = DataLoader(train_tok, batch_size=self.batch_size, shuffle=(train_sampler is None), sampler=train_sampler, num_workers=self.num_workers, pin_memory=True, drop_last=True) eval_loader = None if valid_ds is not None and valid_tok is not None: eval_sampler = DistributedSampler(valid_tok) if self.use_ddp else None eval_loader = DataLoader(valid_tok, batch_size=self.eval_batch_size, shuffle=False, sampler=eval_sampler, num_workers=self.num_workers, pin_memory=True, drop_last=False) # Save shapes summary log(f"[rank {self.rank}] Train samples: {len(train_tok)}; Eval samples: {len(valid_tok) if valid_tok is not None else 0}") return train_loader, eval_loader # ------------------------- # Optimizer / Scheduler setup # ------------------------- def _setup_optimizer_and_scheduler(self, total_training_steps): no_decay = ["bias", "LayerNorm.weight"] params = [ {"params": [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay)], "weight_decay": self.weight_decay}, {"params": [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0}, ] optimizer = AdamW(params, lr=self.learning_rate) scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=self.warmup_steps, num_training_steps=total_training_steps) return optimizer, scheduler # ------------------------- # Training / Eval steps # ------------------------- def _forward(self, batch): # batch: dict with input_ids, attention_mask, (maybe labels) input_ids = batch["input_ids"].to(self.device) attention_mask = batch["attention_mask"].to(self.device) labels = batch.get("labels", None) if labels is not None: labels = torch.tensor(labels).to(self.device) out = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) loss = out.loss logits = out.logits else: out = self.model(input_ids=input_ids, attention_mask=attention_mask) logits = out.logits loss = None return loss, logits, labels def _evaluate(self, eval_loader): self.model.eval() all_preds = [] all_labels = [] total_loss = 0.0 n_batches = 0 with torch.no_grad(): for batch in tqdm(eval_loader, desc=f"Eval (rank {self.rank})", disable=(self.rank != 0)): # transform batch items to tensors if necessary if isinstance(batch, dict) and isinstance(batch.get("labels"), list): # dataset.map produced labels as lists; convert batch = {k: torch.tensor(v) if isinstance(v, list) else v for k,v in batch.items()} loss, logits, labels = self._forward(batch) if loss is not None: total_loss += loss.item() if logits is not None: preds = torch.argmax(logits, dim=-1).cpu().tolist() all_preds.extend(preds) if labels is not None: all_labels.extend(labels.cpu().tolist()) n_batches += 1 # metrics if len(all_labels) == 0: # no labels present, return accuracy 0 return {"loss": total_loss / (n_batches or 1), "accuracy": None} acc = accuracy_score(all_labels, all_preds) precision, recall, f1, _ = precision_recall_fscore_support(all_labels, all_preds, average="weighted", zero_division=0) return {"loss": total_loss / (n_batches or 1), "accuracy": acc, "precision": precision, "recall": recall, "f1": f1} # ------------------------- # Checkpointing helpers # ------------------------- def _save_checkpoint(self, step_or_epoch): ckpt_dir = self.output_dir / f"ckpt-{step_or_epoch}" ckpt_dir.mkdir(parents=True, exist_ok=True) model_to_save = self.model.module if hasattr(self.model, "module") else self.model model_to_save.save_pretrained(ckpt_dir) self.tokenizer.save_pretrained(ckpt_dir) # save training state state = { "global_step": self.global_step, "best_metric": self.best_metric, "cfg": self.cfg } with open(ckpt_dir / "train_state.json", "w") as f: json.dump(state, f) log(f"[rank {self.rank}] Saved checkpoint -> {ckpt_dir}") # ------------------------- # Main train loop # ------------------------- def train(self): log(f"[rank {self.rank}] Starting finetune. device={self.device} fp16={self.fp16} ddp={self.use_ddp}") train_loader, eval_loader = self._build_datasets_and_loaders() # compute total steps steps_per_epoch = math.ceil(len(train_loader) / (1.0 * self.gradient_accumulation_steps)) total_training_steps = int(steps_per_epoch * self.num_epochs) self.total_steps = total_training_steps log(f"[rank {self.rank}] Steps per epoch: {steps_per_epoch}, total training steps: {total_training_steps}") optimizer, scheduler = self._setup_optimizer_and_scheduler(total_training_steps) # DDP wrap if needed if self.use_ddp: # wrap with DistributedDataParallel self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[torch.cuda.current_device()], output_device=torch.cuda.current_device(), find_unused_parameters=False) # training loop self.model.train() optimizer.zero_grad() self.global_step = 0 best_metric = -1.0 for epoch in range(self.num_epochs): if self.use_ddp: train_loader.sampler.set_epoch(epoch) epoch_loss = 0.0 pbar = tqdm(train_loader, desc=f"Train Epoch {epoch} (rank {self.rank})", disable=(self.rank != 0)) for step, batch in enumerate(pbar): # convert list-labels to tensors if needed if isinstance(batch.get("labels"), list): batch["labels"] = torch.tensor(batch["labels"]) # forward with autocast(enabled=(self.fp16 and torch.cuda.is_available())): loss, logits, labels = self._forward(batch) if loss is None: # create dummy loss as mean CE of predicted vs random (shouldn't happen normally) loss = torch.tensor(0.0, device=self.device) loss = loss / self.gradient_accumulation_steps # backward if self.scaler is not None: self.scaler.scale(loss).backward() else: loss.backward() epoch_loss += loss.item() * self.gradient_accumulation_steps # gradient step if (step + 1) % self.gradient_accumulation_steps == 0: # unscale if using scaler if self.scaler is not None: self.scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm) self.scaler.step(optimizer) self.scaler.update() else: torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm) optimizer.step() scheduler.step() optimizer.zero_grad() self.global_step += 1 if self.rank == 0 and (self.global_step % self.logging_steps == 0): pbar.set_postfix({"loss": f"{epoch_loss/((step+1) or 1):.4f}", "step": self.global_step}) # eval and checkpoint if self.rank == 0 and self.eval_steps and (self.global_step % self.eval_steps == 0): if eval_loader is not None: metrics = self._evaluate(eval_loader) log(f"[rank {self.rank}] Eval at step {self.global_step}: {metrics}") # choose metric metric_val = metrics.get("f1") or metrics.get("accuracy") or 0.0 if metric_val > best_metric: best_metric = metric_val self.best_metric = best_metric # save best self._save_checkpoint(f"best-step-{self.global_step}") if self.rank == 0 and self.save_steps and (self.global_step % self.save_steps == 0): self._save_checkpoint(f"step-{self.global_step}") # end epoch log(f"[rank {self.rank}] Epoch {epoch} completed. avg_loss={(epoch_loss/len(train_loader)):.4f}") # epoch-end eval if eval_loader is not None and self.rank == 0: metrics = self._evaluate(eval_loader) log(f"[rank {self.rank}] Epoch {epoch} eval: {metrics}") metric_val = metrics.get("f1") or metrics.get("accuracy") or 0.0 if metric_val > best_metric: best_metric = metric_val self.best_metric = best_metric self._save_checkpoint(f"best-epoch-{epoch}") log(f"[rank {self.rank}] Training complete. Best metric: {self.best_metric}") # Save final if self.rank == 0: self._save_checkpoint("final") # Cleanup DDP if self.use_ddp and torch.distributed.is_initialized(): torch.distributed.barrier() torch.distributed.destroy_process_group()