Upload 7 files

README.md

@@ -0,0 +1,47 @@
+## ChatGCLM-270M
+<<img src="./banner.png" alt="ChatGCLM Hero" width="600">>
+<strong>A high-performance language model architecture.</strong>
+
+---
+
+## Overview
+
+**ChatGCLM** is a generative language model that deviates from the traditional Transformer architecture by utilizing a hybrid approach of **Local** and **Global Convolutions**. By leveraging Fast Fourier Transforms (FFT) for global context, ChatGCLM achieves a massive receptive field with a fraction of the computational overhead associated with standard attention mechanisms.
+
+The architecture is designed for efficiency, speed, and high-quality generation, featuring a custom vocabulary reduction system that optimizes the embedding space for specific datasets.
+
+This repository provides the implementation for training and sampling from the ChatGCLM-270M model, which consists of 270 million parameters.
+
+The model has the full vocabulary of GPT-2, so it can be fine-tuned on any dataset that GPT-2 can be fine-tuned on.
+
+
+## 📦 Installation
+
+Download this repository and extract it.
+
+---
+
+## Usage
+
+### 1. Training the Model
+Place your `.txt` data files in the `data/` directory and run:
+```bash
+python train_chatgclm.py
+```
+This script will build the vocabulary and train the foundation model
+
+### 2. Sample from the model
+Run sample.py to generate text with the model
+```bash
+python sample.py
+```
+
+---
+
+## Fine-tuning
+
+You may fine-tune the model by resuming training from a checkpoint, you may use a different dataset, you may also change parameters such as the learning rate, batch size, etc.
+
+<p align="center">
+  Built with ❤️ by AG 
+</p>

model.py

@@ -0,0 +1,94 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+D_MODEL = 1024 
+N_LAYERS = 22   # Increased to 22 to ensure >270M params
+MAX_SEQ_LEN = 1024  # Reduced from 4096 for 10x speedup
+LOCAL_KERNEL_SIZE = 3 # Reduced from 5
+GLOBAL_KERNEL_SIZE = 512
+USE_GLOBAL_EVERY_N_LAYERS = 2
+FFT_SIZE = 1024 # Match MAX_SEQ_LEN for peak FFT efficiency
+
+class GlobalConv1D(nn.Module):
+    def __init__(self, d_model, kernel_size, fft_size):
+        super().__init__()
+        self.kernel = nn.Parameter(torch.randn(d_model, kernel_size) * 0.01)
+        self.kernel_size = kernel_size
+        self.fft_size = fft_size
+
+    def forward(self, x):
+        B, C, T = x.shape
+        K = min(self.kernel_size, T)
+        overlap = K - 1
+        block = self.fft_size - overlap
+        x = F.pad(x, (overlap, 0))
+        k = self.kernel[:, :K]
+        k = F.pad(k, (0, self.fft_size - K))
+        k_f = torch.fft.rfft(k, n=self.fft_size)
+        outs = []
+        pos = 0
+        while pos < T:
+            seg = x[..., pos:pos+self.fft_size]
+            if seg.shape[-1] < self.fft_size:
+                seg = F.pad(seg, (0, self.fft_size - seg.shape[-1]))
+            y = torch.fft.irfft(torch.fft.rfft(seg, n=self.fft_size) * k_f.unsqueeze(0), n=self.fft_size)
+            outs.append(y[..., overlap:overlap+block])
+            pos += block
+        return torch.cat(outs, dim=-1)[..., :T]
+
+class LocalConv1D(nn.Module):
+    def __init__(self, d_model, k):
+        super().__init__()
+        self.k = k
+        self.dw = nn.Conv1d(d_model, d_model, k, groups=d_model)
+        self.pw = nn.Conv1d(d_model, d_model, 1)
+
+    def forward(self, x):
+        x = F.pad(x, (self.k - 1, 0))
+        return self.pw(F.relu(self.dw(x)))
+
+class Block(nn.Module):
+    def __init__(self, d_model, use_global):
+        super().__init__()
+        self.use_global = use_global
+        self.ln1 = nn.LayerNorm(d_model)
+        self.local = LocalConv1D(d_model, LOCAL_KERNEL_SIZE)
+        if use_global:
+            self.ln2 = nn.LayerNorm(d_model)
+            self.global_conv = GlobalConv1D(d_model, GLOBAL_KERNEL_SIZE, FFT_SIZE)
+        self.ln3 = nn.LayerNorm(d_model)
+        self.ff = nn.Sequential(
+            nn.Linear(d_model, d_model*4),
+            nn.GELU(),
+            nn.Linear(d_model*4, d_model)
+        )
+
+    def forward(self, x):
+        x = x + self.local(self.ln1(x).transpose(1,2)).transpose(1,2)
+        if self.use_global:
+            x = x + self.global_conv(self.ln2(x).transpose(1,2)).transpose(1,2)
+        return x + self.ff(self.ln3(x))
+
+class ChatGCLM(nn.Module):
+    def __init__(self, vocab):
+        super().__init__()
+        self.emb = nn.Embedding(vocab, D_MODEL)
+        self.pos = nn.Embedding(MAX_SEQ_LEN, D_MODEL)
+        self.layers = nn.ModuleList([
+            Block(D_MODEL, i % USE_GLOBAL_EVERY_N_LAYERS == 0)
+            for i in range(N_LAYERS)
+        ])
+        self.ln = nn.LayerNorm(D_MODEL)
+        self.head = nn.Linear(D_MODEL, vocab)
+        self.head.weight = self.emb.weight
+
+    def forward(self, x):
+        T = x.size(1)
+        if T > MAX_SEQ_LEN:
+            x = x[:, -MAX_SEQ_LEN:]
+            T = MAX_SEQ_LEN
+        h = self.emb(x) + self.pos(torch.arange(T, device=x.device))
+        for layer in self.layers:
+            h = layer(h)
+        return self.head(self.ln(h))

