Upload 4 files

chatgclm_base_2.9M.pt

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

sample.py

@@ -0,0 +1,281 @@
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import tiktoken
+
+# ----------------- CONFIG -----------------
+MODEL_PATH = "chatgclm_base_2.9M.pt"
+VOCAB_PATH = "vocab_map.pt"
+TOKENIZER_NAME = "gpt2"
+
+# Defined in training script
+D_MODEL = 256
+N_LAYERS = 4
+MAX_SEQ_LEN = 1024
+LOCAL_KERNEL_SIZE = 5
+GLOBAL_KERNEL_SIZE = 256
+USE_GLOBAL_EVERY_N_LAYERS = 2
+FFT_SIZE = 1024
+
+PAD_ID = 0
+SEP_ID = 1
+EOS_ID = 2
+OFFSET = 3
+# ------------------------------------------
+
+# ----------------- MODEL DEF -----------------
+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 GCLM(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)
+        
+        # Weight tying
+        self.head.weight = self.emb.weight
+
+    def forward(self, x):
+        T = x.size(1)
+        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))
+
+
+# ----------------- UTILS -----------------
+def load_model_and_vocab(device):
+    if not os.path.exists(VOCAB_PATH):
+        print(f"[ERROR] Vocab file not found: {VOCAB_PATH}")
+        return None, None, None
+        
+    vocab_data = torch.load(VOCAB_PATH, map_location="cpu")
+    used_tokens = vocab_data["used_tokens"]
+    id2new = vocab_data["id2new"]
+    vocab_size = len(used_tokens) + OFFSET
+    
+    print(f"[INFO] Vocab loaded. Size: {vocab_size}")
+
+    model = GCLM(vocab_size).to(device)
+    
+    if os.path.exists(MODEL_PATH):
+        print(f"[INFO] Loading model from {MODEL_PATH}...")
+        state_dict = torch.load(MODEL_PATH, map_location=device)
+        model.load_state_dict(state_dict)
+        model.eval()
+    else:
+        print(f"[ERROR] Model file not found: {MODEL_PATH}")
+        return None, None, None
+
+    return model, used_tokens, id2new
+
+@torch.no_grad()
+def generate(model, prompt, tokenizer, id2new, used_tokens, device, max_new_tokens=200, temperature=0.8, top_k=50):
+    model.eval()
+    
+    # Encode prompt
+    raw_ids = tokenizer.encode(prompt)
+    input_ids = []
+    
+    # Map to model IDs
+    for rid in raw_ids:
+        if rid in id2new:
+            input_ids.append(id2new[rid])
+        else:
+            # Skip unknown tokens
+            continue
+            
+    if not input_ids:
+        print("[WARN] No known tokens in prompt.")
+        input_ids = [PAD_ID] # Should not happen ideally
+        
+    x = torch.tensor([input_ids], dtype=torch.long, device=device)
+
+    generated = []
+    
+    for _ in range(max_new_tokens):
+        # Crop to max seq len
+        if x.size(1) > MAX_SEQ_LEN:
+             ctx = x[:, -MAX_SEQ_LEN:]
+        else:
+            ctx = x
+
+        logits = model(ctx)
+        next_token_logits = logits[:, -1, :] / temperature
+        
+        # Optional: Top-k sampling
+        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.append(idx)
+        
+    # Decode result
+    decoded_text = decoder(generated, used_tokens, tokenizer)
+    return decoded_text
+
+def decoder(ids, used_tokens, tokenizer):
+    raw_ids = []
+    for i in ids:
+        if i >= OFFSET:
+            raw_ids.append(used_tokens[i - OFFSET])
+    return tokenizer.decode(raw_ids)
+
+
+
+# ----------------- MAIN -----------------
+if __name__ == "__main__":
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        device = "mps"
+    else:
+        device = "cpu"
+        
+    print(f"Using device: {device}")
+    
+    model, used_tokens, id2new = load_model_and_vocab(device)
+    enc = tiktoken.get_encoding(TOKENIZER_NAME)
+    
+    if model:
+        # Find a good starting token ID (e.g., newline or space)
+        newline_id = id2new.get(enc.encode("\n")[0], OFFSET)
+        
+        while True:
+            print(f"\n--- Generating Sample (Temp=0.8, TopK=50) ---")
+            print("-" * 20)
+            
+            x = torch.tensor([[newline_id]], dtype=torch.long, device=device)
+            generated = []
+            
+            with torch.no_grad():
+                for _ in range(500):
+                    if x.size(1) > MAX_SEQ_LEN:
+                         ctx = x[:, -MAX_SEQ_LEN:]
+                    else:
+                        ctx = x
+
+                    logits = model(ctx)
+                    logits = logits[:, -1, :] / 0.8 # Temperature
+                    
+                    # Top-k
+                    v, _ = torch.topk(logits, min(50, logits.size(-1)))
+                    logits[logits < v[:, [-1]]] = -float('Inf')
+
+                    probs = F.softmax(logits, dim=-1)
+                    next_token = torch.multinomial(probs, num_samples=1)
+                    
+                    idx = next_token.item()
+                    x = torch.cat((x, next_token), dim=1)
+                    generated.append(idx)
+                    
+                    if idx == EOS_ID:
+                        print("[EOS]", end="", flush=True)
+                        break
+                    
+                    if idx >= OFFSET:
+                        raw_id = used_tokens[idx - OFFSET]
+                        token_text = enc.decode([raw_id])
+                        print(token_text, end="", flush=True)
+                    elif idx == PAD_ID:
+                        print("[PAD]", end="", flush=True)
+                    elif idx == SEP_ID:
+                        print("[SEP]", end="", flush=True)
+
+            print("\n" + "-"*20)
+            cont = input("\nPress [Enter] to generate again, or type 'exit': ")
+            if cont.lower() == 'exit':
+                break
+
+

