feat: SAGE v2 upgrade (GQA, CoT, W&B, FineWeb-Edu)

requirements.txt

@@ -6,3 +6,4 @@ tqdm>=4.66.1
 bitsandbytes>=0.41.0
 numpy<2.0.0
 huggingface_hub>=0.20.0
+wandb>=0.16.0

sage/config.py

@@ -7,6 +7,7 @@ class SageConfig:
     # Model dimensions corresponding to T4 (16GB VRAM) fit
     d_model: int = 512
     n_heads: int = 8
+    n_kv_heads: int = 4   # GQA: must divide n_heads
     n_layers: int = 6
     d_ff: int = 2048
     
@@ -32,6 +33,7 @@ class SageConfig:
     
     # Checkpointing and path details
     checkpoint_dir: str = "checkpoints"
+    project_name: str = "sage-v2"
     
     @property
     def device(self):

sage/data.py

@@ -79,7 +79,7 @@ class StreamingTextDataset(IterableDataset):
 
     def __init__(
         self,
-        dataset_name: str = "roneneldan/TinyStories",
+        dataset_name: str = "HuggingFaceFW/fineweb-edu",
         split: str = "train",
         seq_len: int = 512,
         tokenizer: Optional[SageTokenizer] = None,
@@ -97,6 +97,13 @@ class StreamingTextDataset(IterableDataset):
         self.text_field = text_field
         self.min_doc_len = min_doc_len
         self.max_doc_len = max_doc_len
+        
+        # Auto-adjust configuration based on popular datasets
+        if "fineweb-edu" in dataset_name.lower():
+            self.text_field = "text"
+            self.split = "train" if split == "train" else split
+        elif "tinystories" in dataset_name.lower():
+            self.text_field = "text"
 
     def _stream_tokens(self) -> Iterator[int]:
         """Yields individual token IDs from the HuggingFace dataset stream."""
@@ -164,7 +171,7 @@ class StreamingTextDataset(IterableDataset):
 
 def create_dataloader(
     config: SageConfig,
-    dataset_name: str = "roneneldan/TinyStories",
+    dataset_name: str = "HuggingFaceFW/fineweb-edu",
     split: str = "train",
     tokenizer: Optional[SageTokenizer] = None,
 ) -> DataLoader:

sage/finetune.py

@@ -16,6 +16,7 @@ import torch
 import torch.nn as nn
 from torch.amp import GradScaler, autocast
 from tqdm import tqdm
+import wandb
 from typing import Optional, List
 
 from .config import SageConfig
@@ -156,6 +157,13 @@ def finetune_instruction(
     if use_lora:
         model = inject_lora(model, rank=lora_rank)
 
+    # ------- W&B Logging -------
+    wandb.init(
+        project=config.project_name,
+        name=f"finetune-{time.strftime('%Y%m%d-%H%M')}",
+        config=config.__dict__,
+    )
+
     optimizer = create_optimizer(model, config)
 
     # AMP setup
@@ -213,6 +221,10 @@ def finetune_instruction(
             avg = accum_loss / 10
             pbar.set_postfix(loss=f"{avg:.4f}", lr=f"{lr:.2e}")
             logger.info(f"finetune step={step+1} | loss={avg:.4f}")
+            wandb.log({
+                "finetune/loss": avg,
+                "finetune/lr": lr,
+            }, step=step + 1)
             accum_loss = 0.0
 
     # Merge LoRA weights back for clean inference
@@ -221,6 +233,7 @@ def finetune_instruction(
 
     save_checkpoint(model, None, total_steps, config.checkpoint_dir, filename="sage_finetuned.pt")
     logger.info("Instruction fine-tuning complete. Checkpoint saved as sage_finetuned.pt")
+    wandb.finish()
     return model
 
 

sage/memory.py

@@ -178,22 +178,17 @@ class RAGManager:
 # ===================================================================
 
 DEFAULT_SYSTEM_PROMPT = (
-    "You are a high-quality reasoning assistant model.\n"
-    "You must ONLY learn from high-quality instruction and reasoning datasets.\n"
-    "You must IGNORE any previously trained low-quality or repetitive patterns.\n"
-    "\n"
-    "Training preference rules:\n"
-    "1. Prioritize step-by-step reasoning over short or repetitive answers.\n"
-    "2. Always produce structured logical explanations when solving problems.\n"
-    "3. Avoid repetition, filler words, or looped phrases.\n"
-    "4. Prefer datasets with mathematical reasoning and high-quality instruction.\n"
-    "5. Do not imitate noisy conversational or corrupted text patterns.\n"
-    "6. Always prefer clarity, correctness, and structured reasoning.\n"
-    "\n"
-    "Output behavior goal:\n"
-    "- Think in steps.\n"
-    "- Explain logic clearly.\n"
-    "- Produce final answer only after reasoning."
+    "You are SAGE, a high-quality reasoning assistant. "
+    "Your goal is to provide accurate, structured, and deep logical explanations.\n\n"
+    "CRITICAL GUIDELINES:\n"
+    "1. THINKING PHASE: You must ALWAYS start your response with a <thinking> section. "
+    "In this section, break down the user's request, identify key constraints, and plan your logical steps.\n"
+    "2. RESPONSE PHASE: After completing your internal reasoning, provide your final answer within <response> tags.\n"
+    "3. QUALITY: Prioritize step-by-step mathematical or logical derivation over short answers.\n"
+    "4. NO REPETITION: Avoid filler words or circular logic.\n\n"
+    "RESPONSE TEMPLATE:\n"
+    "<thinking>\n[Step-by-step logic here]\n</thinking>\n"
+    "<response>\n[Final clear answer here]\n</response>"
 )
 
 class ConversationHistory:

sage/model.py

@@ -32,17 +32,30 @@ def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor
     
     return xq_out.type_as(xq), xk_out.type_as(xk)
 
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """Repeat Key/Value heads n_rep times to match number of Query heads."""
+    if n_rep == 1:
+        return x
+    B, T, n_kv_heads, head_dim = x.size()
+    return (
+        x[:, :, :, None, :]
+        .expand(B, T, n_kv_heads, n_rep, head_dim)
+        .reshape(B, T, n_kv_heads * n_rep, head_dim)
+    )
+
 class CausalSelfAttention(nn.Module):
     def __init__(self, config: SageConfig):
         super().__init__()
         self.n_heads = config.n_heads
+        self.n_kv_heads = config.n_kv_heads
+        self.n_rep = self.n_heads // self.n_kv_heads
         self.d_model = config.d_model
         assert self.d_model % self.n_heads == 0
         self.head_dim = self.d_model // self.n_heads
         
-        self.wq = nn.Linear(self.d_model, self.d_model, bias=False)
-        self.wk = nn.Linear(self.d_model, self.d_model, bias=False)
-        self.wv = nn.Linear(self.d_model, self.d_model, bias=False)
+        self.wq = nn.Linear(self.d_model, self.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(self.d_model, self.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(self.d_model, self.n_kv_heads * self.head_dim, bias=False)
         self.wo = nn.Linear(self.d_model, self.d_model, bias=False)
         
         self.resid_dropout = nn.Dropout(config.dropout)
@@ -54,14 +67,11 @@ class CausalSelfAttention(nn.Module):
         x: torch.Tensor, 
         freqs_cis: torch.Tensor, 
         kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
-    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
-        B, T, C = x.size() # batch, seq_len, d_model
-        
         q, k, v = self.wq(x), self.wk(x), self.wv(x)
         
         q = q.view(B, T, self.n_heads, self.head_dim)
-        k = k.view(B, T, self.n_heads, self.head_dim)
-        v = v.view(B, T, self.n_heads, self.head_dim)
+        k = k.view(B, T, self.n_kv_heads, self.head_dim)
+        v = v.view(B, T, self.n_kv_heads, self.head_dim)
         
         q, k = apply_rotary_emb(q, k, freqs_cis)
         
@@ -74,6 +84,10 @@ class CausalSelfAttention(nn.Module):
         else:
             new_kv_cache = None
             
+        # Repeat KV heads to match Q heads (GQA)
+        k = repeat_kv(k, self.n_rep)
+        v = repeat_kv(v, self.n_rep)
+
         # Move heads to correct dimension: (B, n_heads, T, head_dim)
         q = q.transpose(1, 2)
         k = k.transpose(1, 2)

sage/train.py

@@ -11,6 +11,7 @@ import torch
 import torch.nn as nn
 from torch.amp import GradScaler, autocast
 from tqdm import tqdm
+import wandb
 from typing import Optional
 
 from .config import SageConfig
@@ -140,6 +141,13 @@ def train(
     loader = create_dataloader(config, dataset_name=dataset_name, tokenizer=tok)
     data_iter = iter(loader)
 
+    # ------- W&B Logging -------
+    wandb.init(
+        project=config.project_name,
+        name=f"pretrain-{time.strftime('%Y%m%d-%H%M')}",
+        config=config.__dict__,
+    )
+
     # ------- gradient checkpointing (saves VRAM) -------
     base_model = getattr(model, "module", model)
     if hasattr(base_model, "layers"):
@@ -220,6 +228,11 @@ def train(
             logger.info(
                 f"step={step+1} | loss={avg_loss:.4f} | ppl={perplexity:.2f} | lr={lr:.2e}"
             )
+            wandb.log({
+                "train/loss": avg_loss,
+                "train/perplexity": perplexity,
+                "train/lr": lr,
+            }, step=step + 1)
             accum_loss = 0.0
 
         # ------- checkpoint every 100 steps -------
@@ -228,4 +241,5 @@ def train(
             logger.info(f"Checkpoint saved at step {step + 1}")
 
     logger.info("Training complete.")
+    wandb.finish()
     return model

sage_single.py

@@ -30,6 +30,7 @@ from torch.amp import GradScaler, autocast
 from torch.utils.data import IterableDataset, DataLoader
 from tqdm import tqdm
 import tiktoken
+import wandb
 
 __version__ = "1.0.0"
 
@@ -42,6 +43,7 @@ __version__ = "1.0.0"
 class SageConfig:
     d_model: int = 512
     n_heads: int = 8
+    n_kv_heads: int = 4
     n_layers: int = 6
     d_ff: int = 2048
     n_experts: int = 4
@@ -57,6 +59,7 @@ class SageConfig:
     weight_decay: float = 0.01
     max_grad_norm: float = 1.0
     checkpoint_dir: str = "checkpoints"
+    project_name: str = "sage-v2"
 
     @property
     def device(self):
@@ -151,15 +154,22 @@ def apply_rotary_emb(xq, xk, freqs_cis):
     xk_out = torch.view_as_real(xk_ * fc).flatten(3)
     return xq_out.type_as(xq), xk_out.type_as(xk)
 
+def repeat_kv(x, n_rep):
+    if n_rep == 1: return x
+    B, T, n_kv_heads, head_dim = x.size()
+    return x[:, :, :, None, :].expand(B, T, n_kv_heads, n_rep, head_dim).reshape(B, T, n_kv_heads * n_rep, head_dim)
+
 class CausalSelfAttention(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.n_heads = config.n_heads
+        self.n_kv_heads = config.n_kv_heads
+        self.n_rep = self.n_heads // self.n_kv_heads
         self.d_model = config.d_model
         self.head_dim = config.d_model // config.n_heads
-        self.wq = nn.Linear(config.d_model, config.d_model, bias=False)
-        self.wk = nn.Linear(config.d_model, config.d_model, bias=False)
-        self.wv = nn.Linear(config.d_model, config.d_model, bias=False)
+        self.wq = nn.Linear(config.d_model, config.n_heads * self.head_dim, bias=False)
+        self.wk = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=False)
+        self.wv = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=False)
         self.wo = nn.Linear(config.d_model, config.d_model, bias=False)
         self.resid_dropout = nn.Dropout(config.dropout)
 
@@ -167,8 +177,8 @@ class CausalSelfAttention(nn.Module):
         B, T, C = x.size()
         q, k, v = self.wq(x), self.wk(x), self.wv(x)
         q = q.view(B, T, self.n_heads, self.head_dim)
-        k = k.view(B, T, self.n_heads, self.head_dim)
-        v = v.view(B, T, self.n_heads, self.head_dim)
+        k = k.view(B, T, self.n_kv_heads, self.head_dim)
+        v = v.view(B, T, self.n_kv_heads, self.head_dim)
         q, k = apply_rotary_emb(q, k, freqs_cis)
         if kv_cache is not None:
             k = torch.cat([kv_cache[0], k], dim=1)
@@ -176,6 +186,7 @@ class CausalSelfAttention(nn.Module):
             new_kv = (k, v)
         else:
             new_kv = None
+        k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
         q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
         is_causal = kv_cache is None and T > 1
         try:
@@ -293,11 +304,13 @@ def clean_text(text):
     return text.strip()
 
 class StreamingTextDataset(IterableDataset):
-    def __init__(self, dataset_name="roneneldan/TinyStories", split="train", seq_len=512, tokenizer=None, buffer_size=1000, text_field="text"):
+    def __init__(self, dataset_name="HuggingFaceFW/fineweb-edu", split="train", seq_len=512, tokenizer=None, buffer_size=1000, text_field="text"):
         super().__init__()
         self.dataset_name, self.split, self.seq_len = dataset_name, split, seq_len
         self.tokenizer = tokenizer or SageTokenizer()
         self.buffer_size, self.text_field = buffer_size, text_field
+        if "fineweb-edu" in dataset_name.lower(): self.text_field = "text"
+        elif "tinystories" in dataset_name.lower(): self.text_field = "text"
 
     def _tokens(self):
         from datasets import load_dataset
@@ -321,7 +334,7 @@ class StreamingTextDataset(IterableDataset):
         random.shuffle(buf)
         yield from buf
 
-def create_dataloader(config, dataset_name="roneneldan/TinyStories", tokenizer=None):
+def create_dataloader(config, dataset_name="HuggingFaceFW/fineweb-edu", tokenizer=None):
     tok = tokenizer or SageTokenizer()
     ds = StreamingTextDataset(dataset_name=dataset_name, seq_len=config.max_seq_len, tokenizer=tok)
     return DataLoader(ds, batch_size=config.batch_size, num_workers=0, pin_memory=True, drop_last=True)
@@ -362,6 +375,7 @@ def train_model(model, config, total_steps=500, dataset_name="roneneldan/TinySto
     scaler = GradScaler("cuda", enabled=use_amp and amp_dtype == torch.float16)
     loader = create_dataloader(config, dataset_name, tok)
     data_iter = iter(loader)
+    wandb.init(project=config.project_name, name=f"pretrain-{time.strftime('%Y%m%d-%H%M')}", config=config.__dict__)
     model.train()
     accum_loss, t0 = 0.0, time.time()
     pbar = tqdm(range(start_step, total_steps), desc="Training")
@@ -387,11 +401,13 @@ def train_model(model, config, total_steps=500, dataset_name="roneneldan/TinySto
         if (step + 1) % 10 == 0:
             avg = accum_loss / 10
             pbar.set_postfix(loss=f"{avg:.4f}", ppl=f"{math.exp(min(avg,20)):.1f}", lr=f"{lr:.2e}")
+            wandb.log({"train/loss": avg, "train/perplexity": math.exp(min(avg, 20)), "train/lr": lr}, step=step + 1)
             accum_loss = 0.0
         if (step + 1) % 100 == 0:
             save_checkpoint(model, opt, step + 1, config.checkpoint_dir)
     save_checkpoint(model, opt, total_steps, config.checkpoint_dir)
     logger.info("Training complete.")
+    wandb.finish()
     return model
 
 
@@ -517,6 +533,7 @@ def finetune(model, config, samples=None, steps=200, use_lora=True, tokenizer=No
     use_amp = device.type == "cuda"
     amp_dtype = torch.bfloat16 if use_amp and torch.cuda.is_bf16_supported() else torch.float16
     scaler = GradScaler("cuda", enabled=use_amp and amp_dtype == torch.float16)
+    wandb.init(project=config.project_name, name=f"finetune-{time.strftime('%Y%m%d-%H%M')}", config=config.__dict__)
     model.train(); accum = 0.0
     for step in tqdm(range(steps), desc="Fine-tuning"):
         lr = get_lr(step, config, steps)
@@ -537,6 +554,7 @@ def finetune(model, config, samples=None, steps=200, use_lora=True, tokenizer=No
     if use_lora: model = merge_lora(model)
     save_checkpoint(model, None, steps, config.checkpoint_dir, "sage_finetuned.pt")
     logger.info("Fine-tuning complete.")
+    wandb.finish()
     return model