prepare_data.py

@@ -0,0 +1,152 @@
+import csv
+import json
+import os
+import sys
+
+# Increase field size limit for large CSV fields
+csv.field_size_limit(sys.maxsize)
+
+def clean_text(text):
+    if not text:
+        return ""
+    return text.strip()
+
+def process_item(system, conversation, output_file):
+    # conversion to text format: 
+    # <system>...
+    # <user>...
+    # <ai>...
+    
+    # Check if we have valid content
+    if not conversation:
+        return
+
+    text_parts = []
+    
+    # Add system if present
+    if system:
+        text_parts.append(f"<system>{clean_text(system)}")
+    
+    # Process conversation turns
+    # conversation is a list of (role, content)
+    # roles: 'user', 'ai' (we map 'assistant'->'ai')
+    
+    for role, content in conversation:
+        content = clean_text(content)
+        if not content:
+            continue
+        
+        if role == 'system':
+            # Handle system in message list if somehow present/overriding
+            text_parts.append(f"<system>{content}")
+        elif role == 'user':
+            text_parts.append(f"<user>{content}")
+        elif role == 'assistant' or role == 'ai':
+            text_parts.append(f"<ai>{content}")
+        else:
+            # Fallback for unknown roles
+            text_parts.append(f"<{role}>{content}")
+            
+    if text_parts:
+        final_str = "\n".join(text_parts)
+        output_file.write(final_str + "\n\n")
+
+def process_csv(filepath, output_path):
+    print(f"Processing CSV: {filepath}")
+    try:
+        with open(filepath, 'r', encoding='utf-8', errors='replace') as f:
+            reader = csv.DictReader(f)
+            with open(output_path, 'a', encoding='utf-8') as out:
+                count = 0
+                for row in reader:
+                    # Mapping logic for this specific CSV structure:
+                    # thread_title -> User context/prompt
+                    # instruction -> System prompt
+                    # message -> AI response
+                    
+                    sys_prompt = row.get('instruction', '')
+                    title = row.get('thread_title', '')
+                    msg = row.get('message', '')
+                    
+                    # If we have mainly 'text' and it looks like it contains everything, we might prefer it?
+                    # But analysis suggested 'text' was just instruction/duplicate in some rows.
+                    # We'll stick to constructing from parts which is safer for structured training.
+                    
+                    conversation = []
+                    if title:
+                        conversation.append(('user', title))
+                    if msg:
+                        conversation.append(('ai', msg))
+                    
+                    if conversation:
+                        process_item(sys_prompt, conversation, out)
+                        count += 1
+                        if count % 10000 == 0:
+                            print(f"CSV Processed {count}...", flush=True)
+            print(f"Finished CSV. Processed {count} rows.")
+    except Exception as e:
+        print(f"Error processing CSV: {e}")
+
+def process_jsonl(filepath, output_path):
+    print(f"Processing JSONL: {filepath}")
+    try:
+        with open(filepath, 'r', encoding='utf-8') as f, open(output_path, 'a', encoding='utf-8') as out:
+            count = 0
+            for line in f:
+                if not line.strip(): continue
+                try:
+                    data = json.loads(line)
+                    messages = data.get('messages', [])
+                    
+                    # Extract system prompt if it exists as a separate field or role
+                    system_prompt = data.get('system', '') 
+                    
+                    conversation = []
+                    
+                    for m in messages:
+                        role = m.get('role', '')
+                        content = m.get('content', '')
+                        
+                        if role == 'system':
+                            # If we hit a system role, treat it as global system or part of flow
+                            # User asked to "add system as <system>"
+                            # I'll just map it directly.
+                            conversation.append(('system', content))
+                        else:
+                            conversation.append((role, content))
+                            
+                    if conversation:
+                        # Pass None for separate system arg since we handle it in loop
+                        process_item(None, conversation, out)
+                        count += 1
+                        
+                    if count % 10000 == 0:
+                        print(f"JSONL Processed {count}...", flush=True)
+                        
+                except json.JSONDecodeError:
+                    continue
+            print(f"Finished JSONL. Processed {count} items.")
+    except Exception as e:
+        print(f"Error processing JSONL: {e}")
+
+def main():
+    data_dir = "data"
+    output_filename = "processed_corpus.txt"
+    output_path = os.path.join(data_dir, output_filename)
+    
+    # Overwrite/Create new
+    with open(output_path, 'w', encoding='utf-8') as f:
+        f.write("")
+        
+    # Process JSONL
+    jsonl_path = os.path.join(data_dir, "dataset.jsonl")
+    if os.path.exists(jsonl_path):
+        process_jsonl(jsonl_path, output_path)
+    
+    # Process CSV
+    csv_path = os.path.join(data_dir, "train.csv")
+    if os.path.exists(csv_path):
+        process_csv(csv_path, output_path)
+
+if __name__ == "__main__":
+    main()