train_gclm_base.py

@@ -0,0 +1,378 @@
+print("Starting...")
+
+###############################################
+# CONFIGURATION — CUSTOMIZE EVERYTHING HERE
+###############################################
+
+# ---- data / vocab ----
+TXT_PATH = "data.txt"
+DATA_PCT = 0.001 # this is small for testing purposes
+TOKENIZER_NAME = "gpt2"
+REDUCE_VOCAB = True
+VOCAB_SAVE_PATH = "vocab_map.pt"
+
+# ---- training ----
+EPOCHS = 25
+MICRO_BATCH_SIZE = 1
+GRAD_ACCUM_STEPS = 8
+LEARNING_RATE = 3e-4
+
+# ---- model ----
+D_MODEL = 256
+N_LAYERS = 4
+MAX_SEQ_LEN = 1024
+
+LOCAL_KERNEL_SIZE = 5
+GLOBAL_KERNEL_SIZE = 256
+USE_GLOBAL_EVERY_N_LAYERS = 2
+
+# ---- FFT conv ----
+FFT_SIZE = 1024   # must be power of 2 and > GLOBAL_KERNEL_SIZE
+
+# ---- checkpointing ----
+SAVE_PATH = "model.pt"
+SAVE_N_EPOCHS = 1
+
+# ---- device ----
+USE_DEVICE = "cuda"
+USE_AMP = True
+USE_ACTIVATION_CHECKPOINTING = False
+
+# ---- torch.compile ----
+COMPILE = False
+COMPILE_MODE = "reduce-overhead"
+COMPILE_BACKEND = "eager"
+
+###############################################
+# END CONFIG
+###############################################
+
+import os
+
+# Windows cannot use expandable_segments — only enable on Linux.
+if os.name != "nt":
+    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+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
+
+# performance settings
+if torch.cuda.is_available():
+    torch.set_float32_matmul_precision("high")
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+###############################################################
+# SPECIAL TOKENS
+###############################################################
+
+PAD_ID = 0
+SEP_ID = 1
+EOS_ID = 2
+OFFSET = 3
+
+###############################################################
+# VOCAB
+###############################################################
+
+def build_dataset_vocab(txt_path, tokenizer, save_path):
+    text = open(txt_path, "r", encoding="utf-8").read()
+    if DATA_PCT < 1.0:
+        text = text[:int(len(text) * DATA_PCT)]
+    token_ids = tokenizer.encode(text)
+    used = sorted(set(token_ids))
+
+    id2new = {tok: i + OFFSET for i, tok in enumerate(used)}
+
+    torch.save({
+        "used_tokens": used,
+        "id2new": id2new,
+        "PAD_ID": PAD_ID,
+        "SEP_ID": SEP_ID,
+        "EOS_ID": EOS_ID,
+    }, save_path)
+
+    print(f"[OK] Vocab size: {len(used) + OFFSET}")
+    return used, id2new
+
+
+###############################################################
+# DATASET
+###############################################################
+
+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) - self.max_len - 1)
+
+    def __getitem__(self, i):
+        x = self.ids[i:i+self.max_len]
+        y = self.ids[i+1:i+self.max_len+1]
+        return torch.tensor(x, dtype=torch.long), torch.tensor(y, dtype=torch.long)
+
+
+###############################################################
+# GLOBAL + LOCAL CONVOLUTION
+###############################################################
+
+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 GCLM(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)
+        
+        # Weight tying: SIGNIFICANTLY reduces parameter count
+        self.head.weight = self.emb.weight
+
+    def forward(self, x):
+        T = x.size(1)
+        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))
+
+
+###############################################################
+# TRAINING LOOP
+###############################################################
+
+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 = []
+    # Check up to 50 batches for validation to save time
+    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():
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        device = "mps"
+    else:
+        device = "cpu"
+    print("[INFO] Device:", device)
+
+    # 1. Prepare Vocab & Data
+    tok = tiktoken.get_encoding(TOKENIZER_NAME)
+    
+    # We call this to generate/load the vocab map
+    used, id2new = build_dataset_vocab(TXT_PATH, tok, VOCAB_SAVE_PATH)
+    vocab = len(used) + OFFSET
+
+    # Load and process full text
+    print("[INFO] Loading and tokenizing text...")
+    text = open(TXT_PATH, "r", encoding="utf-8").read()
+    if DATA_PCT < 1.0:
+        text = text[:int(len(text) * DATA_PCT)]
+    
+    raw_ids = tok.encode(text)
+    # Map to new IDs
+    ids = [id2new.get(i, PAD_ID) for i in raw_ids] + [EOS_ID]
+    
+    # Split Train/Val (90/10)
+    n = len(ids)
+    split_idx = int(n * 0.9)
+    train_ids = ids[:split_idx]
+    val_ids = ids[split_idx:]
+    
+    print(f"[INFO] Tokens: {n} | Train: {len(train_ids)} | Val: {len(val_ids)}")
+
+    train_ds = RemappedTextDataset(train_ids, MAX_SEQ_LEN)
+    val_ds = RemappedTextDataset(val_ids, MAX_SEQ_LEN)
+
+    train_dl = DataLoader(train_ds, batch_size=MICRO_BATCH_SIZE, shuffle=True)
+    val_dl = DataLoader(val_ds, batch_size=MICRO_BATCH_SIZE, shuffle=False)
+
+    model = GCLM(vocab).to(device)
+    
+    # Calculate params
+    num_params = sum(p.numel() for p in model.parameters())
+    param_str = format_params(num_params)
+    save_path = f"chatgclm_base_{param_str}.pt"
+    print(f"[INFO] Model parameters: {num_params:,} ({param_str})")
+    print(f"[INFO] Save path: {save_path}")
+
+    # 🔁 RESUME IF CHECKPOINT EXISTS
+    if os.path.exists(save_path):
+        model.load_state_dict(torch.load(save_path, map_location=device))
+        print(f"[RESUME] Loaded existing checkpoint from {save_path}")
+
+    if device == "cuda" and COMPILE:
+        print("[INFO] Compiling model with torch.compile...")
+        model = torch.compile(
+            model,
+            mode=COMPILE_MODE,
+            fullgraph=False,
+            backend=COMPILE_BACKEND
+        )
+
+    opt = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
+    loss_fn = nn.CrossEntropyLoss(ignore_index=PAD_ID)
+
+    # AMP Context
+    if device == "cuda" and USE_AMP:
+        ctx = torch.amp.autocast(device)
+        scaler = torch.amp.GradScaler(device)
+    else:
+        # Dummy context for cpu/mps
+        import contextlib
+        ctx = contextlib.nullcontext()
+        scaler = None
+
+    for ep in range(EPOCHS):
+        print(f"\nEpoch {ep+1}/{EPOCHS}")
+        opt.zero_grad(set_to_none=True)
+
+        pbar = tqdm(train_dl, desc="Training")
+        running_loss = 0.0
+        
+        for i, (x, y) in enumerate(pbar):
+            x, y = x.to(device), y.to(device)
+
+            with ctx:
+                logits = model(x)
+                loss = loss_fn(logits.reshape(-1, vocab), y.reshape(-1))
+                loss_val = loss.item()
+                loss = loss / GRAD_ACCUM_STEPS
+
+            if scaler:
+                scaler.scale(loss).backward()
+            else:
+                loss.backward()
+
+            if (i+1) % GRAD_ACCUM_STEPS == 0:
+                if scaler:
+                    scaler.step(opt)
+                    scaler.update()
+                else:
+                    opt.step()
+                opt.zero_grad(set_to_none=True)
+
+            # Update progress bar
+            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}")
+
+        # Validate at end of epoch
+        val_loss = estimate_loss(model, val_dl, device, ctx)
+        print(f"Epoch {ep+1} finished. Train Loss: {running_loss:.4f} | Val Loss: {val_loss:.4f}")
+
+        if SAVE_N_EPOCHS and (ep+1) % SAVE_N_EPOCHS == 0:
+            torch.save(model.state_dict(), save_path)
+            print(f"[OK] Saved checkpoint to {save_path}")
+
+    torch.save(model.state_dict(), save_path)
+    print("[DONE] Training complete.")
+
+
+###############################################################
+# ENTRY POINT
+###############################################################
+
+if __name__ == "__main__":
+    train()

vocab_map.pt

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