Module neuralops.trainer.train_lra
Train small Transformer from scratch on LRA tasks, with optional module injection. LRA requires training from scratch, NOT fine-tuning.
Functions
def build_small_transformer(seq_len: int,
num_labels: int,
vocab_size: int,
hidden_size: int = 256,
num_layers: int = 6,
num_heads: int = 8,
intermediate_size: int = 1024)-
Expand source code
def build_small_transformer( seq_len: int, num_labels: int, vocab_size: int, hidden_size: int = 256, num_layers: int = 6, num_heads: int = 8, intermediate_size: int = 1024, ): """Build a standard 6-layer Transformer as used in LRA benchmark literature.""" config = AutoConfig.from_pretrained("bert-base-uncased") config.num_labels = num_labels config.vocab_size = vocab_size config.hidden_size = hidden_size config.num_hidden_layers = num_layers config.num_attention_heads = num_heads config.intermediate_size = intermediate_size config.max_position_embeddings = max(seq_len, 4096) model = AutoModelForSequenceClassification.from_config(config) return modelBuild a standard 6-layer Transformer as used in LRA benchmark literature.
def evaluate(model, dataset, batch_size)-
Expand source code
def evaluate(model, dataset, batch_size): model.eval() loader = DataLoader(dataset, batch_size=batch_size, num_workers=0) correct = 0 total = 0 device = next(model.parameters()).device with torch.no_grad(): for batch in loader: input_ids = batch["input_ids"].to(device) labels = batch["labels"].to(device) outputs = model(input_ids) preds = outputs.logits.argmax(dim=-1) correct += (preds == labels).sum().item() total += labels.size(0) return 100.0 * correct / max(1, total) def train_task(task: str,
module_name: str,
output_dir: str,
epochs: int = None,
max_steps: int = None,
batch_size: int = None,
seq_len: int = None,
lr: float = None,
hidden_size: int = 256,
num_layers: int = 6,
num_heads: int = 8,
intermediate_size: int = 1024,
gradient_checkpointing: bool = False,
warmup_ratio: float = 0.1,
warmup_steps: int = None,
precision: str = 'bf16')-
Expand source code
def train_task( task: str, module_name: str, output_dir: str, epochs: int = None, max_steps: int = None, batch_size: int = None, seq_len: int = None, lr: float = None, hidden_size: int = 256, num_layers: int = 6, num_heads: int = 8, intermediate_size: int = 1024, gradient_checkpointing: bool = False, warmup_ratio: float = 0.1, warmup_steps: int = None, precision: str = "bf16", ): """Orchestrate training of small Transformer on LRA benchmark tasks.""" cfg = TASK_CONFIG[task] num_epochs = epochs if epochs is not None else cfg.get("epochs", 100) current_batch_size = batch_size if batch_size is not None else cfg["batch_size"] current_seq_len = seq_len if seq_len is not None else cfg["seq_len"] current_lr = lr if lr is not None else cfg.get("lr", 5e-4) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") train_ds, test_ds = _setup_datasets(task, current_seq_len, cfg) train_loader = DataLoader(train_ds, batch_size=current_batch_size, shuffle=True, num_workers=0) total_steps = ( num_epochs * len(train_loader) if max_steps is None else min(max_steps, num_epochs * len(train_loader)) ) actual_warmup_steps = ( warmup_steps if warmup_steps is not None else max(1, int(total_steps * warmup_ratio)) ) model, optimizer, scheduler = _setup_model_and_optimizer( module_name=module_name, device=device, current_seq_len=current_seq_len, cfg=cfg, hidden_size=hidden_size, num_layers=num_layers, num_heads=num_heads, intermediate_size=intermediate_size, gradient_checkpointing=gradient_checkpointing, current_lr=current_lr, total_steps=total_steps, actual_warmup_steps=actual_warmup_steps, ) if precision == "fp16": use_amp = torch.cuda.is_available() amp_dtype = torch.float16 elif precision == "bf16": use_amp = torch.cuda.is_available() and torch.cuda.is_bf16_supported() amp_dtype = torch.bfloat16 else: use_amp = False amp_dtype = torch.float32 best_acc = 0.0 global_step = 0 print( f"Starting {task.upper()} training on {module_name}: {num_epochs} epochs, batch_size={current_batch_size}, lr={current_lr}, precision={precision} (amp={use_amp}), warmup_steps={actual_warmup_steps}, total_steps={total_steps}", flush=True, ) for epoch in range(num_epochs): avg_loss, global_step = _train_single_epoch( model=model, train_loader=train_loader, optimizer=optimizer, scheduler=scheduler, device=device, amp_dtype=amp_dtype, use_amp=use_amp, global_step=global_step, max_steps=max_steps, ) # Evaluate every 5 epochs or last epoch if ( (epoch + 1) % 5 == 0 or (epoch + 1) == num_epochs or (max_steps is not None and global_step >= max_steps) ): acc = evaluate(model, test_ds, current_batch_size) print( f"[{task}/{module_name}] Epoch {epoch + 1:03d}/{num_epochs} (step {global_step}) | Loss: {avg_loss:.4f} | Acc: {acc:.2f}% (Best: {max(best_acc, acc):.2f}%)", flush=True, ) if acc > best_acc: best_acc = acc else: print( f"[{task}/{module_name}] Epoch {epoch + 1:03d}/{num_epochs} (step {global_step}) | Loss: {avg_loss:.4f}", flush=True, ) if max_steps is not None and global_step >= max_steps: break metadata = { "task": task, "module": module_name, "best_accuracy": best_acc, "epochs": num_epochs, "seq_len": current_seq_len, "batch_size": current_batch_size, "lr": current_lr, "hidden_size": hidden_size, "num_layers": num_layers, "num_heads": num_heads, "intermediate_size": intermediate_size, } _save_training_results(output_dir, metadata) del model if torch.cuda.is_available(): torch.cuda.empty_cache() return best_accOrchestrate training of small Transformer on LRA benchmark tasks.