sample.py

@@ -0,0 +1,141 @@
+import os
+import torch
+import torch.nn.functional as F
+from collections import OrderedDict
+import tiktoken
+from model import ChatGCLM, MAX_SEQ_LEN
+
+MODEL_PATH = None
+for f in os.listdir("."):
+    if f.startswith("ChatGCLM_") and f.endswith(".pt"):
+        MODEL_PATH = f
+        break
+
+if MODEL_PATH is None:
+    print("Error: No model checkpoint found!")
+    print("Please train the model first with: python3 train_chatgclm.py")
+    exit(1)
+
+TOKENIZER_NAME = "gpt2"
+EOS_ID = 2
+
+def load_model(device):
+    tok = tiktoken.get_encoding(TOKENIZER_NAME)
+    vocab_size = tok.n_vocab
+    
+    model = ChatGCLM(vocab_size).to(device)
+    if os.path.exists(MODEL_PATH) and os.path.getsize(MODEL_PATH) > 0:
+        print(f"Loading model from: {MODEL_PATH}")
+        ckpt = torch.load(MODEL_PATH, map_location=device)
+
+        # Common checkpoint shapes: either a state_dict, or a dict containing
+        # 'model_state_dict' / 'state_dict' keys.
+        if isinstance(ckpt, dict):
+            if 'model_state_dict' in ckpt:
+                state_dict = ckpt['model_state_dict']
+            elif 'state_dict' in ckpt:
+                state_dict = ckpt['state_dict']
+            else:
+                state_dict = ckpt
+        else:
+            state_dict = ckpt
+
+        # If checkpoint was saved from DataParallel/DistributedDataParallel it may
+        # have a 'module.' prefix on every key. Strip it if present.
+        def _strip_module_prefix(sd):
+            keys = list(sd.keys())
+            if any(k.startswith('module.') for k in keys):
+                new_sd = OrderedDict()
+                for k, v in sd.items():
+                    new_key = k[len('module.'): ] if k.startswith('module.') else k
+                    new_sd[new_key] = v
+                return new_sd
+            return sd
+
+        state_dict = _strip_module_prefix(state_dict)
+
+        # Load with strict=False to allow minor mismatches; report missing/unexpected keys.
+        res = model.load_state_dict(state_dict, strict=False)
+        # res is an _IncompatibleKeys object with missing_keys and unexpected_keys
+        missing = getattr(res, 'missing_keys', None)
+        unexpected = getattr(res, 'unexpected_keys', None)
+        if missing:
+            print(f"Warning: missing keys when loading state_dict: {missing}")
+        if unexpected:
+            print(f"Warning: unexpected keys in state_dict: {unexpected}")
+
+        model.eval()
+        return model, tok
+    else:
+        print(f"Error: Could not load model from {MODEL_PATH}")
+        return None, None
+
+@torch.no_grad()
+def generate(model, prompt, tokenizer, device, max_new_tokens=200, temperature=0.8, top_k=50):
+    model.eval()
+    input_ids = tokenizer.encode(prompt)
+    x = torch.tensor([input_ids], dtype=torch.long, device=device)
+    
+    print(f"\n{'='*70}")
+    print(f"PROMPT: {prompt}")
+    print(f"{'='*70}")
+    print("GENERATED TEXT:")
+    print(prompt, end="", flush=True)
+    
+    generated_tokens = []
+    for _ in range(max_new_tokens):
+        ctx = x[:, -MAX_SEQ_LEN:] if x.size(1) > MAX_SEQ_LEN else x
+        logits = model(ctx)
+        next_token_logits = logits[:, -1, :] / temperature
+        
+        if top_k is not None:
+            v, _ = torch.topk(next_token_logits, min(top_k, next_token_logits.size(-1)))
+            next_token_logits[next_token_logits < v[:, [-1]]] = -float('Inf')
+        
+        probs = F.softmax(next_token_logits, dim=-1)
+        next_token = torch.multinomial(probs, num_samples=1)
+        idx = next_token.item()
+        
+        if idx == EOS_ID:
+            break
+        
+        x = torch.cat((x, next_token), dim=1)
+        generated_tokens.append(idx)
+        token_text = tokenizer.decode([idx])
+        print(token_text, end="", flush=True)
+    
+    print(f"\n{'='*70}\n")
+    return tokenizer.decode(generated_tokens)
+
+if __name__ == "__main__":
+    device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+    print(f"Using device: {device}")
+    
+    model, tokenizer = load_model(device)
+    
+    if model is None:
+        exit(1)
+    
+    test_prompts = [
+        "Once upon a time",
+        "The future of AI is",
+        "In a world where",
+    ]
+    
+    print("\n" + "="*70)
+    print("ChatGCLM Text Generation Demo")
+    print("="*70)
+    
+    for prompt in test_prompts:
+        generate(model, prompt, tokenizer, device, max_new_tokens=150, temperature=0.8, top_k=50)
+    
+    print("\n" + "="*70)
+    print("Interactive Mode - Enter your own prompts!")
+    print("="*70)
+    
+    while True:
+        user_prompt = input("\nEnter prompt (or 'exit' to quit): ")
+        if user_prompt.lower() == 'exit':
+            break
+        if user_prompt.strip():
+            generate(model, user_prompt, tokenizer, device, max_new_tokens=200, temperature=0.8, top_k=50)

