Upload 7 files

inference.py

@@ -0,0 +1,244 @@
+# ==============================================================================
+#                          Inference Script
+# ==============================================================================
+# --- 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
+import os
+import glob
+import time
+import datetime
+import traceback
+import dataclasses # Make sure this is imported
+
+# --- Model Configuration ---
+# IMPORTANT: This definition MUST exactly match the one used during training
+#            when the checkpoint was saved.
+@dataclass
+class ModelArgs:
+    # --- ~221M Config used for training step_1200 ---
+    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 # Ensure this matches the value used when saving the checkpoint
+
+    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) ---
+# V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V
+# --- PASTE THE FULL DEFINITIONS OF THE FOLLOWING CLASSES/FUNCTIONS HERE ---
+# --- from your model_architecture.py script:                         ---
+#
+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) # Pass cache=None for non-cached forward
+        h = self.norm(h); output = F.linear(h, self.tok_embeddings.weight); return output # Use tied weights
+
+@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() # CORRECTED: Separate line
+    try:
+        model_device = next(model.parameters()).device
+        model_dtype = next(model.parameters()).dtype
+    except StopIteration:
+        print("Warning: Model has no parameters. Assuming CPU and float32.")
+        model_device = torch.device("cpu")
+        model_dtype = torch.float32
+
+    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
+# --- End Placeholders ---
+
+
+# --- Main Inference Execution ---
+if __name__ == "__main__":
+
+    # --- Configuration for Inference ---
+    # --- !! USE SPECIFIC WINDOWS PATH !! ---
+    raw_checkpoint_path = r".\step_800.pt" # <<< CHANGED to step 1200
+    # --- Normalize the path ---
+    checkpoint_path = os.path.normpath(raw_checkpoint_path)
+    # --- End Adjust ---
+
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    print(f"\n--- Inference Setup ---")
+    print(f"Using device: {device}")
+    print(f"Attempting to load checkpoint: {checkpoint_path}")
+
+    # --- Load Checkpoint and Model Args ---
+    if not os.path.exists(checkpoint_path):
+        # Removed the fallback logic as we are specifying an exact path
+        exit(f"Error: Checkpoint file not found at the specified path: {checkpoint_path}")
+
+    try:
+        # Load checkpoint to CPU first
+        checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False) # weights_only=False needed for ModelArgs
+
+        # Load args dict and instantiate ModelArgs
+        saved_args_data = checkpoint.get('model_args', checkpoint.get('model_args_dict')) # Check both keys
+        if not saved_args_data: exit("Error: model_args not found in checkpoint.")
+        if not isinstance(saved_args_data, dict): saved_args_dict = dataclasses.asdict(saved_args_data)
+        else: saved_args_dict = saved_args_data
+        init_field_names = {f.name for f in dataclasses.fields(ModelArgs) if f.init}
+        filtered_args_dict = {k: v for k, v in saved_args_dict.items() if k in init_field_names}
+        config_inf = ModelArgs(**filtered_args_dict)
+
+        print(f"Loaded model config from checkpoint: {config_inf}")
+
+        # --- Instantiate Model ---
+        model_inf = Llama(config_inf) # Instantiate on CPU
+        print("Model instantiated on CPU.")
+
+        # --- Load Weights ---
+        model_inf.load_state_dict(checkpoint['model_state_dict'])
+        print("Model weights loaded.")
+        model_inf.to(device) # Move model to target device
+        print(f"Model moved to {device}.")
+
+        # --- Prepare for Inference ---
+        model_inf.eval()
+        if device.type == 'cuda':
+            try: model_inf = model_inf.half(); print("Converted loaded model to float16 for inference.")
+            except Exception as e: print(f"Could not convert model to float16: {e}")
+
+    except Exception as e: exit(f"Error loading checkpoint or instantiating model: {e}")
+
+    # --- Load Tokenizer ---
+    tokenizer_name_inf = "deepseek-ai/DeepSeek-R1"
+    print(f"Loading tokenizer: {tokenizer_name_inf}")
+    try:
+        tokenizer_inf = AutoTokenizer.from_pretrained(tokenizer_name_inf, trust_remote_code=True)
+        if tokenizer_inf.pad_token is None:
+            if tokenizer_inf.eos_token: tokenizer_inf.pad_token = tokenizer_inf.eos_token
+            else: tokenizer_inf.add_special_tokens({'pad_token': '[PAD]'})
+        print("Tokenizer loaded.")
+    except Exception as e: exit(f"Error loading tokenizer: {e}")
+
+    # --- Run Generation ---
+    print(f"\n--- Running Generation with Loaded Checkpoint ({os.path.basename(checkpoint_path)}) ---") # Updated print
+    prompt_inf = "Valkyria Chronicles is a tactical role-playing game developed and published by"
+    max_gen_len = 100
+    gen_temperature = 0.7
+    gen_top_k = 50
+    gen_top_p = 0.9
+
+    try:
+        start_time_inf = time.time()
+        _ = generate(
+            model=model_inf, tokenizer=tokenizer_inf, prompt=prompt_inf,
+            max_new_tokens=max_gen_len, temperature=gen_temperature,
+            top_k=gen_top_k, top_p=gen_top_p
+        )
+        end_time_inf = time.time()
+        print(f"\nInference duration: {datetime.timedelta(seconds=int(end_time_inf - start_time_inf))}")
+        print("\n(Output quality depends heavily on limited training. Expect limited coherence.)")
+    except Exception as e: print(f"\nAn error occurred during generation: {e}"); traceback.print_exc()
+
+    print("\n--- Inference Script Section Finished ---")

model_architecture.py

@@ -0,0 +1,519 @@
+# ==============================================================================
+# 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 ---")

step_600.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:530888d8725a62ab23389bea6a1b7d2a116c3f3c1be2b594f5480ab481dbff04
+size 2199232318

step_800.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a2c4b611bfda30d3e87bcc48c1676d4b959da17dd08d00d55f511f1a2b7dd498
+size 2199232318

wikitext2_tokens_128k.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cecec1d0a73cb0418069e202c7890691c82fd8eef1bc7fd8165345cfc8be3e1b
+size 7303192

wikitext2_val_tokens_128k.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:015fabc451eda93a3b3c644a6df177d29234478916d030bba0b86c802e3f0efa
+size 761320