{ "cells": [ { "cell_type": "markdown", "id": "381e1824", "metadata": {}, "source": [ "Imports" ] }, { "cell_type": "code", "execution_count": 1, "id": "5a9e6fb4", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/home/ala-boussoffara/miniconda3/envs/mini-transformer/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", " from .autonotebook import tqdm as notebook_tqdm\n" ] } ], "source": [ "import contextlib\n", "import json\n", "import math\n", "import os\n", "from datetime import datetime\n", "from pathlib import Path\n", "from typing import Any\n", "\n", "import matplotlib.pyplot as plt\n", "import torch\n", "import torch.nn as nn\n", "from datasets import load_from_disk\n", "from hydra import compose, initialize\n", "from hydra.core.global_hydra import GlobalHydra\n", "from hydra.utils import to_absolute_path\n", "from omegaconf import OmegaConf\n", "from torch.nn.utils import clip_grad_norm_\n", "from torch.utils.data import DataLoader\n", "from tqdm import tqdm\n", "from transformers import PreTrainedTokenizerFast\n", "\n", "from mini_transformer import BasicEncoderDecoderTransformer\n", "from mini_transformer.configs import ModelCfg, TokenizerCfg, TrainAppCfg\n", "from mini_transformer.utils import (\n", " check_tokenizer_model_compatibility,\n", " debug_transformer_forward,\n", " make_worker_init_fn,\n", " set_global_seed,\n", ")" ] }, { "cell_type": "markdown", "id": "512d7a41", "metadata": {}, "source": [ "Modifying working directory" ] }, { "cell_type": "code", "execution_count": 2, "id": "d1c043be", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Working directory set to: /mnt/shared/Local/Projets/Mini-Transformer\n" ] } ], "source": [ "project_root = Path(__file__).resolve().parent if \"__file__\" in globals() else Path.cwd()\n", "os.chdir(project_root)\n", "print(f\"Working directory set to: {Path.cwd()}\")" ] }, { "cell_type": "markdown", "id": "9b3f89d1", "metadata": {}, "source": [ "Initiliazation" ] }, { "cell_type": "code", "execution_count": 3, "id": "2bd6dd82", "metadata": {}, "outputs": [], "source": [ "GlobalHydra.instance().clear()\n", "initialize(config_path=\"./configs\", version_base=None)\n", "cfg = compose(\n", " config_name=\"train_mode\",\n", " overrides=[\n", " \"model=small_model\",\n", " \"tokenizer=bpe_8k\",\n", " \"dataset=medium_dataset\",\n", " \"trainer.batch_size=32\",\n", " \"trainer.gradient_accumulation=16\",\n", " \"trainer.epochs=10\",\n", " \"trainer.resume=best\",\n", " \"model.best_checkpoint_path=trained_models/small_model_v1/checkpoints/best/checkpoint.pt\",\n", " ],\n", ")\n", "scfg_temp = OmegaConf.merge(OmegaConf.structured(TrainAppCfg), cfg)\n", "scfg: TrainAppCfg = OmegaConf.to_object(scfg_temp)\n", "set_global_seed(scfg.runtime.seed)\n", "model_cfg = ModelCfg(**vars(scfg.model))\n", "tokenizer_cfg = TokenizerCfg(**vars(scfg.tokenizer))\n", "\n", "# --- Compatibility check ---\n", "check_tokenizer_model_compatibility(model_cfg, tokenizer_cfg)" ] }, { "cell_type": "markdown", "id": "24caf2de", "metadata": {}, "source": [ "Preparing data" ] }, { "cell_type": "code", "execution_count": 4, "id": "8d60e085", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Number of training samples: 1826016\n", "Number of training batches: 57063\n" ] } ], "source": [ "dataset = load_from_disk(cfg.dataset.path)\n", "\n", "train_dataset = dataset[\"train\"]\n", "validation_dataset = dataset[\"validation\"]\n", "test_dataset = dataset[\"test\"]\n", "\n", "\n", "def tuple_collate_fn(batch):\n", " # batch is a list of dicts with 'en' and 'fr' keys\n", " src = [item[\"en\"] for item in batch]\n", " tgt = [item[\"fr\"] for item in batch]\n", " return src, tgt\n", "\n", "\n", "dataloader_generator = torch.Generator().manual_seed(scfg.runtime.seed)\n", "worker_init = make_worker_init_fn(scfg.runtime.seed)\n", "num_workers = int(getattr(scfg.trainer, \"num_workers\", 0))\n", "pin_memory = bool(getattr(scfg.trainer, \"pin_memory\", scfg.runtime.device.startswith(\"cuda\")))\n", "persistent_workers = num_workers > 0\n", "\n", "train_dataloader = DataLoader(\n", " train_dataset,\n", " batch_size=scfg.trainer.batch_size,\n", " shuffle=scfg.trainer.shuffle,\n", " collate_fn=tuple_collate_fn,\n", " generator=dataloader_generator,\n", " worker_init_fn=worker_init,\n", " num_workers=num_workers,\n", " pin_memory=pin_memory,\n", " persistent_workers=persistent_workers,\n", ")\n", "validation_dataloader = DataLoader(\n", " validation_dataset,\n", " batch_size=scfg.trainer.batch_size,\n", " shuffle=False,\n", " collate_fn=tuple_collate_fn,\n", " generator=dataloader_generator,\n", " worker_init_fn=worker_init,\n", " num_workers=num_workers,\n", " pin_memory=pin_memory,\n", " persistent_workers=persistent_workers,\n", ")\n", "test_dataloader = DataLoader(\n", " test_dataset,\n", " batch_size=scfg.trainer.batch_size,\n", " shuffle=False,\n", " collate_fn=tuple_collate_fn,\n", " generator=dataloader_generator,\n", " worker_init_fn=worker_init,\n", " num_workers=num_workers,\n", " pin_memory=pin_memory,\n", " persistent_workers=persistent_workers,\n", ")\n", "\n", "print(\"Number of training samples:\", len(train_dataset))\n", "print(\"Number of training batches:\", len(train_dataloader))" ] }, { "cell_type": "markdown", "id": "96c84adf", "metadata": {}, "source": [ "Model and tokenizer loading" ] }, { "cell_type": "code", "execution_count": 5, "id": "f9da0dab", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Total number of parameters: 25204736\n" ] } ], "source": [ "device = torch.device(scfg.runtime.device)\n", "model = BasicEncoderDecoderTransformer(model_cfg).to(device)\n", "\n", "\n", "# xavier initialization\n", "def init_weights(m):\n", " if isinstance(m, nn.Linear | nn.Embedding):\n", " torch.nn.init.xavier_uniform_(m.weight)\n", " elif isinstance(m, nn.LayerNorm):\n", " m.bias.data.fill_(0)\n", " m.weight.data.fill_(1.0)\n", " if isinstance(m, nn.Linear) and m.bias is not None:\n", " m.bias.data.fill_(0)\n", "\n", "\n", "model.apply(init_weights)\n", "\n", "tokenizer = PreTrainedTokenizerFast(\n", " tokenizer_file=tokenizer_cfg.path,\n", " bos_token=tokenizer_cfg.bos_token,\n", " eos_token=tokenizer_cfg.eos_token,\n", " unk_token=tokenizer_cfg.unk_token,\n", " pad_token=tokenizer_cfg.pad_token,\n", " model_max_length=tokenizer_cfg.max_seq_len,\n", ")\n", "\n", "criterion = nn.CrossEntropyLoss(\n", " ignore_index=tokenizer_cfg.pad_id, label_smoothing=scfg.trainer.label_smoothing\n", ")\n", "\n", "# print total number of parameters\n", "total_params = sum(p.numel() for p in model.parameters())\n", "print(f\"Total number of parameters: {total_params}\")" ] }, { "cell_type": "markdown", "id": "608a205a", "metadata": {}, "source": [ "Optimizer setup" ] }, { "cell_type": "code", "execution_count": 6, "id": "372dea7e", "metadata": {}, "outputs": [], "source": [ "def setup_optimizer(model: nn.Module, cfg: Any) -> torch.optim.Optimizer:\n", " \"\"\"Create an AdamW optimizer with decoupled weight decay groups.\"\"\"\n", "\n", " decay_params: set[str] = set()\n", " no_decay_params: set[str] = set()\n", "\n", " blacklist_weight_modules = (nn.LayerNorm,)\n", "\n", " for module_name, module in model.named_modules():\n", " for param_name, param in module.named_parameters(recurse=False):\n", " if not param.requires_grad:\n", " continue\n", "\n", " full_name = f\"{module_name}.{param_name}\" if module_name else param_name\n", "\n", " if param_name.endswith(\"bias\"):\n", " no_decay_params.add(full_name)\n", " elif isinstance(module, blacklist_weight_modules):\n", " no_decay_params.add(full_name)\n", " elif param.ndim == 1:\n", " no_decay_params.add(full_name)\n", " else:\n", " decay_params.add(full_name)\n", "\n", " for name, param in model.named_parameters():\n", " if not param.requires_grad:\n", " continue\n", " lowered = name.lower()\n", " if \"embedding\" in lowered or (\"pos\" in lowered and \"embed\" in lowered):\n", " decay_params.discard(name)\n", " no_decay_params.add(name)\n", "\n", " param_dict = {name: param for name, param in model.named_parameters() if param.requires_grad}\n", " decay_group = [param_dict[name] for name in sorted(decay_params) if name in param_dict]\n", " no_decay_group = [param_dict[name] for name in sorted(no_decay_params) if name in param_dict]\n", "\n", " optim_cfg = cfg.optim\n", " param_groups = []\n", " if decay_group:\n", " param_groups.append({\"params\": decay_group, \"weight_decay\": optim_cfg.weight_decay})\n", " if no_decay_group:\n", " param_groups.append({\"params\": no_decay_group, \"weight_decay\": 0.0})\n", "\n", " if not param_groups:\n", " raise ValueError(\"No trainable parameters found when setting up the optimizer.\")\n", "\n", " return torch.optim.AdamW(\n", " param_groups,\n", " lr=optim_cfg.lr,\n", " betas=optim_cfg.betas,\n", " eps=optim_cfg.eps,\n", " )\n", "\n", "\n", "optimizer = setup_optimizer(model, scfg)" ] }, { "cell_type": "markdown", "id": "f64a8b38", "metadata": {}, "source": [ "AMP and scaler setup" ] }, { "cell_type": "code", "execution_count": 7, "id": "4ad52841", "metadata": {}, "outputs": [], "source": [ "# 1) Decide whether AMP is on and which dtype to use\n", "use_amp = scfg.trainer.mixed_precision and (scfg.trainer.precision in [\"fp16\", \"bf16\"])\n", "\n", "# 2) Pick the autocast dtype from precision\n", "amp_dtype = torch.float16 if scfg.trainer.precision == \"fp16\" else torch.bfloat16\n", "\n", "# 3) Build the context manager:\n", "# - If AMP is on → use autocast(dtype=...)\n", "# - If AMP is off → use a nullcontext (does nothing)\n", "autocast_ctx = (\n", " torch.amp.autocast(dtype=amp_dtype, device_type=scfg.runtime.device)\n", " if use_amp\n", " else contextlib.nullcontext()\n", ")\n", "\n", "# 4) Create the GradScaler:\n", "# - enabled=True only for fp16 (needs loss scaling)\n", "# - for bf16/fp32 it becomes a no-op automatically\n", "scaler = torch.amp.GradScaler(\n", " enabled=(scfg.trainer.mixed_precision and scfg.trainer.precision == \"fp16\")\n", ")" ] }, { "cell_type": "markdown", "id": "34e29854", "metadata": {}, "source": [ "Scheduler setup" ] }, { "cell_type": "code", "execution_count": 8, "id": "479390a8", "metadata": {}, "outputs": [], "source": [ "def setup_scheduler(\n", " optimizer: torch.optim.Optimizer, trainer_cfg: Any, sched_cfg: Any, steps_per_epoch: int\n", ") -> torch.optim.lr_scheduler.LambdaLR:\n", " \"\"\"Create a cosine LR scheduler driven by warmup/min LR ratios.\"\"\"\n", " total_steps = max(1, int(trainer_cfg.epochs) * max(1, int(steps_per_epoch)))\n", " warmup_ratio = max(0.0, float(getattr(sched_cfg, \"warmup_ratio\", 0.0)))\n", " min_lr_ratio = max(0.0, float(getattr(sched_cfg, \"min_lr_ratio\", 0.0)))\n", " hold_min = bool(getattr(sched_cfg, \"hold_min\", True))\n", " max_lr_ratio = max(0.0, float(getattr(sched_cfg, \"max_lr_ratio\", 1.0)))\n", " warmup_steps = int(round(total_steps * warmup_ratio))\n", " warmup_steps = max(0, min(warmup_steps, max(0, total_steps - 1)))\n", " decay_steps = max(1, total_steps - warmup_steps)\n", " for group in optimizer.param_groups:\n", " base_lr = float(group.get(\"initial_lr\", group[\"lr\"]))\n", " group[\"initial_lr\"] = base_lr\n", " group[\"min_lr\"] = base_lr * min_lr_ratio\n", " group[\"max_lr\"] = base_lr * max_lr_ratio\n", "\n", " def lr_lambda(step: int) -> float:\n", " if warmup_steps > 0 and step < warmup_steps:\n", " warmup_progress = step / max(1, warmup_steps)\n", " warmup_progress = max(0.0, min(1.0, warmup_progress))\n", " return min_lr_ratio + (1.0 - min_lr_ratio) * warmup_progress\n", " if decay_steps <= 0:\n", " return 1.0\n", " progress = (step - warmup_steps) / decay_steps\n", " progress = max(0.0, min(1.0, progress))\n", " cosine = 0.5 * (1.0 + math.cos(math.pi * progress))\n", " factor = min_lr_ratio + (1.0 - min_lr_ratio) * cosine\n", " if progress >= 1.0 and not hold_min:\n", " return 0.0\n", " return factor\n", "\n", " return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n", "\n", "\n", "steps_per_epoch = math.ceil(len(train_dataloader) / max(1, scfg.trainer.gradient_accumulation))\n", "scheduler = setup_scheduler(optimizer, scfg.trainer, scfg.sched, steps_per_epoch=steps_per_epoch)" ] }, { "cell_type": "markdown", "id": "db51c02a", "metadata": {}, "source": [ "Loading checkpoints" ] }, { "cell_type": "code", "execution_count": 9, "id": "a9843809", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Loading checkpoint from /mnt/shared/Local/Projets/Mini-Transformer/trained_models/small_model_v1/checkpoints/best/checkpoint.pt\n" ] } ], "source": [ "output_root = Path(to_absolute_path(scfg.runtime.output_dir))\n", "output_root.mkdir(parents=True, exist_ok=True)\n", "\n", "timestamp = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n", "run_dir = output_root / f\"{model_cfg.name}_{timestamp}\"\n", "suffix = 1\n", "while run_dir.exists():\n", " run_dir = output_root / f\"{model_cfg.name}_{timestamp}_{suffix}\"\n", " suffix += 1\n", "run_dir.mkdir(parents=True, exist_ok=True)\n", "\n", "checkpoints_root = run_dir / \"checkpoints\"\n", "latest_dir = checkpoints_root / \"latest\"\n", "best_dir = checkpoints_root / \"best\"\n", "latest_dir.mkdir(parents=True, exist_ok=True)\n", "best_dir.mkdir(parents=True, exist_ok=True)\n", "\n", "resume_latest_source = getattr(model_cfg, \"latest_checkpoint_path\", \"\")\n", "resume_best_source = getattr(model_cfg, \"best_checkpoint_path\", \"\")\n", "best_checkpoint_export_target = (\n", " Path(to_absolute_path(resume_best_source)) if resume_best_source else None\n", ")\n", "\n", "latest_checkpoint_path = latest_dir / \"checkpoint.pt\"\n", "best_checkpoint_path = best_dir / \"checkpoint.pt\"\n", "\n", "\n", "config_snapshot_path = run_dir / \"config.yaml\"\n", "OmegaConf.save(config=cfg, f=str(config_snapshot_path))\n", "\n", "manifest_path = run_dir / \"manifest.json\"\n", "\n", "resume_mode = getattr(scfg.trainer, \"resume\", \"start_over\")\n", "resume_mode = resume_mode.replace(\" \", \"_\").lower()\n", "if resume_mode not in {\"start_over\", \"best\", \"latest\"}:\n", " print(f\"Unknown resume mode '{resume_mode}', defaulting to 'start_over'.\")\n", " resume_mode = \"start_over\"\n", "\n", "resume_checkpoint_path = None\n", "if resume_mode == \"latest\" and resume_latest_source:\n", " resume_checkpoint_path = Path(to_absolute_path(resume_latest_source))\n", "elif resume_mode == \"best\" and resume_best_source:\n", " resume_checkpoint_path = Path(to_absolute_path(resume_best_source))\n", "\n", "manifest_data = {\n", " \"run_dir\": str(run_dir),\n", " \"model\": model_cfg.name,\n", " \"resume\": {\n", " \"mode\": resume_mode,\n", " \"source\": str(resume_checkpoint_path) if resume_checkpoint_path else None,\n", " \"status\": \"pending\",\n", " },\n", " \"latest\": None,\n", " \"best\": None,\n", "}\n", "manifest_path.write_text(json.dumps(manifest_data, indent=2))\n", "\n", "start_epoch = 0\n", "completed_optimizer_steps = 0\n", "train_losses: list[float] = []\n", "val_losses: list[float] = []\n", "learning_rates: list[float] = []\n", "resume_batch_idx = 0\n", "resume_running_loss = 0.0\n", "best_val_loss = float(\"inf\")\n", "best_val_epoch = -1\n", "\n", "loaded_checkpoint = None\n", "if resume_checkpoint_path and resume_checkpoint_path.is_file():\n", " print(f\"Loading checkpoint from {resume_checkpoint_path}\")\n", " loaded_checkpoint = torch.load(resume_checkpoint_path, map_location=\"cpu\", weights_only=False)\n", " model_state = loaded_checkpoint.get(\"model_state_dict\")\n", " if model_state is not None:\n", " model.load_state_dict(model_state, strict=False)\n", " optim_state = loaded_checkpoint.get(\"optimizer_state_dict\")\n", " if optim_state is not None:\n", " optimizer.load_state_dict(optim_state)\n", " sched_state = loaded_checkpoint.get(\"scheduler_state_dict\")\n", " if sched_state is not None:\n", " scheduler.load_state_dict(sched_state)\n", " scaler_state = loaded_checkpoint.get(\"scaler_state_dict\")\n", " if scaler_state is not None:\n", " scaler.load_state_dict(scaler_state)\n", " train_losses = list(loaded_checkpoint.get(\"train_losses\", train_losses))\n", " val_losses = list(loaded_checkpoint.get(\"val_losses\", val_losses))\n", " learning_rates = list(loaded_checkpoint.get(\"learning_rates\", learning_rates))\n", " start_epoch = int(loaded_checkpoint.get(\"epoch\", start_epoch))\n", " completed_optimizer_steps = int(\n", " loaded_checkpoint.get(\"optimizer_step\", completed_optimizer_steps)\n", " )\n", " resume_batch_idx = int(loaded_checkpoint.get(\"batch_idx\", resume_batch_idx))\n", " resume_running_loss = float(loaded_checkpoint.get(\"running_loss\", resume_running_loss))\n", " best_val_loss = float(\n", " loaded_checkpoint.get(\n", " \"best_val_loss\",\n", " loaded_checkpoint.get(\"best_train_loss\", best_val_loss),\n", " )\n", " )\n", " best_val_epoch = int(\n", " loaded_checkpoint.get(\n", " \"best_val_epoch\",\n", " loaded_checkpoint.get(\"best_train_epoch\", best_val_epoch),\n", " )\n", " )\n", " for state in optimizer.state.values():\n", " for key, value in state.items():\n", " if isinstance(value, torch.Tensor):\n", " state[key] = value.to(device)\n", " model.to(device)\n", "else:\n", " if resume_checkpoint_path:\n", " print(f\"Checkpoint '{resume_checkpoint_path}' not found. Starting fresh.\")\n", " else:\n", " print(\"Starting fresh (no resume checkpoint provided).\")\n", "\n", "if loaded_checkpoint is not None:\n", " manifest_data[\"resume\"][\"status\"] = \"loaded\"\n", "elif resume_checkpoint_path:\n", " manifest_data[\"resume\"][\"status\"] = \"missing\"\n", "else:\n", " manifest_data[\"resume\"][\"status\"] = \"fresh\"\n", "manifest_path.write_text(json.dumps(manifest_data, indent=2))\n", "\n", "\n", "def _update_manifest(section: str, payload: dict) -> None:\n", " manifest_data[section] = payload\n", " manifest_path.write_text(json.dumps(manifest_data, indent=2))\n", "\n", "\n", "if loaded_checkpoint is not None:\n", " torch.save(loaded_checkpoint, latest_checkpoint_path)\n", " _update_manifest(\n", " \"latest\",\n", " {\n", " \"path\": str(latest_checkpoint_path),\n", " \"epoch\": start_epoch,\n", " \"optimizer_step\": completed_optimizer_steps,\n", " \"batch_idx\": resume_batch_idx,\n", " \"running_loss\": resume_running_loss,\n", " \"timestamp\": datetime.now().isoformat(),\n", " },\n", " )\n", " if resume_mode == \"best\":\n", " torch.save(loaded_checkpoint, best_checkpoint_path)\n", " _update_manifest(\n", " \"best\",\n", " {\n", " \"path\": str(best_checkpoint_path),\n", " \"epoch\": start_epoch,\n", " \"optimizer_step\": completed_optimizer_steps,\n", " \"val_loss\": loaded_checkpoint.get(\n", " \"best_val_loss\", loaded_checkpoint.get(\"best_train_loss\")\n", " ),\n", " \"train_loss\": loaded_checkpoint.get(\"metadata\", {}).get(\"train_loss\"),\n", " \"timestamp\": datetime.now().isoformat(),\n", " },\n", " )\n", "\n", "\n", "def save_checkpoint(\n", " epoch_to_resume: int,\n", " optimizer_step: int,\n", " *,\n", " batch_idx: int,\n", " running_loss_value: float,\n", " metadata: dict[str, Any] | None = None,\n", ") -> dict:\n", " checkpoint = {\n", " \"epoch\": epoch_to_resume,\n", " \"optimizer_step\": optimizer_step,\n", " \"batch_idx\": batch_idx,\n", " \"running_loss\": running_loss_value,\n", " \"model_state_dict\": model.state_dict(),\n", " \"optimizer_state_dict\": optimizer.state_dict(),\n", " \"scheduler_state_dict\": scheduler.state_dict(),\n", " \"scaler_state_dict\": scaler.state_dict(),\n", " \"train_losses\": train_losses,\n", " \"val_losses\": val_losses,\n", " \"learning_rates\": learning_rates,\n", " \"best_val_loss\": best_val_loss,\n", " \"best_val_epoch\": best_val_epoch,\n", " \"metadata\": metadata or {},\n", " }\n", " torch.save(checkpoint, latest_checkpoint_path)\n", " _update_manifest(\n", " \"latest\",\n", " {\n", " \"path\": str(latest_checkpoint_path),\n", " \"epoch\": epoch_to_resume,\n", " \"optimizer_step\": optimizer_step,\n", " \"batch_idx\": batch_idx,\n", " \"running_loss\": running_loss_value,\n", " \"metadata\": metadata or {},\n", " \"timestamp\": datetime.now().isoformat(),\n", " },\n", " )\n", " return checkpoint" ] }, { "cell_type": "code", "execution_count": 10, "id": "ab23bf3c", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "'src_ids = torch.randint(0, scfg.tokenizer.vocab_size, (2, 50)).to(device, dtype=torch.long)\\ntgt_ids = torch.randint(0, scfg.tokenizer.vocab_size, (2, 50)).to(device, dtype=torch.long)\\nsrc_padd_mask = (src_ids == scfg.tokenizer.pad_id).to(device, dtype=torch.bool)\\ntgt_padd_mask = (tgt_ids == scfg.tokenizer.pad_id).to(device, dtype=torch.bool)'" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# src,tgt=next(iter(train_dataloader))\n", "src = [\"it can be a very complicated thing, the ocean\"]\n", "tgt = [\"ca peut être très compliqué, l'océan\"]\n", "encoded_src = tokenizer(\n", " src, padding=True, truncation=True, return_tensors=\"pt\", add_special_tokens=False\n", ")\n", "encoded_tgt = tokenizer(tgt, padding=True, truncation=True, return_tensors=\"pt\")\n", "src_ids = encoded_src[\"input_ids\"].to(device)\n", "tgt_ids = encoded_tgt[\"input_ids\"].to(device)\n", "src_padd_mask = (src_ids == scfg.tokenizer.pad_id).to(device, dtype=torch.bool)\n", "tgt_padd_mask = (tgt_ids == scfg.tokenizer.pad_id).to(device, dtype=torch.bool)\n", "# initilaize src_ids and tgt_ids randomly with big lengths\n", "\"\"\"src_ids = torch.randint(0, scfg.tokenizer.vocab_size, (2, 50)).to(device, dtype=torch.long)\n", "tgt_ids = torch.randint(0, scfg.tokenizer.vocab_size, (2, 50)).to(device, dtype=torch.long)\n", "src_padd_mask = (src_ids == scfg.tokenizer.pad_id).to(device, dtype=torch.bool)\n", "tgt_padd_mask = (tgt_ids == scfg.tokenizer.pad_id).to(device, dtype=torch.bool)\"\"\"" ] }, { "cell_type": "code", "execution_count": 11, "id": "689fe371", "metadata": {}, "outputs": [], "source": [ "# Optional transformer debugging / attention capture\n", "debug_forward_options = {\n", " \"enabled\": False, # flip to True to run debug_transformer_forward\n", " \"show_attention\": True, # show attention plots interactively\n", " \"save_attention\": False, # persist attention artefacts\n", " \"attention_layers\": None, # e.g. [0, 1] to restrict encoder/decoder layers\n", " \"attention_heads\": None, # e.g. [0, 3] to filter heads\n", " \"attention_types\": (\"enc_self\", \"dec_self\", \"dec_cross\"),\n", " \"average_heads\": False, # average selected heads into a single map\n", " \"sample_index\": 0, # which row within the batch to inspect\n", " \"attention_figsize\": (2.4, 2.4),\n", " \"skip_special_tokens\": False,\n", " \"save_dir\": None, # override directory; defaults to run_dir/debug/attention\n", " \"return_maps\": False, # set True to receive raw attention tensors\n", "}\n", "\n", "if debug_forward_options[\"save_dir\"]:\n", " debug_default_attention_dir = Path(debug_forward_options[\"save_dir\"])\n", " if not debug_default_attention_dir.is_absolute():\n", " debug_default_attention_dir = run_dir / debug_default_attention_dir\n", "else:\n", " debug_default_attention_dir = run_dir / \"debug\" / \"attention\"\n", "debug_default_attention_dir.mkdir(parents=True, exist_ok=True)" ] }, { "cell_type": "code", "execution_count": 12, "id": "af2e9c6a", "metadata": {}, "outputs": [], "source": [ "if debug_forward_options.get(\"enabled\", False):\n", " sample_index = int(debug_forward_options.get(\"sample_index\", 0))\n", " sample_index = max(0, min(sample_index, src_ids.size(0) - 1))\n", " debug_result = debug_transformer_forward(\n", " model=model,\n", " tokenizer=tokenizer,\n", " src_ids=src_ids,\n", " tgt_ids=tgt_ids,\n", " pad_id=scfg.tokenizer.pad_id,\n", " device=device,\n", " batch_index=0,\n", " sample_index=sample_index,\n", " logits=None,\n", " show_attention=debug_forward_options.get(\"show_attention\", False),\n", " attention_layers=debug_forward_options.get(\"attention_layers\"),\n", " attention_heads=debug_forward_options.get(\"attention_heads\"),\n", " attention_types=debug_forward_options.get(\n", " \"attention_types\", (\"enc_self\", \"dec_self\", \"dec_cross\")\n", " ),\n", " average_heads=debug_forward_options.get(\"average_heads\", False),\n", " save_attention=debug_forward_options.get(\"save_attention\", False),\n", " save_dir=debug_default_attention_dir,\n", " run_dir=run_dir,\n", " attention_figsize=debug_forward_options.get(\"attention_figsize\", (2.4, 2.4)),\n", " skip_special_tokens=debug_forward_options.get(\"skip_special_tokens\", False),\n", " return_maps=debug_forward_options.get(\"return_maps\", False),\n", " )\n", " print(\"[debug] transformer forward summary saved to\", debug_result.get(\"attention\", {}))\n", "else:\n", " debug_result = None" ] }, { "cell_type": "code", "execution_count": 13, "id": "439d4094", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "Epoch 10 [train]: 0%| | 33/57063 [00:03<1:47:43, 8.82it/s]\n", "Epoch 10 [train]: 0%| | 33/57063 [00:03<1:47:43, 8.82it/s]\n" ] }, { "ename": "KeyboardInterrupt", "evalue": "", "output_type": "error", "traceback": [ "\u001b[31m---------------------------------------------------------------------------\u001b[39m", "\u001b[31mKeyboardInterrupt\u001b[39m Traceback (most recent call last)", "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[13]\u001b[39m\u001b[32m, line 59\u001b[39m\n\u001b[32m 57\u001b[39m loss = loss / grad_accum\n\u001b[32m 58\u001b[39m scaled_loss = scaler.scale(loss)\n\u001b[32m---> \u001b[39m\u001b[32m59\u001b[39m \u001b[43mscaled_loss\u001b[49m\u001b[43m.\u001b[49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 60\u001b[39m accum_counter += \u001b[32m1\u001b[39m\n\u001b[32m 62\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m accum_counter == grad_accum:\n", "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/mini-transformer/lib/python3.11/site-packages/torch/_tensor.py:525\u001b[39m, in \u001b[36mTensor.backward\u001b[39m\u001b[34m(self, gradient, retain_graph, create_graph, inputs)\u001b[39m\n\u001b[32m 515\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[32m 516\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[32m 517\u001b[39m Tensor.backward,\n\u001b[32m 518\u001b[39m (\u001b[38;5;28mself\u001b[39m,),\n\u001b[32m (...)\u001b[39m\u001b[32m 523\u001b[39m inputs=inputs,\n\u001b[32m 524\u001b[39m )\n\u001b[32m--> \u001b[39m\u001b[32m525\u001b[39m \u001b[43mtorch\u001b[49m\u001b[43m.\u001b[49m\u001b[43mautograd\u001b[49m\u001b[43m.\u001b[49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m 526\u001b[39m \u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m=\u001b[49m\u001b[43minputs\u001b[49m\n\u001b[32m 527\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n", "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/mini-transformer/lib/python3.11/site-packages/torch/autograd/__init__.py:267\u001b[39m, in \u001b[36mbackward\u001b[39m\u001b[34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[39m\n\u001b[32m 262\u001b[39m retain_graph = create_graph\n\u001b[32m 264\u001b[39m \u001b[38;5;66;03m# The reason we repeat the same comment below is that\u001b[39;00m\n\u001b[32m 265\u001b[39m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[32m 266\u001b[39m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m267\u001b[39m \u001b[43m_engine_run_backward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m 268\u001b[39m \u001b[43m \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 269\u001b[39m \u001b[43m \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 270\u001b[39m \u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 271\u001b[39m \u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 272\u001b[39m \u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 273\u001b[39m \u001b[43m \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[32m 274\u001b[39m \u001b[43m \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[32m 275\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n", "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/mini-transformer/lib/python3.11/site-packages/torch/autograd/graph.py:744\u001b[39m, in \u001b[36m_engine_run_backward\u001b[39m\u001b[34m(t_outputs, *args, **kwargs)\u001b[39m\n\u001b[32m 742\u001b[39m unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)\n\u001b[32m 743\u001b[39m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[32m--> \u001b[39m\u001b[32m744\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mVariable\u001b[49m\u001b[43m.\u001b[49m\u001b[43m_execution_engine\u001b[49m\u001b[43m.\u001b[49m\u001b[43mrun_backward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[32m 745\u001b[39m \u001b[43m \u001b[49m\u001b[43mt_outputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\n\u001b[32m 746\u001b[39m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001b[39;00m\n\u001b[32m 747\u001b[39m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[32m 748\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m attach_logging_hooks:\n", "\u001b[31mKeyboardInterrupt\u001b[39m: " ] } ], "source": [ "model.to(device)\n", "model.enable_gradient_checkpointing(cfg.trainer.gradient_checkpointing)\n", "\n", "grad_accum = max(1, int(scfg.trainer.gradient_accumulation))\n", "save_every_steps = max(1, int(getattr(scfg.trainer, \"save_interval\", 1)))\n", "\n", "optimizer.zero_grad(set_to_none=True)\n", "\n", "optimizer_step = completed_optimizer_steps\n", "current_resume_batch = resume_batch_idx\n", "current_resume_running_loss = resume_running_loss\n", "\n", "for epoch in range(start_epoch, scfg.trainer.epochs):\n", " model.train()\n", " running_loss = current_resume_running_loss if epoch == start_epoch else 0.0\n", " batches_processed = 0\n", " last_batch_idx = -1\n", " accum_counter = 0\n", " for batch_idx, (src, tgt) in enumerate(\n", " tqdm(train_dataloader, desc=f\"Epoch {epoch + 1} [train]\")\n", " ):\n", " if epoch == start_epoch and batch_idx < current_resume_batch:\n", " continue\n", " if epoch == start_epoch and current_resume_batch and batch_idx == current_resume_batch:\n", " current_resume_batch = 0\n", " current_resume_running_loss = 0.0\n", "\n", " batches_processed += 1\n", " last_batch_idx = batch_idx\n", "\n", " encoded_src = tokenizer(\n", " src,\n", " padding=True,\n", " truncation=True,\n", " return_tensors=\"pt\",\n", " add_special_tokens=False,\n", " )\n", " encoded_tgt = tokenizer(\n", " tgt,\n", " padding=True,\n", " truncation=True,\n", " return_tensors=\"pt\",\n", " )\n", " tgt_ids = encoded_tgt[\"input_ids\"].to(device)\n", " decoder_in = tgt_ids[:, :-1]\n", " labels = tgt_ids[:, 1:]\n", "\n", " src_ids = encoded_src[\"input_ids\"].to(device)\n", " src_padd_mask = (encoded_src[\"attention_mask\"] == 0).to(device)\n", " tgt_padd_mask = decoder_in.eq(scfg.tokenizer.pad_id).to(device)\n", "\n", " with autocast_ctx:\n", " logits = model(src_ids, decoder_in, src_padd_mask, tgt_padd_mask)\n", " loss = criterion(logits.reshape(-1, logits.size(-1)), labels.reshape(-1))\n", "\n", " running_loss += loss.item()\n", " loss = loss / grad_accum\n", " scaled_loss = scaler.scale(loss)\n", " scaled_loss.backward()\n", " accum_counter += 1\n", "\n", " if accum_counter == grad_accum:\n", " scaler.unscale_(optimizer)\n", " clip_grad_norm_(model.parameters(), cfg.trainer.clip_grad_norm)\n", " scaler.step(optimizer)\n", " scaler.update()\n", " optimizer.zero_grad(set_to_none=True)\n", " scheduler.step()\n", " optimizer_step += 1\n", " if optimizer_step % save_every_steps == 0:\n", " save_checkpoint(\n", " epoch,\n", " optimizer_step,\n", " batch_idx=batch_idx + 1,\n", " running_loss_value=running_loss,\n", " )\n", " accum_counter = 0\n", "\n", " if accum_counter > 0:\n", " scaler.unscale_(optimizer)\n", " if accum_counter < grad_accum:\n", " adjust_factor = grad_accum / accum_counter\n", " for group in optimizer.param_groups:\n", " for param in group[\"params\"]:\n", " if param.grad is not None:\n", " param.grad.mul_(adjust_factor)\n", " clip_grad_norm_(model.parameters(), cfg.trainer.clip_grad_norm)\n", " scaler.step(optimizer)\n", " scaler.update()\n", " optimizer.zero_grad(set_to_none=True)\n", " scheduler.step()\n", " optimizer_step += 1\n", " if optimizer_step % save_every_steps == 0:\n", " save_checkpoint(\n", " epoch,\n", " optimizer_step,\n", " batch_idx=last_batch_idx + 1,\n", " running_loss_value=running_loss,\n", " )\n", " accum_counter = 0\n", "\n", " avg_train_loss = running_loss / len(train_dataloader)\n", " print(f\"Epoch [{epoch + 1}/{scfg.trainer.epochs}] Train Loss: {avg_train_loss:.4f}\")\n", "\n", " train_losses.append(avg_train_loss)\n", " last_lr = scheduler.get_last_lr()\n", " learning_rates.append(float(last_lr[0]) if last_lr else optimizer.param_groups[0][\"lr\"])\n", " latest_val_loss: float | None = None\n", " if epoch % scfg.trainer.eval_interval == 0:\n", " model.eval()\n", " val_loss = 0.0\n", " with torch.no_grad(), autocast_ctx:\n", " for src, tgt in tqdm(validation_dataloader, desc=f\"Epoch {epoch + 1} [val]\"):\n", " encoded_src = tokenizer(\n", " src,\n", " padding=True,\n", " truncation=True,\n", " return_tensors=\"pt\",\n", " add_special_tokens=False,\n", " )\n", " encoded_tgt = tokenizer(\n", " tgt,\n", " padding=True,\n", " truncation=True,\n", " return_tensors=\"pt\",\n", " )\n", "\n", " src_ids = encoded_src[\"input_ids\"].to(device)\n", " src_padd_mask = (encoded_src[\"attention_mask\"] == 0).to(device)\n", " tgt_ids = encoded_tgt[\"input_ids\"].to(device)\n", " decoder_in = tgt_ids[:, :-1]\n", " labels = tgt_ids[:, 1:]\n", " tgt_padd_mask = decoder_in.eq(scfg.tokenizer.pad_id).to(device)\n", "\n", " logits = model(src_ids, decoder_in, src_padd_mask, tgt_padd_mask)\n", " loss = criterion(logits.reshape(-1, logits.size(-1)), labels.reshape(-1))\n", " val_loss += loss.item()\n", "\n", " avg_val_loss = val_loss / max(1, len(validation_dataloader))\n", " print(f\"Epoch [{epoch + 1}/{scfg.trainer.epochs}] Val Loss: {avg_val_loss:.4f}\")\n", " val_losses.append(avg_val_loss)\n", " latest_val_loss = avg_val_loss\n", "\n", " model.train()\n", "\n", " is_new_best = latest_val_loss is not None and latest_val_loss < best_val_loss\n", " if is_new_best:\n", " best_val_loss = latest_val_loss\n", " best_val_epoch = epoch + 1\n", "\n", " epoch_metadata = {\"train_loss\": avg_train_loss}\n", " if latest_val_loss is not None:\n", " epoch_metadata[\"val_loss\"] = latest_val_loss\n", "\n", " latest_payload = save_checkpoint(\n", " epoch + 1,\n", " optimizer_step,\n", " batch_idx=0,\n", " running_loss_value=0.0,\n", " metadata=epoch_metadata,\n", " )\n", "\n", " if is_new_best and best_checkpoint_path is not None:\n", " torch.save(latest_payload, best_checkpoint_path)\n", " _update_manifest(\n", " \"best\",\n", " {\n", " \"path\": str(best_checkpoint_path),\n", " \"epoch\": best_val_epoch,\n", " \"optimizer_step\": optimizer_step,\n", " \"train_loss\": avg_train_loss,\n", " \"val_loss\": best_val_loss,\n", " \"metadata\": latest_payload.get(\"metadata\", {}),\n", " \"timestamp\": datetime.now().isoformat(),\n", " },\n", " )\n", " print(\n", " f\"New best checkpoint saved at {best_checkpoint_path} \"\n", " f\"(epoch {best_val_epoch}, val {best_val_loss:.4f}, train {avg_train_loss:.4f})\"\n", " )\n", "\n", " current_resume_batch = 0\n", " current_resume_running_loss = 0.0" ] }, { "cell_type": "markdown", "id": "d55a574f", "metadata": {}, "source": [ "Exporting" ] }, { "cell_type": "code", "execution_count": null, "id": "ad75105a", "metadata": {}, "outputs": [], "source": [ "\"\"\"if scfg.trainer.export_best_to_model_path and best_checkpoint_export_target:\n", " if best_checkpoint_path.exists():\n", " best_checkpoint_export_target.parent.mkdir(parents=True, exist_ok=True)\n", " shutil.copy2(best_checkpoint_path, best_checkpoint_export_target)\n", " print(f\"[export] Best checkpoint copied to {best_checkpoint_export_target}\")\n", " else:\n", " print(f\"[warn] No best checkpoint found at {best_checkpoint_path} to export.\")\"\"\"" ] }, { "cell_type": "markdown", "id": "4b21a8a3", "metadata": {}, "source": [ "Plotting" ] }, { "cell_type": "code", "execution_count": null, "id": "6875ddbe", "metadata": {}, "outputs": [], "source": [ "plots_dir = run_dir / \"plots\"\n", "plots_dir.mkdir(parents=True, exist_ok=True)\n", "\n", "if train_losses:\n", " epochs = list(range(1, len(train_losses) + 1))\n", " fig, ax = plt.subplots(figsize=(8, 5))\n", " ax.plot(epochs, train_losses, label=\"Training Loss\")\n", " if val_losses:\n", " val_epochs = list(range(1, len(val_losses) + 1))\n", " ax.plot(val_epochs, val_losses, label=\"Validation Loss\")\n", " ax.set_xlabel(\"Epochs\")\n", " ax.set_ylabel(\"Loss\")\n", " ax.set_title(\"Training, Validation Loss & Learning Rate\")\n", " lines = ax.get_lines()\n", " labels = [line.get_label() for line in lines]\n", " if learning_rates:\n", " lr_epochs = epochs[: len(learning_rates)]\n", " ax2 = ax.twinx()\n", " ax2.plot(\n", " lr_epochs, learning_rates, label=\"Learning Rate\", color=\"tab:green\", linestyle=\"--\"\n", " )\n", " ax2.set_ylabel(\"Learning Rate\")\n", " lr_lines = ax2.get_lines()\n", " lines = list(lines) + list(lr_lines)\n", " labels += [line.get_label() for line in lr_lines]\n", " ax.legend(lines, labels, loc=\"upper right\")\n", " plot_path = plots_dir / \"loss_curves.png\"\n", " fig.savefig(plot_path, dpi=150, bbox_inches=\"tight\")\n", " print(f\"[saved] Loss curves -> {plot_path}\")\n", " plt.show()\n", "else:\n", " print(\"No training history available to plot.\")" ] }, { "cell_type": "markdown", "id": "17bf7bc8", "metadata": {}, "source": [ "Test Dataset" ] }, { "cell_type": "code", "execution_count": null, "id": "39f61da8", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "Testing: 32%|███▏ | 101/313 [02:04<03:10, 1.11it/s]" ] } ], "source": [ "# test the model (extended: compute BLEU/chrF/ROUGE, perplexity, exact/token accuracy)\n", "import time\n", "\n", "import evaluate\n", "import torch.nn.functional as F\n", "from tqdm import tqdm\n", "\n", "# select some samples to test\n", "test_dataset = test_dataset.select(range(10000))\n", "test_dataloader = DataLoader(\n", " test_dataset,\n", " batch_size=scfg.trainer.batch_size,\n", " shuffle=False,\n", " collate_fn=tuple_collate_fn,\n", " generator=dataloader_generator,\n", " worker_init_fn=worker_init,\n", " num_workers=num_workers,\n", " pin_memory=pin_memory,\n", " persistent_workers=persistent_workers,\n", ")\n", "model.eval()\n", "test_loss = 0.0\n", "all_hyps = []\n", "all_refs = []\n", "latencies = []\n", "total_nll = 0.0\n", "total_tokens = 0\n", "exact_matches = 0\n", "token_matches = 0\n", "token_total = 0\n", "\n", "bleu_metric = evaluate.load(\"sacrebleu\")\n", "chrf_metric = evaluate.load(\"chrf\")\n", "rouge_metric = evaluate.load(\"rouge\")\n", "\n", "pad_id = (\n", " getattr(tokenizer_cfg, \"pad_id\", None)\n", " if \"tokenizer_cfg\" in globals()\n", " else getattr(model_cfg, \"pad_id\", None)\n", ")\n", "\n", "with torch.no_grad(), autocast_ctx:\n", " for src, tgt in tqdm(test_dataloader, desc=\"Testing\"):\n", " # tokenization (src and tgt are lists of strings per tuple_collate_fn)\n", " encoded_src = tokenizer(\n", " src,\n", " padding=True,\n", " truncation=True,\n", " return_tensors=\"pt\",\n", " add_special_tokens=False,\n", " )\n", " encoded_tgt = tokenizer(\n", " tgt,\n", " padding=True,\n", " truncation=True,\n", " return_tensors=\"pt\",\n", " )\n", "\n", " src_ids = encoded_src[\"input_ids\"].to(device)\n", " src_padd_mask = (encoded_src[\"attention_mask\"] == 0).to(device)\n", " tgt_ids = encoded_tgt[\"input_ids\"].to(device)\n", " decoder_in = tgt_ids[:, :-1]\n", " labels = tgt_ids[:, 1:]\n", " tgt_padd_mask = decoder_in.eq(pad_id)\n", "\n", " # forward for test loss (uses configured criterion)\n", " logits = model(src_ids, decoder_in, src_padd_mask, tgt_padd_mask)\n", " loss = criterion(logits.reshape(-1, logits.size(-1)), labels.reshape(-1))\n", " test_loss += loss.item()\n", "\n", " # compute token-summed NLL for perplexity\n", " V = logits.size(-1)\n", " logits_flat = logits.reshape(-1, V) # Changed from view to reshape\n", " labels_flat = labels.reshape(-1) # Changed from view to reshape\n", " nll_sum = F.cross_entropy(logits_flat, labels_flat, ignore_index=pad_id, reduction=\"sum\")\n", " non_pad = (labels_flat != pad_id).sum().item()\n", " total_nll += nll_sum.item()\n", " total_tokens += non_pad\n", "\n", " # generation and latency\n", " t0 = time.perf_counter()\n", " out_ids = model.generate(\n", " src_ids, src_padd_mask, max_new_tokens=128, temperature=1.0, do_sample=False\n", " )\n", " t1 = time.perf_counter()\n", " latencies.append((t1 - t0) / max(1, src_ids.size(0)))\n", "\n", " hyps = tokenizer.batch_decode(out_ids.cpu(), skip_special_tokens=True)\n", " all_hyps.extend(hyps)\n", " all_refs.extend(tgt)\n", "\n", " # exact match and token-level accuracy\n", " for h, r in zip(hyps, tgt, strict=False):\n", " if h.strip().lower() == r.strip().lower():\n", " exact_matches += 1\n", "\n", " hyp_tok = tokenizer(hyps, padding=True, truncation=True, return_tensors=\"pt\")[\"input_ids\"]\n", " ref_tok = tgt_ids[:, : hyp_tok.size(1)].cpu()\n", " matches = (hyp_tok[:, : ref_tok.size(1)] == ref_tok).cpu()\n", " token_matches += int(matches.sum().item())\n", " token_total += int(ref_tok.numel())\n", "\n", "# aggregate and report\n", "avg_test_loss = test_loss / max(1, len(test_dataloader))\n", "avg_latency = sum(latencies) / len(latencies) if latencies else None\n", "p95_latency = sorted(latencies)[int(0.95 * len(latencies))] if latencies else None\n", "avg_nll = (total_nll / total_tokens) if total_tokens > 0 else float(\"nan\")\n", "perplexity = math.exp(avg_nll) if total_tokens > 0 and avg_nll < 100 else float(\"inf\")\n", "exact_match_rate = exact_matches / len(all_hyps) if all_hyps else 0.0\n", "token_accuracy = token_matches / token_total if token_total > 0 else 0.0\n", "\n", "# corpus-level metrics\n", "refs_for_metric = [[r] for r in all_refs]\n", "bleu_res = bleu_metric.compute(predictions=all_hyps, references=refs_for_metric)\n", "chrf_res = chrf_metric.compute(predictions=all_hyps, references=refs_for_metric)\n", "rouge_res = rouge_metric.compute(predictions=all_hyps, references=[r[0] for r in refs_for_metric])\n", "\n", "print(f\"Test Loss: {avg_test_loss:.4f}\")\n", "print(f\"Generated {len(all_hyps)} hypotheses\")\n", "print(f\"Avg latency per sentence: {avg_latency:.4f} s\")\n", "print(f\"p95 latency per sentence: {p95_latency:.4f} s\")\n", "print(f\"Perplexity (corpus): {perplexity:.3f}\")\n", "print(f\"Exact match rate: {exact_match_rate:.3%}\")\n", "print(f\"Token accuracy: {token_accuracy:.3%}\")\n", "print(\"BLEU:\", bleu_res)\n", "print(\"chrF:\", chrf_res)\n", "print(\"ROUGE:\", rouge_res)\n", "\n", "# Prepare test metrics\n", "test_metrics = {\n", " \"test_loss\": avg_test_loss,\n", " \"perplexity\": perplexity,\n", " \"avg_latency_s\": avg_latency,\n", " \"p95_latency_s\": p95_latency,\n", " \"exact_match_rate\": exact_match_rate,\n", " \"token_accuracy\": token_accuracy,\n", " \"bleu\": bleu_res,\n", " \"chrf\": chrf_res,\n", " \"rouge\": rouge_res,\n", " \"timestamp\": datetime.now().isoformat(),\n", " \"model_name\": model_cfg.name,\n", " \"dataset_size\": len(test_dataset),\n", " \"checkpoint_path\": str(resume_checkpoint_path) if resume_checkpoint_path else None,\n", "}\n", "\n", "# Save to manifest\n", "manifest_data[\"test\"] = test_metrics\n", "\n", "# Save metrics to a separate file\n", "metrics_file = run_dir / \"test_metrics.json\"\n", "metrics_file.write_text(json.dumps(test_metrics, indent=2))\n", "print(f\"[saved] Test metrics -> {metrics_file}\")\n", "\n", "manifest_path.write_text(json.dumps(manifest_data, indent=2))" ] }, { "cell_type": "markdown", "id": "f8652b1c", "metadata": {}, "source": [ "Inference" ] }, { "cell_type": "code", "execution_count": null, "id": "45a55eaf", "metadata": {}, "outputs": [], "source": [ "\"\"\"from scripts.inference import inference\n", "GlobalHydra.instance().clear()\n", "initialize(config_path=\"./configs\", version_base=None)\n", "cfg = compose(config_name=\"infer_mode\")\n", "print(inference(cfg))\"\"\"" ] } ], "metadata": { "kernelspec": { "display_name": "mini-transformer", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.13" } }, "nbformat": 4, "nbformat_minor": 5 }