train_chatgclm.py

@@ -0,0 +1,272 @@
+import math
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from tqdm import tqdm
+import tiktoken
+import contextlib
+from model import ChatGCLM, MAX_SEQ_LEN
+
+if os.name != "nt":
+    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+if torch.cuda.is_available():
+    torch.set_float32_matmul_precision("high")
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+DATA_DIR = "data"
+DATA_PCT = 0.002
+MPS_SEQ_LEN = 512
+MPS_STEPS_PER_EPOCH = 18
+CPU_SEQ_LEN = 512
+CPU_STEPS_PER_EPOCH = 48
+TOKENIZER_NAME = "gpt2"
+VOCAB_SAVE_PATH = "vocab_map.pt"
+FINE_TUNE = True
+FINE_TUNE_FILE = "chat_data.txt"
+
+EPOCHS = 50
+MICRO_BATCH_SIZE = 8 # Increased for better GPU utilization
+GRAD_ACCUM_STEPS = 4  # Total batch size 32
+STEPS_PER_EPOCH = 500 # Reduced to 500 for ~5-7 min epochs
+LEARNING_RATE = 5e-4
+MIN_LR = 1e-5
+
+SAVE_N_EPOCHS = 1
+
+PAD_ID = 0
+SEP_ID = 1
+EOS_ID = 2
+OFFSET = 3
+
+def build_dataset_vocab(data_dir, tokenizer, save_path):
+    vocab_size = tokenizer.n_vocab
+    torch.save({
+        "vocab_size": vocab_size,
+        "PAD_ID": PAD_ID,
+        "SEP_ID": SEP_ID,
+        "EOS_ID": EOS_ID,
+    }, save_path)
+    return vocab_size
+
+class RemappedTextDataset(Dataset):
+    def __init__(self, ids, max_len):
+        self.ids = ids
+        self.max_len = max_len
+
+    def __len__(self):
+        return max(0, (len(self.ids) - 1) // self.max_len)
+
+    def __getitem__(self, i):
+        start = i * self.max_len
+        x = self.ids[start : start + self.max_len]
+        y = self.ids[start + 1 : start + self.max_len + 1]
+        
+        if len(x) < self.max_len:
+            x = x + [0] * (self.max_len - len(x))
+        if len(y) < self.max_len:
+            y = y + [0] * (self.max_len - len(y))
+            
+        return torch.tensor(x, dtype=torch.long), torch.tensor(y, dtype=torch.long)
+
+def format_params(num):
+    if num >= 1_000_000_000:
+        return f"{num/1_000_000_000:.1f}B"
+    elif num >= 1_000_000:
+        return f"{num/1_000_000:.1f}M"
+    else:
+        return f"{num/1_000:.1f}K"
+
+@torch.no_grad()
+def estimate_loss(model, dl, device, ctx):
+    model.eval()
+    losses = []
+    limit = 50
+    for i, (x, y) in enumerate(dl):
+        if i >= limit: break
+        x, y = x.to(device), y.to(device)
+        with ctx:
+            logits = model(x)
+            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1), ignore_index=PAD_ID)
+        losses.append(loss.item())
+    model.train()
+    return sum(losses) / len(losses) if losses else 0.0
+
+def train():
+    device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+    tok = tiktoken.get_encoding(TOKENIZER_NAME)
+    effective_batch_target = MICRO_BATCH_SIZE * GRAD_ACCUM_STEPS
+    micro_batch_size = MICRO_BATCH_SIZE
+    grad_accum_steps = GRAD_ACCUM_STEPS
+    train_seq_len = MAX_SEQ_LEN
+    steps_per_epoch = STEPS_PER_EPOCH
+
+    if device == "mps":
+        if hasattr(torch, "mps"):
+            torch.mps.empty_cache()
+        micro_batch_size = 1
+        grad_accum_steps = max(1, math.ceil(effective_batch_target / micro_batch_size))
+        train_seq_len = min(MAX_SEQ_LEN, MPS_SEQ_LEN)
+        steps_per_epoch = min(STEPS_PER_EPOCH, MPS_STEPS_PER_EPOCH)
+    elif device == "cpu":
+        micro_batch_size = min(4, MICRO_BATCH_SIZE)
+        grad_accum_steps = max(1, math.ceil(effective_batch_target / micro_batch_size))
+        train_seq_len = min(MAX_SEQ_LEN, CPU_SEQ_LEN)
+        steps_per_epoch = min(STEPS_PER_EPOCH, CPU_STEPS_PER_EPOCH)
+
+    steps_per_epoch = max(1, steps_per_epoch)
+    effective_batch_size = micro_batch_size * grad_accum_steps
+    vocab = build_dataset_vocab(DATA_DIR, tok, VOCAB_SAVE_PATH)
+
+    full_text = ""
+    # Read target text files
+    if FINE_TUNE:
+        candidate = os.path.join(DATA_DIR, FINE_TUNE_FILE)
+        target_files = [FINE_TUNE_FILE] if os.path.isfile(candidate) else []
+        print(f"Fine-tune mode ON -> using {FINE_TUNE_FILE} only")
+        if not target_files:
+            raise FileNotFoundError(f"Expected fine-tune file '{FINE_TUNE_FILE}' in {DATA_DIR}")
+    else:
+        target_files = [f for f in os.listdir(DATA_DIR) if f.endswith(".txt")]
+        target_files.sort()
+    print(f"Loading {len(target_files)} text file(s) from {DATA_DIR}...")
+    for f in target_files:
+        fpath = os.path.join(DATA_DIR, f)
+        print(f" - Reading {f}...")
+        try:
+            with open(fpath, "r", encoding="utf-8") as file:
+                content = file.read()
+                full_text += content + "\n"
+        except Exception as e:
+            print(f"Error reading {f}: {e}")
+
+    print(f"Total dataset size: {len(full_text):,} characters")
+    ids = tok.encode(full_text) + [EOS_ID]
+    if 0 < DATA_PCT < 1.0:
+        target_tokens = max(MAX_SEQ_LEN + 1, int(len(ids) * DATA_PCT))
+        ids = ids[:target_tokens]
+        print(f"Using {DATA_PCT*100:.2f}% of tokens -> {len(ids):,} tokens")
+    else:
+        print(f"Tokenized dataset -> {len(ids):,} tokens")
+    
+    n = len(ids)
+    split_idx = int(n * 0.95) # 95% train, 5% val
+    train_ids = ids[:split_idx]
+    val_ids = ids[split_idx:]
+    
+    train_ds = RemappedTextDataset(train_ids, train_seq_len)
+    val_ds = RemappedTextDataset(val_ids, train_seq_len)
+    
+    # Accelerated DataLoader settings
+    kwargs = {'num_workers': 4, 'pin_memory': True} if device == "cuda" else {}
+    train_dl = DataLoader(train_ds, batch_size=micro_batch_size, shuffle=True, **kwargs)
+    val_dl = DataLoader(val_ds, batch_size=micro_batch_size, shuffle=False, **kwargs)
+
+    model = ChatGCLM(vocab).to(device)
+    
+
+    # Multi-GPU support
+    if torch.cuda.device_count() > 1:
+        print(f"Using {torch.cuda.device_count()} GPUs!")
+        model = nn.DataParallel(model)
+        
+    num_params = sum(p.numel() for p in model.parameters())
+    param_str = format_params(num_params)
+    save_path = f"ChatGCLM_{param_str}.pt"
+    
+    print("-" * 30)
+    print(f"ChatGCLM TRAINING START")
+    print(f"Model ID:    {save_path}")
+    print(f"Parameters:  {num_params:,}")
+    print(f"Device:      {device}")
+    print(f"Vocab Size:  {vocab}")
+    print(f"Learning Rate: {LEARNING_RATE}")
+    print(f"Micro Batch: {micro_batch_size}")
+    print(f"Grad Accum:  {grad_accum_steps}")
+    print(f"Effective Batch: {effective_batch_size}")
+    print(f"Train Seq:  {train_seq_len}")
+    print(f"Epoch Steps: {steps_per_epoch}")
+    print(f"Fine-tune:  {FINE_TUNE}")
+    print(f"Epochs:      {EPOCHS}")
+    print("-" * 30)
+
+    if os.path.exists(save_path) and os.path.getsize(save_path) > 0:
+        print(f"⏳ Found checkpoint at {save_path}, loading...")
+        state_dict = torch.load(save_path, map_location=device)
+        # Handle DataParallel/Module prefix mismatch
+        if isinstance(model, nn.DataParallel):
+            if "module." not in list(state_dict.keys())[0]:
+                 new_state_dict = {f"module.{k}": v for k, v in state_dict.items()}
+                 state_dict = new_state_dict
+        elif "module." in list(state_dict.keys())[0]:
+             # If loading DP checkpoint into non-DP model
+             new_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+             state_dict = new_state_dict
+             
+        model.load_state_dict(state_dict)
+        print("✓ Model weights loaded successfully! Resuming training.")
+    else:
+        print("ℹ No checkpoint found. Starting training from scratch.")
+
+    # Use fused AdamW if available
+    opt_kwargs = {"lr": LEARNING_RATE}
+    if device == "cuda":
+        opt_kwargs["fused"] = True
+    opt = torch.optim.AdamW(model.parameters(), **opt_kwargs)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=EPOCHS, eta_min=MIN_LR)
+    loss_fn = nn.CrossEntropyLoss(ignore_index=PAD_ID)
+    if device == "cuda":
+        ctx = torch.amp.autocast(device_type="cuda")
+        scaler = torch.amp.GradScaler()
+    else:
+        ctx = contextlib.nullcontext()
+        scaler = None
+
+    for ep in range(EPOCHS):
+        model.train()
+        opt.zero_grad(set_to_none=True)
+        total_steps = min(len(train_dl), steps_per_epoch)
+        pbar = tqdm(train_dl, desc=f"Epoch {ep+1}/{EPOCHS}", total=total_steps)
+        running_loss = 0.0
+        steps_since_update = 0
+        for step_idx, (x, y) in enumerate(pbar):
+            if step_idx >= total_steps:
+                break
+            x, y = x.to(device), y.to(device)
+            steps_since_update += 1
+            is_last_batch = (step_idx + 1) == total_steps
+            accum_divisor = grad_accum_steps if not is_last_batch else steps_since_update
+            with ctx:
+                logits = model(x)
+                loss = loss_fn(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
+                loss_val = loss.item()
+                loss = loss / accum_divisor
+            if scaler:
+                scaler.scale(loss).backward()
+            else:
+                loss.backward()
+            should_step = steps_since_update == grad_accum_steps or is_last_batch
+            if should_step:
+                if scaler:
+                    scaler.step(opt)
+                    scaler.update()
+                else:
+                    opt.step()
+                opt.zero_grad(set_to_none=True)
+                if device == "mps" and hasattr(torch, "mps"):
+                    torch.mps.empty_cache()
+                steps_since_update = 0
+            running_loss = 0.9 * running_loss + 0.1 * loss_val if running_loss > 0 else loss_val
+            pbar.set_postfix(loss=f"{running_loss:.4f}")
+        val_loss = estimate_loss(model, val_dl, device, ctx)
+        current_lr = scheduler.get_last_lr()[0]
+        print(f"Epoch {ep+1} | Train Loss: {running_loss:.4f} | Val Loss: {val_loss:.4f} | LR: {current_lr:.6f}")
+        torch.save(model.state_dict(), save_path)
+        print(f"✓ Model saved successfully after epoch {ep+1} to {save_path}")
+        scheduler.step()
+
+if __name__ == "__main__":
+    train()

vocab_map.pt

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