File size: 6,619 Bytes
3b2d368 | 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 | # main.py (updated fragment to add bert_finetune mode)
import os
import re
import shutil
import csv
import sys
import json
from pathlib import Path
import torch
import torch.distributed as dist
from dotenv import load_dotenv
import hydra
from omegaconf import OmegaConf
# existing imports from your repo
from lmr.config import initialize_config
from lmr.tokenizer import Tokenizer
from lmr.models import get_model
from lmr.data import get_dataset_splits
from lmr.checkpointing import Checkpointing
from lmr.utils.seed import set_seed
from lmr.training import Bert_Trainer, Trainer
from lmr.generation import Generator
from lmr.benchmark import Benchmark
from lmr.ddp import unwrap_model
# NEW: import the BertFineTuneTrainer implementation you created earlier.
# Adjust this import to where you saved the class. Example: lmr.training.bert_finetune_trainer
from lmr.training.bert_finetune_trainer import BertFineTuneTrainer
# transformers
from transformers import AutoTokenizer, AutoConfig, AutoModelForSequenceClassification, BertForNextSentencePrediction
DATASET_DIR = Path("datasets")
CHECKPOINT_DIR = Path("/work/jf381/checkpoints")
BENCHMARK_DIR = Path("output")
# -------------------------
# Helper: create model+tokenizer for finetune mode
# -------------------------
def create_tokenizer_and_model_for_finetune(model_name_or_path: str, task: str, num_labels: int = None):
"""
Returns (tokenizer, model).
- task == "sentence_pair" -> AutoModelForSequenceClassification (num_labels required or inferred)
- task == "next_sentence_prediction" -> BertForNextSentencePrediction
"""
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
if task == "sentence_pair":
# If user didn't pass num_labels, default to 2 (binary)
nlab = num_labels if num_labels is not None else 2
model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path, num_labels=nlab)
elif task == "next_sentence_prediction":
# load BERT NSP model which exposes next sentence prediction head
model = BertForNextSentencePrediction.from_pretrained(model_name_or_path)
else:
raise ValueError(f"Unsupported finetune task: {task}")
return tokenizer, model
# =============================================================================
# MAIN
# =============================================================================
@hydra.main(config_path="config", config_name="config", version_base="1.3")
def main(config):
# common init
load_dotenv()
set_seed(config)
initialize_config(config)
mode = config.mode
maybe_print = print # keep simple; you may swap for Logger if desired
# existing modes (train/generate/...)
if mode == "train":
# unchanged: legacy pretrain flow
if torch.cuda.is_available() and torch.cuda.device_count() > 1:
if not dist.is_initialized():
dist.init_process_group(backend="nccl")
torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count())
tokenizer, model = setup_model_and_tokenizer(config)
tokenized_dataset_dir = DATASET_DIR / config.tokenizer_base
splits = get_dataset_splits(config.dataset, 1024, tokenized_dataset_dir)
checkpointing = Checkpointing(model, CHECKPOINT_DIR / config.checkpoint_name)
if "bert" in str(config.model).lower():
from lmr.training import Bert_Trainer
trainer = Bert_Trainer(config.training, model, tokenizer, splits, checkpointing, None)
else:
trainer = Trainer(config.training, model, tokenizer, splits, checkpointing)
trainer.train()
elif mode == "generate":
generate(config)
elif mode == "generate_all":
generate_all(config)
elif mode == "benchmark":
tokenizer, model = setup_model_and_tokenizer(config)
checkpointing = Checkpointing(model, CHECKPOINT_DIR / config.checkpoint_name)
benchmarking = Benchmark(config.benchmark, model, tokenizer, checkpointing, BENCHMARK_DIR / config.checkpoint_name)
benchmarking.run_benchmarks()
# ---------------------------
# NEW: bert fine-tune mode using HF datasets directly
# ---------------------------
elif mode == "bert_finetune":
# config should contain fields:
# model_name_or_path: pretrained model (e.g., bert-base-uncased or local folder)
# task: "sentence_pair" | "next_sentence_prediction"
# dataset: huggingface dataset id (e.g., "glue/mrpc" or "glue", config "mrpc")
# batch_size, num_epochs, lr, save_dir, use_ddp, etc
maybe_print("=== BERT FINETUNE MODE ===")
maybe_print(f"Config: task={config.finetune.task}, dataset={config.finetune.dataset}, model={config.finetune.model_name_or_path}")
# create tokenizer + model appropriate for finetune task
tokenizer, model = create_tokenizer_and_model_for_finetune(
model_name_or_path=config.finetune.model_name_or_path,
task=config.finetune.task,
num_labels=getattr(config.finetune, "num_labels", None)
)
# device
use_ddp = getattr(config.finetune, "use_ddp", False)
if use_ddp and torch.cuda.is_available():
# torchrun should provide local_rank; hydra may not, so rely on env var for rank
local_rank = int(os.environ.get("LOCAL_RANK", os.environ.get("RANK", 0)))
world_size = int(os.environ.get("WORLD_SIZE", 1))
device = torch.device(f"cuda:{local_rank}")
maybe_print(f"[finetune] DDP mode: local_rank={local_rank} world_size={world_size} device={device}")
cfg_namespace = argparse.Namespace(**config.finetune) # convert to simple namespace
# instantiate trainer with DDP flags
finetune_trainer = BertFineTuneTrainer(model=model, tokenizer=tokenizer, cfg=cfg_namespace, device=device)
# caller should run via torchrun --nproc_per_node=N python main.py mode=bert_finetune ...
finetune_trainer.train()
else:
# single-GPU / CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
maybe_print(f"[finetune] Single-device mode: device={device}")
cfg_namespace = argparse.Namespace(**config.finetune)
finetune_trainer = BertFineTuneTrainer(model=model, tokenizer=tokenizer, cfg=cfg_namespace, device=device)
finetune_trainer.train()
else:
print(f"❌ Unknown mode: {mode}")
if __name__ == "__main__":
main()
|