Uploaded Training / Testing File / Eval

Make sure to update file paths if intend to use this.

<3

baseline_eval.py

@@ -0,0 +1,73 @@
+"""Baseline evaluation - compare Loop model to standard Qwen3."""
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from datasets import load_dataset
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+MODEL_PATH = "/content/Qwen3-0.6B"
+BATCH_SIZE = 8
+MAX_LENGTH = 256
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=torch.float16).to(device)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+tokenizer.pad_token = tokenizer.eos_token
+
+print("\n2. Loading validation data...")
+dataset = load_dataset("wikitext", "wikitext-2-raw-v1")
+
+def tokenize_fn(examples):
+    return tokenizer(examples["text"], truncation=True, max_length=MAX_LENGTH, padding="max_length")
+
+tokenized = dataset.map(tokenize_fn, batched=True, remove_columns=["text"])
+tokenized = tokenized.filter(lambda x: sum(1 for t in x["input_ids"] if t != tokenizer.pad_token_id) > 10)
+
+val_data = tokenized["validation"]
+print(f"   Validation samples: {len(val_data)}")
+
+def collate_fn(batch):
+    input_ids = torch.tensor([x["input_ids"] for x in batch])
+    attention_mask = torch.tensor([x["attention_mask"] for x in batch])
+    labels = input_ids.clone()
+    labels[attention_mask == 0] = -100
+    return {"input_ids": input_ids.to(device), "attention_mask": attention_mask.to(device), "labels": labels.to(device)}
+
+loader = DataLoader(val_data, batch_size=BATCH_SIZE, collate_fn=collate_fn)
+
+print("\n3. Calculating Baseline Loss...")
+model.eval()
+total_loss = 0
+steps = 0
+
+with torch.no_grad():
+    for batch in tqdm(loader, desc="Evaluating"):
+        with torch.amp.autocast('cuda', dtype=torch.float16):
+            outputs = model(**batch)
+        total_loss += outputs.loss.item()
+        steps += 1
+
+baseline_loss = total_loss / steps
+baseline_ppl = torch.exp(torch.tensor(baseline_loss)).item()
+
+print("\n" + "=" * 60)
+print("RESULTS")
+print("=" * 60)
+print(f"Baseline Qwen3-0.6B Loss: {baseline_loss:.4f}")
+print(f"Baseline Qwen3-0.6B PPL:  {baseline_ppl:.2f}")
+print(f"")
+print(f"Loop Attention Loss:      3.5549 (Epoch 3)")
+print(f"Loop Attention PPL:       35.01")
+print(f"")
+
+if baseline_loss > 3.5549:
+    delta = baseline_loss - 3.5549
+    print(f"✅ SUCCESS: Loop Attention beats baseline by {delta:.4f}!")
+elif abs(baseline_loss - 3.5549) < 0.05:
+    print("📊 NEUTRAL: Loop Attention matches baseline (within noise).")
+else:
+    delta = 3.5549 - baseline_loss
+    print(f"📉 Loop Attention is {delta:.4f} behind baseline.")
+print("=" * 60)

