""" Training and evaluation pipeline for Part 1 (MoE Transformer Ablations). Evaluates perplexity, BLEU, expert utilization heatmaps, and saves artifacts. """ import os import json import logging from typing import Dict, Any, List, Optional, Tuple import torch import torch.nn as nn from torch.utils.data import DataLoader from tqdm import tqdm import matplotlib.pyplot as plt import numpy as np from .config import TransformerConfig from .model import Transformer from .moe import MoE from .dataset import build_or_load_tokenizer, TranslationDataset, collate_fn from ..utils import set_seed, compute_perplexity, compute_bleu, WandbLogger, count_parameters logger = logging.getLogger("Part1_Train") def train_epoch( model: Transformer, dataloader: DataLoader, optimizer: torch.optim.Optimizer, device: torch.device, wandb_logger: Optional[WandbLogger] = None, step_offset: int = 0 ) -> Tuple[float, int]: model.train() total_loss = 0.0 total_tokens = 0 step = step_offset for batch in tqdm(dataloader, desc="Training", leave=False): input_ids = batch["input_ids"].to(device) targets = batch["targets"].to(device) optimizer.zero_grad() logits, loss, aux_loss = model(input_ids, targets=targets) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() step += 1 num_valid_tokens = (targets != -100).sum().item() total_loss += loss.item() * num_valid_tokens total_tokens += max(num_valid_tokens, 1) if wandb_logger and step % 20 == 0: wandb_logger.log({ "train/loss": loss.item(), "train/aux_loss": aux_loss.item() if isinstance(aux_loss, torch.Tensor) else aux_loss, "train/step": step }, step=step) avg_loss = total_loss / max(total_tokens, 1) return avg_loss, step @torch.no_grad() def evaluate_and_extract_heatmaps( model: Transformer, dataloader: DataLoader, tokenizer: Any, device: torch.device, max_eval_samples: int = 1000 ) -> Dict[str, Any]: """ Evaluates perplexity and BLEU, while recording layer-wise expert activations split by input language ('vi' vs 'ja'). """ model.eval() total_loss = 0.0 total_tokens = 0 predictions: List[str] = [] references: List[List[str]] = [] # Track expert activation counts: lang -> layer_idx -> [count_per_expert] has_moe = model.config.use_moe num_layers = len(model.blocks) routed_experts_count = ( model.config.num_experts if model.config.ablation_variant != "shared_routed" else 3 ) expert_activity = { "vi": np.zeros((num_layers, routed_experts_count), dtype=np.float64), "ja": np.zeros((num_layers, routed_experts_count), dtype=np.float64) } lang_token_counts = {"vi": 0, "ja": 0} eval_samples_processed = 0 for batch in tqdm(dataloader, desc="Evaluating", leave=False): input_ids = batch["input_ids"].to(device) targets = batch["targets"].to(device) src_langs = batch["src_langs"] prompts = batch["prompts"] tgt_texts = batch["tgt_texts"] logits, loss, _ = model(input_ids, targets=targets) num_valid_tokens = (targets != -100).sum().item() total_loss += loss.item() * num_valid_tokens total_tokens += max(num_valid_tokens, 1) # Record expert routing decisions if MoE if has_moe: for l_idx, block in enumerate(model.blocks): if isinstance(block.ffn, MoE) and block.ffn.last_expert_indices is not None: # last_expert_indices: (B * N_ctx, top_k) indices = block.ffn.last_expert_indices.view(input_ids.shape[0], input_ids.shape[1], -1) for b_idx, lang in enumerate(src_langs): if lang in expert_activity: sample_indices = indices[b_idx].cpu().numpy().flatten() for exp_id in sample_indices: if 0 <= exp_id < routed_experts_count: expert_activity[lang][l_idx, exp_id] += 1 lang_token_counts[lang] += 1 # Greedy decoding for BLEU evaluation on a subset of samples if eval_samples_processed < max_eval_samples: for b_idx in range(len(prompts)): prompt = prompts[b_idx] target_text = tgt_texts[b_idx] prompt_enc = tokenizer(prompt, return_tensors="pt").input_ids.to(device) # Generate up to 50 tokens curr_ids = prompt_enc.clone() for _ in range(50): if curr_ids.shape[1] >= model.config.N_ctx: break out_logits, _, _ = model(curr_ids) next_tok = out_logits[:, -1, :].argmax(dim=-1, keepdim=True) curr_ids = torch.cat([curr_ids, next_tok], dim=1) if next_tok.item() == tokenizer.eos_token_id: break gen_tokens = curr_ids[0, prompt_enc.shape[1]:] pred_text = tokenizer.decode(gen_tokens.cpu(), skip_special_tokens=True) predictions.append(pred_text) references.append([target_text]) eval_samples_processed += 1 if eval_samples_processed >= max_eval_samples: break avg_loss = total_loss / max(total_tokens, 1) ppl = compute_perplexity(avg_loss) bleu = compute_bleu(predictions, references) # Normalize expert activity to distribution per language for lang in ["vi", "ja"]: total_act = expert_activity[lang].sum() if total_act > 0: expert_activity[lang] = expert_activity[lang] / total_act return { "loss": avg_loss, "perplexity": ppl, "bleu": bleu, "expert_activity": expert_activity, "sample_predictions": predictions[:5], "sample_references": references[:5] } def save_expert_heatmap( expert_activity: Dict[str, np.ndarray], save_path: str = "assets/expert_heatmap.png", variant_name: str = "MoE" ) -> None: """Plots and saves expert activation heatmap comparing Vietnamese vs Japanese.""" os.makedirs(os.path.dirname(os.path.abspath(save_path)), exist_ok=True) fig, axes = plt.subplots(1, 2, figsize=(14, 5)) for idx, lang in enumerate(["vi", "ja"]): ax = axes[idx] data = expert_activity[lang] if not isinstance(data, np.ndarray): data = np.asarray(data, dtype=float) im = ax.imshow(data, cmap="Blues", aspect="auto") ax.set_title(f"Expert Activation: {lang.upper()} ({variant_name})") ax.set_xlabel("Expert Index") ax.set_ylabel("Transformer Layer") fig.colorbar(im, ax=ax) plt.tight_layout() plt.savefig(save_path, dpi=200) plt.close() logger.info(f"Expert activation heatmap saved to {save_path}") def run_part1_training( ablation_variant: str = "dense", config: Optional[TransformerConfig] = None, dataset_name: str = "belumind/en-vi-ja-curated-500k-triplets", max_samples: Optional[int] = 10000, batch_size: int = 32, epochs: int = 1, max_steps: Optional[int] = None, device: Optional[torch.device] = None, output_dir: str = "assets/part1", save_checkpoint_dir: str = "checkpoints/part1" ) -> Dict[str, Any]: """ End-to-end training and evaluation pipeline for Part 1 MoE ablations. Trains model, extracts expert language specialization, and saves heatmaps. """ os.makedirs(output_dir, exist_ok=True) os.makedirs(save_checkpoint_dir, exist_ok=True) if device is None: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if config is None: config = TransformerConfig(ablation_variant=ablation_variant) logger.info(f"Initializing model for ablation: {ablation_variant.upper()} on device {device}") model = Transformer(config).to(device) # Load or build multilingual byte-level BPE tokenizer tokenizer = build_or_load_tokenizer(vocab_size=config.N_vocab) # Load translation dataset logger.info(f"Loading dataset {dataset_name} (max_samples={max_samples})...") try: from datasets import load_dataset ds = load_dataset(dataset_name, split="train") raw_samples = [] for i, item in enumerate(ds): if max_samples is not None and i >= max_samples: break raw_samples.append(item) except Exception as e: logger.warning(f"Could not load Hugging Face dataset ({e}). Using demonstration seed corpus.") raw_samples = [ {"en": "Hello, how are you today?", "vi": "Xin chào, hôm nay bạn thế nào?", "ja": "こんにちは、今日はお元気ですか？"}, {"en": "Machine translation with neural networks is fascinating.", "vi": "Dịch máy với mạng nơ-ron rất hấp dẫn.", "ja": "ニューラルネットワークによる機械翻訳は非常に興味深い。"}, {"en": "Mixture of experts scales model capacity efficiently.", "vi": "Mô hình chuyên gia kết hợp mở rộng dung lượng hiệu quả.", "ja": "エキスパートの混合はモデル容量を効率的に拡張します。"}, {"en": "Artificial intelligence continues to transform language technology.", "vi": "Trí tuệ nhân tạo tiếp tục chuyển đổi công nghệ ngôn ngữ.", "ja": "人工知能は言語技術を変革し続けています。"} ] * 200 # Split train and validation split_idx = int(len(raw_samples) * 0.9) train_samples = raw_samples[:split_idx] val_samples = raw_samples[split_idx:] if split_idx < len(raw_samples) else raw_samples[:10] train_ds = TranslationDataset(train_samples, tokenizer, max_seq_len=config.N_ctx) val_ds = TranslationDataset(val_samples, tokenizer, max_seq_len=config.N_ctx) train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, collate_fn=collate_fn) val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, collate_fn=collate_fn) optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay) wandb_logger = None if os.environ.get("WANDB_API_KEY") and os.environ.get("WANDB_MODE") != "offline": try: wandb_logger = WandbLogger( project="ANLP-Assignment-2", name=f"Part1-MoE-{ablation_variant}", config={"ablation": ablation_variant, "epochs": epochs, "samples": max_samples} ) except Exception as e: logger.warning(f"Could not initialize WandB ({e}).") logger.info(f"Starting training for {epochs} epoch(s)...") step_offset = 0 for epoch in range(epochs): avg_loss, step_offset = train_epoch( model, train_loader, optimizer, device, wandb_logger=wandb_logger, step_offset=step_offset ) logger.info(f"Epoch {epoch + 1}/{epochs} | Average Training Loss: {avg_loss:.4f}") # Evaluate and extract expert heatmaps logger.info("Evaluating model and extracting language routing decisions...") eval_results = evaluate_and_extract_heatmaps(model, val_loader, tokenizer, device) if wandb_logger: wandb_logger.log({ "val/loss": eval_results["loss"], "val/perplexity": eval_results["perplexity"], "val/bleu": eval_results["bleu"] }) wandb_logger.finish() # Save heatmap if MoE if config.use_moe: heatmap_path = os.path.join(output_dir, f"expert_heatmap_{ablation_variant}.png") save_expert_heatmap(eval_results["expert_activity"], save_path=heatmap_path, variant_name=ablation_variant) # Also copy to root assets save_expert_heatmap(eval_results["expert_activity"], save_path="assets/expert_heatmap.png", variant_name=ablation_variant) # Save checkpoint ckpt_path = os.path.join(save_checkpoint_dir, f"model_{ablation_variant}.pt") torch.save({ "model_state_dict": model.state_dict(), "config": config, "eval_results": {k: v for k, v in eval_results.items() if k != "expert_activity"} }, ckpt_path) logger.info(f"Checkpoint saved to {ckpt_path}") # Save evaluation summary summary_path = os.path.join(output_dir, f"summary_{ablation_variant}.json") with open(summary_path, "w") as f: json.dump({ "variant": ablation_variant, "loss": eval_results["loss"], "perplexity": eval_results["perplexity"], "bleu": eval_results["bleu"], "sample_predictions": eval_results["sample_predictions"], "sample_references": eval_results["sample_references"] }, f, indent=2) return eval_results