From 4a515453a8c5dad59d97a9e8b0a3476629836277 Mon Sep 17 00:00:00 2001 From: Harikrishnan V Date: Thu, 2 Jul 2026 01:02:22 +0530 Subject: [PATCH 1/2] add error handling, logging, and fix critical import bugs --- README.md | 4 +- configs/model.py | 34 +++ configs/{config.py => training.py} | 19 +- ...ained_prediction.py => sample_pretrain.py} | 49 ++-- scripts/train_model.py | 76 ++++--- src/data_preparation/data_loader.py | 57 +++-- src/data_preparation/make_shards.py | 139 ++++++------ src/logger.py | 15 ++ src/{model_files => models}/__init__.py | 0 src/{model_files => models}/attention.py | 7 +- .../layer_normalization.py | 0 src/{model_files => models}/model.py | 77 +++++-- .../positional_embeddings.py | 0 .../swiglu_feed_forward.py | 0 .../transformer_block.py | 8 +- src/pre_training/model_evaluation.py | 18 -- src/pre_training/model_training.py | 107 --------- src/{pre_training => pretraining}/__init__.py | 0 src/pretraining/evaluator.py | 32 +++ src/pretraining/trainer.py | 138 ++++++++++++ src/tokenizer_train.py | 213 +++++++++++++++--- 21 files changed, 654 insertions(+), 339 deletions(-) create mode 100644 configs/model.py rename configs/{config.py => training.py} (51%) rename scripts/{check_pre_trained_prediction.py => sample_pretrain.py} (63%) create mode 100644 src/logger.py rename src/{model_files => models}/__init__.py (100%) rename src/{model_files => models}/attention.py (85%) rename src/{model_files => models}/layer_normalization.py (100%) rename src/{model_files => models}/model.py (60%) rename src/{model_files => models}/positional_embeddings.py (100%) rename src/{model_files => models}/swiglu_feed_forward.py (100%) rename src/{model_files => models}/transformer_block.py (78%) delete mode 100644 src/pre_training/model_evaluation.py delete mode 100644 src/pre_training/model_training.py rename src/{pre_training => pretraining}/__init__.py (100%) create mode 100644 src/pretraining/evaluator.py create mode 100644 src/pretraining/trainer.py 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/configs/model.py b/configs/model.py new file mode 100644 index 0000000..085f3da --- /dev/null +++ b/configs/model.py @@ -0,0 +1,34 @@ +from dataclasses import dataclass + + +@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 + + 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/config.py b/configs/training.py similarity index 51% rename from configs/config.py rename to configs/training.py index 016271d..2631bf2 100644 --- a/configs/config.py +++ b/configs/training.py @@ -1,18 +1,10 @@ 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: @@ -23,4 +15,11 @@ class TrainingConfig: 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 + + 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/sample_pretrain.py similarity index 63% rename from scripts/check_pre_trained_prediction.py rename to scripts/sample_pretrain.py index 46a1f4f..4eaba6f 100644 --- a/scripts/check_pre_trained_prediction.py +++ b/scripts/sample_pretrain.py @@ -2,19 +2,37 @@ import torch.nn.functional as F from tokenizers import Tokenizer -from configs.config import ModelConfig, TrainingConfig -from src.model_files.model import BetterGPT +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): + 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) + 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] @@ -24,7 +42,7 @@ def load_and_predict(text, model_path, tokenizer_path, device, 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 + ids = ids[:-1] # drop trailing eos token seqs.append(ids) lengths = {len(s) for s in seqs} @@ -33,38 +51,35 @@ def load_and_predict(text, model_path, tokenizer_path, device, f"mask the model lacks. Pass equal-length prompts, or generate one at a time." ) - idx = torch.tensor(seqs, device=device) + 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 + 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 = TrainingConfig().device - + + 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", type=bool, default=True) + 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, @@ -73,6 +88,6 @@ def load_and_predict(text, model_path, tokenizer_path, device, max_tokens=args.max_tokens, temp=args.temp, top_k=args.top_k, - stop_at_eos=args.stop_at_eos + stop_at_eos=args.stop_at_eos, ) - print(output) \ No newline at end of file + 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/make_shards.py b/src/data_preparation/make_shards.py index 0bf7794..6d4b7cd 100644 --- a/src/data_preparation/make_shards.py +++ b/src/data_preparation/make_shards.py @@ -1,78 +1,77 @@ 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= 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): + 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): + 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..9050e7f --- /dev/null +++ b/src/logger.py @@ -0,0 +1,15 @@ +import logging +import os + + +def get_logger(name: str) -> logging.Logger: + 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/__init__.py b/src/models/__init__.py similarity index 100% rename from src/model_files/__init__.py rename to src/models/__init__.py diff --git a/src/model_files/attention.py b/src/models/attention.py similarity index 85% rename from src/model_files/attention.py rename to src/models/attention.py index ad59d41..1f80972 100644 --- a/src/model_files/attention.py +++ b/src/models/attention.py @@ -3,7 +3,7 @@ import torch.nn as nn from torch.nn.functional import scaled_dot_product_attention -class RopeAttention(nn.Module): +class MHAttention(nn.Module): def __init__(self,emb_dim,head_count,head_dim,seq_length,rope): super().__init__() @@ -28,7 +28,6 @@ def forward(self,x,attention_mask): 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) @@ -36,13 +35,10 @@ def forward(self,x,attention_mask): 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) @@ -55,5 +51,4 @@ def forward(self,x,attention_mask): 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/models/layer_normalization.py similarity index 100% rename from src/model_files/layer_normalization.py rename to src/models/layer_normalization.py diff --git a/src/model_files/model.py b/src/models/model.py similarity index 60% rename from src/model_files/model.py rename to src/models/model.py index 5c67438..4640ab9 100644 --- a/src/model_files/model.py +++ b/src/models/model.py @@ -4,11 +4,13 @@ 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 +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): @@ -16,40 +18,42 @@ class BetterGPT(nn.Module): 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): + 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 + 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}" + 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.lm_head = nn.Linear(config.emb_dim, config.vocab_size, bias=False) - self.transformer_block = nn.ModuleList([TransformerBlock( + 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, + seq_length=self.seq_length, rope=rope - ) for _ in range(config.num_blocks)]) + ) for _ in range(config.num_blocks)]) self.apply(self._init_weights) @@ -57,24 +61,20 @@ def __init__(self,config:ModelConfig): 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 + 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): + 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 = block(x, attention_mask) x = self.rmsnorm(x) x = self.lm_head(x) @@ -83,14 +83,31 @@ def forward(self,input_dim,attention_mask=None): @torch.no_grad() def generate(self, idx, max_tokens, temp, top_k=None, eos_id=None): + 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: @@ -117,11 +134,23 @@ def generate(self, idx, max_tokens, temp, top_k=None, eos_id=None): @torch.no_grad() def generate_single(self, idx, max_tokens, temp, top_k=None, eos_id=None): + 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: @@ -138,4 +167,4 @@ def generate_single(self, idx, max_tokens, temp, top_k=None, eos_id=None): break idx = torch.cat([idx, torch.tensor([out_idx], device=idx.device)], dim=0) - return idx \ No newline at end of file + return idx diff --git a/src/model_files/positional_embeddings.py b/src/models/positional_embeddings.py similarity index 100% rename from src/model_files/positional_embeddings.py rename to src/models/positional_embeddings.py diff --git a/src/model_files/swiglu_feed_forward.py b/src/models/swiglu_feed_forward.py similarity index 100% rename from src/model_files/swiglu_feed_forward.py rename to src/models/swiglu_feed_forward.py diff --git a/src/model_files/transformer_block.py b/src/models/transformer_block.py similarity index 78% rename from src/model_files/transformer_block.py rename to src/models/transformer_block.py index 54f6a9a..0aa1c05 100644 --- a/src/model_files/transformer_block.py +++ b/src/models/transformer_block.py @@ -2,9 +2,9 @@ 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 +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): def __init__( @@ -20,7 +20,7 @@ def __init__( super().__init__() self.pre_attn_norm = RMSNorm(emb_dim) self.pre_ffn_norm = RMSNorm(emb_dim) - self.attention = RopeAttention( + self.attention = MHAttention( emb_dim=emb_dim, head_dim=head_dim, head_count=head_count, 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/pre_training/__init__.py b/src/pretraining/__init__.py similarity index 100% rename from src/pre_training/__init__.py rename to src/pretraining/__init__.py diff --git a/src/pretraining/evaluator.py b/src/pretraining/evaluator.py new file mode 100644 index 0000000..c54aafe --- /dev/null +++ b/src/pretraining/evaluator.py @@ -0,0 +1,32 @@ +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): + 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..cd5c6ff --- /dev/null +++ b/src/pretraining/trainer.py @@ -0,0 +1,138 @@ +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, +): + 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 From 7c4237c832b8d155cf0c46f5b6f4640ada1e54bc Mon Sep 17 00:00:00 2001 From: Harikrishnan V Date: Thu, 2 Jul 2026 01:10:39 +0530 Subject: [PATCH 2/2] added docstring --- action | 0 configs/model.py | 5 +++ configs/training.py | 1 + scripts/sample_pretrain.py | 19 +++++++++ src/data_preparation/dataset.py | 37 +++++++++++++++++ src/data_preparation/make_shards.py | 24 +++++++++++- src/logger.py | 4 ++ src/models/attention.py | 61 +++++++++++++++++++---------- src/models/layer_normalization.py | 16 +++++++- src/models/model.py | 37 +++++++++++++++++ src/models/positional_embeddings.py | 46 +++++++++++++++------- src/models/swiglu_feed_forward.py | 27 +++++++++---- src/models/transformer_block.py | 47 +++++++++++++++------- src/pretraining/evaluator.py | 11 ++++++ src/pretraining/trainer.py | 22 +++++++++++ 15 files changed, 299 insertions(+), 58 deletions(-) create mode 100644 action diff --git a/action b/action new file mode 100644 index 0000000..e69de29 diff --git a/configs/model.py b/configs/model.py index 085f3da..87511c4 100644 --- a/configs/model.py +++ b/configs/model.py @@ -3,6 +3,11 @@ @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 diff --git a/configs/training.py b/configs/training.py index 2631bf2..c877838 100644 --- a/configs/training.py +++ b/configs/training.py @@ -8,6 +8,7 @@ @dataclass class TrainingConfig: + """Hyperparameters and runtime settings for pretraining.""" token_count: int = 470_000_000 batch_size: int = 64 learning_rate: float = 3e-4 diff --git a/scripts/sample_pretrain.py b/scripts/sample_pretrain.py index 4eaba6f..4c36f61 100644 --- a/scripts/sample_pretrain.py +++ b/scripts/sample_pretrain.py @@ -12,6 +12,25 @@ 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: 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 6d4b7cd..f04172f 100644 --- a/src/data_preparation/make_shards.py +++ b/src/data_preparation/make_shards.py @@ -9,6 +9,13 @@ class ShardDataset: + """Tokenize a HuggingFace streaming dataset and write it to binary shard files. + + Token IDs are stored as uint16, limiting vocab_size to < 65 000. + Each shard holds buffer_size tokens packed contiguously; the final shard may + be smaller. Shards that already exist on disk are skipped. + """ + def __init__( self, dataset_name: str, @@ -16,8 +23,17 @@ def __init__( out_dir: str, split: str = "train", data_column_name: str = "text", - buffer_size: int = 20_000_000 + buffer_size: int = 20_000_000, ): + """ + Args: + dataset_name: HuggingFace dataset identifier (e.g. 'roneneldan/TinyStories'). + tokenizer: Trained tokenizers.Tokenizer instance. + out_dir: Root output directory; shards are written under /_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" @@ -34,12 +50,18 @@ def __init__( ) 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 diff --git a/src/logger.py b/src/logger.py index 9050e7f..c4050d7 100644 --- a/src/logger.py +++ b/src/logger.py @@ -3,6 +3,10 @@ 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() diff --git a/src/models/attention.py b/src/models/attention.py index 1f80972..ac43931 100644 --- a/src/models/attention.py +++ b/src/models/attention.py @@ -3,52 +3,73 @@ import torch.nn as nn from torch.nn.functional import scaled_dot_product_attention + class MHAttention(nn.Module): - def __init__(self,emb_dim,head_count,head_dim,seq_length,rope): + """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.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", + self.register_buffer( + "mask", torch.tril(torch.ones(seq_length, seq_length)), - persistent=False + persistent=False, ) + def forward(self, x, attention_mask): + """Apply multi-head attention with RoPE and an optional padding mask. - def forward(self,x,attention_mask): + Args: + x: Input tensor of shape (B, T, emb_dim). + attention_mask: Optional bool tensor (B, T); True = attend, False = mask out. - batch,seq_length,emb_dim = x.shape + 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, 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) + 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() + 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) + 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 = scaled_dot_product_attention(rotated_q, rotated_k, v, is_causal=True) - y = y.view(batch,seq_length,emb_dim) + y = y.transpose(1, 2).contiguous() + y = y.view(batch, seq_length, emb_dim) y = self.out_proj(y) - return y \ No newline at end of file + return y diff --git a/src/models/layer_normalization.py b/src/models/layer_normalization.py index 18f650d..627a032 100644 --- a/src/models/layer_normalization.py +++ b/src/models/layer_normalization.py @@ -1,13 +1,27 @@ import torch import torch.nn as nn + class RMSNorm(nn.Module): - def __init__(self,emb_dim,eps:float=1e-6): + """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 index 4640ab9..f049c6b 100644 --- a/src/models/model.py +++ b/src/models/model.py @@ -64,6 +64,7 @@ def __init__(self, config: ModelConfig): 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: @@ -72,6 +73,15 @@ def _init_weights(self, module): 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) @@ -83,6 +93,21 @@ def forward(self, input_dim, attention_mask=None): @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: @@ -134,6 +159,18 @@ def generate(self, idx, max_tokens, temp, top_k=None, eos_id=None): @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: diff --git a/src/models/positional_embeddings.py b/src/models/positional_embeddings.py index 96946ac..7cb5885 100644 --- a/src/models/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 index 27556ec..8749adb 100644 --- a/src/models/swiglu_feed_forward.py +++ b/src/models/swiglu_feed_forward.py @@ -2,16 +2,29 @@ import torch.nn as nn import torch.nn.functional as F + class SwiGLU_FFN(nn.Module): - def __init__(self, emb_dim,hid_dim): + """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) + 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): + 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) - + down = self.down_proj(gate * up) return down diff --git a/src/models/transformer_block.py b/src/models/transformer_block.py index 0aa1c05..7c627c6 100644 --- a/src/models/transformer_block.py +++ b/src/models/transformer_block.py @@ -6,17 +6,25 @@ from src.models.attention import MHAttention from src.models.layer_normalization import RMSNorm + class TransformerBlock(nn.Module): - def __init__( - self, - emb_dim, - hid_dim, - seq_length, - rope, - head_count, - head_dim - ): + """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) @@ -25,11 +33,20 @@ def __init__( 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) + 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. - def forward(self,x, attention_mask): - x = x + self.attention(self.pre_attn_norm(x),attention_mask) + 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 \ No newline at end of file + return x diff --git a/src/pretraining/evaluator.py b/src/pretraining/evaluator.py index c54aafe..b2cfe23 100644 --- a/src/pretraining/evaluator.py +++ b/src/pretraining/evaluator.py @@ -8,6 +8,17 @@ @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 = [] diff --git a/src/pretraining/trainer.py b/src/pretraining/trainer.py index cd5c6ff..57ce835 100644 --- a/src/pretraining/trainer.py +++ b/src/pretraining/trainer.py @@ -25,6 +25,28 @@ def training( 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 = []