test_loop_generation.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python3
+"""Test generation with Loop Attention (use_cache=False)."""
+
+import sys
+import torch
+sys.path.insert(0, '/content')
+from modeling_qwen_loop import Qwen3LoopForCausalLM
+from transformers import AutoTokenizer
+
+MODEL_PATH = "/content/Qwen3-0.6B"
+GATE_PATH = "/content/Qwen3-0.6B-looped/checkpoints/gate_projections_epoch_3.pt"
+
+print("\n1. Loading model...")
+model = Qwen3LoopForCausalLM.from_pretrained(MODEL_PATH)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+
+print("2. Loading trained gates...")
+gate_state = torch.load(GATE_PATH, map_location=device)
+for key, value in gate_state.items():
+    parts = key.split('.')
+    layer_idx = int(parts[1])
+    param_name = parts[-1]
+    if param_name == 'weight':
+        model.model.layers[layer_idx].self_attn.gate.weight.data = value.to(device)
+    elif param_name == 'bias':
+        model.model.layers[layer_idx].self_attn.gate.bias.data = value.to(device)
+print("   Gates loaded!")
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+tokenizer.pad_token = tokenizer.eos_token
+
+model.eval()
+
+prompts = [
+    "The capital of France is",
+    "def fibonacci(n):",
+    "In the year 2050,",
+    "The quick brown fox",
+    "Explain quantum computing in simple terms:"
+]
+
+
+for prompt in prompts:
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+ 
+        out = model.generate(
+            input_ids=inputs.input_ids,
+            max_new_tokens=50,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            use_cache=False, 
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+    text = tokenizer.decode(out[0], skip_special_tokens=True)
+    print(f"\nPrompt: {prompt}")
+    print(f"Output: {text}")
+

train.py

@@ -0,0 +1,179 @@
+import os
+import sys
+import time
+import json
+import torch
+import glob
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+torch.set_float32_matmul_precision('high')
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer
+from datasets import load_dataset
+from tqdm import tqdm
+
+sys.path.insert(0, '/content/Qwen3-0.6B-looped')
+from modeling_qwen_loop import Qwen3LoopForCausalLM
+
+MODEL_PATH = "/content/Qwen3-0.6B"
+OUTPUT_DIR = "/content/Qwen3-0.6B-looped/checkpoints"
+BATCH_SIZE = 20
+GRADIENT_ACCUMULATION_STEPS = 4
+LEARNING_RATE = 1e-4
+MAX_LENGTH = 1024
+NUM_EPOCHS = 3
+NUM_WORKERS = 8
+PIN_MEMORY = True
+
+print("=" * 60)
+print("TRAINING v3: Optimized (Compile + Workers + Checkpointing)")
+print("=" * 60)
+
+print("\n1. Loading model...")
+checkpoints = sorted(glob.glob(f"{OUTPUT_DIR}/epoch_*"))
+start_epoch = 0
+
+if checkpoints:
+    latest_checkpoint = checkpoints[-1]
+    print(f"   Resuming from checkpoint: {latest_checkpoint}")
+    model = Qwen3LoopForCausalLM.from_pretrained(MODEL_PATH)
+    state_path = os.path.join(latest_checkpoint, "pytorch_model.bin")
+    if os.path.exists(state_path):
+        model.load_state_dict(torch.load(state_path))
+    else:
+        print("   Warning: Checkpoint found but pytorch_model.bin missing. Starting fresh.")
+
+    try:
+        start_epoch = int(latest_checkpoint.split("_")[-1])
+        print(f"   Resuming at Epoch {start_epoch + 1}")
+    except:
+        start_epoch = 0
+else:
+    model = Qwen3LoopForCausalLM.from_pretrained(MODEL_PATH)
+
+device = torch.device("cuda")
+model = model.to(device)
+
+print("\n2. Unfreezing gates + layer norms...")
+model.enable_gate_and_layernorm_training()
+
+print("   Compiling model with torch.compile()...")
+try:
+    model = torch.compile(model)
+except Exception as e:
+    print(f"   Warning: torch.compile failed (ignoring): {e}")
+
+print("\n3. Loading WikiText-2...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+tokenizer.pad_token = tokenizer.eos_token
+
+dataset = load_dataset("wikitext", "wikitext-2-raw-v1")
+
+def tokenize_fn(examples):
+    return tokenizer(examples["text"], truncation=True, max_length=MAX_LENGTH, padding="max_length")
+
+tokenized = dataset.map(tokenize_fn, batched=True, remove_columns=["text"])
+tokenized = tokenized.filter(lambda x: sum(1 for t in x["input_ids"] if t != tokenizer.pad_token_id) > 10)
+
+print(f"   Train samples: {len(tokenized['train'])}")
+print(f"   Val samples: {len(tokenized['validation'])}")
+
+def collate_fn(batch):
+    input_ids = torch.tensor([x["input_ids"] for x in batch])
+    attention_mask = torch.tensor([x["attention_mask"] for x in batch])
+    labels = input_ids.clone()
+    labels[attention_mask == 0] = -100
+    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}
+
+train_loader = DataLoader(
+    tokenized["train"], 
+    batch_size=BATCH_SIZE, 
+    shuffle=True, 
+    collate_fn=collate_fn,
+    num_workers=NUM_WORKERS,
+    pin_memory=PIN_MEMORY
+)
+val_loader = DataLoader(
+    tokenized["validation"], 
+    batch_size=BATCH_SIZE, 
+    shuffle=False, 
+    collate_fn=collate_fn,
+    num_workers=NUM_WORKERS,
+    pin_memory=PIN_MEMORY
+)
+
+optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=0.01)
+total_steps = len(train_loader) * NUM_EPOCHS // GRADIENT_ACCUMULATION_STEPS
+warmup_steps = total_steps // 10
+
+def get_lr(step):
+    if step < warmup_steps:
+        return step / warmup_steps
+    return max(0.1, 1.0 - (step - warmup_steps) / (total_steps - warmup_steps))
+
+scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, get_lr)
+
+print("\n4. Training Configuration:")
+print(f"   Context length: {MAX_LENGTH}")
+print(f"   Batch size: {BATCH_SIZE} (Effective: {BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS})")
+print(f"   Workers: {NUM_WORKERS}")
+print(f"   Total steps: {total_steps}")
+
+print("\n" + "=" * 60)
+print("Starting Training...")
+print("=" * 60)
+
+scaler = torch.amp.GradScaler('cuda')
+model.train()
+global_step = 0
+start_time = time.time()
+
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+for epoch in range(start_epoch, NUM_EPOCHS):
+    epoch_loss = 0
+    epoch_steps = 0
+    progress = tqdm(train_loader, desc=f"Epoch {epoch+1}/{NUM_EPOCHS}")
+
+    for step, batch in enumerate(progress):
+        batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
+
+        with torch.amp.autocast('cuda', dtype=torch.bfloat16):
+            outputs = model(**batch, use_cache=False)
+            loss = outputs.loss / GRADIENT_ACCUMULATION_STEPS
+
+        scaler.scale(loss).backward()
+        epoch_loss += loss.item() * GRADIENT_ACCUMULATION_STEPS
+        epoch_steps += 1
+
+        if (step + 1) % GRADIENT_ACCUMULATION_STEPS == 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            scaler.step(optimizer)
+            scaler.update()
+            scheduler.step()
+            optimizer.zero_grad()
+            global_step += 1
+
+            current_lr = scheduler.get_last_lr()[0]
+            mem_usage = torch.cuda.memory_allocated() / 1024**3
+            progress.set_postfix(loss=loss.item() * GRADIENT_ACCUMULATION_STEPS, lr=current_lr, mem=f"{mem_usage:.1f}GB")
+
+
+    print(f"Saving checkpoint for Epoch {epoch+1}...")
+
+    model_to_save = model._orig_mod if hasattr(model, '_orig_mod') else model
+    model_to_save.save_pretrained(f"{OUTPUT_DIR}/epoch_{epoch+1}")
+    
+    gate_state_dict = {k: v for k, v in model_to_save.state_dict().items() if 'gate' in k}
+    torch.save(gate_state_dict, f"{OUTPUT_DIR}/gate_projections.pt")
+    torch.save(gate_state_dict, f"{OUTPUT_DIR}/gate_projections_epoch_{epoch+1}.pt")
+
+print("Training complete.")
+
+model_to_save = model._orig_mod if hasattr(model, '_orig_mod') else model
+model_to_save.save_pretrained(f"{OUTPUT_DIR}/final")
+gate_state_dict = {k: v for k, v in model_to_save.state_dict().items() if 'gate' in k}
+torch.save(gate_state_dict, f"{OUTPUT_DIR}/gate_projections.pt")