|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.optim as optim |
|
|
from torch.utils.data import Dataset, DataLoader |
|
|
import torch.nn.functional as F |
|
|
from torch.cuda.amp import GradScaler, autocast |
|
|
import os |
|
|
import json |
|
|
import argparse |
|
|
import time |
|
|
import math |
|
|
import glob |
|
|
from typing import Dict, List, Optional |
|
|
from tqdm import tqdm |
|
|
import numpy as np |
|
|
import gc |
|
|
import logging |
|
|
from collections import defaultdict |
|
|
import multiprocessing |
|
|
|
|
|
|
|
|
try: |
|
|
from model_slm import MixtureOfRecursions, count_parameters, TextGenerator |
|
|
from custom_tokenizer import TechnicalTokenizer |
|
|
except ImportError as e: |
|
|
raise ImportError(f"Failed to import custom modules: {e}") |
|
|
|
|
|
|
|
|
DEFAULT_MAX_LENGTH = 128 |
|
|
DEFAULT_MAX_EXAMPLES = 50000 |
|
|
DEFAULT_D_MODEL = 384 |
|
|
DEFAULT_N_LAYERS = 6 |
|
|
DEFAULT_N_HEADS = 6 |
|
|
DEFAULT_EPOCHS = 15 |
|
|
DEFAULT_BATCH_SIZE = 16 |
|
|
DEFAULT_LEARNING_RATE = 5e-4 |
|
|
DEFAULT_GRAD_ACCUM_STEPS = 1 |
|
|
DEFAULT_EVAL_EVERY = 500 |
|
|
DEFAULT_WARMUP_RATIO = 0.05 |
|
|
DEFAULT_CHECKPOINT_DIR = "checkpoints" |
|
|
DEFAULT_LOG_LEVEL = "INFO" |
|
|
|
|
|
|
|
|
logging.basicConfig( |
|
|
level=DEFAULT_LOG_LEVEL, |
|
|
format="%(asctime)s [%(levelname)s] %(message)s", |
|
|
handlers=[ |
|
|
logging.StreamHandler(), |
|
|
logging.FileHandler("training.log") |
|
|
] |
|
|
) |
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
class FastTechnicalTextDataset(Dataset): |
|
|
"""Optimized dataset for fast loading and processing of technical text.""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
data_file: str, |
|
|
tokenizer: TechnicalTokenizer, |
|
|
max_length: int = DEFAULT_MAX_LENGTH, |
|
|
max_examples: int = DEFAULT_MAX_EXAMPLES |
|
|
): |
|
|
""" |
|
|
Initialize the dataset with optimized loading. |
|
|
|
|
|
Args: |
|
|
data_file (str): Path to the training data file. |
|
|
tokenizer (TechnicalTokenizer): Tokenizer for encoding text. |
|
|
max_length (int): Maximum sequence length. |
|
|
max_examples (int): Maximum number of examples to load. |
|
|
|
|
|
Raises: |
|
|
FileNotFoundError: If the data file does not exist. |
|
|
ValueError: If max_length or max_examples is invalid. |
|
|
""" |
|
|
if not os.path.exists(data_file): |
|
|
raise FileNotFoundError(f"Data file not found: {data_file}") |
|
|
if max_length <= 0 or max_examples <= 0: |
|
|
raise ValueError("max_length and max_examples must be positive") |
|
|
|
|
|
self.tokenizer = tokenizer |
|
|
self.max_length = max_length |
|
|
self.pad_token_id = tokenizer.vocab.get('<pad>', 0) |
|
|
self.max_examples = max_examples |
|
|
self.examples = [] |
|
|
|
|
|
logger.info(f"Loading dataset from {data_file} with max_length={max_length}, max_examples={max_examples}") |
|
|
start_time = time.time() |
|
|
self._fast_load_data(data_file) |
|
|
self._tensorize_data() |
|
|
logger.info(f"Loaded {len(self.examples)} examples in {time.time() - start_time:.1f}s") |
|
|
|
|
|
if torch.cuda.is_available(): |
|
|
torch.cuda.empty_cache() |
|
|
gc.collect() |
|
|
|
|
|
def _fast_load_data(self, data_file: str) -> None: |
|
|
"""Load and filter data efficiently.""" |
|
|
logger.info("Reading and filtering data...") |
|
|
with open(data_file, 'r', encoding='utf-8') as f: |
|
|
lines = f.readlines() |
|
|
|
|
|
logger.info(f"File contains {len(lines)} lines") |
|
|
good_examples = [] |
|
|
seen_hashes = set() |
|
|
|
|
|
for line in lines[:self.max_examples * 3]: |
|
|
line = line.strip() |
|
|
if ( |
|
|
50 <= len(line) <= 400 and |
|
|
line.count(' ') >= 8 and |
|
|
not line.lower().startswith(('http', 'www', 'ftp')) and |
|
|
line.count('.') <= len(line) * 0.1 |
|
|
): |
|
|
line_hash = hash(line[:100]) |
|
|
if line_hash not in seen_hashes: |
|
|
seen_hashes.add(line_hash) |
|
|
good_examples.append(line) |
|
|
if len(good_examples) >= self.max_examples: |
|
|
break |
|
|
|
|
|
logger.info(f"Filtered to {len(good_examples)} quality examples") |
|
|
|
|
|
batch_size = 1000 |
|
|
for i in range(0, len(good_examples), batch_size): |
|
|
batch = good_examples[i:i + batch_size] |
|
|
for line in batch: |
|
|
try: |
|
|
if not line.endswith('<|endoftext|>'): |
|
|
line += ' <|endoftext|>' |
|
|
tokens = self.tokenizer.encode_ids(line, add_special_tokens=True) |
|
|
if 30 <= len(tokens) <= self.max_length: |
|
|
if len(tokens) < self.max_length: |
|
|
tokens.extend([self.pad_token_id] * (self.max_length - len(tokens))) |
|
|
self.examples.append(tokens) |
|
|
except Exception as e: |
|
|
logger.warning(f"Failed to process line: {e}") |
|
|
continue |
|
|
if i % 5000 == 0: |
|
|
logger.info(f"Processed {len(self.examples)} examples...") |
|
|
|
|
|
logger.info(f"Final dataset size: {len(self.examples)} examples") |
|
|
|
|
|
def _tensorize_data(self) -> None: |
|
|
"""Pre-tensorize data for faster training.""" |
|
|
logger.info("Pre-tensorizing data...") |
|
|
seq_len = self.max_length - 1 |
|
|
tensorized_examples = [] |
|
|
|
|
|
for tokens in self.examples: |
|
|
if len(tokens) != self.max_length: |
|
|
continue |
|
|
input_ids = torch.tensor(tokens[:-1], dtype=torch.long) |
|
|
targets = torch.tensor(tokens[1:], dtype=torch.long) |
|
|
original_len = next((i for i, x in enumerate(tokens) if x == self.pad_token_id), self.max_length) |
|
|
mask_len = min(original_len, seq_len) |
|
|
attention_mask = torch.zeros(seq_len, dtype=torch.long) |
|
|
attention_mask[:mask_len] = 1 |
|
|
tensorized_examples.append({ |
|
|
'input_ids': input_ids, |
|
|
'targets': targets, |
|
|
'attention_mask': attention_mask |
|
|
}) |
|
|
|
|
|
self.examples = tensorized_examples |
|
|
logger.info("Data pre-tensorized successfully") |
|
|
|
|
|
def __len__(self) -> int: |
|
|
"""Return the number of examples in the dataset.""" |
|
|
return len(self.examples) |
|
|
|
|
|
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: |
|
|
"""Return a single example from the dataset.""" |
|
|
return self.examples[idx] |
|
|
|
|
|
class FastCosineScheduler: |
|
|
"""Cosine learning rate scheduler with warmup.""" |
|
|
|
|
|
def __init__(self, optimizer: optim.Optimizer, total_steps: int, warmup_ratio: float = DEFAULT_WARMUP_RATIO): |
|
|
""" |
|
|
Initialize the cosine scheduler. |
|
|
|
|
|
Args: |
|
|
optimizer (optim.Optimizer): Optimizer to schedule. |
|
|
total_steps (int): Total training steps. |
|
|
warmup_ratio (float): Ratio of steps for warmup phase. |
|
|
|
|
|
Raises: |
|
|
ValueError: If total_steps or warmup_ratio is invalid. |
|
|
""" |
|
|
if total_steps <= 0 or not 0 <= warmup_ratio <= 1: |
|
|
raise ValueError("total_steps must be positive and warmup_ratio must be in [0, 1]") |
|
|
|
|
|
self.optimizer = optimizer |
|
|
self.total_steps = total_steps |
|
|
self.warmup_steps = int(total_steps * warmup_ratio) |
|
|
self.base_lr = optimizer.param_groups[0]['lr'] |
|
|
self.step_count = 0 |
|
|
|
|
|
def step(self) -> float: |
|
|
""" |
|
|
Update the learning rate. |
|
|
|
|
|
Returns: |
|
|
float: Current learning rate. |
|
|
""" |
|
|
self.step_count += 1 |
|
|
if self.step_count <= self.warmup_steps: |
|
|
lr = self.base_lr * self.step_count / self.warmup_steps |
|
|
else: |
|
|
progress = (self.step_count - self.warmup_steps) / (self.total_steps - self.warmup_steps) |
|
|
lr = self.base_lr * 0.5 * (1 + math.cos(math.pi * progress)) |
|
|
|
|
|
for param_group in self.optimizer.param_groups: |
|
|
param_group['lr'] = lr |
|
|
return lr |
|
|
|
|
|
class UltraFastTrainer: |
|
|
"""Trainer optimized for fast training of transformer models.""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
model: nn.Module, |
|
|
tokenizer: TechnicalTokenizer, |
|
|
train_dataset: FastTechnicalTextDataset, |
|
|
val_dataset: Optional[FastTechnicalTextDataset] = None, |
|
|
config: Optional[Dict] = None |
|
|
): |
|
|
""" |
|
|
Initialize the trainer. |
|
|
|
|
|
Args: |
|
|
model (nn.Module): The transformer model to train. |
|
|
tokenizer (TechnicalTokenizer): Tokenizer for encoding/decoding. |
|
|
train_dataset (FastTechnicalTextDataset): Training dataset. |
|
|
val_dataset (Optional[FastTechnicalTextDataset]): Validation dataset. |
|
|
config (Optional[Dict]): Training configuration. |
|
|
|
|
|
Raises: |
|
|
ValueError: If config contains invalid parameters. |
|
|
""" |
|
|
self.model = model |
|
|
self.tokenizer = tokenizer |
|
|
self.train_dataset = train_dataset |
|
|
self.val_dataset = val_dataset |
|
|
self.config = config or {} |
|
|
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
self.model.to(self.device) |
|
|
|
|
|
self._validate_config() |
|
|
self._fast_init_weights() |
|
|
self._setup_fast_optimizer() |
|
|
|
|
|
epochs = self.config.get('epochs', DEFAULT_EPOCHS) |
|
|
batch_size = self.config.get('batch_size', DEFAULT_BATCH_SIZE) |
|
|
total_steps = len(train_dataset) // batch_size * epochs |
|
|
self.scheduler = FastCosineScheduler(self.optimizer, total_steps) |
|
|
self.scaler = GradScaler() if self.device.type == 'cuda' else None |
|
|
self.global_step = 0 |
|
|
self.best_loss = float('inf') |
|
|
self.grad_accum_steps = self.config.get('gradient_accumulation_steps', DEFAULT_GRAD_ACCUM_STEPS) |
|
|
self.eval_every = self.config.get('eval_every', DEFAULT_EVAL_EVERY) |
|
|
|
|
|
def _validate_config(self) -> None: |
|
|
"""Validate training configuration.""" |
|
|
if self.config.get('batch_size', DEFAULT_BATCH_SIZE) <= 0: |
|
|
raise ValueError("batch_size must be positive") |
|
|
if self.config.get('epochs', DEFAULT_EPOCHS) <= 0: |
|
|
raise ValueError("epochs must be positive") |
|
|
if self.config.get('learning_rate', DEFAULT_LEARNING_RATE) <= 0: |
|
|
raise ValueError("learning_rate must be positive") |
|
|
if self.config.get('gradient_accumulation_steps', DEFAULT_GRAD_ACCUM_STEPS) <= 0: |
|
|
raise ValueError("gradient_accumulation_steps must be positive") |
|
|
|
|
|
def _fast_init_weights(self) -> None: |
|
|
"""Initialize model weights.""" |
|
|
def fast_init(module: nn.Module) -> None: |
|
|
if isinstance(module, nn.Linear): |
|
|
nn.init.normal_(module.weight, std=0.02) |
|
|
if module.bias is not None: |
|
|
nn.init.zeros_(module.bias) |
|
|
elif isinstance(module, nn.Embedding): |
|
|
nn.init.normal_(module.weight, std=0.02) |
|
|
self.model.apply(fast_init) |
|
|
logger.info("Model weights initialized") |
|
|
|
|
|
def _setup_fast_optimizer(self) -> None: |
|
|
"""Set up AdamW optimizer.""" |
|
|
lr = self.config.get('learning_rate', DEFAULT_LEARNING_RATE) |
|
|
params = [p for p in self.model.parameters() if p.requires_grad] |
|
|
self.optimizer = optim.AdamW(params, lr=lr, betas=(0.9, 0.99), weight_decay=0.01, eps=1e-6) |
|
|
logger.info(f"Optimizer initialized with learning rate: {lr}") |
|
|
|
|
|
def compute_fast_loss(self, logits: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: |
|
|
""" |
|
|
Compute masked cross-entropy loss. |
|
|
|
|
|
Args: |
|
|
logits (torch.Tensor): Model output logits of shape (batch_size, seq_len, vocab_size). |
|
|
targets (torch.Tensor): Target token IDs of shape (batch_size, seq_len). |
|
|
mask (torch.Tensor): Attention mask of shape (batch_size, seq_len). |
|
|
|
|
|
Returns: |
|
|
torch.Tensor: Computed loss. |
|
|
""" |
|
|
logits_flat = logits.view(-1, logits.size(-1)) |
|
|
targets_flat = targets.view(-1) |
|
|
mask_flat = mask.view(-1).bool() |
|
|
|
|
|
if not mask_flat.any(): |
|
|
return torch.tensor(0.0, device=logits.device, requires_grad=True) |
|
|
|
|
|
return F.cross_entropy(logits_flat[mask_flat], targets_flat[mask_flat]) |
|
|
|
|
|
def train_epoch_fast(self, epoch: int, dataloader: DataLoader) -> Dict[str, float]: |
|
|
""" |
|
|
Train for one epoch. |
|
|
|
|
|
Args: |
|
|
epoch (int): Current epoch number. |
|
|
dataloader (DataLoader): Training data loader. |
|
|
|
|
|
Returns: |
|
|
Dict[str, float]: Training metrics (loss, perplexity, epoch_time_min). |
|
|
""" |
|
|
self.model.train() |
|
|
total_loss = 0 |
|
|
num_batches = 0 |
|
|
start_time = time.time() |
|
|
|
|
|
progress_bar = tqdm(dataloader, desc=f"Epoch {epoch}", leave=False, miniters=50) |
|
|
for batch_idx, batch in enumerate(progress_bar): |
|
|
input_ids = batch['input_ids'].to(self.device, non_blocking=True) |
|
|
targets = batch['targets'].to(self.device, non_blocking=True) |
|
|
mask = batch['attention_mask'].to(self.device, non_blocking=True) |
|
|
|
|
|
with autocast(enabled=self.device.type == 'cuda'): |
|
|
logits, comp_loss = self.model(input_ids, mask) |
|
|
lm_loss = self.compute_fast_loss(logits, targets, mask) |
|
|
total_loss_step = lm_loss + 0.0001 * comp_loss |
|
|
if self.grad_accum_steps > 1: |
|
|
total_loss_step = total_loss_step / self.grad_accum_steps |
|
|
|
|
|
if self.scaler: |
|
|
self.scaler.scale(total_loss_step).backward() |
|
|
if (batch_idx + 1) % self.grad_accum_steps == 0: |
|
|
self.scaler.unscale_(self.optimizer) |
|
|
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) |
|
|
self.scaler.step(self.optimizer) |
|
|
self.scaler.update() |
|
|
self.optimizer.zero_grad(set_to_none=True) |
|
|
self.scheduler.step() |
|
|
self.global_step += 1 |
|
|
else: |
|
|
total_loss_step.backward() |
|
|
if (batch_idx + 1) % self.grad_accum_steps == 0: |
|
|
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) |
|
|
self.optimizer.step() |
|
|
self.optimizer.zero_grad(set_to_none=True) |
|
|
self.scheduler.step() |
|
|
self.global_step += 1 |
|
|
|
|
|
total_loss += lm_loss.item() |
|
|
num_batches += 1 |
|
|
|
|
|
if batch_idx % 100 == 0: |
|
|
current_loss = total_loss / num_batches |
|
|
progress_bar.set_postfix({'loss': f"{current_loss:.3f}", 'ppl': f"{math.exp(min(current_loss, 10)):.1f}"}) |
|
|
|
|
|
if batch_idx % 200 == 0 and batch_idx > 0 and self.device.type == 'cuda': |
|
|
torch.cuda.empty_cache() |
|
|
|
|
|
avg_loss = total_loss / max(num_batches, 1) |
|
|
return { |
|
|
'loss': avg_loss, |
|
|
'perplexity': math.exp(min(avg_loss, 10)), |
|
|
'epoch_time_min': (time.time() - start_time) / 60 |
|
|
} |
|
|
|
|
|
def validate_fast(self, dataloader: DataLoader) -> Dict[str, float]: |
|
|
""" |
|
|
Validate the model on the validation dataset. |
|
|
|
|
|
Args: |
|
|
dataloader (DataLoader): Validation data loader. |
|
|
|
|
|
Returns: |
|
|
Dict[str, float]: Validation metrics (loss, perplexity). |
|
|
""" |
|
|
self.model.eval() |
|
|
total_loss = 0 |
|
|
num_batches = 0 |
|
|
max_val_batches = min(100, len(dataloader)) |
|
|
|
|
|
with torch.no_grad(): |
|
|
for batch_idx, batch in enumerate(dataloader): |
|
|
if batch_idx >= max_val_batches: |
|
|
break |
|
|
input_ids = batch['input_ids'].to(self.device, non_blocking=True) |
|
|
targets = batch['targets'].to(self.device, non_blocking=True) |
|
|
mask = batch['attention_mask'].to(self.device, non_blocking=True) |
|
|
|
|
|
with autocast(enabled=self.device.type == 'cuda'): |
|
|
logits, _ = self.model(input_ids, mask) |
|
|
loss = self.compute_fast_loss(logits, targets, mask) |
|
|
|
|
|
total_loss += loss.item() |
|
|
num_batches += 1 |
|
|
|
|
|
avg_loss = total_loss / max(num_batches, 1) |
|
|
return {'loss': avg_loss, 'perplexity': math.exp(min(avg_loss, 10))} |
|
|
|
|
|
def save_checkpoint_fast(self, epoch: int, metrics: Dict, save_dir: str = DEFAULT_CHECKPOINT_DIR) -> Optional[str]: |
|
|
""" |
|
|
Save a checkpoint if the loss improves. |
|
|
|
|
|
Args: |
|
|
epoch (int): Current epoch number. |
|
|
metrics (Dict): Training and validation metrics. |
|
|
save_dir (str): Directory to save checkpoints. |
|
|
|
|
|
Returns: |
|
|
Optional[str]: Path to the saved checkpoint or None. |
|
|
""" |
|
|
os.makedirs(save_dir, exist_ok=True) |
|
|
val_loss = metrics.get('val_loss', metrics.get('loss', float('inf'))) |
|
|
|
|
|
if val_loss < self.best_loss: |
|
|
self.best_loss = val_loss |
|
|
checkpoint = { |
|
|
'epoch': epoch, |
|
|
'model_state_dict': self.model.state_dict(), |
|
|
'optimizer_state_dict': self.optimizer.state_dict(), |
|
|
'metrics': metrics, |
|
|
'scaler_state_dict': self.scaler.state_dict() if self.scaler else None |
|
|
} |
|
|
best_path = os.path.join(save_dir, "best_model.pt") |
|
|
torch.save(checkpoint, best_path) |
|
|
logger.info(f"New best checkpoint saved: {best_path}, Loss: {val_loss:.4f}") |
|
|
return best_path |
|
|
return None |
|
|
|
|
|
def train_ultra_fast(self, num_epochs: int = DEFAULT_EPOCHS, batch_size: int = DEFAULT_BATCH_SIZE) -> List[Dict]: |
|
|
""" |
|
|
Train the model with optimized settings. |
|
|
|
|
|
Args: |
|
|
num_epochs (int): Number of training epochs. |
|
|
batch_size (int): Batch size for training. |
|
|
|
|
|
Returns: |
|
|
List[Dict]: Training history with metrics for each epoch. |
|
|
""" |
|
|
logger.info(f"Starting ultra-fast training: {num_epochs} epochs, batch_size={batch_size}") |
|
|
logger.info("Target: Loss < 2.0, PPL < 12, Time: 4-5 hours") |
|
|
|
|
|
train_loader = DataLoader( |
|
|
self.train_dataset, |
|
|
batch_size=batch_size, |
|
|
shuffle=True, |
|
|
num_workers=min(multiprocessing.cpu_count(), 4), |
|
|
pin_memory=self.device.type == 'cuda', |
|
|
persistent_workers=True, |
|
|
drop_last=True |
|
|
) |
|
|
|
|
|
val_loader = None |
|
|
if self.val_dataset: |
|
|
val_loader = DataLoader( |
|
|
self.val_dataset, |
|
|
batch_size=batch_size * 2, |
|
|
shuffle=False, |
|
|
num_workers=min(multiprocessing.cpu_count() // 2, 2), |
|
|
pin_memory=self.device.type == 'cuda' |
|
|
) |
|
|
|
|
|
total_start_time = time.time() |
|
|
history = [] |
|
|
|
|
|
for epoch in range(1, num_epochs + 1): |
|
|
logger.info(f"Starting epoch {epoch}/{num_epochs}") |
|
|
train_metrics = self.train_epoch_fast(epoch, train_loader) |
|
|
|
|
|
val_metrics = {} |
|
|
if val_loader and (epoch % 2 == 0 or epoch == num_epochs): |
|
|
val_metrics = self.validate_fast(val_loader) |
|
|
|
|
|
epoch_time = train_metrics['epoch_time_min'] * 60 |
|
|
epoch_info = { |
|
|
'epoch': epoch, |
|
|
'train_loss': train_metrics['loss'], |
|
|
'train_ppl': train_metrics['perplexity'], |
|
|
'epoch_time_min': train_metrics['epoch_time_min'] |
|
|
} |
|
|
if val_metrics: |
|
|
epoch_info.update({'val_loss': val_metrics['loss'], 'val_ppl': val_metrics['perplexity']}) |
|
|
|
|
|
history.append(epoch_info) |
|
|
|
|
|
elapsed_hours = (time.time() - total_start_time) / 3600 |
|
|
remaining_hours = elapsed_hours * (num_epochs - epoch) / max(epoch, 1) |
|
|
|
|
|
logger.info(f"Epoch {epoch} results:") |
|
|
logger.info(f" Epoch time: {epoch_time/60:.1f} min") |
|
|
logger.info(f" Total elapsed: {elapsed_hours:.1f}h") |
|
|
logger.info(f" Est. remaining: {remaining_hours:.1f}h") |
|
|
logger.info(f" Train Loss: {train_metrics['loss']:.4f}") |
|
|
logger.info(f" Train PPL: {train_metrics['perplexity']:.1f}") |
|
|
if val_metrics: |
|
|
logger.info(f" Val Loss: {val_metrics['loss']:.4f}") |
|
|
logger.info(f" Val PPL: {val_metrics['perplexity']:.1f}") |
|
|
|
|
|
current_loss = val_metrics.get('loss', train_metrics['loss']) |
|
|
current_ppl = val_metrics.get('perplexity', train_metrics['perplexity']) |
|
|
if current_loss < 2.0 and current_ppl < 12: |
|
|
logger.info(f"Targets achieved: Loss={current_loss:.4f} < 2.0, PPL={current_ppl:.1f} < 12") |
|
|
|
|
|
combined_metrics = {**train_metrics} |
|
|
if val_metrics: |
|
|
combined_metrics.update({f"val_{k}": v for k, v in val_metrics.items()}) |
|
|
self.save_checkpoint_fast(epoch, combined_metrics) |
|
|
|
|
|
if self.device.type == 'cuda': |
|
|
torch.cuda.empty_cache() |
|
|
gc.collect() |
|
|
|
|
|
if current_loss < 1.8 and current_ppl < 10: |
|
|
logger.info("Early stopping: Excellent performance achieved!") |
|
|
break |
|
|
|
|
|
total_time = (time.time() - total_start_time) / 3600 |
|
|
logger.info(f"Training completed in {total_time:.1f} hours") |
|
|
logger.info(f"Best loss: {self.best_loss:.4f}") |
|
|
return history |
|
|
|
|
|
def run_ultra_fast_training() -> int: |
|
|
""" |
|
|
Run the ultra-fast training pipeline. |
|
|
|
|
|
Returns: |
|
|
int: Exit code (0 for success, 1 for failure). |
|
|
""" |
|
|
parser = argparse.ArgumentParser(description="Ultra-Fast Training for MixtureOfRecursions Model") |
|
|
parser.add_argument("--train_file", default=None, help="Path to training data file") |
|
|
parser.add_argument("--val_file", default=None, help="Path to validation data file") |
|
|
parser.add_argument("--tokenizer_dir", default="tokenizer", help="Directory for tokenizer files") |
|
|
parser.add_argument("--max_examples", type=int, default=DEFAULT_MAX_EXAMPLES, help="Maximum number of training examples") |
|
|
parser.add_argument("--d_model", type=int, default=DEFAULT_D_MODEL, help="Model embedding dimension") |
|
|
parser.add_argument("--n_layers", type=int, default=DEFAULT_N_LAYERS, help="Number of transformer layers") |
|
|
parser.add_argument("--n_heads", type=int, default=DEFAULT_N_HEADS, help="Number of attention heads") |
|
|
parser.add_argument("--max_seq_len", type=int, default=DEFAULT_MAX_LENGTH, help="Maximum sequence length") |
|
|
parser.add_argument("--epochs", type=int, default=DEFAULT_EPOCHS, help="Number of training epochs") |
|
|
parser.add_argument("--batch_size", type=int, default=DEFAULT_BATCH_SIZE, help="Batch size for training") |
|
|
parser.add_argument("--learning_rate", type=float, default=DEFAULT_LEARNING_RATE, help="Learning rate") |
|
|
parser.add_argument("--gradient_accumulation_steps", type=int, default=DEFAULT_GRAD_ACCUM_STEPS, help="Gradient accumulation steps") |
|
|
parser.add_argument("--eval_every", type=int, default=DEFAULT_EVAL_EVERY, help="Evaluate every N steps") |
|
|
|
|
|
args = parser.parse_args() |
|
|
|
|
|
torch.manual_seed(42) |
|
|
np.random.seed(42) |
|
|
|
|
|
logger.info("Starting ultra-fast training pipeline") |
|
|
|
|
|
if args.train_file is None: |
|
|
patterns = ["*train*.txt", "*_train.txt"] |
|
|
files = [] |
|
|
for pattern in patterns: |
|
|
files.extend(glob.glob(pattern)) |
|
|
files.extend(glob.glob(os.path.join("split_data", pattern))) |
|
|
files.extend(glob.glob(os.path.join("data", pattern))) |
|
|
if files: |
|
|
args.train_file = files[0] |
|
|
logger.info(f"Found training file: {args.train_file}") |
|
|
else: |
|
|
logger.error("No training files found!") |
|
|
return 1 |
|
|
|
|
|
try: |
|
|
tokenizer = TechnicalTokenizer() |
|
|
tokenizer.load(args.tokenizer_dir) |
|
|
logger.info(f"Tokenizer loaded with vocab size: {tokenizer.get_vocab_size()}") |
|
|
except Exception as e: |
|
|
logger.error(f"Failed to load tokenizer: {e}") |
|
|
return 1 |
|
|
|
|
|
logger.info("Creating training dataset...") |
|
|
try: |
|
|
train_dataset = FastTechnicalTextDataset( |
|
|
args.train_file, tokenizer, args.max_seq_len, args.max_examples |
|
|
) |
|
|
except Exception as e: |
|
|
logger.error(f"Failed to create training dataset: {e}") |
|
|
return 1 |
|
|
|
|
|
val_dataset = None |
|
|
if args.val_file and os.path.exists(args.val_file): |
|
|
try: |
|
|
val_dataset = FastTechnicalTextDataset( |
|
|
args.val_file, tokenizer, args.max_seq_len, max_examples=5000 |
|
|
) |
|
|
logger.info("Validation dataset created") |
|
|
except Exception as e: |
|
|
logger.warning(f"Failed to create validation dataset: {e}") |
|
|
|
|
|
try: |
|
|
model = MixtureOfRecursions( |
|
|
vocab_size=tokenizer.get_vocab_size(), |
|
|
d_model=args.d_model, |
|
|
n_layers=args.n_layers, |
|
|
n_heads=args.n_heads, |
|
|
max_seq_len=args.max_seq_len - 1, |
|
|
padding_idx=tokenizer.vocab.get('<pad>', 0) |
|
|
) |
|
|
logger.info("Model initialized") |
|
|
except Exception as e: |
|
|
logger.error(f"Failed to initialize model: {e}") |
|
|
return 1 |
|
|
|
|
|
config = { |
|
|
'learning_rate': args.learning_rate, |
|
|
'gradient_accumulation_steps': args.gradient_accumulation_steps, |
|
|
'eval_every': args.eval_every, |
|
|
'batch_size': args.batch_size, |
|
|
'epochs': args.epochs |
|
|
} |
|
|
|
|
|
try: |
|
|
trainer = UltraFastTrainer(model, tokenizer, train_dataset, val_dataset, config) |
|
|
results = trainer.train_ultra_fast(args.epochs, args.batch_size) |
|
|
|
|
|
with open('ultra_fast_results.json', 'w') as f: |
|
|
json.dump(results, f, indent=2) |
|
|
logger.info("Training results saved to ultra_fast_results.json") |
|
|
|
|
|
return 0 |
|
|
except Exception as e: |
|
|
logger.error(f"Training failed: {e}") |
|
|
return 1 |
|
|
|
|
|
if __name__ == "__main__": |
|
|
exit(run_ultra_fast_training()) |
