diff --git a/README.md b/README.md index 7a7ebb1..76005ac 100644 --- a/README.md +++ b/README.md @@ -83,8 +83,8 @@ scripts/ src/ │ ├── data_preparation/ -├── model_files/ -├── pre_training/ +├── models/ +├── pretraining/ ├── finetuning/ └── paths.py ``` diff --git a/src/model_files/__init__.py b/action similarity index 100% rename from src/model_files/__init__.py rename to action diff --git a/configs/config.py b/configs/config.py deleted file mode 100644 index 016271d..0000000 --- a/configs/config.py +++ /dev/null @@ -1,26 +0,0 @@ -import torch -from dataclasses import dataclass -from src.paths import CHECKPOINT_DIR - -checkpoint_dir = CHECKPOINT_DIR - -@dataclass -class ModelConfig: - vocab_size: int = 8192 - emb_dim: int = 512 - num_blocks: int = 8 - head_count: int = 8 - seq_length: int = 512 - ffn_multiple: int = 128 - - -@dataclass -class TrainingConfig: - token_count: int = 470_000_000 - batch_size: int = 64 - learning_rate: float = 3e-4 - betas: tuple = (0.9, 0.95) - eps: float = 1e-8 - device: str = "cuda" if torch.cuda.is_available() else "cpu" - checkpoint_dir: str = checkpoint_dir - \ No newline at end of file diff --git a/configs/model.py b/configs/model.py new file mode 100644 index 0000000..87511c4 --- /dev/null +++ b/configs/model.py @@ -0,0 +1,39 @@ +from dataclasses import dataclass + + +@dataclass +class ModelConfig: + """Hyperparameters for the BetterGPT model architecture. + + Validated at construction time: emb_dim must be divisible by head_count, + and the resulting head_dim must be even (required by RoPE). + """ + vocab_size: int = 8192 + emb_dim: int = 512 + num_blocks: int = 8 + head_count: int = 8 + seq_length: int = 512 + ffn_multiple: int = 128 + + def __post_init__(self): + if self.vocab_size <= 0: + raise ValueError(f"vocab_size must be > 0, got {self.vocab_size}") + if self.emb_dim <= 0: + raise ValueError(f"emb_dim must be > 0, got {self.emb_dim}") + if self.head_count <= 0: + raise ValueError(f"head_count must be > 0, got {self.head_count}") + if self.emb_dim % self.head_count != 0: + raise ValueError( + f"emb_dim ({self.emb_dim}) must be divisible by head_count ({self.head_count})" + ) + head_dim = self.emb_dim // self.head_count + if head_dim % 2 != 0: + raise ValueError( + f"head_dim ({head_dim}) must be even for RoPE; adjust emb_dim or head_count" + ) + if self.num_blocks <= 0: + raise ValueError(f"num_blocks must be > 0, got {self.num_blocks}") + if self.seq_length <= 0: + raise ValueError(f"seq_length must be > 0, got {self.seq_length}") + if self.ffn_multiple <= 0: + raise ValueError(f"ffn_multiple must be > 0, got {self.ffn_multiple}") diff --git a/configs/training.py b/configs/training.py new file mode 100644 index 0000000..c877838 --- /dev/null +++ b/configs/training.py @@ -0,0 +1,26 @@ +import torch +from dataclasses import dataclass + +from src.paths import CHECKPOINT_DIR + +checkpoint_dir = CHECKPOINT_DIR + + +@dataclass +class TrainingConfig: + """Hyperparameters and runtime settings for pretraining.""" + token_count: int = 470_000_000 + batch_size: int = 64 + learning_rate: float = 3e-4 + betas: tuple = (0.9, 0.95) + eps: float = 1e-8 + device: str = "cuda" if torch.cuda.is_available() else "cpu" + checkpoint_dir: str = checkpoint_dir + + def __post_init__(self): + if self.token_count <= 0: + raise ValueError(f"token_count must be > 0, got {self.token_count}") + if self.batch_size < 1: + raise ValueError(f"batch_size must be >= 1, got {self.batch_size}") + if self.learning_rate <= 0: + raise ValueError(f"learning_rate must be > 0, got {self.learning_rate}") diff --git a/scripts/check_pre_trained_prediction.py b/scripts/check_pre_trained_prediction.py deleted file mode 100644 index 46a1f4f..0000000 --- a/scripts/check_pre_trained_prediction.py +++ /dev/null @@ -1,78 +0,0 @@ -import torch -import torch.nn.functional as F -from tokenizers import Tokenizer - -from configs.config import ModelConfig, TrainingConfig -from src.model_files.model import BetterGPT - - -def load_and_predict(text, model_path, tokenizer_path, device, - max_tokens=100, temp=0.4, top_k=7, stop_at_eos=True): - tokenizer = Tokenizer.from_file(tokenizer_path) - eos_id = tokenizer.token_to_id("") if stop_at_eos else None - - ckpt = torch.load(model_path, map_location=device, weights_only=True) - model = BetterGPT(ModelConfig(**ckpt["model_config"])) - model.load_state_dict(ckpt["model_state"]) - model = model.to(device) - - if isinstance(text, str): - text = [text] - enc_batch = tokenizer.encode_batch(text) - - seqs = [] - for enc in enc_batch: - ids = enc.ids - if eos_id is not None and ids and ids[-1] == eos_id: - ids = ids[:-1] #drop eos token id - seqs.append(ids) - - lengths = {len(s) for s in seqs} - assert len(lengths) == 1, ( - f"prompts differ in length {lengths}; batched gen needs a key-padding " - f"mask the model lacks. Pass equal-length prompts, or generate one at a time." - ) - - idx = torch.tensor(seqs, device=device) - out = model.generate(idx, max_tokens=max_tokens, temp=temp, - top_k=top_k, eos_id=eos_id) - return tokenizer.decode_batch(out.tolist(), skip_special_tokens=True) - - - -if __name__ == "__main__": - import argparse - import os - - path = TrainingConfig().checkpoint_dir - model_path = os.path.join(path, "best_model.pt") - - if not os.path.exists(model_path): - raise FileNotFoundError(f"Model checkpoint not found at {model_path}") - - - tokenizer_path = "./tokenizer.json" - device = TrainingConfig().device - - parser = argparse.ArgumentParser() - parser.add_argument("--text", type=str, required=True) - parser.add_argument("--max_tokens", type=int, default=500) - parser.add_argument("--temp", type=float, default=0.4) - parser.add_argument("--top_k", type=int, default=7) - parser.add_argument("--stop_at_eos", type=bool, default=True) - - args = parser.parse_args() - - - - output = load_and_predict( - text=args.text, - model_path=model_path, - tokenizer_path=tokenizer_path, - device=device, - max_tokens=args.max_tokens, - temp=args.temp, - top_k=args.top_k, - stop_at_eos=args.stop_at_eos - ) - print(output) \ No newline at end of file diff --git a/scripts/sample_pretrain.py b/scripts/sample_pretrain.py new file mode 100644 index 0000000..4c36f61 --- /dev/null +++ b/scripts/sample_pretrain.py @@ -0,0 +1,112 @@ +import torch +import torch.nn.functional as F +from tokenizers import Tokenizer + +from configs.model import ModelConfig +from configs.training import TrainingConfig +from src.models.model import BetterGPT +from src.logger import get_logger + +logger = get_logger("sample") + + +def load_and_predict(text, model_path, tokenizer_path, device, + max_tokens=100, temp=0.4, top_k=7, stop_at_eos=True): + """Load a trained checkpoint and generate text for one or more prompts. + + All prompts in `text` must encode to the same token length because the model + uses batched generation without a key-padding mask. For prompts of different + lengths, call this function once per prompt. + + Args: + text: A single prompt string or a list of equal-length prompt strings. + model_path: Path to the .pt checkpoint file produced by the trainer. + tokenizer_path: Path to the tokenizer JSON file. + device: Torch device string ('cpu', 'cuda', etc.). + max_tokens: Maximum number of new tokens to generate per prompt. + temp: Sampling temperature; 0 for greedy decoding. + top_k: Restrict sampling to the top-k logits; None for unrestricted. + stop_at_eos: Stop generation when the token is produced. + + Returns: + List of decoded output strings, one per prompt. + """ + if max_tokens <= 0: + raise ValueError(f"max_tokens must be > 0, got {max_tokens}") + if temp < 0: + raise ValueError(f"temp must be >= 0, got {temp}") + if top_k is not None and top_k <= 0: + raise ValueError(f"top_k must be > 0, got {top_k}") + + tokenizer = Tokenizer.from_file(tokenizer_path) + eos_id = tokenizer.token_to_id("") if stop_at_eos else None + + ckpt = torch.load(model_path, map_location=device, weights_only=True) + model = BetterGPT(ModelConfig(**ckpt["model_config"])) + model.load_state_dict(ckpt["model_state"]) + model = model.to(device) + logger.info( + "Model loaded from %s | vocab_size=%d | emb_dim=%d | num_blocks=%d", + model_path, + ckpt["model_config"].get("vocab_size"), + ckpt["model_config"].get("emb_dim"), + ckpt["model_config"].get("num_blocks"), + ) + + if isinstance(text, str): + text = [text] + enc_batch = tokenizer.encode_batch(text) + + seqs = [] + for enc in enc_batch: + ids = enc.ids + if eos_id is not None and ids and ids[-1] == eos_id: + ids = ids[:-1] # drop trailing eos token + seqs.append(ids) + + lengths = {len(s) for s in seqs} + assert len(lengths) == 1, ( + f"prompts differ in length {lengths}; batched gen needs a key-padding " + f"mask the model lacks. Pass equal-length prompts, or generate one at a time." + ) + + idx = torch.tensor(seqs, device=device) + out = model.generate(idx, max_tokens=max_tokens, temp=temp, + top_k=top_k, eos_id=eos_id) + return tokenizer.decode_batch(out.tolist(), skip_special_tokens=True) + + +if __name__ == "__main__": + import argparse + import os + + training_config = TrainingConfig() + path = training_config.checkpoint_dir + model_path = os.path.join(path, "best_model.pt") + + if not os.path.exists(model_path): + raise FileNotFoundError(f"Model checkpoint not found at {model_path}") + + tokenizer_path = "./tokenizer.json" + device = training_config.device + + parser = argparse.ArgumentParser() + parser.add_argument("--text", type=str, required=True) + parser.add_argument("--max_tokens", type=int, default=500) + parser.add_argument("--temp", type=float, default=0.4) + parser.add_argument("--top_k", type=int, default=7) + parser.add_argument("--stop_at_eos", action="store_true", default=True) + + args = parser.parse_args() + + output = load_and_predict( + text=args.text, + model_path=model_path, + tokenizer_path=tokenizer_path, + device=device, + max_tokens=args.max_tokens, + temp=args.temp, + top_k=args.top_k, + stop_at_eos=args.stop_at_eos, + ) + print(output) diff --git a/scripts/train_model.py b/scripts/train_model.py index 9381459..9de56df 100644 --- a/scripts/train_model.py +++ b/scripts/train_model.py @@ -1,31 +1,43 @@ import os +import sys import torch from transformers import get_cosine_schedule_with_warmup -from src.model_files.model import BetterGPT -from configs.config import ModelConfig, TrainingConfig +from src.models.model import BetterGPT +from configs.model import ModelConfig +from configs.training import TrainingConfig from src.data_preparation.data_loader import train_loader, val_loader -from src.pre_training.model_training import training +from src.pretraining.trainer import training +from src.logger import get_logger -model = BetterGPT(ModelConfig()) +logger = get_logger("train") +model_config = ModelConfig() +training_config = TrainingConfig() -path = TrainingConfig().checkpoint_dir +model = BetterGPT(model_config) +total_params = sum(p.numel() for p in model.parameters()) +logger.info("Model initialized | parameters=%dM", total_params // 1_000_000) + +path = training_config.checkpoint_dir if not os.path.exists(path): os.mkdir(path) + logger.info("Created checkpoint directory: %s", path) -#param groups +# param groups: weight decay for matrices, no decay for 1-D params decay, no_decay = [], [] for name, p in model.named_parameters(): if not p.requires_grad: continue - if p.dim() >= 2: # linear/embedding weight matrices + if p.dim() >= 2: decay.append(p) - else: # RMSNorm, gamma (1D), biases + else: no_decay.append(p) +logger.info("Optimizer groups | decay=%d params | no_decay=%d params", len(decay), len(no_decay)) + optim_groups = [ {"params": decay, "weight_decay": 0.1}, {"params": no_decay, "weight_decay": 0.0}, @@ -33,17 +45,17 @@ optimizer = torch.optim.AdamW( optim_groups, - lr=TrainingConfig().learning_rate, - betas=TrainingConfig().betas, - eps=TrainingConfig().eps, + lr=training_config.learning_rate, + betas=training_config.betas, + eps=training_config.eps, ) -# total_steps = num_epochs * len(val_loader) if using dataset class, -# here I am using IterableDataset class, which doesn't have __len__ method, -# so calculating total steps based on total tokens and batch size and seq length +# total_steps = num_epochs * len(val_loader) if using dataset class, +# here using IterableDataset which has no __len__, so derive from token budget +total_steps = int(training_config.token_count / (training_config.batch_size * model_config.seq_length)) +warmup_steps = int(0.02 * total_steps) # ~2% -total_steps = TrainingConfig().token_count / (TrainingConfig().batch_size * ModelConfig().seq_length) #total token/(batch*seq) -warmup_steps = int(0.02 * total_steps) # ~2% +logger.info("LR schedule | total_steps=%d | warmup_steps=%d", total_steps, warmup_steps) lr_scheduler = get_cosine_schedule_with_warmup( optimizer, @@ -53,19 +65,21 @@ if __name__ == "__main__": - - - best_val_loss = training( - model, - ModelConfig(), - total_steps, - train_loader, - val_loader, - optimizer, - lr_scheduler, - device = TrainingConfig().device, - save_path=path, - resume_checkpoint=None, - eval_every=500 + try: + best_val_loss = training( + model, + model_config, + total_steps, + train_loader, + val_loader, + optimizer, + lr_scheduler, + device=training_config.device, + save_path=path, + resume_checkpoint=None, + eval_every=500, ) - print(f"model trained, with best validation loss of {best_val_loss}") \ No newline at end of file + logger.info("Training complete | best_val_loss=%.4f", best_val_loss) + except Exception: + logger.exception("Training failed with an unhandled exception") + sys.exit(1) diff --git a/src/data_preparation/data_loader.py b/src/data_preparation/data_loader.py index ed06707..2a30d34 100644 --- a/src/data_preparation/data_loader.py +++ b/src/data_preparation/data_loader.py @@ -1,43 +1,56 @@ import torch from src.data_preparation.dataset import TinyDataset from torch.utils.data import DataLoader -from configs.config import TrainingConfig, ModelConfig +from configs.model import ModelConfig +from configs.training import TrainingConfig from src.paths import DATA_DIR +from src.logger import get_logger -path = DATA_DIR +logger = get_logger(__name__) +path = DATA_DIR pin_memory = torch.cuda.is_available() - -train_dataset = TinyDataset( - path=path, - seq_length=ModelConfig().seq_length, - split="train", - infinite=True +num_workers = 4 if torch.cuda.is_available() else 0 + +try: + train_dataset = TinyDataset( + path=path, + seq_length=ModelConfig().seq_length, + split="train", + infinite=True ) -val_dataset = TinyDataset( - path=path, - seq_length=ModelConfig().seq_length, - split="validation", - infinite=False + val_dataset = TinyDataset( + path=path, + seq_length=ModelConfig().seq_length, + split="validation", + infinite=False ) +except FileNotFoundError as e: + raise FileNotFoundError( + f"{e}\n" + "Data shards not found. Run `python scripts/create_data_shards.py` first." + ) from e train_loader = DataLoader( dataset=train_dataset, batch_size=TrainingConfig().batch_size, - num_workers=4 if torch.cuda.is_available() else 0, - pin_memory=True if torch.cuda.is_available() else False, - persistent_workers=True if torch.cuda.is_available() else False, + num_workers=num_workers, + pin_memory=pin_memory, + persistent_workers=pin_memory, drop_last=True - ) +) val_loader = DataLoader( dataset=val_dataset, batch_size=TrainingConfig().batch_size, - num_workers=4 if torch.cuda.is_available() else 0, - pin_memory=True if torch.cuda.is_available() else False, - persistent_workers=True if torch.cuda.is_available() else False, + num_workers=num_workers, + pin_memory=pin_memory, + persistent_workers=pin_memory, drop_last=True - ) - +) +logger.info( + "DataLoaders ready | batch_size=%d | seq_length=%d | workers=%d", + TrainingConfig().batch_size, ModelConfig().seq_length, num_workers, +) diff --git a/src/data_preparation/dataset.py b/src/data_preparation/dataset.py index 2ca0dd1..c037bcd 100644 --- a/src/data_preparation/dataset.py +++ b/src/data_preparation/dataset.py @@ -23,6 +23,16 @@ class TinyDataset(IterableDataset): def __init__(self, path, seq_length, split, shuffle_buffer=4096, infinite=True, seed=1337): + """ + Args: + path: Directory containing the binary shard files. + seq_length: Number of tokens per training window (input length). + split: Shard file prefix to glob for (e.g. 'train', 'validation'). + shuffle_buffer: Reservoir size for the shuffle buffer; larger = more random. + infinite: If True, repeat indefinitely across epochs (for training). + If False, yield each window exactly once (for evaluation). + seed: Base random seed; advanced per worker and per epoch to vary shuffling. + """ super().__init__() self.data_files = sorted(glob(f"{path}/{split}*.bin")) self.seq_length = seq_length @@ -33,6 +43,12 @@ def __init__(self, path, seq_length, split, shuffle_buffer=4096, raise FileNotFoundError(f"No files found for split '{split}' in {path}") def _worker_files(self): + """Return the subset of shard files assigned to this DataLoader worker. + + Returns: + Tuple of (files list, worker_id int). Single-process mode returns + all files with worker_id=0. + """ info = get_worker_info() if info is None: return self.data_files, 0 @@ -40,6 +56,15 @@ def _worker_files(self): return files, info.id def _chunks_from_file(self, file, rng): + """Memory-map a shard file and yield shuffled (x, y) token windows. + + Uses a try/finally to ensure the mmap is closed even if iteration is + interrupted. Windows smaller than seq_length+1 (too few tokens) are skipped. + + Args: + file: Path to a binary uint16 shard file. + rng: random.Random instance for shuffling chunk order within the file. + """ try: data = np.memmap(file, dtype=np.uint16, mode="r") n = len(data) @@ -58,12 +83,24 @@ def _chunks_from_file(self, file, rng): data._mmap.close() def _stream(self, files, rng): + """Yield chunks from all files in a randomly shuffled order. + + Args: + files: List of shard file paths to stream from. + rng: random.Random instance controlling file-level shuffle order. + """ file_order = list(files) rng.shuffle(file_order) for file in file_order: yield from self._chunks_from_file(file, rng) def __iter__(self): + """Iterate over token windows with reservoir-sampled shuffling. + + In infinite mode, repeats across epochs with a fresh seed each epoch so + the shuffle order differs. In finite mode, flushes the remaining buffer + after exhausting the shards. + """ files, worker_id = self._worker_files() if not files: return diff --git a/src/data_preparation/make_shards.py b/src/data_preparation/make_shards.py index 0bf7794..f04172f 100644 --- a/src/data_preparation/make_shards.py +++ b/src/data_preparation/make_shards.py @@ -1,78 +1,99 @@ import os import numpy as np -from datasets import Dataset,load_dataset +from datasets import Dataset, load_dataset from tokenizers import Tokenizer -class ShardDataset: - def __init__( - self, - dataset_name:str, - tokenizer: Tokenizer, - out_dir:str, - split: str="train", - data_column_name:str = "text", - buffer_size:int = 2_0000_000 - ): - - self.tokenizer = tokenizer - # self.split = split - # self.data = load_dataset(dataset_name,split=self.split,streaming=True).shuffle(seed=42, buffer_size=10_000) - self.dataset_name = dataset_name - data_folder_name = f"{dataset_name.split("/")[-1]}_data" - self.data_column_name = data_column_name - self.out_dir = os.path.join(out_dir,data_folder_name) - self.buffer_size = buffer_size - - self.vocab_size = int(tokenizer.get_vocab_size()) - os.makedirs(self.out_dir,exist_ok = True) - - assert self.vocab_size<65000, "vocab size more, can't store data in uint16" - - def save_shards(self, data_ids, shard_name): - shard_path = os.path.join(self.out_dir, shard_name) - - if not os.path.exists(shard_path): - data_ids.tofile(shard_path) - print(f"{shard_name} saved") - - - def run(self,split): - shard_id = 1 - buffer = np.empty(self.buffer_size, dtype=np.uint16) - buffer_idx = 0 - - ds = load_dataset(self.dataset_name,split=split,streaming=True).shuffle(seed=42, buffer_size=10_000) - - for samples in ds.iter(1000): - encoded = self.tokenizer.encode_batch(samples[self.data_column_name]) +from src.logger import get_logger - for item in encoded: - ids = np.array(item.ids,dtype=np.uint16) - start=0 +logger = get_logger(__name__) - while start/_data/. + split: Dataset split to process ('train', 'validation', etc.). + data_column_name: Column name that contains the raw text strings. + buffer_size: Number of tokens to accumulate before flushing one shard file. + """ + self.tokenizer = tokenizer + self.dataset_name = dataset_name + data_folder_name = f"{dataset_name.split('/')[-1]}_data" + self.data_column_name = data_column_name + self.out_dir = os.path.join(out_dir, data_folder_name) + self.buffer_size = buffer_size + + self.vocab_size = int(tokenizer.get_vocab_size()) + os.makedirs(self.out_dir, exist_ok=True) + + if self.vocab_size >= 65000: + raise ValueError( + f"vocab_size={self.vocab_size} exceeds 64999; cannot store token IDs in uint16." + ) + + def save_shards(self, data_ids, shard_name): + """Write data_ids to a binary shard file, skipping if the file already exists.""" + shard_path = os.path.join(self.out_dir, shard_name) + if not os.path.exists(shard_path): + data_ids.tofile(shard_path) + logger.info("Saved shard: %s", shard_name) + + def run(self, split): + """Tokenize the given dataset split and write token IDs to shard files. + + Args: + split: Dataset split name to process ('train', 'validation', etc.). + """ + shard_id = 1 + buffer = np.empty(self.buffer_size, dtype=np.uint16) + buffer_idx = 0 + + ds = load_dataset(self.dataset_name, split=split, streaming=True).shuffle(seed=42, buffer_size=10_000) + + for samples in ds.iter(1000): + encoded = self.tokenizer.encode_batch(samples[self.data_column_name]) + + for item in encoded: + ids = np.array(item.ids, dtype=np.uint16) + start = 0 + + while start < len(ids): + remaining = self.buffer_size - buffer_idx + take = min(remaining, len(ids) - start) + + buffer[buffer_idx:buffer_idx + take] = ids[start:start + take] + + buffer_idx += take + start += take + + if buffer_idx == self.buffer_size: + shard_name = f"{split}_shard{shard_id:04d}.bin" + self.save_shards(buffer, shard_name) + + if shard_id % 5 == 0: + logger.info("Progress: %d shards created", shard_id) + + shard_id += 1 + buffer_idx = 0 + + if buffer_idx > 0: shard_name = f"{split}_shard{shard_id:04d}.bin" - - self.save_shards(buffer, shard_name) - - if shard_id % 5 == 0: - print(f"Created {shard_id} shards") - - shard_id += 1 - buffer_idx = 0 - - if buffer_idx>0: - shard_name = f"{split}_shard{shard_id:04d}.bin" - self.save_shards(buffer[:buffer_idx], shard_name) - - + self.save_shards(buffer[:buffer_idx], shard_name) diff --git a/src/logger.py b/src/logger.py new file mode 100644 index 0000000..c4050d7 --- /dev/null +++ b/src/logger.py @@ -0,0 +1,19 @@ +import logging +import os + + +def get_logger(name: str) -> logging.Logger: + """Return a named logger with a StreamHandler, creating handlers only once. + + Level defaults to INFO; override with the LOG_LEVEL environment variable. + """ + logger = logging.getLogger(name) + if not logger.handlers: + handler = logging.StreamHandler() + handler.setFormatter(logging.Formatter( + "%(asctime)s | %(levelname)-8s | %(name)s — %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + )) + logger.addHandler(handler) + logger.setLevel(os.getenv("LOG_LEVEL", "INFO").upper()) + return logger diff --git a/src/model_files/attention.py b/src/model_files/attention.py deleted file mode 100644 index ad59d41..0000000 --- a/src/model_files/attention.py +++ /dev/null @@ -1,59 +0,0 @@ - -import torch -import torch.nn as nn -from torch.nn.functional import scaled_dot_product_attention - -class RopeAttention(nn.Module): - def __init__(self,emb_dim,head_count,head_dim,seq_length,rope): - super().__init__() - - self.head_count = head_count - self.head_dim = head_dim - self.rope = rope - - - self.qkv_proj = nn.Linear(emb_dim,3*emb_dim,bias=False) - - self.out_proj = nn.Linear(emb_dim,emb_dim,bias=False) - - self.register_buffer("mask", - torch.tril(torch.ones(seq_length, seq_length)), - persistent=False - ) - - - def forward(self,x,attention_mask): - - batch,seq_length,emb_dim = x.shape - qkv = self.qkv_proj(x) - - q,k,v = qkv.chunk(3,dim=-1) - # print(f"q,k,v shape after chunk {q.shape}") - - q = q.view(batch,seq_length,self.head_count,self.head_dim).transpose(1,2) - k = k.view(batch,seq_length,self.head_count,self.head_dim).transpose(1,2) - v = v.view(batch,seq_length,self.head_count,self.head_dim).transpose(1,2) - - rotated_q = self.rope(q) - rotated_k = self.rope(k) - # print(f"rotated shape q {rotated_q.shape},k {rotated_k.shape}") - # drop_out_p = self.attn_dropout if self.training else 0.0 - - mask = self.mask[:seq_length,:seq_length].bool() - - if attention_mask is not None: - # print("attention mask") - attention_mask = attention_mask.bool().unsqueeze(1).unsqueeze(1) - combined_mask = attention_mask & mask - attn_bias = torch.zeros_like(combined_mask, dtype=q.dtype) - attn_bias.masked_fill_(~combined_mask, torch.finfo(q.dtype).min) - y = scaled_dot_product_attention(rotated_q,rotated_k,v,attn_mask=attn_bias,is_causal=False) - else: - y = scaled_dot_product_attention(rotated_q,rotated_k,v,is_causal=True) - - y = y.transpose(1,2).contiguous() - - y = y.view(batch,seq_length,emb_dim) - y = self.out_proj(y) - # print(y.shape) - return y \ No newline at end of file diff --git a/src/model_files/layer_normalization.py b/src/model_files/layer_normalization.py deleted file mode 100644 index 18f650d..0000000 --- a/src/model_files/layer_normalization.py +++ /dev/null @@ -1,13 +0,0 @@ -import torch -import torch.nn as nn - -class RMSNorm(nn.Module): - def __init__(self,emb_dim,eps:float=1e-6): - super().__init__() - self.eps = eps - self.gamma = nn.Parameter(torch.ones(emb_dim)) - - def forward(self, x): - ms = x.float().pow(2).mean(dim=-1, keepdim=True) - x_normed = x.float() * torch.rsqrt(ms + self.eps) - return (x_normed * self.gamma.float()).type_as(x) diff --git a/src/model_files/model.py b/src/model_files/model.py deleted file mode 100644 index 5c67438..0000000 --- a/src/model_files/model.py +++ /dev/null @@ -1,141 +0,0 @@ -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from configs.config import ModelConfig -from src.model_files.transformer_block import TransformerBlock -from src.model_files.positional_embeddings import RoPESplitHalf -from src.model_files.layer_normalization import RMSNorm - - - -class BetterGPT(nn.Module): - """A small decoder only langauge model inspired from architeture of GPT and LLama models, - packed with RoPE positional embedding, flash attention for optimized attention computation with - grouped query attention, and used fused kernals and weight tying. - """ - def __init__(self,config:ModelConfig): - super().__init__() - - hid = int((8*config.emb_dim)//3) - hid_dim = config.ffn_multiple*((hid+config.ffn_multiple-1)//config.ffn_multiple) - head_dim = config.emb_dim//config.head_count - self.seq_length = config.seq_length - - assert head_dim % 2 == 0, ( - f"RoPE requires even head_dim, got {head_dim}" - ) - - rope = RoPESplitHalf( - head_dim=head_dim, - max_seq_len=config.seq_length - ) - - self.emb_layer = nn.Embedding( - num_embeddings=config.vocab_size, - embedding_dim=config.emb_dim - ) - - self.rmsnorm = RMSNorm(config.emb_dim) - - self.lm_head = nn.Linear(config.emb_dim,config.vocab_size,bias=False) - - self.transformer_block = nn.ModuleList([TransformerBlock( - head_count=config.head_count, - head_dim=head_dim, - emb_dim=config.emb_dim, - hid_dim=hid_dim, - seq_length = self.seq_length, - rope=rope - ) for _ in range(config.num_blocks)]) - - self.apply(self._init_weights) - - for name, p in self.named_parameters(): - if name.endswith("out_proj.weight") or name.endswith("down_proj.weight"): - nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.num_blocks)) - - self.lm_head.weight = self.emb_layer.weight #weight tying - - def _init_weights(self, module): - - """weight initialization for Linear and embedding layers""" - - if isinstance(module, nn.Linear): - nn.init.normal_(module.weight, mean=0.0, std=0.02) - if module.bias is not None: - nn.init.zeros_(module.bias) - - elif isinstance(module, nn.Embedding): - nn.init.normal_(module.weight, mean=0.0, std=0.02) - - def forward(self,input_dim,attention_mask=None): - x = self.emb_layer(input_dim) - for block in self.transformer_block: - x = block(x,attention_mask) - - x = self.rmsnorm(x) - x = self.lm_head(x) - - return x - - @torch.no_grad() - def generate(self, idx, max_tokens, temp, top_k=None, eos_id=None): - self.eval() - B = idx.size(0) - finished = torch.zeros(B, 1, dtype=torch.bool, device=idx.device) - - for _ in range(max_tokens): - idx_cond = idx[:, -self.seq_length:] - logits = self(idx_cond)[:, -1, :] - - if temp == 0: - out_idx = torch.argmax(logits, dim=-1, keepdim=True) - else: - logits = logits / temp - if top_k is not None: - k = min(top_k, logits.size(-1)) - val, _ = torch.topk(logits, k) - min_val = val[:, -1].unsqueeze(-1) - logits = logits.masked_fill(logits < min_val, float('-inf')) - probs = F.softmax(logits, dim=-1) - out_idx = torch.multinomial(probs, num_samples=1) - - if eos_id is not None: - # rows already done keep emitting EOS (clean padding, not real tokens) - out_idx = torch.where(finished, torch.full_like(out_idx, eos_id), out_idx) - finished = finished | (out_idx == eos_id) - - idx = torch.cat([idx, out_idx], dim=1) - - if eos_id is not None and finished.all(): - break - - return idx - - @torch.no_grad() - def generate_single(self, idx, max_tokens, temp, top_k=None, eos_id=None): - self.eval() - for _ in range(max_tokens): - idx_cond = idx[-self.seq_length:] - logits = self(idx_cond.unsqueeze(0))[:, -1, :] - - if temp == 0: - out_idx = torch.argmax(logits, dim=-1).item() - else: - logits = logits / temp - if top_k is not None: - k = min(top_k, logits.size(-1)) - val, _ = torch.topk(logits, k) - min_val = val[:, -1].unsqueeze(-1) - logits = logits.masked_fill(logits < min_val, float('-inf')) - probs = F.softmax(logits, dim=-1) - out_idx = torch.multinomial(probs, num_samples=1).item() - - if eos_id is not None and out_idx == eos_id: - break - - idx = torch.cat([idx, torch.tensor([out_idx], device=idx.device)], dim=0) - return idx \ No newline at end of file diff --git a/src/model_files/swiglu_feed_forward.py b/src/model_files/swiglu_feed_forward.py deleted file mode 100644 index 27556ec..0000000 --- a/src/model_files/swiglu_feed_forward.py +++ /dev/null @@ -1,17 +0,0 @@ - -import torch.nn as nn -import torch.nn.functional as F - -class SwiGLU_FFN(nn.Module): - def __init__(self, emb_dim,hid_dim): - super().__init__() - self.gate_proj = nn.Linear(emb_dim,hid_dim,bias=False) - self.up_proj = nn.Linear(emb_dim,hid_dim,bias=False) - self.down_proj = nn.Linear(hid_dim,emb_dim,bias=False) - - def forward(self,x): - gate = F.silu(self.gate_proj(x)) - up = self.up_proj(x) - down = self.down_proj(gate*up) - - return down diff --git a/src/model_files/transformer_block.py b/src/model_files/transformer_block.py deleted file mode 100644 index 54f6a9a..0000000 --- a/src/model_files/transformer_block.py +++ /dev/null @@ -1,35 +0,0 @@ - -import torch -import torch.nn as nn - -from src.model_files.swiglu_feed_forward import SwiGLU_FFN -from src.model_files.attention import RopeAttention -from src.model_files.layer_normalization import RMSNorm - -class TransformerBlock(nn.Module): - def __init__( - self, - emb_dim, - hid_dim, - seq_length, - rope, - head_count, - head_dim - ): - - super().__init__() - self.pre_attn_norm = RMSNorm(emb_dim) - self.pre_ffn_norm = RMSNorm(emb_dim) - self.attention = RopeAttention( - emb_dim=emb_dim, - head_dim=head_dim, - head_count=head_count, - seq_length=seq_length, - rope=rope - ) - self.ffn = SwiGLU_FFN(emb_dim=emb_dim,hid_dim=hid_dim) - - def forward(self,x, attention_mask): - x = x + self.attention(self.pre_attn_norm(x),attention_mask) - x = x + self.ffn(self.pre_ffn_norm(x)) - return x \ No newline at end of file diff --git a/src/pre_training/__init__.py b/src/models/__init__.py similarity index 100% rename from src/pre_training/__init__.py rename to src/models/__init__.py diff --git a/src/models/attention.py b/src/models/attention.py new file mode 100644 index 0000000..ac43931 --- /dev/null +++ b/src/models/attention.py @@ -0,0 +1,75 @@ + +import torch +import torch.nn as nn +from torch.nn.functional import scaled_dot_product_attention + + +class MHAttention(nn.Module): + """Multi-head self-attention with Rotary Position Embeddings (RoPE). + + Uses PyTorch's scaled_dot_product_attention (flash-attention kernel when available). + A lower-triangular causal mask is pre-computed at init time; an optional per-token + attention_mask is AND-ed with it at forward time. + """ + + def __init__(self, emb_dim, head_count, head_dim, seq_length, rope): + """ + Args: + emb_dim: Model embedding dimension. + head_count: Number of attention heads. + head_dim: Dimension per head (emb_dim // head_count). + seq_length: Maximum sequence length for pre-computing the causal mask. + rope: RoPE module applied to queries and keys before attention. + """ + super().__init__() + + self.head_count = head_count + self.head_dim = head_dim + self.rope = rope + + self.qkv_proj = nn.Linear(emb_dim, 3 * emb_dim, bias=False) + self.out_proj = nn.Linear(emb_dim, emb_dim, bias=False) + + self.register_buffer( + "mask", + torch.tril(torch.ones(seq_length, seq_length)), + persistent=False, + ) + + def forward(self, x, attention_mask): + """Apply multi-head attention with RoPE and an optional padding mask. + + Args: + x: Input tensor of shape (B, T, emb_dim). + attention_mask: Optional bool tensor (B, T); True = attend, False = mask out. + + Returns: + Output tensor of shape (B, T, emb_dim). + """ + batch, seq_length, emb_dim = x.shape + qkv = self.qkv_proj(x) + + q, k, v = qkv.chunk(3, dim=-1) + + q = q.view(batch, seq_length, self.head_count, self.head_dim).transpose(1, 2) + k = k.view(batch, seq_length, self.head_count, self.head_dim).transpose(1, 2) + v = v.view(batch, seq_length, self.head_count, self.head_dim).transpose(1, 2) + + rotated_q = self.rope(q) + rotated_k = self.rope(k) + + mask = self.mask[:seq_length, :seq_length].bool() + + if attention_mask is not None: + attention_mask = attention_mask.bool().unsqueeze(1).unsqueeze(1) + combined_mask = attention_mask & mask + attn_bias = torch.zeros_like(combined_mask, dtype=q.dtype) + attn_bias.masked_fill_(~combined_mask, torch.finfo(q.dtype).min) + y = scaled_dot_product_attention(rotated_q, rotated_k, v, attn_mask=attn_bias, is_causal=False) + else: + y = scaled_dot_product_attention(rotated_q, rotated_k, v, is_causal=True) + + y = y.transpose(1, 2).contiguous() + y = y.view(batch, seq_length, emb_dim) + y = self.out_proj(y) + return y diff --git a/src/models/layer_normalization.py b/src/models/layer_normalization.py new file mode 100644 index 0000000..627a032 --- /dev/null +++ b/src/models/layer_normalization.py @@ -0,0 +1,27 @@ +import torch +import torch.nn as nn + + +class RMSNorm(nn.Module): + """Root Mean Square Layer Normalization (https://arxiv.org/abs/1910.07467). + + More efficient than LayerNorm — omits mean-centering and bias terms. + Computation is promoted to float32 then cast back to the input dtype to + avoid precision loss with bfloat16 inputs. + """ + + def __init__(self, emb_dim, eps: float = 1e-6): + """ + Args: + emb_dim: Size of the last dimension to normalize over. + eps: Small constant added inside rsqrt for numerical stability. + """ + super().__init__() + self.eps = eps + self.gamma = nn.Parameter(torch.ones(emb_dim)) + + def forward(self, x): + """Normalize x by its RMS and scale by the learnable gamma parameter.""" + ms = x.float().pow(2).mean(dim=-1, keepdim=True) + x_normed = x.float() * torch.rsqrt(ms + self.eps) + return (x_normed * self.gamma.float()).type_as(x) diff --git a/src/models/model.py b/src/models/model.py new file mode 100644 index 0000000..f049c6b --- /dev/null +++ b/src/models/model.py @@ -0,0 +1,207 @@ +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from configs.model import ModelConfig +from src.models.transformer_block import TransformerBlock +from src.models.positional_embeddings import RoPESplitHalf +from src.models.layer_normalization import RMSNorm +from src.logger import get_logger + +logger = get_logger(__name__) + + +class BetterGPT(nn.Module): + """A small decoder only langauge model inspired from architeture of GPT and LLama models, + packed with RoPE positional embedding, flash attention for optimized attention computation with + grouped query attention, and used fused kernals and weight tying. + """ + def __init__(self, config: ModelConfig): + super().__init__() + + hid = int((8 * config.emb_dim) // 3) + hid_dim = config.ffn_multiple * ((hid + config.ffn_multiple - 1) // config.ffn_multiple) + head_dim = config.emb_dim // config.head_count + self.seq_length = config.seq_length + + if head_dim % 2 != 0: + raise ValueError( + f"RoPE requires even head_dim, got {head_dim} " + f"(emb_dim={config.emb_dim}, head_count={config.head_count})" + ) + + rope = RoPESplitHalf( + head_dim=head_dim, + max_seq_len=config.seq_length + ) + + self.emb_layer = nn.Embedding( + num_embeddings=config.vocab_size, + embedding_dim=config.emb_dim + ) + + self.rmsnorm = RMSNorm(config.emb_dim) + + self.lm_head = nn.Linear(config.emb_dim, config.vocab_size, bias=False) + + self.transformer_block = nn.ModuleList([TransformerBlock( + head_count=config.head_count, + head_dim=head_dim, + emb_dim=config.emb_dim, + hid_dim=hid_dim, + seq_length=self.seq_length, + rope=rope + ) for _ in range(config.num_blocks)]) + + self.apply(self._init_weights) + + for name, p in self.named_parameters(): + if name.endswith("out_proj.weight") or name.endswith("down_proj.weight"): + nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.num_blocks)) + + self.lm_head.weight = self.emb_layer.weight # weight tying + + def _init_weights(self, module): + """Initialize Linear weights with N(0, 0.02) and zero biases; Embedding weights with N(0, 0.02).""" + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + + def forward(self, input_dim, attention_mask=None): + """Run a forward pass through embedding → transformer blocks → RMSNorm → LM head. + + Args: + input_dim: Token index tensor of shape (B, T). + attention_mask: Optional padding mask of shape (B, T); True = attend. + + Returns: + Logits tensor of shape (B, T, vocab_size). + """ + x = self.emb_layer(input_dim) + for block in self.transformer_block: + x = block(x, attention_mask) + + x = self.rmsnorm(x) + x = self.lm_head(x) + + return x + + @torch.no_grad() + def generate(self, idx, max_tokens, temp, top_k=None, eos_id=None): + """Generate tokens autoregressively for a batch of prompts. + + Rows that produce eos_id keep emitting it as padding so the batch stays + rectangular; generation stops early when every row has finished. + + Args: + idx: Prompt token indices of shape (B, T). + max_tokens: Maximum number of new tokens to append. + temp: Sampling temperature; 0 for greedy decoding. + top_k: If set, restrict sampling to the top-k logits. + eos_id: Token ID that signals end of sequence; None to disable. + + Returns: + Token index tensor of shape (B, T + generated_length). + """ + if idx.dim() != 2: + raise ValueError(f"generate() expects a 2-D input (B, T), got shape {tuple(idx.shape)}") + if temp < 0: + raise ValueError(f"temp must be >= 0, got {temp}") + if top_k is not None and top_k <= 0: + raise ValueError(f"top_k must be > 0, got {top_k}") + if max_tokens <= 0: + raise ValueError(f"max_tokens must be > 0, got {max_tokens}") + + self.eval() + B = idx.size(0) + finished = torch.zeros(B, 1, dtype=torch.bool, device=idx.device) + + for _ in range(max_tokens): + if idx.size(1) > self.seq_length: + logger.warning( + "Input length %d exceeds seq_length %d; truncating to last %d tokens.", + idx.size(1), self.seq_length, self.seq_length, + ) + idx_cond = idx[:, -self.seq_length:] + logits = self(idx_cond)[:, -1, :] + + if torch.isnan(logits).any() or torch.isinf(logits).any(): + raise RuntimeError("NaN or Inf detected in model logits — check model weights.") + + if temp == 0: + out_idx = torch.argmax(logits, dim=-1, keepdim=True) + else: + logits = logits / temp + if top_k is not None: + k = min(top_k, logits.size(-1)) + val, _ = torch.topk(logits, k) + min_val = val[:, -1].unsqueeze(-1) + logits = logits.masked_fill(logits < min_val, float('-inf')) + probs = F.softmax(logits, dim=-1) + out_idx = torch.multinomial(probs, num_samples=1) + + if eos_id is not None: + # rows already done keep emitting EOS (clean padding, not real tokens) + out_idx = torch.where(finished, torch.full_like(out_idx, eos_id), out_idx) + finished = finished | (out_idx == eos_id) + + idx = torch.cat([idx, out_idx], dim=1) + + if eos_id is not None and finished.all(): + break + + return idx + + @torch.no_grad() + def generate_single(self, idx, max_tokens, temp, top_k=None, eos_id=None): + """Generate tokens autoregressively for a single prompt (no batch dimension). + + Args: + idx: Prompt token indices of shape (T,). + max_tokens: Maximum number of new tokens to append. + temp: Sampling temperature; 0 for greedy decoding. + top_k: If set, restrict sampling to the top-k logits. + eos_id: Token ID that signals end of sequence; None to disable. + + Returns: + Token index tensor of shape (T + generated_length,). + """ + if idx.dim() != 1: + raise ValueError(f"generate_single() expects a 1-D input (T,), got shape {tuple(idx.shape)}") + if temp < 0: + raise ValueError(f"temp must be >= 0, got {temp}") + if top_k is not None and top_k <= 0: + raise ValueError(f"top_k must be > 0, got {top_k}") + if max_tokens <= 0: + raise ValueError(f"max_tokens must be > 0, got {max_tokens}") + + self.eval() + for _ in range(max_tokens): + idx_cond = idx[-self.seq_length:] + logits = self(idx_cond.unsqueeze(0))[:, -1, :] + + if torch.isnan(logits).any() or torch.isinf(logits).any(): + raise RuntimeError("NaN or Inf detected in model logits — check model weights.") + + if temp == 0: + out_idx = torch.argmax(logits, dim=-1).item() + else: + logits = logits / temp + if top_k is not None: + k = min(top_k, logits.size(-1)) + val, _ = torch.topk(logits, k) + min_val = val[:, -1].unsqueeze(-1) + logits = logits.masked_fill(logits < min_val, float('-inf')) + probs = F.softmax(logits, dim=-1) + out_idx = torch.multinomial(probs, num_samples=1).item() + + if eos_id is not None and out_idx == eos_id: + break + + idx = torch.cat([idx, torch.tensor([out_idx], device=idx.device)], dim=0) + return idx diff --git a/src/model_files/positional_embeddings.py b/src/models/positional_embeddings.py similarity index 51% rename from src/model_files/positional_embeddings.py rename to src/models/positional_embeddings.py index 96946ac..7cb5885 100644 --- a/src/model_files/positional_embeddings.py +++ b/src/models/positional_embeddings.py @@ -3,10 +3,21 @@ import torch.nn as nn -#two ways to create vectorized rope - class RoPE_Interleave(nn.Module): + """Rotary Position Embedding — interleaved variant. + + Sin/cos values are expanded with repeat_interleave so each pair of adjacent + elements (even, odd) in the head dimension shares a frequency. Pre-computes + sin/cos tables up to max_seq_len at init time. + """ + def __init__(self, max_seq_len, head_dim, base=10000): + """ + Args: + max_seq_len: Maximum sequence length to pre-compute tables for. + head_dim: Per-head dimension; must be even. + base: Frequency base for the geometric sequence of inv-frequencies. + """ super().__init__() inv_freq = 1.0 / ( @@ -14,7 +25,6 @@ def __init__(self, max_seq_len, head_dim, base=10000): ) pos = torch.arange(max_seq_len).float() - # angles = pos[:,None]@inv_freq[None,:] angles = torch.outer(pos, inv_freq) sin = angles.sin().repeat_interleave(2, dim=-1) @@ -22,27 +32,36 @@ def __init__(self, max_seq_len, head_dim, base=10000): self.register_buffer("sin", sin) self.register_buffer("cos", cos) - def rotate_half(self, x): + """Swap adjacent pairs with a sign flip: [-x1, x0, -x3, x2, ...].""" even = x[..., 0::2] odd = x[..., 1::2] - return torch.stack((-odd, even), dim=-1).flatten(-2) def forward(self, x): + """Apply RoPE rotation to query or key tensor x of shape (B, H, T, head_dim).""" T = x.size(-2) - sin = self.sin[:T] cos = self.cos[:T] - # print(x.shape) - return x * cos + self.rotate_half(x) * sin class RoPESplitHalf(nn.Module): - - def __init__(self,head_dim,max_seq_len,base=10000): + """Rotary Position Embedding — split-half variant (used by LLaMA). + + The head dimension is split into two halves; the rotation mixes the first + half with the negated second half. Pre-computes sin/cos tables up to + max_seq_len at init time. Buffers are non-persistent (not saved to checkpoints). + """ + + def __init__(self, head_dim, max_seq_len, base=10000): + """ + Args: + head_dim: Per-head dimension; must be even. + max_seq_len: Maximum sequence length to pre-compute tables for. + base: Frequency base for the geometric sequence of inv-frequencies. + """ super().__init__() inv_freq = 1.0 / ( @@ -58,15 +77,14 @@ def __init__(self,head_dim,max_seq_len,base=10000): self.register_buffer("sin", sin, persistent=False) self.register_buffer("cos", cos, persistent=False) - def rotate_half(self,x): + def rotate_half(self, x): + """Return [-x2 | x1] where x is split into equal halves x1 and x2.""" x1, x2 = x.chunk(2, dim=-1) return torch.cat((-x2, x1), dim=-1) def forward(self, x): + """Apply RoPE rotation to query or key tensor x of shape (B, H, T, head_dim).""" T = x.size(-2) - sin = self.sin[:T] cos = self.cos[:T] - # print(x.shape) - return x * cos + self.rotate_half(x) * sin diff --git a/src/models/swiglu_feed_forward.py b/src/models/swiglu_feed_forward.py new file mode 100644 index 0000000..8749adb --- /dev/null +++ b/src/models/swiglu_feed_forward.py @@ -0,0 +1,30 @@ + +import torch.nn as nn +import torch.nn.functional as F + + +class SwiGLU_FFN(nn.Module): + """SwiGLU feed-forward network from PaLM / LLaMA. + + Computes: down_proj(silu(gate_proj(x)) * up_proj(x)). + All three linear layers are bias-free. + """ + + def __init__(self, emb_dim, hid_dim): + """ + Args: + emb_dim: Input and output dimension. + hid_dim: Hidden (expanded) dimension; typically ~(8/3)*emb_dim rounded up + to the nearest ffn_multiple. + """ + super().__init__() + self.gate_proj = nn.Linear(emb_dim, hid_dim, bias=False) + self.up_proj = nn.Linear(emb_dim, hid_dim, bias=False) + self.down_proj = nn.Linear(hid_dim, emb_dim, bias=False) + + def forward(self, x): + """Apply the SwiGLU gating: silu(gate) * up, then project back down.""" + gate = F.silu(self.gate_proj(x)) + up = self.up_proj(x) + down = self.down_proj(gate * up) + return down diff --git a/src/models/transformer_block.py b/src/models/transformer_block.py new file mode 100644 index 0000000..7c627c6 --- /dev/null +++ b/src/models/transformer_block.py @@ -0,0 +1,52 @@ + +import torch +import torch.nn as nn + +from src.models.swiglu_feed_forward import SwiGLU_FFN +from src.models.attention import MHAttention +from src.models.layer_normalization import RMSNorm + + +class TransformerBlock(nn.Module): + """Single pre-norm transformer block. + + Applies: norm → attention → residual, then norm → FFN → residual. + Pre-normalization (norm before the sub-layer rather than after) improves + training stability for deep models. + """ + + def __init__(self, emb_dim, hid_dim, seq_length, rope, head_count, head_dim): + """ + Args: + emb_dim: Model embedding dimension. + hid_dim: FFN hidden (expanded) dimension. + seq_length: Maximum sequence length; passed to MHAttention for mask pre-computation. + rope: Shared RoPE module applied inside attention. + head_count: Number of attention heads. + head_dim: Per-head dimension (emb_dim // head_count). + """ + super().__init__() + self.pre_attn_norm = RMSNorm(emb_dim) + self.pre_ffn_norm = RMSNorm(emb_dim) + self.attention = MHAttention( + emb_dim=emb_dim, + head_dim=head_dim, + head_count=head_count, + seq_length=seq_length, + rope=rope, + ) + self.ffn = SwiGLU_FFN(emb_dim=emb_dim, hid_dim=hid_dim) + + def forward(self, x, attention_mask): + """Apply one transformer block with residual connections. + + Args: + x: Input tensor of shape (B, T, emb_dim). + attention_mask: Optional padding mask forwarded to MHAttention. + + Returns: + Output tensor of shape (B, T, emb_dim). + """ + x = x + self.attention(self.pre_attn_norm(x), attention_mask) + x = x + self.ffn(self.pre_ffn_norm(x)) + return x diff --git a/src/pre_training/model_evaluation.py b/src/pre_training/model_evaluation.py deleted file mode 100644 index 83173c7..0000000 --- a/src/pre_training/model_evaluation.py +++ /dev/null @@ -1,18 +0,0 @@ -import torch -import torch.nn.functional as F - -@torch.no_grad() -def evaluate(model, val_loader, device, max_batches=50): - model.eval() - losses = [] - for i, (input_seq, tar_seq) in enumerate(val_loader): - if i >= max_batches: - break - input_seq = input_seq.to(device, non_blocking=True) - tar_seq = tar_seq.to(device, non_blocking=True) - with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): - logits = model(input_seq) - B, T, V = logits.shape - loss = F.cross_entropy(logits.view(B * T, V), tar_seq.view(B * T)) - losses.append(loss.item()) - return sum(losses) / len(losses) \ No newline at end of file diff --git a/src/pre_training/model_training.py b/src/pre_training/model_training.py deleted file mode 100644 index 55eba16..0000000 --- a/src/pre_training/model_training.py +++ /dev/null @@ -1,107 +0,0 @@ -import math -import torch -import os - -from dataclasses import asdict -import torch.nn.functional as F -from tqdm import tqdm - -from src.pre_training.model_evaluation import evaluate - - - -def training( - model, - model_config, - max_steps, - train_loader, - val_loader, - optimizer, - lr_scheduler, - device, - save_path, - resume_checkpoint=None, - eval_every=500, -): - step = 0 - best_val_loss = float("inf") - window_loss = [] - print(f"Training on {device} for {max_steps} steps...") - model = model.to(device) - - checkpoint_path=os.path.join(save_path,"checkpoint.pt") - best_path=os.path.join(save_path,"best_model.pt") - - #resume: weights + optimizer + scheduler + step. stream restarts. - if resume_checkpoint is not None and os.path.exists(resume_checkpoint): - ckpt = torch.load(resume_checkpoint, map_location=device, weights_only=True) - model.load_state_dict(ckpt["model_state"]) - optimizer.load_state_dict(ckpt["optimizer_state"]) - lr_scheduler.load_state_dict(ckpt["scheduler_state"]) - step = ckpt["step"] - best_val_loss = ckpt["best_val_loss"] - print(f"Resumed from step {step} (data stream restarts from beginning)") - - def save_ckpt(path): - tmp = f"{path}.tmp" - torch.save({ - "step": step, - "model_config": asdict(model_config), - "model_state": model.state_dict(), - "optimizer_state": optimizer.state_dict(), - "scheduler_state": lr_scheduler.state_dict(), - "best_val_loss": best_val_loss, - }, tmp) - os.replace(tmp, path) - - model.train() - train_iter = iter(train_loader) - pbar = tqdm(total=max_steps, initial=step, desc="Training") - - while step < max_steps: - input_seq, tar_seq = next(train_iter) - input_seq, tar_seq = input_seq.to(device, non_blocking=True), tar_seq.to(device, non_blocking=True) - - optimizer.zero_grad(set_to_none=True) - with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): - logits = model(input_seq) - B, T, V = logits.shape - loss = F.cross_entropy(logits.view(B * T, V), tar_seq.view(B * T)) - - loss.backward() - grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) - optimizer.step() - lr_scheduler.step() - - window_loss.append(loss.item()) - step += 1 - pbar.update(1) - - if step % eval_every == 0: - val_loss = evaluate(model, val_loader, device) - avg_train = sum(window_loss) / len(window_loss) - print(f"\nStep {step} | train {avg_train:.4f} | val {val_loss:.4f} | " - f"grad_norm {grad_norm:.4f} | lr {optimizer.param_groups[0]['lr']:.2e}") - - save_ckpt(checkpoint_path) - if val_loss < best_val_loss: - best_val_loss = val_loss - save_ckpt(best_path) - print(f" best updated: {best_val_loss:.4f}") - - window_loss.clear() - model.train() - - - if window_loss: - val_loss = evaluate(model, val_loader, device) - avg_train = sum(window_loss) / len(window_loss) - print(f"\nFinal step {step} | train {avg_train:.4f} | val {val_loss:.4f}") - save_ckpt(checkpoint_path) - if val_loss < best_val_loss: - best_val_loss = val_loss - save_ckpt(best_path) - - pbar.close() - return best_val_loss - diff --git a/src/pretraining/__init__.py b/src/pretraining/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/pretraining/evaluator.py b/src/pretraining/evaluator.py new file mode 100644 index 0000000..b2cfe23 --- /dev/null +++ b/src/pretraining/evaluator.py @@ -0,0 +1,43 @@ +import torch +import torch.nn.functional as F + +from src.logger import get_logger + +logger = get_logger(__name__) + + +@torch.no_grad() +def evaluate(model, val_loader, device, max_batches=50): + """Compute average cross-entropy loss on the validation set. + + Args: + model: BetterGPT model instance. + val_loader: DataLoader for the validation split. + device: Torch device string or torch.device. + max_batches: Maximum number of batches to evaluate; caps evaluation time. + + Returns: + Average loss over evaluated batches, or float('inf') if the loader is empty. + """ + device_str = device if isinstance(device, str) else device.type + model.eval() + losses = [] + + for i, (input_seq, tar_seq) in enumerate(val_loader): + if i >= max_batches: + break + input_seq = input_seq.to(device, non_blocking=True) + tar_seq = tar_seq.to(device, non_blocking=True) + with torch.amp.autocast(device_type=device_str, dtype=torch.bfloat16, enabled=(device_str == "cuda")): + logits = model(input_seq) + B, T, V = logits.shape + loss = F.cross_entropy(logits.view(B * T, V), tar_seq.view(B * T)) + losses.append(loss.item()) + + if not losses: + logger.warning("Evaluator received zero batches — validation set may be empty.") + return float("inf") + + avg_loss = sum(losses) / len(losses) + logger.debug("Evaluated %d batches | avg_val_loss=%.4f", len(losses), avg_loss) + return avg_loss diff --git a/src/pretraining/trainer.py b/src/pretraining/trainer.py new file mode 100644 index 0000000..57ce835 --- /dev/null +++ b/src/pretraining/trainer.py @@ -0,0 +1,160 @@ +import math +import torch +import os + +from dataclasses import asdict +import torch.nn.functional as F +from tqdm import tqdm + +from src.pretraining.evaluator import evaluate +from src.logger import get_logger + +logger = get_logger(__name__) + + +def training( + model, + model_config, + max_steps, + train_loader, + val_loader, + optimizer, + lr_scheduler, + device, + save_path, + resume_checkpoint=None, + eval_every=500, +): + """Run the pretraining loop with periodic evaluation and checkpointing. + + Checkpoints are written atomically (write to .tmp then rename). If a + resume_checkpoint is provided and exists, training continues from the saved + step; the data stream always restarts from the beginning. + + Args: + model: BetterGPT model instance. + model_config: ModelConfig dataclass (serialized into checkpoints). + max_steps: Total number of optimizer steps to run. + train_loader: DataLoader for the training split (infinite IterableDataset). + val_loader: DataLoader for the validation split. + optimizer: Configured AdamW optimizer. + lr_scheduler: Learning rate scheduler (stepped every optimizer step). + device: Torch device string or torch.device. + save_path: Directory where checkpoint.pt and best_model.pt are written. + resume_checkpoint: Optional path to an existing checkpoint to resume from. + eval_every: Number of steps between validation evaluations. + + Returns: + best_val_loss: Lowest validation loss recorded during training. + """ + step = 0 + best_val_loss = float("inf") + window_loss = [] + + device_str = device if isinstance(device, str) else device.type + logger.info("Training on %s for %d steps...", device_str, max_steps) + model = model.to(device) + + checkpoint_path = os.path.join(save_path, "checkpoint.pt") + best_path = os.path.join(save_path, "best_model.pt") + + # resume: weights + optimizer + scheduler + step. stream restarts. + if resume_checkpoint is not None and os.path.exists(resume_checkpoint): + ckpt = torch.load(resume_checkpoint, map_location=device, weights_only=True) + model.load_state_dict(ckpt["model_state"]) + optimizer.load_state_dict(ckpt["optimizer_state"]) + lr_scheduler.load_state_dict(ckpt["scheduler_state"]) + step = ckpt["step"] + best_val_loss = ckpt["best_val_loss"] + logger.info("Resumed from step %d (data stream restarts from beginning)", step) + else: + logger.info("Starting fresh training run") + + def save_ckpt(path): + tmp = f"{path}.tmp" + try: + torch.save({ + "step": step, + "model_config": asdict(model_config), + "model_state": model.state_dict(), + "optimizer_state": optimizer.state_dict(), + "scheduler_state": lr_scheduler.state_dict(), + "best_val_loss": best_val_loss, + }, tmp) + os.replace(tmp, path) + except OSError as e: + logger.error("Failed to save checkpoint to %s: %s", path, e) + + model.train() + train_iter = iter(train_loader) + pbar = tqdm(total=max_steps, initial=step, desc="Training") + + while step < max_steps: + try: + input_seq, tar_seq = next(train_iter) + except StopIteration: + raise RuntimeError( + f"Training dataloader exhausted at step {step}/{max_steps}. " + "Increase token_count or reduce batch_size/seq_length." + ) + + input_seq = input_seq.to(device, non_blocking=True) + tar_seq = tar_seq.to(device, non_blocking=True) + + optimizer.zero_grad(set_to_none=True) + try: + with torch.amp.autocast(device_type=device_str, dtype=torch.bfloat16, enabled=(device_str == "cuda")): + logits = model(input_seq) + B, T, V = logits.shape + loss = F.cross_entropy(logits.view(B * T, V), tar_seq.view(B * T)) + + if loss.isnan() or loss.isinf(): + raise RuntimeError(f"Loss is {loss.item()} at step {step} — training has diverged.") + + loss.backward() + except torch.cuda.OutOfMemoryError: + logger.error( + "CUDA out of memory at step %d. Consider reducing batch_size or seq_length.", step + ) + raise + + grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + optimizer.step() + lr_scheduler.step() + + window_loss.append(loss.item()) + step += 1 + pbar.update(1) + + if step % eval_every == 0: + val_loss = evaluate(model, val_loader, device) + avg_train = sum(window_loss) / len(window_loss) + logger.info( + "Step %d/%d | train_loss=%.4f | val_loss=%.4f | grad_norm=%.4f | lr=%.2e", + step, max_steps, avg_train, val_loss, + grad_norm, optimizer.param_groups[0]["lr"], + ) + + save_ckpt(checkpoint_path) + if val_loss < best_val_loss: + best_val_loss = val_loss + save_ckpt(best_path) + logger.info("New best model saved (val_loss=%.4f)", best_val_loss) + + window_loss.clear() + model.train() + + if window_loss: + val_loss = evaluate(model, val_loader, device) + avg_train = sum(window_loss) / len(window_loss) + logger.info( + "Final step %d | train_loss=%.4f | val_loss=%.4f", + step, avg_train, val_loss, + ) + save_ckpt(checkpoint_path) + if val_loss < best_val_loss: + best_val_loss = val_loss + save_ckpt(best_path) + + pbar.close() + return best_val_loss diff --git a/src/tokenizer_train.py b/src/tokenizer_train.py index 2ca5166..e79dfdd 100644 --- a/src/tokenizer_train.py +++ b/src/tokenizer_train.py @@ -3,7 +3,6 @@ import os import shutil import tempfile -from dataclasses import dataclass, field from typing import Any, Callable, Iterable, List, Optional from tokenizers import Regex, Tokenizer @@ -42,17 +41,56 @@ def __init__( config: TokenizerConfig, text_extractor: Optional[Callable[[Any], str]] = None, ): + """Initialize the trainer with a dataset, config, and optional text extractor. + + Args: + dataset: Iterable of samples to train the tokenizer on. + config: Tokenizer configuration (vocab size, special tokens, etc.). + text_extractor: Optional callable that maps a sample to its text. + Defaults to :meth:`_default_extractor` when not provided. + """ + if dataset is None: + raise ValueError("dataset must not be None") + if config is None: + raise ValueError("config must not be None") + if text_extractor is not None and not callable(text_extractor): + raise TypeError("text_extractor must be callable") self.dataset = dataset self.config = config self.text_extractor = text_extractor or self._default_extractor self.tokenizer: Optional[Tokenizer] = None def _default_extractor(self, sample: Any) -> str: - if isinstance(sample, dict): - return str(sample.get(self.config.text_field, "")).strip() - return str(sample).strip() + """Extract text from a sample. + + For dict samples, reads the configured ``text_field``; otherwise casts + the sample to a string. The result is stripped of surrounding whitespace. + + Args: + sample: A single dataset sample. + + Returns: + The extracted, whitespace-stripped text. + """ + try: + if isinstance(sample, dict): + return str(sample.get(self.config.text_field, "")).strip() + return str(sample).strip() + except (TypeError, ValueError, AttributeError) as e: + logger.warning("failed to extract text from sample: %s", e) + return "" def _build_tokenizer(self) -> Tokenizer: + """Construct an untrained byte-level BPE tokenizer. + + Wires up a GPT-style regex split pre-tokenizer followed by byte-level + encoding, and a matching byte-level decoder for reversible round trips. + + Returns: + A configured but untrained :class:`Tokenizer` instance. + """ + logger.debug("building byte-level BPE tokenizer (unk_token=%r)", + self.config.unk_token) tok = Tokenizer(BPE(unk_token=self.config.unk_token)) tok.pre_tokenizer = Sequence([ Split(Regex(SPLIT_PATTERN), behavior="isolated"), @@ -62,6 +100,20 @@ def _build_tokenizer(self) -> Tokenizer: return tok def _build_trainer(self) -> BpeTrainer: + """Create the BPE trainer from the configuration. + + Seeds the initial alphabet with the full byte-level alphabet so every + byte is representable and training stays reversible. + + Returns: + A configured :class:`BpeTrainer` instance. + """ + logger.debug( + "building BPE trainer (vocab_size=%d, min_frequency=%d, special_tokens=%d)", + self.config.vocab_size, + self.config.min_frequency, + len(self.config.special_tokens), + ) return BpeTrainer( vocab_size=self.config.vocab_size, min_frequency=self.config.min_frequency, @@ -71,9 +123,26 @@ def _build_trainer(self) -> BpeTrainer: ) def _batches(self) -> Iterable[List[str]]: - buf, n_yielded = [], 0 + """Yield batches of extracted text for streaming training. + + Iterates the dataset, extracts and skips empty text, and yields lists of + up to ``batch_size`` strings, logging progress periodically. A final + partial batch is yielded if any leftover texts remain. + + Yields: + Lists of non-empty text strings. + """ + buf, n_yielded, n_skipped = [], 0, 0 for sample in self.dataset: - text = self.text_extractor(sample) + try: + text = self.text_extractor(sample) + except Exception as e: + # text_extractor is arbitrary/user-supplied; no specific + # exception type to catch. Skip the bad sample rather than + # aborting a long training run. + n_skipped += 1 + logger.warning("text_extractor raised, skipping sample: %s", e) + continue if not text: continue buf.append(text) @@ -85,9 +154,20 @@ def _batches(self) -> Iterable[List[str]]: buf = [] if buf: yield buf + if n_skipped: + logger.warning("skipped %d samples due to extraction errors", n_skipped) def _attach_post_processor(self) -> None: - """Look up real IDs from the trained tokenizer, not input order.""" + """Attach a BOS/EOS template post-processor to the trained tokenizer. + + Resolves the actual token IDs from the trained vocab (rather than + assuming input order) and installs a :class:`TemplateProcessing` that + wraps single and paired sequences with BOS/EOS. No-op when BOS/EOS are + not configured. + + Raises: + RuntimeError: If a configured BOS/EOS token is absent from the vocab. + """ bos, eos = self.config.bos_token, self.config.eos_token if not (bos and eos): logger.info("bos/eos not configured; skipping post-processor") @@ -104,9 +184,19 @@ def _attach_post_processor(self) -> None: pair=f"{bos} $A {eos} {bos} $B:1 {eos}:1", special_tokens=[(bos, bos_id), (eos, eos_id)], ) + logger.info("attached post-processor (bos=%s:%d, eos=%s:%d)", + bos, bos_id, eos, eos_id) def _verify_round_trip(self) -> None: - """Strict round-trip: no strip, exact match required.""" + """Verify encode/decode is lossless on a set of edge-case phrases. + + Encodes then decodes each phrase in :data:`ROUND_TRIP_TESTS` and requires + an exact match (no stripping). Aggregates and logs all mismatches or + exceptions before failing. + + Raises: + RuntimeError: If any phrase fails to round-trip exactly. + """ failures = [] for phrase in ROUND_TRIP_TESTS: try: @@ -123,11 +213,28 @@ def _verify_round_trip(self) -> None: logger.info("round-trip OK on %d phrases", len(ROUND_TRIP_TESTS)) def _verify_special_tokens(self) -> None: + """Ensure every configured special token exists in the trained vocab. + + Raises: + RuntimeError: If any configured special token is missing. + """ for t in self.config.special_tokens: if self.tokenizer.token_to_id(t) is None: raise RuntimeError(f"special token {t!r} missing from vocab") + logger.info("verified %d special tokens present in vocab", + len(self.config.special_tokens)) def _save_hf(self, raw_tokenizer_path: str, save_dir: str) -> None: + """Wrap the raw tokenizer as a HuggingFace tokenizer and save it. + + Builds a :class:`PreTrainedTokenizerFast` from the raw tokenizer file, + carrying over max length, special tokens, and any configured + unk/pad/bos/eos tokens, then writes it to ``save_dir``. + + Args: + raw_tokenizer_path: Path to the saved raw ``tokenizer.json``. + save_dir: Directory to write the HuggingFace tokenizer files into. + """ kwargs = { "tokenizer_file": raw_tokenizer_path, "model_max_length": self.config.model_max_length, @@ -139,8 +246,17 @@ def _save_hf(self, raw_tokenizer_path: str, save_dir: str) -> None: kwargs[name] = val hf = PreTrainedTokenizerFast(**kwargs) hf.save_pretrained(save_dir) + logger.debug("saved HuggingFace tokenizer to %s", save_dir) def _write_metadata(self, save_dir: str) -> None: + """Write training metadata as JSON alongside the tokenizer. + + Records target vs. actual vocab size, min frequency, special tokens, and + the split pattern to ``training_metadata.json`` in ``save_dir``. + + Args: + save_dir: Directory to write the metadata file into. + """ meta = { "vocab_size_target": self.config.vocab_size, "vocab_size_actual": self.tokenizer.get_vocab_size(), @@ -148,10 +264,34 @@ def _write_metadata(self, save_dir: str) -> None: "special_tokens": self.config.special_tokens, "split_pattern": SPLIT_PATTERN, } - with open(os.path.join(save_dir, "training_metadata.json"), "w") as f: - json.dump(meta, f, indent=2) + try: + with open(os.path.join(save_dir, "training_metadata.json"), "w") as f: + json.dump(meta, f, indent=2) + except OSError as e: + raise RuntimeError(f"failed to write training metadata: {e}") from e + logger.debug("wrote training metadata to %s", save_dir) def train(self, output_dir: str) -> Tokenizer: + """Train the tokenizer end to end and save it to ``output_dir``. + + Builds the tokenizer and trainer, trains from the batched dataset, + attaches the post-processor, verifies special tokens and round trips, + then atomically writes the raw tokenizer, HuggingFace wrapper, and + metadata into ``output_dir`` via a temporary staging directory. + + Args: + output_dir: Directory to save the trained tokenizer artifacts into. + + Returns: + The trained :class:`Tokenizer` instance. + + Raises: + PermissionError: If ``output_dir`` is not writable. + RuntimeError: If training fails, produces an empty vocab, verification + of special tokens or round trips fails, or artifacts cannot be + saved. + """ + logger.info("starting tokenizer training run -> %s", output_dir) os.makedirs(output_dir, exist_ok=True) if not os.access(output_dir, os.W_OK): raise PermissionError(f"output_dir not writable: {output_dir}") @@ -160,29 +300,46 @@ def train(self, output_dir: str) -> Tokenizer: trainer = self._build_trainer() logger.info("training: target_vocab=%d", self.config.vocab_size) - self.tokenizer.train_from_iterator(self._batches(), trainer=trainer) - logger.info("training done: actual_vocab=%d", - self.tokenizer.get_vocab_size()) + try: + self.tokenizer.train_from_iterator(self._batches(), trainer=trainer) + except Exception as e: + raise RuntimeError(f"tokenizer training failed: {e}") from e + + actual_vocab = self.tokenizer.get_vocab_size() + if actual_vocab <= len(self.config.special_tokens): + raise RuntimeError( + f"training produced an empty vocab (size={actual_vocab}); " + f"the dataset may be empty or yielded no usable text" + ) + logger.info("training done: actual_vocab=%d", actual_vocab) self._attach_post_processor() self._verify_special_tokens() self._verify_round_trip() - with tempfile.TemporaryDirectory(dir=os.path.dirname(output_dir) or ".") as tmp: - raw_path = os.path.join(tmp, "tokenizer.json") - self.tokenizer.save(raw_path) - self._save_hf(raw_path, tmp) - self._write_metadata(tmp) - - for name in os.listdir(tmp): - src = os.path.join(tmp, name) - dst = os.path.join(output_dir, name) - if os.path.exists(dst): - if os.path.isdir(dst): - shutil.rmtree(dst) - else: - os.remove(dst) - shutil.move(src, dst) + logger.info("saving tokenizer artifacts to %s", output_dir) + try: + with tempfile.TemporaryDirectory(dir=os.path.dirname(output_dir) or ".") as tmp: + logger.debug("staging artifacts in temp dir %s", tmp) + raw_path = os.path.join(tmp, "tokenizer.json") + self.tokenizer.save(raw_path) + self._save_hf(raw_path, tmp) + self._write_metadata(tmp) + + for name in os.listdir(tmp): + src = os.path.join(tmp, name) + dst = os.path.join(output_dir, name) + if os.path.exists(dst): + if os.path.isdir(dst): + shutil.rmtree(dst) + else: + os.remove(dst) + shutil.move(src, dst) + logger.debug("moved artifact %s -> %s", name, dst) + except OSError as e: + raise RuntimeError( + f"failed to save tokenizer artifacts to {output_dir}: {e}" + ) from e logger.info("saved to %s", output_dir) return self.tokenizer