# ============================================================================== # Single-File Script ~221M Model - Resume Training for ~4 Hours # ============================================================================== # --- Necessary Imports --- import torch import torch.nn as nn from dataclasses import dataclass, field import math import torch.nn.functional as F from transformers import AutoTokenizer, get_cosine_schedule_with_warmup import os from tqdm import tqdm import traceback # Corrected import: Added IterableDataset AND Dataset from torch.utils.data import Dataset, DataLoader, IterableDataset import torch.optim as optim # Use torch.amp imports (recommended over torch.cuda.amp) from torch.amp import GradScaler, autocast from datasets import load_dataset, IterableDataset as HFIterableDataset import datetime import random import matplotlib.pyplot as plt import glob import time import dataclasses # Make sure this is imported # --- Model Configuration --- @dataclass class ModelArgs: # --- ~221M Config for 4GB VRAM --- hidden_size: int = 768; num_hidden_layers: int = 12; num_attention_heads: int = 12 num_key_value_heads: int = 12; intermediate_size: int = 2048; vocab_size: int = 128000 rms_norm_eps: float = 1e-5; rope_theta: float = 500000.0; max_position_embeddings: int = 4096 head_dim: int = field(init=False) add_recency_bias: bool = False # Keep this option if desired def __post_init__(self): self.head_dim = self.hidden_size // self.num_attention_heads if self.hidden_size % self.num_attention_heads != 0: raise ValueError("hidden_size % num_attention_heads != 0") if self.num_attention_heads % self.num_key_value_heads != 0: raise ValueError("num_attention_heads % num_key_value_heads != 0") # --- Model Components (RMSNorm, RoPE funcs, Attention, FeedForward, TransformerBlock, Llama) --- class RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__(); self.eps = eps; self.weight = nn.Parameter(torch.ones(dim)) def _norm(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) def forward(self, x): original_dtype = x.dtype; output = self._norm(x.float()).to(original_dtype); return output * self.weight def precompute_theta_pos_frequencies(head_dim: int, seq_len: int, device: str | torch.device, theta: float = 10000.0): if head_dim % 2 != 0: raise ValueError("head_dim must be even for RoPE") theta_indices = torch.arange(0, head_dim, 2).float(); theta_freqs = 1.0 / (theta**(theta_indices / head_dim)) target_device = torch.device(device) if isinstance(device, str) else device; theta_freqs = theta_freqs.to(target_device) positions = torch.arange(seq_len, device=target_device).float(); freqs = torch.outer(positions, theta_freqs).float(); return freqs, positions def apply_rotary_embeddings(x: torch.Tensor, freqs_cis_full: torch.Tensor, positions: torch.Tensor): positions = positions.long(); max_pos = freqs_cis_full.shape[0] if torch.max(positions) >= max_pos: positions = torch.clamp(positions, max=max_pos - 1) freqs = freqs_cis_full[positions]; freqs = freqs.unsqueeze(0).unsqueeze(2) bsz, seq_len, n_part_heads, head_dim = x.shape; x1 = x[..., : head_dim // 2]; x2 = x[..., head_dim // 2 :] cos_freqs = torch.cos(freqs).type_as(x); sin_freqs = torch.sin(freqs).type_as(x) rotated_x1 = x1 * cos_freqs - x2 * sin_freqs; rotated_x2 = x1 * sin_freqs + x2 * cos_freqs rotated_x = torch.cat([rotated_x1, rotated_x2], dim=-1); return rotated_x.type_as(x) class Attention(nn.Module): def __init__(self, args: ModelArgs): super().__init__(); self.args = args; self.num_heads = args.num_attention_heads; self.num_kv_heads = args.num_key_value_heads self.head_dim = args.head_dim; self.repeats = self.num_heads // self.num_kv_heads self.wq = nn.Linear(args.hidden_size, args.num_attention_heads * args.head_dim, bias=False) self.wk = nn.Linear(args.hidden_size, args.num_key_value_heads * args.head_dim, bias=False) self.wv = nn.Linear(args.hidden_size, args.num_key_value_heads * args.head_dim, bias=False) self.wo = nn.Linear(args.num_attention_heads * args.head_dim, args.hidden_size, bias=False) def _repeat_kv(self, x: torch.Tensor, n_rep: int) -> torch.Tensor: bsz, n_kv_heads, seqlen, head_dim = x.shape; if n_rep == 1: return x return (x[:, :, None, :, :].expand(bsz, n_kv_heads, n_rep, seqlen, head_dim).reshape(bsz, n_kv_heads * n_rep, seqlen, head_dim)) def _create_recency_bias(self, seqlen, full_seqlen, device, dtype, bias_strength=0.1, decay_rate=0.9): bias = torch.zeros((1, 1, seqlen, full_seqlen), device=device, dtype=dtype); indices = torch.arange(full_seqlen, device=device) rel_pos = torch.arange(seqlen, device=device).unsqueeze(1) - indices.unsqueeze(0); mask = rel_pos >= 0 decaying_bias = bias_strength * (decay_rate ** (-rel_pos[mask])); bias[:, :, mask] = decaying_bias.type_as(bias); return bias def forward(self, x: torch.Tensor, freqs_cis_full: torch.Tensor, positions: torch.Tensor, mask: torch.Tensor | None = None, cache: tuple[torch.Tensor, torch.Tensor] | None = None) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: bsz, seqlen, _ = x.shape; xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) xq = xq.view(bsz, seqlen, self.num_heads, self.head_dim); xk = xk.view(bsz, seqlen, self.num_kv_heads, self.head_dim); xv = xv.view(bsz, seqlen, self.num_kv_heads, self.head_dim) xq = apply_rotary_embeddings(xq, freqs_cis_full, positions); xk = apply_rotary_embeddings(xk, freqs_cis_full, positions) xk = xk.transpose(1, 2); xv = xv.transpose(1, 2) if cache is not None: cache_k, cache_v = cache; keys = torch.cat((cache_k.to(xk.device), xk), dim=2); values = torch.cat((cache_v.to(xv.device), xv), dim=2) else: keys = xk; values = xv updated_cache = (keys.detach(), values.detach()); keys_repeated = self._repeat_kv(keys, self.repeats); values_repeated = self._repeat_kv(values, self.repeats) xq = xq.transpose(1, 2); scores = torch.matmul(xq.float(), keys_repeated.transpose(-2, -1).float()) / math.sqrt(self.head_dim) if self.args.add_recency_bias: full_seqlen = keys_repeated.shape[-2]; recency_bias = self._create_recency_bias(seqlen, full_seqlen, device=scores.device, dtype=scores.dtype); scores = scores + recency_bias if mask is not None: full_seqlen = keys_repeated.shape[-2]; expected_mask_shape_end = (seqlen, full_seqlen) if mask.shape[-2:] != expected_mask_shape_end: try: mask_slice = mask[:, :, -seqlen:, :full_seqlen]; scores = scores + mask_slice.float() except Exception: pass else: scores = scores + mask.float() scores = nn.functional.softmax(scores, dim=-1).type_as(xq); output = torch.matmul(scores, values_repeated) output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1); output = self.wo(output); return output, updated_cache class FeedForward(nn.Module): def __init__(self, args: ModelArgs): super().__init__(); self.w1 = nn.Linear(args.hidden_size, args.intermediate_size, bias=False); self.w3 = nn.Linear(args.hidden_size, args.intermediate_size, bias=False); self.w2 = nn.Linear(args.intermediate_size, args.hidden_size, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x)) class TransformerBlock(nn.Module): def __init__(self, args: ModelArgs): super().__init__(); self.args = args; self.attention_norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps); self.attention = Attention(args); self.ffn_norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps); self.feed_forward = FeedForward(args) def forward(self, x: torch.Tensor, freqs_cis_full: torch.Tensor, positions: torch.Tensor, mask: torch.Tensor | None = None, cache: tuple[torch.Tensor, torch.Tensor] | None = None) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: r, cache = self.attention(self.attention_norm(x), freqs_cis_full, positions, mask, cache); h = x + r; r = self.feed_forward(self.ffn_norm(h)); out = h + r; return out, cache class Llama(nn.Module): def __init__(self, args: ModelArgs): super().__init__(); self.args = args; self.tok_embeddings = nn.Embedding(args.vocab_size, args.hidden_size); self.layers = nn.ModuleList([TransformerBlock(args) for _ in range(args.num_hidden_layers)]) self.norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps); self.tok_embeddings.weight.requires_grad = True freqs_cis, _ = precompute_theta_pos_frequencies(args.head_dim, args.max_position_embeddings, device='cpu', theta=args.rope_theta) self.register_buffer("freqs_cis", freqs_cis, persistent=False) def forward(self, tokens: torch.Tensor, positions: torch.Tensor): bsz, seqlen = tokens.shape; h = self.tok_embeddings(tokens); freqs_cis_full = self.freqs_cis.to(h.device); mask = None if seqlen > 1: mask = torch.full((1, 1, seqlen, seqlen), float("-inf"), device=h.device); mask = torch.triu(mask, diagonal=1).type_as(h) positions = positions.to(h.device) for layer in self.layers: h, _ = layer(h, freqs_cis_full, positions, mask, cache=None) h = self.norm(h); output = F.linear(h, self.tok_embeddings.weight); return output # --- Generate function (Added Top-P Sampling) --- @torch.no_grad() def generate(model: Llama, tokenizer: AutoTokenizer, prompt: str, max_new_tokens: int, temperature: float = 1.0, top_k: int | None = None, top_p: float | None = None): model.eval() try: model_device = next(model.parameters()).device; model_dtype = next(model.parameters()).dtype except StopIteration: model_device = torch.device("cpu"); model_dtype = torch.float32; print("Warning: Model has no parameters.") prompt_ids = tokenizer.encode(prompt, add_special_tokens=True); tokens = torch.tensor([prompt_ids], dtype=torch.long, device=model_device) cache = [(torch.zeros((1, model.args.num_key_value_heads, 0, model.args.head_dim), device=model_device, dtype=model_dtype), torch.zeros((1, model.args.num_key_value_heads, 0, model.args.head_dim), device=model_device, dtype=model_dtype)) for _ in range(model.args.num_hidden_layers)] generated_token_ids = []; current_tokens = tokens; print(f"Generating {max_new_tokens} tokens from prompt: '{prompt}'"); print("Output: ", end='') full_freqs_cis = model.freqs_cis.to(model_device) for i in range(max_new_tokens): current_seq_len = current_tokens.shape[1]; start_pos = cache[0][0].shape[2]; positions = torch.arange(start_pos, start_pos + current_seq_len, device=model_device) current_mask = None; if i == 0 and current_seq_len > 1: current_mask = torch.full((1, 1, current_seq_len, current_seq_len), float("-inf"), device=model_device); current_mask = torch.triu(current_mask, diagonal=1).type(model_dtype) h = model.tok_embeddings(current_tokens); updated_cache_list = [] for layer_idx, layer in enumerate(model.layers): h, updated_layer_cache = layer(h, full_freqs_cis, positions, current_mask, cache[layer_idx]); updated_cache_list.append(updated_layer_cache) cache = updated_cache_list; h = model.norm(h); logits = F.linear(h, model.tok_embeddings.weight) next_token_logits = logits[:, -1, :] if temperature == 0: next_token_id = torch.argmax(next_token_logits, dim=-1, keepdim=True) else: next_token_logits = next_token_logits / temperature if top_k is not None and top_k > 0: v, _ = torch.topk(next_token_logits, min(top_k, next_token_logits.size(-1))); next_token_logits[next_token_logits < v[:, [-1]]] = float('-inf') if top_p is not None and 0.0 < top_p < 1.0: probs_for_filter = F.softmax(next_token_logits, dim=-1); probs_sort, probs_idx = torch.sort(probs_for_filter, descending=True); probs_sum = torch.cumsum(probs_sort, dim=-1) mask_top_p = probs_sum > top_p; mask_top_p[..., 0] = False; mask_top_p[..., 1:] = mask_top_p[..., :-1].clone(); indices_to_remove = mask_top_p.scatter(1, probs_idx, mask_top_p); next_token_logits[indices_to_remove] = float('-inf') probs = F.softmax(next_token_logits, dim=-1); next_token_id = torch.multinomial(probs, num_samples=1) if tokenizer.eos_token_id is not None and next_token_id.item() == tokenizer.eos_token_id: print("\n[EOS token reached]"); break next_token_id_item = next_token_id.item(); generated_token_ids.append(next_token_id_item); current_tokens = next_token_id.clone() print(tokenizer.decode([next_token_id_item]), end='', flush=True) if len(generated_token_ids) >= max_new_tokens: break print("\n--- Generation Complete ---"); final_token_ids = prompt_ids + generated_token_ids; full_generated_text = tokenizer.decode(final_token_ids, skip_special_tokens=False) print(f"\nFull generated text:\n{full_generated_text}"); return full_generated_text # --- Dataset Class (Map Style for WikiText) --- class SimpleLMDataset(Dataset): def __init__(self, token_ids: list[int], sequence_length: int): self.token_ids = token_ids; self.sequence_length = sequence_length self.num_sequences = max(0, len(token_ids) - sequence_length) if self.num_sequences == 0: raise ValueError(f"Dataset token count ({len(token_ids)}) not > sequence length ({sequence_length}).") def __len__(self): return self.num_sequences def __getitem__(self, idx): chunk = self.token_ids[idx : idx + self.sequence_length + 1] if len(chunk) < self.sequence_length + 1: last_valid_idx = len(self.token_ids) - self.sequence_length - 1 chunk = self.token_ids[last_valid_idx : last_valid_idx + self.sequence_length + 1] input_ids = torch.tensor(chunk[:-1], dtype=torch.long); target_ids = torch.tensor(chunk[1:], dtype=torch.long) return input_ids, target_ids # --- Dataset Class (Iterable for SlimPajama - Kept for reference/fallback) --- class TokenizedSequenceDataset(IterableDataset): def __init__(self, dataset_name, dataset_config, split, tokenizer, sequence_length, buffer_size=10000): try: self.dataset = load_dataset(dataset_name, dataset_config, split=split, streaming=True); print(f"Successfully loaded streaming dataset: {dataset_name} ({split})") except Exception as e: raise RuntimeError(f"Failed to load streaming dataset {dataset_name} ({split}): {e}") from e self.tokenizer = tokenizer; self.sequence_length = sequence_length; self.buffer_size = buffer_size; self.buffer = [] try: self.iter_dataset = iter(self.dataset) except Exception as e: raise RuntimeError(f"Failed to create iterator for dataset {dataset_name} ({split}): {e}") from e def __iter__(self): while True: while len(self.buffer) < self.sequence_length + 1: try: item = next(self.iter_dataset); text = item.get('text', '') if text and text.strip(): token_ids = self.tokenizer.encode(text, add_special_tokens=False); self.buffer.extend(token_ids) except StopIteration: if len(self.buffer) < self.sequence_length + 1: return else: break if len(self.buffer) < self.sequence_length + 1: return chunk = self.buffer[:self.sequence_length + 1]; input_ids = torch.tensor(chunk[:-1], dtype=torch.long); target_ids = torch.tensor(chunk[1:], dtype=torch.long) yield input_ids, target_ids; self.buffer = self.buffer[1:] # --- Checkpoint Loading Function --- def load_checkpoint(checkpoint_dir: str, model: Llama, optimizer, scaler, scheduler, device): latest_checkpoint_path = None; highest_step = -1 if os.path.isdir(checkpoint_dir): checkpoints = glob.glob(os.path.join(checkpoint_dir, "step_*.pt")) for ckpt_path in checkpoints: try: step = int(os.path.basename(ckpt_path).split('_')[-1].split('.')[0]); except ValueError: continue if step > highest_step: highest_step = step; latest_checkpoint_path = ckpt_path if latest_checkpoint_path: print(f"Loading checkpoint from: {latest_checkpoint_path}") try: checkpoint = torch.load(latest_checkpoint_path, map_location='cpu', weights_only=False) # Use False for safety current_args_dict = model.args.__dict__ saved_args_data = checkpoint.get('model_args', checkpoint.get('model_args_dict')) if not saved_args_data: print("Warning: Checkpoint missing model_args."); saved_args_dict=None; args_match=False elif not isinstance(saved_args_data, dict): saved_args_dict = dataclasses.asdict(saved_args_data) # Use imported module else: saved_args_dict = saved_args_data args_match = True if saved_args_dict: for f in dataclasses.fields(ModelArgs): # Use dataclasses.fields if f.init and f.name != 'head_dim': current_val = current_args_dict.get(f.name); saved_val = saved_args_dict.get(f.name) if current_val != saved_val: print(f"Mismatch in arg '{f.name}': Current={current_val}, Saved={saved_val}"); args_match = False; break else: args_match = False if not args_match: print("ERROR: Model args mismatch. Cannot load checkpoint."); return 0 model.load_state_dict(checkpoint['model_state_dict']); model.to(device) if optimizer is not None: try: optimizer.load_state_dict(checkpoint['optimizer_state_dict']) except Exception as e: print(f"Warning: Could not load optimizer state dict: {e}") for state in optimizer.state.values(): for k, v in state.items(): if isinstance(v, torch.Tensor): state[k] = v.to(device) if scaler is not None: try: scaler.load_state_dict(checkpoint['scaler_state_dict']) except Exception as e: print(f"Warning: Could not load scaler state dict: {e}") if scheduler is not None: try: scheduler.load_state_dict(checkpoint['scheduler_state_dict']) except Exception as e: print(f"Warning: Could not load scheduler state dict: {e}") start_step = checkpoint['step']; print(f"Resuming training from step {start_step + 1}"); return start_step except Exception as e: print(f"Error loading checkpoint {latest_checkpoint_path}: {e}"); traceback.print_exc(); return 0 else: print("No checkpoint found. Starting training from scratch."); return 0 # --- Plotting Function --- def plot_loss(train_losses, val_losses, val_steps_list, checkpoint_dir, start_step=0): plt.figure(figsize=(12, 6)); smoothing_window = 10 train_steps = list(range(start_step + 1, start_step + len(train_losses) + 1)) plt.plot(train_steps, train_losses, label='Training Loss (Raw)', alpha=0.3) if len(train_losses) > smoothing_window: train_losses_smoothed = [sum(train_losses[max(0, i-smoothing_window):i+1])/min(i+1, smoothing_window) for i in range(len(train_losses))] plt.plot(train_steps, train_losses_smoothed, label=f'Training Loss (Smoothed {smoothing_window} steps)', alpha=0.9) if val_losses and val_steps_list: plt.plot(val_steps_list, val_losses, label='Validation Loss', marker='o', linestyle='--') plt.xlabel("Optimizer Steps"); plt.ylabel("Loss"); plt.yscale('log'); plt.title("Training and Validation Loss Over Steps") plt.legend(); plt.grid(True); plot_filename = f"loss_plot_step_{start_step}_to_{start_step+len(train_losses)}.png" plot_path = os.path.join(checkpoint_dir, plot_filename); plt.savefig(plot_path) print(f"Loss plot saved to {plot_path}") # --- Basic Training Function (Single GPU, AMP, LR Schedule, Validation, Checkpointing, Plotting) --- def simple_train( model: Llama, tokenizer: AutoTokenizer, train_dataset: IterableDataset | Dataset, val_dataset: IterableDataset | Dataset | None, optimizer: torch.optim.Optimizer, criterion: nn.Module, scheduler, num_epochs: int, device: torch.device, gradient_accumulation_steps: int = 1, use_amp: bool = False, max_train_steps: int | None = None, start_step: int = 0, save_interval: int = 1000, checkpoint_dir: str = ".", validation_interval: int = 500, val_steps: int = 50 ): model.train(); total_steps = start_step; global_step_this_run = 0 scaler = GradScaler(enabled=use_amp and device.type == 'cuda') os.makedirs(checkpoint_dir, exist_ok=True) train_loss_history = []; val_loss_history = []; val_steps_history = [] print(f"\n--- Starting Training (Resuming from step {start_step}, Target Steps: {max_train_steps if max_train_steps else 'N/A'}) ---") print(f"--- (AMP: {use_amp and device.type == 'cuda'}) ---") is_iterable = isinstance(train_dataset, IterableDataset) train_loader = DataLoader(train_dataset, batch_size=1, num_workers=0, shuffle=(not is_iterable)) if val_dataset: val_loader = DataLoader(val_dataset, batch_size=1, num_workers=0) training_complete = False # Adjust tqdm total based on remaining steps tqdm_total = (max_train_steps - start_step) if max_train_steps is not None else None print(f"Starting loop, aiming for {max_train_steps} total steps...") # Use total=None for iterable datasets if max_steps not set, as length is unknown pbar = tqdm(total=tqdm_total, desc=f"Optim Steps ({start_step}...)") # Need to manually track iterations vs optimizer steps data_iterator = iter(train_loader) accum_count = 0 # Counter for gradient accumulation steps while not training_complete: # Check if we need to stop before starting the next optimizer step if max_train_steps is not None and total_steps >= max_train_steps: training_complete = True; break # --- Accumulation Loop --- accum_loss = 0.0 optimizer.zero_grad() # Zero gradients at start of accumulation cycle for _ in range(gradient_accumulation_steps): try: input_ids, target_ids = next(data_iterator) except StopIteration: print("\nDataLoader exhausted within accumulation cycle or epoch.") # If loader exhausted before completing max_steps, stop training training_complete = True; break # Break inner accum loop input_ids = input_ids.to(device); target_ids = target_ids.to(device) seqlen = input_ids.shape[1]; positions = torch.arange(seqlen, device=device) # Use torch.amp.autocast with autocast(device_type=device.type, dtype=torch.float16, enabled=use_amp and device.type == 'cuda'): logits = model(input_ids, positions) loss = criterion(logits.view(-1, logits.size(-1)).float(), target_ids.view(-1)) loss = loss / gradient_accumulation_steps # Normalize loss for accumulation scaler.scale(loss).backward() accum_loss += loss.item() # Accumulate *normalized* loss item if training_complete: break # Exit outer loop if data exhausted # --- Optimizer Step --- scaler.unscale_(optimizer) # Optional: torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) scaler.step(optimizer); scaler.update(); scheduler.step(); optimizer.zero_grad(set_to_none=True) total_steps += 1; global_step_this_run += 1 pbar.update(1) # Update progress bar by one optimizer step # --- Logging --- current_loss = accum_loss * gradient_accumulation_steps # Log un-normalized loss for this step train_loss_history.append(current_loss) # Note: epoch_loss calculation might be less meaningful with iterable dataset and max_steps # avg_loss_so_far = sum(train_loss_history[-50:]) / min(len(train_loss_history), 50) # Example: rolling average pbar.set_postfix({"Loss": f"{current_loss:.4f}", "LR": f"{scheduler.get_last_lr()[0]:.6f}", "Steps": total_steps}) # --- Validation --- if val_dataset and total_steps % validation_interval == 0 and total_steps > 0: model.eval(); val_loss = 0.0; val_batches = 0; print(f"\nRunning validation at step {total_steps}...") val_pbar = tqdm(enumerate(val_loader), total=val_steps, desc="Validation") with torch.no_grad(): val_iter = iter(val_loader) for val_step in range(val_steps): try: val_input_ids, val_target_ids = next(val_iter) val_input_ids = val_input_ids.to(device); val_target_ids = val_target_ids.to(device) val_seqlen = val_input_ids.shape[1]; val_positions = torch.arange(val_seqlen, device=device) with autocast(device_type=device.type, dtype=torch.float16, enabled=use_amp and device.type == 'cuda'): val_logits = model(val_input_ids, val_positions) v_loss = criterion(val_logits.view(-1, val_logits.size(-1)).float(), val_target_ids.view(-1)) val_loss += v_loss.item(); val_batches += 1; val_pbar.update(1); val_pbar.set_postfix({"Val Loss": f"{val_loss/val_batches:.4f}"}) except StopIteration: print("Validation loader exhausted early."); break val_pbar.close() avg_val_loss = val_loss / val_batches if val_batches > 0 else float('inf') val_loss_history.append(avg_val_loss); val_steps_history.append(total_steps) print(f"Validation finished. Average Val Loss: {avg_val_loss:.4f}"); model.train() # --- Checkpointing --- if total_steps % save_interval == 0 and total_steps > 0: save_path = os.path.join(checkpoint_dir, f"step_{total_steps}.pt") try: model_args_dict = dataclasses.asdict(model.args) save_content = { 'step': total_steps, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'scaler_state_dict': scaler.state_dict(), 'scheduler_state_dict': scheduler.state_dict(), 'model_args_dict': model_args_dict } torch.save(save_content, save_path); print(f"\nCheckpoint saved to {save_path}") except Exception as e: print(f"\nError saving checkpoint: {e}") # --- Check Max Steps --- if max_train_steps is not None and total_steps >= max_train_steps: print(f"\nReached max_train_steps ({max_train_steps}). Stopping training."); training_complete = True; break pbar.close() # Close pbar if loop finishes naturally print("--- Training Finished ---") return train_loss_history, val_loss_history, val_steps_history # --- Main Execution Block --- if __name__ == "__main__": # --- Configuration --- config = ModelArgs(add_recency_bias=False) # Use ~221M config device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(f"Model Configuration:\n{config}") print(f"Calculated Head Dimension: {config.head_dim}") print(f"\nUsing device: {device}") # --- Component Tests (Commented out) --- """ """ # --- Tokenizer --- print("\n--- Tokenizer Loading ---") tokenizer_name = "deepseek-ai/DeepSeek-R1" print(f"Loading tokenizer: {tokenizer_name}") try: tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True) print("Tokenizer loaded successfully.") if tokenizer.vocab_size != config.vocab_size: exit(f"FATAL: Tokenizer vocab size mismatch!") else: print(f"Tokenizer vocab size ({tokenizer.vocab_size}) matches model config.") if tokenizer.pad_token is None: if tokenizer.eos_token: tokenizer.pad_token = tokenizer.eos_token; print(f"Set PAD token to EOS token: {tokenizer.pad_token}") else: tokenizer.add_special_tokens({'pad_token': '[PAD]'}); print("Added a generic [PAD] token.") except Exception as e: exit(f"Error loading tokenizer '{tokenizer_name}': {e}") # --- Training Setup --- print("\n--- Training Setup ---") train_batch_size = 1 train_seq_len = 256 grad_accum_steps = 16 use_amp_training = True if device.type == 'cuda' else False learning_rate = 1e-4 # Lower LR num_epochs = 1 # --- ADJUSTED MAX STEPS for ~4 hour run --- max_steps_for_run = 1200 # Absolute target step for this run (start_step + new_steps) # --- ADJUSTED Total Scheduler Steps (longer term goal) --- total_scheduler_steps = 10000 # Example longer goal warmup_steps = 100 # --- Save to current directory --- checkpoint_dir = "." save_interval = 200 # Save less frequently validation_interval = 100 # Validate less frequently val_steps = 20 # --- Dataset --- print("\nLoading and preparing WikiText-2 dataset...") train_dataset, val_dataset = None, None try: # Using WikiText-2 directly token_file = "./wikitext2_tokens_128k.pt" val_token_file = "./wikitext2_val_tokens_128k.pt" force_remake_dataset = False if os.path.exists(token_file) and os.path.exists(val_token_file) and not force_remake_dataset: print(f"Loading tokenized data from {token_file} and {val_token_file}...") all_token_ids = torch.load(token_file) all_val_token_ids = torch.load(val_token_file) print("Tokenized data loaded.") else: print("Token files not found or remake forced, processing WikiText-2...") print("Processing train split...") train_raw_dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train") train_full_text = "\n".join([item['text'] for item in train_raw_dataset if item['text'].strip()]) all_token_ids = tokenizer.encode(train_full_text) torch.save(all_token_ids, token_file) print(f"Saved tokenized train data ({len(all_token_ids)} tokens) to {token_file}") print("Processing validation split...") val_raw_dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="validation") val_full_text = "\n".join([item['text'] for item in val_raw_dataset if item['text'].strip()]) all_val_token_ids = tokenizer.encode(val_full_text) torch.save(all_val_token_ids, val_token_file) print(f"Saved tokenized validation data ({len(all_val_token_ids)} tokens) to {val_token_file}") if len(all_token_ids) <= train_seq_len: exit("Train dataset too short.") if len(all_val_token_ids) <= train_seq_len: exit("Validation dataset too short.") train_dataset = SimpleLMDataset(all_token_ids, sequence_length=train_seq_len) val_dataset = SimpleLMDataset(all_val_token_ids, sequence_length=train_seq_len) print("Using WikiText-2 dataset.") except Exception as e: exit(f"Dataset error: {e}") # DataLoaders train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True, num_workers=2, pin_memory=True) val_loader = DataLoader(val_dataset, batch_size=train_batch_size, num_workers=2, pin_memory=True) if val_dataset else None print(f"DataLoaders created. Training Seq Len: {train_seq_len}") print(f"Train sequences: {len(train_dataset)}, Val sequences: {len(val_dataset) if val_dataset else 0}") # --- Model, Optimizer, Scheduler, Loss --- train_model = Llama(config).to(device) print(f"Training model instantiated ({'float32' if not use_amp_training else 'mixed precision'}). Recency Bias: {config.add_recency_bias}") total_params_train = sum(p.numel() for p in train_model.parameters() if p.requires_grad) print(f"Total Trainable Parameters: {total_params_train / 1e6:.2f} Million") optimizer = optim.AdamW(train_model.parameters(), lr=learning_rate, weight_decay=0.1) criterion = nn.CrossEntropyLoss() scheduler = get_cosine_schedule_with_warmup( optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_scheduler_steps ) scaler = GradScaler(enabled=use_amp_training and device.type == 'cuda') print(f"Optimizer: AdamW, Base LR: {learning_rate}") print(f"Scheduler: Cosine with {warmup_steps} warmup steps up to {total_scheduler_steps} steps") print(f"Loss Function: CrossEntropyLoss") # --- Load Checkpoint --- # Pass optimizer, scaler, scheduler to be loaded start_step = load_checkpoint(checkpoint_dir, train_model, optimizer, scaler, scheduler, device) # Calculate steps to run in this session steps_to_run_this_session = max(0, max_steps_for_run - start_step) # The absolute step number to stop at in this run current_run_target_step = start_step + steps_to_run_this_session if steps_to_run_this_session <= 0: print(f"Already completed or exceeded target steps ({max_steps_for_run}). Exiting.") exit() # --- Run Training --- print(f"\n--- Running Training (Will run for {steps_to_run_this_session} steps in this session, target total: {max_steps_for_run}) ---") start_time = time.time() train_loss_hist, val_loss_hist, val_steps_hist = [], [], [] try: # Pass the absolute target step for this run train_loss_hist, val_loss_hist, val_steps_hist = simple_train( model=train_model, tokenizer=tokenizer, train_dataset=train_dataset, val_dataset=val_dataset, optimizer=optimizer, criterion=criterion, scheduler=scheduler, num_epochs=num_epochs, device=device, gradient_accumulation_steps=grad_accum_steps, use_amp=use_amp_training, max_train_steps=current_run_target_step, start_step=start_step, save_interval=save_interval, checkpoint_dir=checkpoint_dir, # Pass "." validation_interval=validation_interval, val_steps=val_steps ) print("\nTraining loop finished.") end_time = time.time(); print(f"Training duration for this session: {datetime.timedelta(seconds=int(end_time - start_time))}") # --- Plotting --- if train_loss_hist: # Adjust steps for plotting if resuming plot_train_steps = list(range(start_step + 1, start_step + len(train_loss_hist) + 1)) # Filter validation steps/losses that occurred *during this run* plot_val_steps = [s for s in val_steps_history if s >= start_step] plot_val_loss = [val_loss_history[i] for i, s in enumerate(val_steps_history) if s >= start_step] plot_loss(train_loss_hist, plot_val_loss, plot_val_steps, checkpoint_dir, start_step=start_step) # Pass "." # --- Generation After Training --- print("\n--- Generation After Training ---") train_model.eval() if device.type == 'cuda': try: train_model = train_model.half(); print("Trained model converted to float16 for generation.") except Exception as e: print(f"Could not convert trained model to float16: {e}.") test_prompt_after = "The meaning of life is" _ = generate(model=train_model, tokenizer=tokenizer, prompt=test_prompt_after, max_new_tokens=60, temperature=0.7, top_k=50, top_p=0.9) print("\n(Check if output shows more structure than random)") except torch.cuda.OutOfMemoryError: print("\n--- CUDA Out of Memory during Training ---"); print("Try reducing train_seq_len or gradient_accumulation_steps further.") except Exception as e: print(f"\nAn error occurred during training: {e}"); traceback.print_exc() print("\n--- Script Finished ---")